Merge pull request #603 from xylar/add-job-script-to-steps

Add job script to steps

compass/setup.py

@@ -161,10 +161,7 @@ def setup_cases(tests=None, numbers=None, config_file=None, machine=None,  # noq
     with open(pickle_file, 'wb') as handle:
         pickle.dump(test_suite, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
-    if 'LOAD_COMPASS_ENV' in os.environ:
-        script_filename = os.environ['LOAD_COMPASS_ENV']
-        # make a symlink to the script for loading the compass conda env.
-        symlink(script_filename, os.path.join(work_dir, 'load_compass_env.sh'))
+    _symlink_load_script(work_dir)
 
     max_cores, max_of_min_cores = _get_required_cores(test_cases)
 
@@ -311,6 +308,14 @@ def setup_case(path, test_case, config_file, machine, work_dir, baseline_dir,
             pickle.dump((test_case, step), handle,
                         protocol=pickle.HIGHEST_PROTOCOL)
 
+        _symlink_load_script(step.work_dir)
+
+        if machine is not None:
+            cores = step.cpus_per_task * step.ntasks
+            min_cores = step.min_cpus_per_task * step.min_tasks
+            write_job_script(config, machine, cores, min_cores,
+                             step.work_dir)
+
     # pickle the test case and step for use at runtime
     pickle_filename = os.path.join(test_case.work_dir, 'test_case.pickle')
     with open(pickle_filename, 'wb') as handle:
@@ -319,11 +324,7 @@ def setup_case(path, test_case, config_file, machine, work_dir, baseline_dir,
                       'work_dir': test_case.work_dir}
         pickle.dump(test_suite, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
-    if 'LOAD_COMPASS_ENV' in os.environ:
-        script_filename = os.environ['LOAD_COMPASS_ENV']
-        # make a symlink to the script for loading the compass conda env.
-        symlink(script_filename, os.path.join(test_case_dir,
-                                              'load_compass_env.sh'))
+    _symlink_load_script(test_case_dir)
 
     if machine is not None:
         max_cores, max_of_min_cores = _get_required_cores({path: test_case})
@@ -457,3 +458,11 @@ def _get_basic_config(config_file, machine, mpas_model_path, mpas_core):
         config.set('paths', 'mpas_model', mpas_model_path, user=True)
 
     return config
+
+
+def _symlink_load_script(work_dir):
+    """ make a symlink to the script for loading the compass conda env. """
+    if 'LOAD_COMPASS_ENV' in os.environ:
+        script_filename = os.environ['LOAD_COMPASS_ENV']
+        symlink(script_filename,
+                os.path.join(work_dir, 'load_compass_env.sh'))