From f8beddb298d22b787012664c0023bf873316a604 Mon Sep 17 00:00:00 2001
From: Sai Sankalp Reddy TV <isankalp@amazon.com>
Date: Mon, 10 Mar 2025 19:46:52 +0530
Subject: [PATCH 1/5] feat: add pre-scripts and post-scripts args for start-job
 command

---
 README.md                        |  4 +++-
 src/hyperpod_cli/commands/job.py | 25 ++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index d4e63793..816c3b14 100644
--- a/README.md
+++ b/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`.
diff --git a/src/hyperpod_cli/commands/job.py b/src/hyperpod_cli/commands/job.py
index bf41b6fb..89bf9f29 100644
--- a/src/hyperpod_cli/commands/job.py
+++ b/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:

From eca509a9eabf04da1f6b31a93298543e530a8230 Mon Sep 17 00:00:00 2001
From: Sai Sankalp Reddy TV <isankalp@amazon.com>
Date: Mon, 10 Mar 2025 21:09:03 +0530
Subject: [PATCH 2/5] fix: corrected delimiter in documentation and updatetd
 config file instead of using override-parameters

---
 README.md                        |  4 ++--
 src/hyperpod_cli/commands/job.py | 30 +++++++++++++++++++-----------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 816c3b14..a2dd358b 100644
--- a/README.md
+++ b/README.md
@@ -148,8 +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.
+* `pre-script` (list[string]) - Optional. Commands to run before the job starts. Multiple commands should be separated by comma.
+* `post-script` (list[string]) - Optional. Commands to run after the job completes. Multiple commands should be separated by comma.
 * `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`.
diff --git a/src/hyperpod_cli/commands/job.py b/src/hyperpod_cli/commands/job.py
index 89bf9f29..bf65a08b 100644
--- a/src/hyperpod_cli/commands/job.py
+++ b/src/hyperpod_cli/commands/job.py
@@ -720,6 +720,23 @@ def start_job(
                 custom_labels[KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY] = priority
                 priority = None
 
+            # Handle pre_script
+            if pre_script:
+                _override_or_remove(
+                    config["cluster"]["cluster_config"], 
+                    "pre_script", 
+                    pre_script.split(',')
+                )
+            
+            # Handle post_script
+            if post_script:
+                _override_or_remove(
+                    config["cluster"]["cluster_config"], 
+                    "post_script", 
+                    post_script.split(',')
+                )
+
+
             _override_or_remove(
                 config["cluster"]["cluster_config"],
                 "custom_labels",
@@ -806,9 +823,7 @@ def start_job(
         auto_resume=auto_resume,
         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,
+        deep_health_check_passed_nodes_only=deep_health_check_passed_nodes_only
     )
     # TODO: Unblock this after fixing customer using EKS cluster.
     console_link = utils.get_cluster_console_url()
@@ -989,8 +1004,7 @@ 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,
-                      pre_script=None, post_script=None):
+                      auto_resume=None, label_selector=None, max_retry=None, deep_health_check_passed_nodes_only=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()
@@ -1052,12 +1066,6 @@ 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:

From 441cd0f2ca2993ca6e6c0039dc9d1a67aebf6974 Mon Sep 17 00:00:00 2001
From: Sai Sankalp Reddy TV <isankalp@amazon.com>
Date: Mon, 10 Mar 2025 21:49:59 +0530
Subject: [PATCH 3/5] test: added unit test cases

---
 test/unit_tests/test_job.py | 45 +++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/test/unit_tests/test_job.py b/test/unit_tests/test_job.py
index 098aff85..6466f717 100644
--- a/test/unit_tests/test_job.py
+++ b/test/unit_tests/test_job.py
@@ -723,6 +723,51 @@ def test_start_job_with_cli_args_label_selection_not_json_str(
         )
         self.assertEqual(result.exit_code, 1)
 
+    @mock.patch("yaml.dump")
+    @mock.patch("hyperpod_cli.clients.kubernetes_client.KubernetesClient.__new__")
+    @mock.patch("hyperpod_cli.commands.job.JobValidator")
+    @mock.patch("boto3.Session")
+    def test_start_job_with_cli_args_pre_script_and_post_script(
+        self,
+        mock_boto3,
+        mock_validator_cls,
+        mock_kubernetes_client,
+        mock_yaml_dump,
+    ):
+        mock_validator = mock_validator_cls.return_value
+        mock_validator.validate_aws_credential.return_value = True
+        mock_kubernetes_client.get_current_context_namespace.return_value = "kubeflow"
+        mock_yaml_dump.return_value = None
+        result = self.runner.invoke(
+            start_job,
+            [
+                "--job-name",
+                "test-job",
+                "--instance-type",
+                "ml.c5.xlarge",
+                "--image",
+                "pytorch:1.9.0-cuda11.1-cudnn8-runtime",
+                "--node-count",
+                "2",
+                "--label-selector",
+                "{NonJsonStr",
+                "--entry-script",
+                "/opt/train/src/train.py",
+                "--pre-script",
+                "echo 'test', echo 'test 1'",
+                "--post-script",
+                "echo 'test 1', echo 'test 2'"
+            ],
+        )
+
+        # Assert that yaml.dump was called with the correct configuration
+        mock_yaml_dump.assert_called_once()
+        call_args = mock_yaml_dump.call_args[0]
+        self.assertEqual(call_args[0]['cluster']['cluster_config']['pre_script'], ["echo 'test'", " echo 'test 1'"])
+        self.assertEqual(call_args[0]['cluster']['cluster_config']['post_script'], ["echo 'test 1'", " echo 'test 2'"])
+        
+        self.assertEqual(result.exit_code, 1)
+
     @mock.patch("yaml.dump")
     @mock.patch("hyperpod_cli.clients.kubernetes_client.KubernetesClient.__new__")
     @mock.patch("hyperpod_cli.commands.job.JobValidator")

From beb20a376ecbbfeb0b412e22a23b40a5065372ec Mon Sep 17 00:00:00 2001
From: Sai Sankalp Reddy TV <isankalp@amazon.com>
Date: Wed, 19 Mar 2025 22:50:56 +0530
Subject: [PATCH 4/5] fix: updated config field used to training_cfg

---
 src/hyperpod_cli/commands/job.py | 4 ++--
 test/unit_tests/test_job.py      | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/hyperpod_cli/commands/job.py b/src/hyperpod_cli/commands/job.py
index bf65a08b..7cb52093 100644
--- a/src/hyperpod_cli/commands/job.py
+++ b/src/hyperpod_cli/commands/job.py
@@ -723,7 +723,7 @@ def start_job(
             # Handle pre_script
             if pre_script:
                 _override_or_remove(
-                    config["cluster"]["cluster_config"], 
+                    config["training_cfg"],
                     "pre_script", 
                     pre_script.split(',')
                 )
@@ -731,7 +731,7 @@ def start_job(
             # Handle post_script
             if post_script:
                 _override_or_remove(
-                    config["cluster"]["cluster_config"], 
+                    config["training_cfg"],
                     "post_script", 
                     post_script.split(',')
                 )
diff --git a/test/unit_tests/test_job.py b/test/unit_tests/test_job.py
index 6466f717..d42dac04 100644
--- a/test/unit_tests/test_job.py
+++ b/test/unit_tests/test_job.py
@@ -763,8 +763,8 @@ def test_start_job_with_cli_args_pre_script_and_post_script(
         # Assert that yaml.dump was called with the correct configuration
         mock_yaml_dump.assert_called_once()
         call_args = mock_yaml_dump.call_args[0]
-        self.assertEqual(call_args[0]['cluster']['cluster_config']['pre_script'], ["echo 'test'", " echo 'test 1'"])
-        self.assertEqual(call_args[0]['cluster']['cluster_config']['post_script'], ["echo 'test 1'", " echo 'test 2'"])
+        self.assertEqual(call_args[0]['training_cfg']['pre_script'], ["echo 'test'", " echo 'test 1'"])
+        self.assertEqual(call_args[0]['training_cfg']['post_script'], ["echo 'test 1'", " echo 'test 2'"])
         
         self.assertEqual(result.exit_code, 1)
 

From 03e4d5b1ccff807727ffb6716f2d7df4cef42b9e Mon Sep 17 00:00:00 2001
From: Sai Sankalp Reddy TV <isankalp@amazon.com>
Date: Thu, 20 Mar 2025 20:30:51 +0530
Subject: [PATCH 5/5] chore: found bug in cli adding extra parm in test to
 unblock PR

---
 test/unit_tests/test_job.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/unit_tests/test_job.py b/test/unit_tests/test_job.py
index 4773e2b3..e4aca64b 100644
--- a/test/unit_tests/test_job.py
+++ b/test/unit_tests/test_job.py
@@ -919,7 +919,9 @@ def test_start_job_with_cli_args_pre_script_and_post_script(
                 "--pre-script",
                 "echo 'test', echo 'test 1'",
                 "--post-script",
-                "echo 'test 1', echo 'test 2'"
+                "echo 'test 1', echo 'test 2'",
+                "--label-selector", 
+                '{"preferred": {"node.kubernetes.io/instance-type": ["ml.c5.xlarge"]}}'
             ],
         )