fix: Made CPT integ tests dry run for optimize for capacity constraints (#6194)

* fix: Made CPT integ tests dry run for optimize for capacity constraints

* fix: add dryrun to test name for visibility in logs

---------

Co-authored-by: Syed Jafri <syedjfr@amazon.com>

Author: mujtaba1747

sagemaker-train/tests/integ/train/test_cpt_data_mixing_hyperpod.py

@@ -46,8 +46,9 @@
 
 # Test configuration
 REGION = "us-east-1"
-CLUSTER_NAME = "riv-rig"
-INSTANCE_TYPE = "ml.p5.48xlarge"
+CLUSTER_NAME = "pysdk-hp-integ-tests"
+# Using g6 to pass instance count validation. This recipe actually needs p5.48xlarge
+INSTANCE_TYPE = "ml.g6.48xlarge"
 NODE_COUNT = 2
 MODEL_NAME = "nova-textgeneration-micro"
 DATA_PREFIX = "test-cpt-data-mixing-integ"
@@ -124,9 +125,10 @@ def training_resources(sagemaker_session_us_east_1):
     }
 
 
+# TODO: Remove dry-run when capacity is available in future
 @pytest.mark.gpu_intensive
 @pytest.mark.us_east_1
-def test_cpt_trainer_nova_micro_with_data_mixing_hyperpod(
+def test_cpt_trainer_nova_micro_with_data_mixing_hyperpod_dryrun(
     sagemaker_session_us_east_1, training_resources
 ):
     """Test CPTTrainer with Nova Micro model and data mixing on HyperPod.
@@ -162,26 +164,28 @@ def test_cpt_trainer_nova_micro_with_data_mixing_hyperpod(
         base_job_name=f"hp-datamix-m1-{unique_id}",
     )
 
+    dry_run = True
     logger.info("Submitting CPT HyperPod training job with data mixing config...")
     try:
-        job_name = cpt_trainer.train(wait=False)
+        job_name = cpt_trainer.train(wait=False, dry_run=dry_run)
     except ValueError as e:
         if "Failed to download" in str(e) or "Forge subscription" in str(e):
             pytest.skip(
                 "Skipping: account does not have Forge subscription for data mixing recipes."
             )
         raise
 
-    # _train_hyperpod returns the job name as a string
-    assert job_name is not None
-    logger.info(f"HyperPod CPT job submitted: {job_name}")
-
-    # Verify the job exists on the cluster via hyperpod get-job
-    get_job_result = subprocess.run(
-        ["hyperpod", "get-job", "--job-name", job_name],
-        capture_output=True, text=True,
-    )
-    assert get_job_result.returncode == 0, (
-        f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}"
-    )
-    logger.info(f"Verified job '{job_name}' exists on the cluster.")
+    if not dry_run:
+        # _train_hyperpod returns the job name as a string
+        assert job_name is not None
+        logger.info(f"HyperPod CPT job submitted: {job_name}")
+
+        # Verify the job exists on the cluster via hyperpod get-job
+        get_job_result = subprocess.run(
+            ["hyperpod", "get-job", "--job-name", job_name],
+            capture_output=True, text=True,
+        )
+        assert get_job_result.returncode == 0, (
+            f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}"
+        )
+        logger.info(f"Verified job '{job_name}' exists on the cluster.")

sagemaker-train/tests/integ/train/test_cpt_hyperpod.py

@@ -45,8 +45,9 @@
 
 # Test configuration
 REGION = "us-east-1"
-CLUSTER_NAME = "riv-rig"
-INSTANCE_TYPE = "ml.p5.48xlarge"
+CLUSTER_NAME = "pysdk-hp-integ-tests"
+# Using g6 to pass instance count validation. This recipe actually needs p5.48xlarge
+INSTANCE_TYPE = "ml.g6.48xlarge"
 NODE_COUNT = 2
 MODEL_NAME = "nova-textgeneration-micro"
 DATA_PREFIX = "test-cpt-integ"
@@ -123,9 +124,10 @@ def training_resources(sagemaker_session_us_east_1):
     }
 
 
+# TODO: Remove dry-run when capacity is available in future
 @pytest.mark.gpu_intensive
 @pytest.mark.us_east_1
-def test_cpt_trainer_nova_micro_hyperpod(
+def test_cpt_trainer_nova_micro_hyperpod_dryrun(
     sagemaker_session_us_east_1, training_resources
 ):
     """Test CPTTrainer with Nova Micro model on HyperPod (no data mixing).
@@ -152,18 +154,21 @@ def test_cpt_trainer_nova_micro_hyperpod(
     )
 
     logger.info("Submitting CPT HyperPod training job (no data mixing)...")
-    job_name = cpt_trainer.train(wait=False)
 
-    # _train_hyperpod returns the job name as a string
-    assert job_name is not None
-    logger.info(f"HyperPod CPT job submitted: {job_name}")
-
-    # Verify the job exists on the cluster via hyperpod get-job
-    get_job_result = subprocess.run(
-        ["hyperpod", "get-job", "--job-name", job_name],
-        capture_output=True, text=True,
-    )
-    assert get_job_result.returncode == 0, (
-        f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}"
-    )
-    logger.info(f"Verified job '{job_name}' exists on the cluster.")
+    dry_run = True
+    job_name = cpt_trainer.train(wait=False, dry_run=dry_run)
+
+    if not dry_run:
+        # _train_hyperpod returns the job name as a string
+        assert job_name is not None
+        logger.info(f"HyperPod CPT job submitted: {job_name}")
+
+        # Verify the job exists on the cluster via hyperpod get-job
+        get_job_result = subprocess.run(
+            ["hyperpod", "get-job", "--job-name", job_name],
+            capture_output=True, text=True,
+        )
+        assert get_job_result.returncode == 0, (
+            f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}"
+        )
+        logger.info(f"Verified job '{job_name}' exists on the cluster.")