Merge pull request #68 from isi-vista/allow-pypy-in-conda-script

Allow Conda Job Script to Specify pypy3 rather than python
isi-vista · Nov 13, 2020 · 55f57b7 · 55f57b7
2 parents f623ea3 + cac5fd4
commit 55f57b7
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 3 deletions.
diff --git a/pegasus_wrapper/__init__.py b/pegasus_wrapper/__init__.py
@@ -61,6 +61,7 @@ def run_python_on_parameters(
     resource_request: Optional[ResourceRequest] = None,
     override_conda_config: Optional[CondaConfiguration] = None,
     category: Optional[str] = None,
+    use_pypy: bool = False,
 ) -> DependencyNode:
     """
     Schedule a job to run the given *python_module* on the given *parameters*.
@@ -80,6 +81,7 @@ def run_python_on_parameters(
         resource_request=resource_request,
         override_conda_config=override_conda_config,
         category=category,
+        use_pypy=use_pypy,
     )
 
 

diff --git a/pegasus_wrapper/conda_job_script.py b/pegasus_wrapper/conda_job_script.py
@@ -95,6 +95,7 @@ def generate_shell_script(
         stdout_file: Path,
         ckpt_path: Optional[Path] = None,
         override_conda_config: Optional[CondaConfiguration] = None,
+        python: str = "python",
     ) -> str:
         """
         Returns the content of a shell script to run the given Python entry point
@@ -118,6 +119,7 @@ def generate_shell_script(
             param_file=param_file,
             ckpt_line=f"touch {ckpt_path.absolute()}" if ckpt_path else "",
             stdout_file=stdout_file,
+            python=python,
         )
 
     def write_shell_script_to(
@@ -131,6 +133,7 @@ def write_shell_script_to(
         stdout_file: Optional[Path] = None,
         ckpt_path: Optional[Path] = None,
         override_conda_config: Optional[CondaConfiguration] = None,
+        python: str = "python",
     ) -> None:
         if isinstance(parameters, Path):
             if params_path:
@@ -161,6 +164,7 @@ def write_shell_script_to(
                 working_directory=working_directory,
                 ckpt_path=ckpt_path,
                 override_conda_config=override_conda_config,
+                python=python,
             ),
             encoding="utf-8",
         )
@@ -185,8 +189,8 @@ def write_shell_script_to(
 {spack_lines}
 cd {working_directory}
 echo `which python`
-echo python -m {entry_point} {param_file}
-python -m {entry_point} {param_file} 2>&1 | tee {stdout_file}
+echo {python} -m {entry_point} {param_file}
+{python} -m {entry_point} {param_file} 2>&1 | tee {stdout_file}
 {ckpt_line}
 """
 

diff --git a/pegasus_wrapper/workflow.py b/pegasus_wrapper/workflow.py
@@ -127,6 +127,7 @@ def run_python_on_parameters(
         resource_request: Optional[ResourceRequest] = None,
         override_conda_config: Optional[CondaConfiguration] = None,
         category: Optional[str] = None,
+        use_pypy: bool = False,
     ) -> DependencyNode:
         """
         Schedule a job to run the given *python_module* on the given *parameters*.
@@ -173,6 +174,7 @@ def run_python_on_parameters(
             stdout_file=stdout_path,
             ckpt_path=checkpoint_path,
             override_conda_config=override_conda_config,
+            python="pypy3" if use_pypy else "python",
         )
         script_executable = Executable(
             namespace=self._namespace,

diff --git a/setup.py b/setup.py
@@ -23,7 +23,7 @@
     python_requires="~=3.6",
     install_requires=[
         "importlib-resources==1.4.0",
-	"vistautils>=0.21.0"
+	    "vistautils>=0.21.0",
     ],
     classifiers=[
         "Programming Language :: Python :: 3",