Fix Slurm handling of comma-separated env vars

smartsim/launcher/step/slurmStep.py

@@ -130,13 +130,17 @@ def get_launch_cmd(self):
 
         srun_cmd = [srun, "--output", output, "--error", error, "--job-name", self.name]
 
+
         if self.alloc:
             srun_cmd += ["--jobid", str(self.alloc)]
 
         if self.run_settings.env_vars:
-            env_var_str = self.run_settings.format_env_vars()
+            env_var_str, comma_separated_env_vars = self.run_settings.format_env_vars()
             srun_cmd += ["--export", env_var_str]
 
+            if comma_separated_env_vars:
+                srun_cmd = ["env"] + comma_separated_env_vars + srun_cmd
+
         srun_cmd += self.run_settings.format_run_args()
         srun_cmd += self._build_exe()
         return srun_cmd

smartsim/settings/slurmSettings.py

@@ -180,6 +180,8 @@ def format_env_vars(self):
         # TODO make these overridable by user
         presets = ["PATH", "LD_LIBRARY_PATH", "PYTHONPATH"]
 
+        comma_separated_format_str = []
+
         def add_env_var(var, format_str):
             try:
                 value = os.environ[var]
@@ -196,8 +198,12 @@ def add_env_var(var, format_str):
 
         # add user supplied variables
         for k, v in self.env_vars.items():
-            format_str += "=".join((k, str(v))) + ","
-        return format_str.rstrip(",")
+            if "," in str(v):
+                comma_separated_format_str += ["=".join((k, str(v)))]
+                format_str += k + ","
+            else:
+                format_str += "=".join((k, str(v))) + ","
+        return format_str.rstrip(","), comma_separated_format_str
 
 
 class SbatchSettings(BatchSettings):

tests/on_wlm/test_launch_orc_slurm.py

@@ -64,3 +64,40 @@ def test_launch_slurm_cluster_orc(fileutils, wlmutils):
     exp.stop(orc)
     status = exp.get_status(orc)
     assert all([stat == constants.STATUS_CANCELLED for stat in status])
+
+
+def test_incoming_entities(fileutils, wlmutils):
+    """Mirroring of how SmartSim generates SSKEYIN"""
+
+    exp_name = "test-incoming-entities"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    network_interface = wlmutils.get_test_interface()
+    orc = SlurmOrchestrator(6780, db_nodes=1, batch=False, interface=network_interface)
+    orc.set_path(test_dir)
+
+    sleep = fileutils.get_test_conf_path("sleep.py")
+    sleep_settings = wlmutils.get_run_settings("python", f"{sleep} --time=3")
+
+    sleep_ensemble = exp.create_ensemble("sleep-ensemble",
+                                         run_settings=sleep_settings,
+                                         replicas=2)
+
+    sskeyin_reader = fileutils.get_test_conf_path("incoming_entities_reader.py")
+    sskeyin_reader_settings = wlmutils.get_run_settings("python", f"{sskeyin_reader}")
+    sskeyin_reader_settings.env_vars["NAME_0"] = sleep_ensemble.entities[0].name
+    sskeyin_reader_settings.env_vars["NAME_1"] = sleep_ensemble.entities[1].name
+    sskeyin_reader = exp.create_model("sskeyin_reader", path=test_dir, run_settings=sskeyin_reader_settings)
+    sskeyin_reader.register_incoming_entity(sleep_ensemble.entities[0])
+    sskeyin_reader.register_incoming_entity(sleep_ensemble.entities[1])
+
+    exp.start(orc, block=False)
+    try:
+        exp.start(sskeyin_reader, block=True)
+        assert exp.get_status(sskeyin_reader)[0] == constants.STATUS_COMPLETED
+    except Exception as e:
+        exp.stop(orc)
+        raise e
+
+    exp.stop(orc)

tests/test_configs/incoming_entities_reader.py

@@ -0,0 +1,25 @@
+from os import environ
+
+"""This script checks that the names of multiple
+   incoming entities is correctly passed through
+   a WLM.
+   SSKEYIN is the environment variable passed by
+   SmartSim and it should contain two names:
+   sleep_0 and sleep_1. For this test, we
+   also pass them in as separate variables 
+   NAME_0 and NAME_1. This program will fail if
+   SSKEYIN does not contain NAME_0 and NAME_1 in.
+   the right order. 
+"""
+
+sskeyin = environ["SSKEYIN"]
+name_0 = environ["NAME_0"]
+name_1 = environ["NAME_1"]
+
+smartsim_separator = ","  # this is the current choice
+
+expected_sskeyin = smartsim_separator.join((name_0, name_1))
+
+if sskeyin != expected_sskeyin:
+    raise ValueError(f"SSKEYIN expected to be {expected_sskeyin}, " 
+                     f"but was {sskeyin}")

tests/test_slurm_settings.py

@@ -46,11 +46,13 @@ def test_update_env():
 
 
 def test_format_env():
-    env_vars = {"OMP_NUM_THREADS": 20, "LOGGING": "verbose"}
+    env_vars = {"OMP_NUM_THREADS": 20, "LOGGING": "verbose", "SSKEYIN": "name_0,name_1"}
     settings = SrunSettings("python", env_vars=env_vars)
-    formatted = settings.format_env_vars()
+    formatted, comma_separated_formatted = settings.format_env_vars()
     assert "OMP_NUM_THREADS" in formatted
     assert "LOGGING" in formatted
+    assert "SSKEYIN" in formatted
+    assert "SSKEYIN=name_0,name_1" in comma_separated_formatted
 
 
 # ---- Sbatch ---------------------------------------------------