test(train): Fail fast in HyperPod SFT integ tests (#6449)

The two HyperPod Nova Micro SFT integ tests polled only for the job's
S3 manifest, for up to 6 hours. When the cluster cannot schedule the
job, or the job succeeds but its outputs are never synced to S3, the
manifest never appears, so the tests held the GPU integ build until
CodeBuild timed out, and every run left another PyTorchJob queued in
the cluster.

Add tests/integ/train/hyperpod_job_wait.py, which also watches the
PyTorchJob through the Kubernetes API:
- fail immediately if the job fails;
- skip if the job does not start running within 30 minutes;
- fail if it runs for more than 75 minutes;
- fail with a pointer to the output sync if the job succeeded but no
  manifest appears within 15 minutes.
The job is removed from the cluster when the wait ends (helm uninstall
of the HyperPod CLI release, falling back to deleting the PyTorchJob).

Author: lucasjia-aws

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

@@ -0,0 +1,214 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Bounded wait for deep integ tests that train on a HyperPod (EKS) cluster.
+
+``SFTTrainer.train()`` on HyperPod submits a Kubeflow PyTorchJob through the
+HyperPod CLI and returns its name. Completion is signalled by
+``{s3_output_path}/{job_name}/manifest.json``, which a service-managed component
+copies from the cluster to S3 after training. Polling only for that manifest
+cannot tell "still training" from "never scheduled" or "succeeded but the
+outputs were never synced", so a broken cluster made such tests poll for hours
+until the CodeBuild build timed out, and every run left another job queued in
+the cluster.
+
+``wait_for_hyperpod_job`` also watches the PyTorchJob itself:
+
+* the job failed -> fail immediately with its condition message;
+* the job never starts running within ``schedule_timeout`` (no status, or pods
+  that stay Pending) -> skip as a cluster capacity/scheduling problem;
+* the job runs longer than ``run_timeout`` -> fail;
+* the job succeeded but no manifest appears within ``manifest_timeout`` -> fail
+  with a message pointing at the output sync.
+
+Whatever the outcome, the job is removed from the cluster when the wait ends
+(``helm uninstall`` of the release the HyperPod CLI created, falling back to
+deleting the PyTorchJob), so tests no longer leave jobs behind.
+"""
+
+from __future__ import absolute_import
+
+import logging
+import subprocess
+import time
+
+import pytest
+
+from sagemaker.train.base_trainer import BaseTrainer
+
+logger = logging.getLogger(__name__)
+
+PYTORCHJOB_GROUP = "kubeflow.org"
+PYTORCHJOB_VERSION = "v1"
+PYTORCHJOB_PLURAL = "pytorchjobs"
+
+# Healthy HyperPod Nova Micro SFT jobs in this suite were scheduled within
+# seconds and finished in ~25-55 minutes, and the manifest followed within a
+# few minutes. Worst case for one test: 30 + 75 + 15 minutes, well inside the
+# CI build budget.
+DEFAULT_SCHEDULE_TIMEOUT_SECONDS = 30 * 60
+DEFAULT_RUN_TIMEOUT_SECONDS = 75 * 60
+DEFAULT_MANIFEST_TIMEOUT_SECONDS = 15 * 60
+DEFAULT_POLL_INTERVAL_SECONDS = 60
+
+
+def _load_k8s_clients():
+    """Return (CustomObjectsApi, CoreV1Api) for the kubeconfig written by
+    ``hyperpod connect-cluster`` during ``train()``.
+
+    ``kubernetes`` is a dependency of the HyperPod CLI, which these tests need
+    anyway, so it is imported lazily.
+    """
+    from kubernetes import client, config
+
+    config.load_kube_config()
+    return client.CustomObjectsApi(), client.CoreV1Api()
+
+
+def _true_condition(job_obj, condition_type):
+    for cond in (job_obj.get("status") or {}).get("conditions") or []:
+        if cond.get("type") == condition_type and str(cond.get("status")) == "True":
+            return cond
+    return None
+
+
+def _pod_phases(core_api, namespace, job_name):
+    pods = core_api.list_namespaced_pod(
+        namespace, label_selector=f"training.kubeflow.org/job-name={job_name}"
+    )
+    return [p.status.phase for p in pods.items]
+
+
+def delete_hyperpod_job(job_name, namespace="kubeflow", custom_api=None):
+    """Best-effort removal of a HyperPod CLI job from the cluster.
+
+    The CLI installs each job as a Helm release named after the job (one
+    ConfigMap plus one PyTorchJob), so ``helm uninstall`` removes both. If helm is
+    unavailable or fails, delete the PyTorchJob directly.
+    """
+    try:
+        result = subprocess.run(
+            ["helm", "uninstall", job_name, "--namespace", namespace],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode == 0:
+            logger.info("Removed HyperPod job %s (helm uninstall)", job_name)
+            return
+        logger.warning("helm uninstall %s failed: %s", job_name, result.stderr.strip())
+    except FileNotFoundError:
+        logger.warning("helm not on PATH; deleting PyTorchJob %s directly", job_name)
+    try:
+        if custom_api is None:
+            custom_api, _ = _load_k8s_clients()
+        custom_api.delete_namespaced_custom_object(
+            PYTORCHJOB_GROUP, PYTORCHJOB_VERSION, namespace, PYTORCHJOB_PLURAL, job_name
+        )
+        logger.info("Deleted PyTorchJob %s", job_name)
+    except Exception as e:  # pylint: disable=broad-except
+        logger.warning("Failed to delete PyTorchJob %s: %s", job_name, e)
+
+
+def wait_for_hyperpod_job(
+    job_name,
+    output_s3_path,
+    sagemaker_session=None,
+    namespace="kubeflow",
+    schedule_timeout=DEFAULT_SCHEDULE_TIMEOUT_SECONDS,
+    run_timeout=DEFAULT_RUN_TIMEOUT_SECONDS,
+    manifest_timeout=DEFAULT_MANIFEST_TIMEOUT_SECONDS,
+    poll_interval=DEFAULT_POLL_INTERVAL_SECONDS,
+    cleanup=True,
+):
+    """Wait for a HyperPod PyTorchJob and its S3 manifest, failing fast.
+
+    Args:
+        job_name (str): Job name returned by ``train(wait=False)``.
+        output_s3_path (str): The trainer's ``s3_output_path``.
+        sagemaker_session: Session used to read the manifest from S3.
+        namespace (str): Kubernetes namespace the job was submitted to.
+        schedule_timeout (int): Seconds the job may wait to start running before
+            the test is skipped.
+        run_timeout (int): Seconds the job may run before the test fails.
+        manifest_timeout (int): Seconds after the job succeeded that the manifest
+            may take to appear before the test fails.
+        poll_interval (int): Seconds between checks.
+        cleanup (bool): Remove the job from the cluster when the wait ends.
+
+    Returns:
+        str: The checkpoint S3 URI from the manifest.
+    """
+    custom_api, core_api = _load_k8s_clients()
+    start = time.time()
+    running_since = None
+    succeeded_since = None
+    # Hard cap so persistent API errors cannot keep the loop alive past the phases.
+    max_wait = schedule_timeout + run_timeout + manifest_timeout
+    try:
+        while time.time() - start < max_wait:
+            checkpoint = BaseTrainer._resolve_checkpoint_from_manifest(
+                job_name=job_name,
+                output_s3_path=output_s3_path,
+                sagemaker_session=sagemaker_session,
+            )
+            if checkpoint:
+                logger.info("Job %s produced checkpoint %s", job_name, checkpoint)
+                return checkpoint
+
+            now = time.time()
+            try:
+                job_obj = custom_api.get_namespaced_custom_object(
+                    PYTORCHJOB_GROUP, PYTORCHJOB_VERSION, namespace, PYTORCHJOB_PLURAL, job_name
+                )
+                phases = _pod_phases(core_api, namespace, job_name)
+            except Exception as e:  # pylint: disable=broad-except
+                if getattr(e, "status", None) == 404:
+                    pytest.fail(f"PyTorchJob {job_name} no longer exists in {namespace}.")
+                logger.warning("Could not read PyTorchJob %s (will retry): %s", job_name, e)
+                time.sleep(poll_interval)
+                continue
+
+            failed = _true_condition(job_obj, "Failed")
+            if failed:
+                pytest.fail(
+                    f"HyperPod job {job_name} failed: {failed.get('reason')}: "
+                    f"{failed.get('message')}"
+                )
+
+            if _true_condition(job_obj, "Succeeded"):
+                succeeded_since = succeeded_since or now
+                if now - succeeded_since > manifest_timeout:
+                    pytest.fail(
+                        f"HyperPod job {job_name} succeeded but no manifest appeared at "
+                        f"{output_s3_path.rstrip('/')}/{job_name}/manifest.json within "
+                        f"{manifest_timeout}s; the cluster's output sync to S3 is not "
+                        f"delivering job outputs."
+                    )
+            else:
+                if _true_condition(job_obj, "Running") or "Running" in phases:
+                    running_since = running_since or now
+                    if now - running_since > run_timeout:
+                        pytest.fail(
+                            f"HyperPod job {job_name} still running after {run_timeout}s "
+                            f"(pod phases: {phases})."
+                        )
+                elif running_since is None and now - start > schedule_timeout:
+                    pytest.skip(
+                        f"HyperPod job {job_name} did not start within {schedule_timeout}s "
+                        f"(job conditions: {(job_obj.get('status') or {}).get('conditions')}, "
+                        f"pod phases: {phases}); the cluster has no schedulable capacity."
+                    )
+            time.sleep(poll_interval)
+        pytest.fail(f"HyperPod job {job_name} did not complete within {max_wait}s.")
+    finally:
+        if cleanup:
+            delete_hyperpod_job(job_name, namespace=namespace, custom_api=custom_api)

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

@@ -25,6 +25,8 @@
 from sagemaker.core.training.configs import HyperPodCompute
 from sagemaker.core.shapes import StoppingCondition
 
+from .hyperpod_job_wait import wait_for_hyperpod_job
+
 logger = logging.getLogger(__name__)
 
 S3_BUCKET_PREFIX = "mc-flows-sdk-testing"
@@ -114,28 +116,12 @@ def test_sft_trainer_nova_micro_hyperpod_lora(
 
     logger.info(f"HyperPod SFT job submitted: {job_name}")
 
-    # Poll for job completion by checking for the manifest in S3
-    from sagemaker.train.base_trainer import BaseTrainer
-
-    max_wait_time = 21600  # 6 hour timeout (HyperPod jobs can take longer)
-    poll_interval = 60  # Check every 60 seconds
-    start_time = time.time()
-    checkpoint_path = None
-
-    while time.time() - start_time < max_wait_time:
-        checkpoint_path = BaseTrainer._resolve_checkpoint_from_manifest(
-            job_name=job_name,
-            output_s3_path=s3_output_path,
-        )
-        if checkpoint_path:
-            logger.info(f"Checkpoint resolved: {checkpoint_path}")
-            break
-
-        elapsed = int(time.time() - start_time)
-        logger.info(f"Waiting for manifest... ({elapsed}s elapsed)")
-        time.sleep(poll_interval)
-
-    assert (
-        checkpoint_path is not None
-    ), f"Job {job_name} did not produce a manifest within {max_wait_time}s"
+    # Wait for the PyTorchJob and its S3 manifest, failing fast on a stuck cluster,
+    # and remove the job from the cluster afterwards (see hyperpod_job_wait.py).
+    checkpoint_path = wait_for_hyperpod_job(
+        job_name=job_name,
+        output_s3_path=s3_output_path,
+        namespace=hyperpod_compute.namespace,
+    )
+    assert checkpoint_path is not None
     logger.info(f"Training complete. Checkpoint: {checkpoint_path}")

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

@@ -33,14 +33,16 @@
 import logging
 import time
 import random
+import subprocess
 
 import pytest
 from sagemaker.train.sft_trainer import SFTTrainer
 from sagemaker.train.common import TrainingType
-from sagemaker.train.base_trainer import BaseTrainer
 from sagemaker.train.data_mixing_config import DataMixingConfig
 from sagemaker.core.training.configs import HyperPodCompute
 
+from .hyperpod_job_wait import delete_hyperpod_job, wait_for_hyperpod_job
+
 logging.basicConfig(
     level=logging.INFO,
     format="%(levelname)s - %(name)s - %(message)s",
@@ -179,43 +181,26 @@ def test_sft_trainer_nova_micro_data_mixing_hyperpod(
     logger.info(f"HyperPod training job submitted: {job_name}")
 
     # Verify the job exists on the cluster via hyperpod get-job
-    import subprocess
-
+    namespace = compute.namespace or "kubeflow"
     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}"
+    if get_job_result.returncode != 0:
+        delete_hyperpod_job(job_name, namespace=namespace)
+        pytest.fail(f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}")
     logger.info(f"Verified job '{job_name}' exists on the cluster using hp-cli.")
 
-    # Poll for job completion by checking for the manifest in S3.
-    # The manifest is written under {output_s3_path}/{job_name}/manifest.json
-    # once training finishes, so its presence confirms end-to-end completion.
-
-    output_s3_path = training_resources["s3_output_path"]
-    max_wait_time = 21600  # 6 hour timeout (HyperPod jobs can take longer)
-    poll_interval = 60  # Check every 60 seconds
-    start_time = time.time()
-    checkpoint_path = None
-
-    while time.time() - start_time < max_wait_time:
-        checkpoint_path = BaseTrainer._resolve_checkpoint_from_manifest(
-            job_name=job_name,
-            output_s3_path=output_s3_path,
-            sagemaker_session=sagemaker_session_us_east_1,
-        )
-        if checkpoint_path:
-            logger.info(f"Checkpoint resolved: {checkpoint_path}")
-            break
-
-        elapsed = int(time.time() - start_time)
-        logger.info(f"Waiting for manifest... ({elapsed}s elapsed)")
-        time.sleep(poll_interval)
-
-    assert (
-        checkpoint_path is not None
-    ), f"Job {job_name} did not produce a manifest within {max_wait_time}s"
+    # Wait for the PyTorchJob and its S3 manifest
+    # ({output_s3_path}/{job_name}/manifest.json), failing fast on a stuck
+    # cluster, and remove the job from the cluster afterwards
+    # (see hyperpod_job_wait.py).
+    checkpoint_path = wait_for_hyperpod_job(
+        job_name=job_name,
+        output_s3_path=training_resources["s3_output_path"],
+        sagemaker_session=sagemaker_session_us_east_1,
+        namespace=namespace,
+    )
+    assert checkpoint_path is not None
     logger.info(f"Training complete. Checkpoint: {checkpoint_path}")