feat: add pre-scripts and post-scripts args for start-job command

Author: sai-sankalp-reddy

README.md

@@ -134,7 +134,7 @@ hyperpod connect-cluster --cluster-name <cluster-name> [--region <region>] [--na
 This command submits a new training job to the connected SageMaker HyperPod cluster.
 
 ```
-hyperpod start-job --job-name <job-name> [--namespace <namespace>] [--job-kind <kubeflow/PyTorchJob>] [--image <image>] [--command <command>] [--entry-script <script>] [--script-args <arg1 arg2>] [--environment <key=value>] [--pull-policy <Always|IfNotPresent|Never>] [--instance-type <instance-type>] [--node-count <count>] [--tasks-per-node <count>] [--label-selector <key=value>] [--deep-health-check-passed-nodes-only] [--scheduler-type <Kueue SageMaker None>] [--queue-name <queue-name>] [--priority <priority>] [--auto-resume] [--max-retry <count>] [--restart-policy <Always|OnFailure|Never|ExitCode>] [--volumes <volume1,volume2>] [--persistent-volume-claims <claim1:/mount/path,claim2:/mount/path>] [--results-dir <dir>] [--service-account-name <account>]
+hyperpod start-job --job-name <job-name> [--namespace <namespace>] [--job-kind <kubeflow/PyTorchJob>] [--image <image>] [--command <command>] [--entry-script <script>] [--script-args <arg1 arg2>] [--environment <key=value>] [--pull-policy <Always|IfNotPresent|Never>] [--instance-type <instance-type>] [--node-count <count>] [--tasks-per-node <count>] [--label-selector <key=value>] [--deep-health-check-passed-nodes-only] [--scheduler-type <Kueue SageMaker None>] [--queue-name <queue-name>] [--priority <priority>] [--auto-resume] [--max-retry <count>] [--restart-policy <Always|OnFailure|Never|ExitCode>] [--volumes <volume1,volume2>] [--persistent-volume-claims <claim1:/mount/path,claim2:/mount/path>] [--results-dir <dir>] [--service-account-name <account>] [--pre-script <cmd1 cmd2>] [--post-script <cmd1 cmd2>]
 ```
 
 * `job-name` (string) - Required. The name of the job.
@@ -148,6 +148,8 @@ hyperpod start-job --job-name <job-name> [--namespace <namespace>] [--job-kind <
 * `environment` (dict[string, string]) - Optional. The environment variables (key-value pairs) to set in the containers.
 * `node-count` (int) - Required. The number of nodes (instances) to launch the jobs on.
 * `instance-type` (string) - Required. The instance type to launch the job on. Note that the instance types you can use are the available instances within your SageMaker quotas for instances prefixed with `ml`.
+* `pre-script` (string) - Optional. Commands to run before the job starts. Multiple commands should be separated by semicolons.
+* `post-script` (string) - Optional. Commands to run after the job completes. Multiple commands should be separated by semicolons.
 * `tasks-per-node` (int) - Optional. The number of devices to use per instance.
 * `label-selector` (dict[string, list[string]]) - Optional. A dictionary of labels and their values that will override the predefined node selection rules based on the SageMaker HyperPod `node-health-status` label and values. If users provide this field, the CLI will launch the job with this customized label selection.
 * `deep-health-check-passed-nodes-only` (bool) - Optional. If set to `true`, the job will be launched only on nodes that have the `deep-health-check-status` label with the value `passed`.

src/hyperpod_cli/commands/job.py

@@ -430,6 +430,18 @@ def cancel_job(
     help="Optional. Add a temp directory for containers to store data in the hosts."
     " <volume_name>:</host/mount/path>:</container/mount/path>,<volume_name>:</host/mount/path1>:</container/mount/path1>",
 )
+@click.option(
+    "--pre-script",
+    type=click.STRING,
+    required=False,
+    help="Optional. Commands to run before the job starts. Multiple commands should be separated by semicolons.",
+)
+@click.option(
+    "--post-script",
+    type=click.STRING,
+    required=False,
+    help="Optional. Commands to run after the job completes. Multiple commands should be separated by semicolons.",
+)
 @click.option(
     "--recipe",
     type=click.STRING,
@@ -548,6 +560,8 @@ def start_job(
     service_account_name: Optional[str],
     persistent_volume_claims: Optional[str],
     volumes: Optional[str],
+    pre_script: Optional[str],
+    post_script: Optional[str],
     recipe: Optional[str],
     override_parameters: Optional[str],
     debug: bool,
@@ -793,6 +807,8 @@ def start_job(
         label_selector=label_selector,
         max_retry=max_retry,
         deep_health_check_passed_nodes_only=deep_health_check_passed_nodes_only,
+        pre_script=pre_script,
+        post_script=post_script,
     )
     # TODO: Unblock this after fixing customer using EKS cluster.
     console_link = utils.get_cluster_console_url()
@@ -973,7 +989,8 @@ def execute_command(cmd, env=None):
 def start_training_job(recipe, override_parameters, job_name, config_file, launcher_config_path=None, launcher_config_file_name=None,
                       pull_policy=None, restart_policy=None, namespace=None,
                       service_account_name=None, priority_class_name=None, volumes=None, persistent_volume_claims=None,
-                      auto_resume=None, label_selector=None, max_retry=None, deep_health_check_passed_nodes_only=None):
+                      auto_resume=None, label_selector=None, max_retry=None, deep_health_check_passed_nodes_only=None,
+                      pre_script=None, post_script=None):
 
     logger.info(f"recipe: {recipe}, override_parameters: {override_parameters}, job_name: {job_name}, config_file: {config_file}, launcher_config_path: {launcher_config_path}, launcher_config_file_name: {launcher_config_file_name}")
     env = os.environ.copy()
@@ -1035,6 +1052,12 @@ def start_training_job(recipe, override_parameters, job_name, config_file, launc
                 cmd.append(f'+cluster.persistent_volume_claims.{idx}.claimName="{claim_name}"')
                 cmd.append(f'+cluster.persistent_volume_claims.{idx}.mountPath="{mount_path}"')
 
+        if pre_script:
+            cmd.append(f'+cluster.pre_script="{pre_script}"')
+
+        if post_script:
+            cmd.append(f'+cluster.post_script="{post_script}"')
+
         if label_selector:
             cmd.append(f'+cluster.label_selector={label_selector}')
         elif deep_health_check_passed_nodes_only: