test: fail fast on Trainium capacity queueing in neuron slow test

test_gated_model_training_v2_neuron submits a training job on
ml.trn1.32xlarge, which can sit in "Pending - waiting for capacity"
for hours without ever failing, hanging the slow-tests build until it
times out and masking other tests' results.

Add fit_estimator_with_capacity_xfail() helper that submits the job
non-blocking, polls its secondary status, and if it stays queued for
capacity beyond a timeout (default 30m), stops the job and raises a
CapacityError so the existing x_fail_if_ice decorator marks the test
xfail instead of hanging. Capacity is a transient region-level
condition, not an SDK defect, and quota is sufficient (trn1.32xlarge
training quota=20, test needs 1).

X-AI-Prompt: Implement non-blocking fit with capacity-queue timeout that triggers xfail via x_fail_if_ice for the neuron slow test
X-AI-Tool: kiro-cli

Author: lucasjia-aws

tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py

@@ -28,6 +28,7 @@
     get_sm_session,
     get_training_dataset_for_model_and_version,
     x_fail_if_ice,
+    fit_estimator_with_capacity_xfail,
 )
 
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
@@ -221,11 +222,12 @@ def test_gated_model_training_v2_neuron(setup):
     )
 
     # uses ml.trn1.32xlarge instance
-    estimator.fit(
+    fit_estimator_with_capacity_xfail(
+        estimator,
         {
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, '*')}",
-        }
+        },
     )
 
     # uses ml.inf2.xlarge instance

tests/integ/sagemaker/jumpstart/utils.py

@@ -92,6 +92,82 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+# Default time (in seconds) we allow a training job to sit in
+# "Pending - waiting for capacity" before treating it as a capacity shortage.
+# Scarce accelerators (e.g. Trainium ml.trn1.32xlarge) can queue for hours
+# without ever failing, which would otherwise exhaust the build timeout.
+DEFAULT_CAPACITY_WAIT_TIMEOUT_SECONDS = 30 * 60
+_CAPACITY_POLL_INTERVAL_SECONDS = 60
+
+
+def fit_estimator_with_capacity_xfail(
+    estimator,
+    inputs,
+    capacity_wait_timeout_seconds: int = DEFAULT_CAPACITY_WAIT_TIMEOUT_SECONDS,
+):
+    """Submit a training job and wait for it, but fail fast on capacity queueing.
+
+    Unlike ``estimator.fit()`` (which blocks until the job reaches a terminal
+    state), this submits the job non-blocking and then polls its secondary
+    status. Insufficient capacity for scarce instances (e.g. Trainium) does not
+    surface as an exception; SageMaker simply parks the job in
+    ``Pending`` / "waiting for capacity" indefinitely. If the job stays in that
+    state longer than ``capacity_wait_timeout_seconds`` before training starts,
+    we raise a ``RuntimeError`` whose message contains ``"CapacityError"`` so the
+    ``x_fail_if_ice`` decorator marks the test as xfail instead of letting it hang
+    until the build times out.
+
+    Args:
+        estimator: A SageMaker estimator (e.g. ``JumpStartEstimator``).
+        inputs: The training inputs passed through to ``estimator.fit()``.
+        capacity_wait_timeout_seconds: Max time to allow the job to remain in
+            "waiting for capacity" before declaring a capacity shortage.
+    """
+    estimator.fit(inputs, wait=False)
+
+    training_job_name = estimator.latest_training_job.name
+    sagemaker_client = estimator.sagemaker_session.sagemaker_client
+
+    capacity_wait_deadline = time.time() + capacity_wait_timeout_seconds
+    training_started = False
+
+    while True:
+        desc = sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
+        status = desc["TrainingJobStatus"]
+        secondary_status = desc.get("SecondaryStatus", "")
+        secondary_message = desc.get("SecondaryStatusMessage", "")
+
+        # Once the job leaves the pre-capacity phases, capacity has been granted
+        # and normal training-time limits (MaxRuntimeInSeconds) take over.
+        if secondary_status not in ("Starting", "Pending"):
+            training_started = True
+
+        if status in ("Completed", "Failed", "Stopped"):
+            if status != "Completed":
+                raise RuntimeError(
+                    f"Training job {training_job_name} ended with status {status}: "
+                    f"{desc.get('FailureReason', secondary_message)}"
+                )
+            return desc
+
+        if not training_started and time.time() > capacity_wait_deadline:
+            # Stop the queued job so it does not keep holding the request, then
+            # raise a CapacityError so x_fail_if_ice converts this to an xfail.
+            try:
+                sagemaker_client.stop_training_job(TrainingJobName=training_job_name)
+            except ClientError:
+                pass
+            raise RuntimeError(
+                "CapacityError: training job "
+                f"{training_job_name} stayed in '{secondary_status}' "
+                f"({secondary_message!r}) for over "
+                f"{capacity_wait_timeout_seconds}s without acquiring "
+                f"{estimator.instance_type} capacity."
+            )
+
+        time.sleep(_CAPACITY_POLL_INTERVAL_SECONDS)
+
+
 def download_inference_assets():
 
     if not os.path.exists(TMP_DIRECTORY_PATH):