From 1d6838038951346cd417c0213b22183a259fe541 Mon Sep 17 00:00:00 2001
From: Lucas Jia <lucasjia@amazon.com>
Date: Fri, 9 Oct 2026 15:15:43 -0700
Subject: [PATCH] test(train): Bound capacity wait in SMTJ SFT integ test

test_sft_trainer_serverful_smtj trains on one ml.g6.12xlarge in
us-east-1. When there is no capacity the job stays in Pending, and the
test polled for up to 3 hours, which held the GPU integ build until
CodeBuild timed out. The job was then left queued behind the killed
build and trained hours later with nobody waiting for it.

Poll with the shared bounded capacity wait from tests/integ/capacity.py:
after 90 minutes in Pending the job is stopped and the test is skipped,
and the job is stopped whenever the wait ends with it still running.
The total wait is 2 hours, covering the capacity wait plus the job's
~21 minute run time, so the test stays well inside the build timeout.
---
 .../train/test_sft_trainer_serverful_smtj.py  | 46 ++++++++++---------
 1 file changed, 25 insertions(+), 21 deletions(-)

diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py b/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py
index 8ff87dfc3c..5324cb9ce9 100644
--- a/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py
+++ b/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py
@@ -41,9 +41,20 @@
 from sagemaker.train.common import TrainingType
 from sagemaker.core.training.configs import TrainingJobCompute
 
+from ..capacity import wait_for_training_job_with_capacity_timeout
+
 logger = logging.getLogger(__name__)
 
 # Test configuration
+SMTJ_INSTANCE_TYPE = "ml.g6.12xlarge"
+# In us-east-1, SMTJ jobs for this test waited 0 min to 48 h for an ml.g6.12xlarge
+# (median ~2 h); once the instance is acquired, download + training + upload take
+# <= ~21 min. A job still waiting after 90 min is treated as a capacity shortage.
+# The total bound covers that wait plus the job's own run time, and keeps the test
+# well inside the CI build timeout.
+SMTJ_PENDING_TIMEOUT_SECONDS = 90 * 60
+SMTJ_MAX_WAIT_SECONDS = 2 * 3600
+
 DATA_PREFIX = "sft-smtj-integ"
 DATA_S3_KEY = f"{DATA_PREFIX}/sft_smtj_sample_data.jsonl"
 
@@ -116,7 +127,7 @@ def test_sft_trainer_serverful_smtj(sagemaker_session_us_east_1, training_resour
         training_dataset=training_resources["training_dataset"],
         s3_output_path=training_resources["s3_output_path"],
         compute=TrainingJobCompute(
-            instance_type="ml.g6.12xlarge",
+            instance_type=SMTJ_INSTANCE_TYPE,
             instance_count=1,
         ),
         sagemaker_session=sagemaker_session_us_east_1,
@@ -140,27 +151,20 @@ def test_sft_trainer_serverful_smtj(sagemaker_session_us_east_1, training_resour
     assert training_job.training_job_name is not None
     logger.info(f"Training job submitted: {training_job.training_job_name}")
 
-    # Poll for completion
-    max_wait_time = 10800  # 3 hour timeout (Nova training can be slow)
-    poll_interval = 60
-    start_time = time.time()
-
-    while time.time() - start_time < max_wait_time:
-        training_job.refresh()
-        status = training_job.training_job_status
-
-        if status in ["Completed", "Failed", "Stopped"]:
-            break
-
-        logger.info(
-            f"Job {training_job.training_job_name} status: {status} "
-            f"({int(time.time() - start_time)}s elapsed)"
-        )
-        time.sleep(poll_interval)
+    # Poll for completion with a bounded wait for capacity. On a capacity shortage
+    # the job is stopped and the test is skipped, so it neither holds the build
+    # until CodeBuild times out nor leaves the job queued behind a killed build.
+    status = wait_for_training_job_with_capacity_timeout(
+        training_job,
+        max_wait_time=SMTJ_MAX_WAIT_SECONDS,
+        poll_interval=60,
+        pending_timeout=SMTJ_PENDING_TIMEOUT_SECONDS,
+        capacity_label=SMTJ_INSTANCE_TYPE,
+    )
 
-    assert training_job.training_job_status == "Completed", (
-        f"Training job {training_job.training_job_name} ended with status: "
-        f"{training_job.training_job_status}"
+    assert status == "Completed", (
+        f"Training job {training_job.training_job_name} ended with status {status}: "
+        f"{training_job.failure_reason}"
     )
     logger.info(f"Training job completed successfully: {training_job.training_job_name}")