Add --time-limit and --dependency (#48)

* Added support for setting the job time limit from the CLI as well as
the jobs dependency.

* Fix to pass through dependency arg

* Fixed the matrix specific launch args to only be on run commands.
Added support for dependency chaining from the CLI.

* Added support for providing the job-name field from the CLI.

* Removed unused file.

---------

Co-authored-by: Peggy Li <li54@llnl.gov>

Author: bvanessen

hpc_launcher/cli/common_args.py

@@ -91,6 +91,14 @@ def setup_arguments(parser: argparse.ArgumentParser):
 
     group.add_argument("-q", "--queue", default=None, help="Specifies the queue to use")
 
+    group.add_argument(
+        "-t",
+        "--time-limit",
+        type=int,
+        default=None,
+        help="Set a time limit for the job in minutes",
+    )
+
     # Constraints
     group.add_argument(
         "-g",
@@ -225,6 +233,19 @@ def setup_arguments(parser: argparse.ArgumentParser):
         help="Specify the account (or bank) to use fo the job",
     )
 
+    group.add_argument(
+        "--dependency",
+        default=None,
+        help="Specify a scheduler dependency of the submitted job.",
+    )
+
+    group.add_argument(
+        "-J",
+        "--job-name",
+        default=None,
+        help="Specify a name to use fo the job",
+    )
+
     group.add_argument(
         "--reservation",
         default=None,

hpc_launcher/schedulers/flux.py

@@ -73,7 +73,27 @@ def build_scheduler_specific_arguments(
             self.common_launch_args['--env=LD_PRELOAD'] = f'{",".join(self.ld_preloads)}'
 
         if self.time_limit is not None:
-            self.common_launch_args["--time"] = f"{self.time_limit}m"
+            if blocking:
+                self.common_launch_args["--time"] = f"{self.time_limit}m"
+            else:
+                self.submit_only_args["--time"] = f"{self.time_limit}m"
+
+        if self.dependency is not None:
+            self.common_launch_args["--dependency"] = f"{self.dependency}"
+        dependency = self.common_launch_args.get('--dependency', None)
+        if self.override_launch_args and self.override_launch_args.get('--dependency', None):
+            dependency = self.override_launch_args['--dependency']
+        if dependency and not blocking:
+            try:
+                del self.common_launch_args['--dependency']
+            except KeyError:
+                pass
+            try:
+                if self.override_launch_args:
+                    del self.override_launch_args['--dependency']
+            except KeyError:
+                pass
+            self.submit_only_args["--dependency"] = dependency
 
         if self.job_name:
             self.common_launch_args["--job-name"] = f"{self.job_name}"

hpc_launcher/schedulers/scheduler.py

@@ -66,6 +66,8 @@ class Scheduler:
     account: Optional[str] = None
     # The reservation to use for the scheduler
     reservation: Optional[str] = None
+    # Dependency str
+    dependency: Optional[str] = None
     # Hijack preload commands into a scheduler
     ld_preloads: Optional[list[str]] = None
     # Capture the original command so that it can be added to the launch script
@@ -88,6 +90,7 @@ def build_scheduler_specific_arguments(
 
     def build_command_string_and_batch_script(
             self, system: "System", blocking: bool = True, cli_env_only: bool = False,
+            for_launch_cmd: bool = True,
     ) -> (str, list[str]):
         """
         Returns the strings used for a launch command as well as a batch script
@@ -98,6 +101,7 @@ def build_command_string_and_batch_script(
         :param system: The system to use.
         :param blocking: Is the job interactive of blocking
         :param cli_env_only: Append environment variables to CLI not a launch script
+        :param for_launch_cmd:  Some args should not be in both the header and launch cmnd. Ex: flux --dependency=afterany:XXX
         :return: A tuple of (shell script as a string, list of command-line arguments).
         """
         env_vars = system.environment_variables()
@@ -149,11 +153,15 @@ def build_command_string_and_batch_script(
         if not blocking: # Only add batch script header items on non-blocking calls
             prefix = self.batch_script_prefix()
             for k,v in self.submit_only_args.items():
+                if not for_launch_cmd and k == '--dependency':
+                    continue
                 if not v:
                     header.write(f"{prefix} {k}\n")
                 else:
                     header.write(f"{prefix} {k}={v}\n")
             for k,v in self.common_launch_args.items():
+                if not for_launch_cmd and k == '--dependency':
+                    continue
                 if not v:
                     header.write(f"{prefix} {k}\n")
                 else:
@@ -300,7 +308,7 @@ def launcher_script(
         script = ""
         # Launch command only use the cmd_args to construct the shell script to be launched
         (header_lines, cmd_args) = self.build_command_string_and_batch_script(
-            system, blocking, False,
+            system, blocking, False, for_launch_cmd=False
         )
         # For batch jobs add any common args to the internal command
         if not blocking:

hpc_launcher/schedulers/slurm.py

@@ -83,6 +83,23 @@ def build_scheduler_specific_arguments(
         if self.time_limit is not None:
             self.common_launch_args["--time"] = f"{_time_string(self.time_limit)}"
 
+        if self.dependency is not None:
+            self.common_launch_args["--dependency"] = f"{self.dependency}"
+        dependency = self.common_launch_args.get('--dependency', None)
+        if self.override_launch_args and self.override_launch_args.get('--dependency', None):
+            dependency = self.override_launch_args['--dependency']
+        if dependency and not blocking:
+            try:
+                del self.common_launch_args['--dependency']
+            except KeyError:
+                pass
+            try:
+                if self.override_launch_args:
+                    del self.override_launch_args['--dependency']
+            except KeyError:
+                pass
+            self.submit_only_args["--dependency"] = dependency
+
         if self.job_name:
             self.common_launch_args["--job-name"] = f"{self.job_name}"
 

hpc_launcher/systems/lc/cts2.py

@@ -59,8 +59,8 @@ def environment_variables(self) -> list[tuple[str, str]]:
 
     def customize_scheduler(self, scheduler):
         if self.system_name == "matrix" and type(scheduler) is SlurmScheduler:
-            scheduler.common_launch_args["--mpibind"] = "off"
-            scheduler.common_launch_args["--gpu-bind"] = "none"
+            scheduler.run_only_args["--mpibind"] = "off"
+            scheduler.run_only_args["--gpu-bind"] = "none"
 
         return