min viable remote submission function

Author: paxtonfitzpatrick

cluster_scripts/config.py

@@ -8,7 +8,6 @@
 job_config['datadir'] = opj(job_config['startir'], 'data')
 job_config['workingdir'] = opj(job_config['startir'], 'scripts')
 job_config['template'] = opj(dirname(realpath(__file__)), 'run_job_cluster.sh')
-
 job_config['scriptdir'] = opj(job_config['workingdir'], 'scripts')
 job_config['lockdir'] = opj(job_config['workingdir'], 'locks')
 

remote_submit.py

@@ -2,11 +2,11 @@
 from os.path import dirname, realpath, join as opj
 from spurplus import connect_with_retries
 from .upload_scripts import upload_scripts
-from ._helpers import attempt_load_config
+from ._helpers import attempt_load_config, parse_config
 from .cluster_scripts.config import job_config
 
 
-def remote_submit(config_path=None, sync_changes=False, view_output=False):
+def remote_submit(config_path=None, sync_changes=False, await_output=False):
     """
     main function that handles submitting jobs on the cluster from your local
     machine
@@ -16,7 +16,7 @@ def remote_submit(config_path=None, sync_changes=False, view_output=False):
     configs/template_config.txt, you can simply leave this empty
     :param sync_changes: (bool, default: False) if True, upload any local
     changes to cluster scripts before submitting jobs
-    :param view_output: (bool, default: False) if True, keep the connection with
+    :param await_output: (bool, default: False) if True, keep the connection with
     the remote open until your submit script is finished creating jobs.
     Otherwise, terminate the connection after callin the submit script and allow
     job submission to happen in the background.
@@ -25,18 +25,46 @@ def remote_submit(config_path=None, sync_changes=False, view_output=False):
     the ssh connection may fail before job submission is finished
     :return: None (other than some hopefully some results, eventually!)
     """
-    config = attempt_load_config()
+    if config_path is None:
+        config = attempt_load_config()
+    else:
+        config = parse_config(config_path)
+
     hostname = config['hostname']
     username = config['username']
     password = config['password']
+    confirm_overwrite = config['confirm_overwrite_on_upload']
 
     modules = job_config['modules']
-    submit_cmd = job_config['cmd_wrapper']
+    env_type = job_config['env_type']
+    env_name = job_config['env_name']
+    submit_cmd_wrapper = job_config['cmd_wrapper']
+    # TODO: ability to handle custom-named submission script
+    submit_script_path = opj(job_config['workingdir'], 'submit.py')
 
+    # pre-submission commands to be concatenated and run together in remote shell
+    remote_cmds = ['sh', '-c']
+    # command for loading module(s)
+    module_load_cmd = f'module load {modules}'
+    # command activating virtual environment
+    if env_type == 'conda':
+        activate_cmd = 'source activate'
+    else:
+        # TODO: add commands for venv & virtualenv activation
+        raise ValueError("Only conda environments are currently supported")
+    env_activate_cmd = f'{activate_cmd} {env_name}'
+    # command for calling submit script
+    submit_cmd = f'{submit_cmd_wrapper} {submit_script_path}'
+
+    full_submission_cmd = ' && '.join([
+        module_load_cmd,
+        env_activate_cmd,
+        submit_cmd
+    ])
+
+    remote_cmds.append(full_submission_cmd)
 
-    confirm_overwrite = config['confirm_overwrite_on_upload']
 
-    script_dir = opj(dirname(realpath(__file__)), 'cluster_scripts')
 
 
     with connect_with_retries(
@@ -45,10 +73,16 @@ def remote_submit(config_path=None, sync_changes=False, view_output=False):
             password=password
     ) as cluster:
         if sync_changes:
+            script_dir = opj(dirname(realpath(__file__)), 'cluster_scripts')
             upload_scripts(
                 cluster,
                 script_dir,
                 job_config,
                 confirm_overwrite=confirm_overwrite
             )
 
+        if await_output:
+            output = cluster.check_output(remote_cmds)
+            print(output)
+        else:
+            cluster.run(remote_cmds)