Updated to version 0.9.5

* Now --config in slaunch and py-srun can take multiple values and be stacked.
* Fixed the order of stacking of slurm_params and profile params.

Author: maharshi95

bin/py-srun

@@ -6,11 +6,12 @@ import subprocess
 import argparse
 
 from submititnow.umiacs.handlers import profile_handlers
+from rich import print as rich_print
 
 parser = argparse.ArgumentParser()
-parser.add_argument("config", type=str)
-parser.add_argument("shell", nargs="+", default="zsh")
-args = parser.parse_args()
+parser.add_argument("--config", default=[], action="append")
+parser.add_argument("--job-name", default="dev")
+parser.add_argument("shell", nargs="+", default=["bash"])
 
 
 def removeprefix(var: str, prefix: str):
@@ -30,17 +31,41 @@ def load_config(config_filename: str):
     }
 
 
-cmd_args = load_config(args.config)
+def parse_arguments():
+    """Returns command args for SRUN and the downstream command to execute on the execution node."""
+    args, downstream_args = parser.parse_known_args()
 
+    # Create args dict from downstream args
+    cli_args = {}
+    for arg in downstream_args:
+        if arg.startswith("--"):
+            key, value = arg[2:].split("=")
+            cli_args[key] = value
 
-# Make Bash command
-cmd = "srun"
-for key, value in cmd_args.items():
-    cmd += f" --{key}={value}"
-cmd += " --job-name=llms"
-shell_cmd = " ".join(args.shell)
-cmd += f" --pty {shell_cmd}"
+    # Populate arguments from config files
+    config_args = {}
+    for config_filename in args.config:
+        config_args = {**config_args, **load_config(config_filename)}
 
-print(cmd)
+    # Merge arguments from config files and cli
+    cmd_args = {**config_args, **cli_args, "job-name": args.job_name}
 
-subprocess.run(cmd, shell=True)
+    return cmd_args, " ".join(args.shell)
+
+
+def make_shell_command(cmd_args: dict, downstream_cmd: str):
+    """Returns the full SRUN command to execute on the shell."""
+    shell_cmd = "srun"
+    for key, value in cmd_args.items():
+        shell_cmd += f" --{key}={value}"
+    shell_cmd += f" --pty {downstream_cmd}"
+    return shell_cmd
+
+
+if __name__ == "__main__":
+    cmd_args, downstream_cmd = parse_arguments()
+    cmd = make_shell_command(cmd_args, downstream_cmd)
+
+    rich_print(f"[bold]{cmd}")
+
+    subprocess.run(cmd, shell=True)

setup.py

@@ -7,7 +7,7 @@
 
 setuptools.setup(
     name="submititnow",
-    version="0.9.4.1",
+    version="0.9.5",
     author="Maharshi Gor",
     author_email="maharshigor@gmail.com",
     description="A package to make submitit easier to use",

submititnow/options.py

@@ -36,7 +36,8 @@ def add_slurm_arguments(parser: argparse.ArgumentParser):
     slurm_group = parser.add_argument_group("SLURM parameters")
     slurm_group.add_argument(
         "--config",
-        default=None,
+        default=[],
+        action="append",
         help="SubmititNow config file for SLURM.",
         dest="slurm_config",
     )
@@ -134,11 +135,16 @@ def load_slurm_config(config_filename: str) -> Dict[str, Any]:
 
 def get_slurm_params(args: argparse.Namespace) -> Dict[str, Any]:
     # Grabs all SLURM arguments from the parser that are explicitly set to a value
-    slurm_args = {
+    slurm_cli_args = {
         k: v for k, v in vars(args).items() if k.startswith("slurm_") and v is not None
     }
-    if slurm_args.get("slurm_config") is not None:
-        config_filename = slurm_args.pop("slurm_config")
-        default_args = load_slurm_config(config_filename)
-        slurm_args = {**default_args, **slurm_args}
-    return slurm_args
+    slurm_config_args = {}
+
+    config_filenames = slurm_cli_args.pop("slurm_config")
+    for config_filename in config_filenames:
+        slurm_config_args = {
+            **slurm_config_args,
+            **load_slurm_config(config_filename),
+        }
+    slurm_cli_args = {**slurm_config_args, **slurm_cli_args}
+    return slurm_cli_args

submititnow/umiacs/handlers.py

@@ -3,46 +3,50 @@
 
 def clip_profile_handler(slurm_params: Dict[str, Any]):
     return {
-        **slurm_params,
         "slurm_account": "clip",
         "slurm_partition": "clip",
+        **slurm_params,
     }
 
 
 def scavenger_profile_handler(slurm_params: Dict[str, Any]):
     return {
-        **slurm_params,
         "slurm_account": "scavenger",
         "slurm_partition": "scavenger",
         "slurm_qos": "scavenger",
+        **slurm_params,
     }
 
+
 def cml_zhou_profile_handler(slurm_params: Dict[str, Any]):
     return {
-        **slurm_params,
         "slurm_account": "cml-zhou",
         "slurm_partition": "cml-dpart",
+        **slurm_params,
     }
-    
+
+
 def cml_profile_handler(slurm_params: Dict[str, Any]):
     return {
-        **slurm_params,
         "slurm_account": "cml",
         "slurm_partition": "cml-dpart",
+        **slurm_params,
     }
 
+
 def cml_scavenger_profile_handler(slurm_params: Dict[str, Any]):
     return {
-        **slurm_params,
         "slurm_account": "cml-scavenger",
         "slurm_partition": "cml-scavenger",
         "slurm_qos": "cml-scavenger",
+        **slurm_params,
     }
-    
+
+
 profile_handlers = {
     "clip": clip_profile_handler,
     "scavenger": scavenger_profile_handler,
     "cml": cml_profile_handler,
     "cml-zhou": cml_zhou_profile_handler,
     "cml-scavenger": cml_scavenger_profile_handler,
-}
+}