From 59aafc07cbbdf49117d6131f4736979b2dcbb38c Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Mon, 12 Oct 2026 12:06:37 -0700 Subject: [PATCH] test(train): Fail fast in HyperPod SFT integ tests 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). --- .../tests/integ/train/hyperpod_job_wait.py | 214 ++++++++++++++++++ .../integ/train/test_nova_sft_hyperpod.py | 34 +-- .../train/test_sft_data_mixing_hyperpod.py | 51 ++--- 3 files changed, 242 insertions(+), 57 deletions(-) create mode 100644 sagemaker-train/tests/integ/train/hyperpod_job_wait.py diff --git a/sagemaker-train/tests/integ/train/hyperpod_job_wait.py b/sagemaker-train/tests/integ/train/hyperpod_job_wait.py new file mode 100644 index 0000000000..f13d2278f8 --- /dev/null +++ b/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) diff --git a/sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py b/sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py index 802e52d5c8..28dd13ed6a 100644 --- a/sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py +++ b/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}") diff --git a/sagemaker-train/tests/integ/train/test_sft_data_mixing_hyperpod.py b/sagemaker-train/tests/integ/train/test_sft_data_mixing_hyperpod.py index d6aef79d15..a92f42b1b4 100644 --- a/sagemaker-train/tests/integ/train/test_sft_data_mixing_hyperpod.py +++ b/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}")