Merge pull request #3 from outerbounds/custom-python-path

use `/usr/bin/python3` by default

README.md

@@ -85,6 +85,13 @@ PS -- If you are on the [Outerbounds](https://outerbounds.com/) platform, the au
 
 - The extension runs workloads via shell scripts and `sbatch` in a linux native environment
     - i.e. the workloads are NOT run inside docker containers
-    - As such, the compute instances should not have `python2` installed and both `python` and `python3` should refer to a python version above 3.8 preferrably.
+    - As such, the compute instances should have `python3` installed (above 3.8 preferrably)
+    - If the default `python` points to `python2`, one can use the `path_to_python3` argument of the decorator i.e.
+
+```py
+@slurm(
+    path_to_python3="/usr/bin/python3",
+)
+```
 
 ### Fin.

metaflow_extensions/slurm_ext/plugins/slurm/slurm.py

@@ -57,6 +57,7 @@ def __init__(
         self.datastore = datastore
         self.metadata = metadata
         self.environment = environment
+        self.python3_executable = slurm_access_params.pop("path_to_python3", None) or "/usr/bin/python3"
         self.slurm_client = SlurmClient(**slurm_access_params)
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
@@ -115,6 +116,7 @@ def _command(self, environment, code_package_url, step_name, step_cmds, task_spe
             '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
             % cmd_str
         )
+        cmd_str = 'python() { %s \\"$@\\"; }; export -f python && %s' % (self.python3_executable, cmd_str)
 
         return shlex.split('bash -c "%s"' % cmd_str)
 

metaflow_extensions/slurm_ext/plugins/slurm/slurm_cli.py

@@ -41,6 +41,7 @@ def slurm():
 @click.option("--ssh-key-file", help="SSH key file for login node for Slurm.")
 @click.option("--cert-file", help="Certificate file for login node for Slurm.")
 @click.option("--remote-workdir", help="Remote working directory for Slurm.")
+@click.option("--path-to-python3", help="Path to python3 executable on Slurm.")
 @click.option(
     "--cleanup", help="Cleanup created artifacts on Slurm.", is_flag=True, default=False
 )
@@ -88,6 +89,7 @@ def step(
     ssh_key_file=None,
     cert_file=None,
     remote_workdir=None,
+    path_to_python3=None,
     cleanup=False,
     partition=None,
     nodes=None,
@@ -206,6 +208,7 @@ def _sync_metadata():
                 "ssh_key_file": ssh_key_file,
                 "cert_file": cert_file,
                 "remote_workdir": remote_workdir,
+                "path_to_python3": path_to_python3,
                 "cleanup": cleanup,
             },
         )

metaflow_extensions/slurm_ext/plugins/slurm/slurm_decorator.py

@@ -41,6 +41,8 @@ class SlurmDecorator(StepDecorator):
     remote_workdir : str, optional, default None
         Working directory on the remote Slurm cluster where job files will be stored.
         If not specified, defaults to METAFLOW_SLURM_REMOTE_WORKDIR environment variable.
+    path_to_python3: str, optional, default None
+        The path to the python3 executable on the Slurm cluster.
     cleanup : bool, default False
         If True, cleanup created artifacts on Slurm after job completion. This doesn't
         delete the log files.
@@ -63,6 +65,7 @@ class SlurmDecorator(StepDecorator):
         "ssh_key_file": None,
         "cert_file": None,
         "remote_workdir": None,
+        "path_to_python3": None,
         "cleanup": False,
         "partition": None,
         "nodes": None,

metaflow_extensions/slurm_ext/toplevel/toplevel.py

@@ -1,2 +1,2 @@
 __mf_extensions__ = "slurm"
-__version__ = "0.0.1"
+__version__ = "0.0.2"

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_namespace_packages
 
-version = "0.0.1"
+version = "0.0.2"
 
 setup(
     name="metaflow-slurm",