From 5953f68621c89fc352fdf3182af66450403bbb9e Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 6 Oct 2026 14:36:20 -0700 Subject: [PATCH] test(train): Serialize Nova serverless integ jobs Nova serverless fine-tuning jobs (nova-textgeneration-lite-v2 SFT, SFT data mixing, RLVR) run on a shared p5 pool that customers also use, and each job takes 4 p5 instances. The deep integ suite runs under pytest-xdist, so these three tests submitted their jobs at the same time and could occupy most of the pool, leaving no capacity for customer jobs. Add tests/integ/train/nova_capacity.py: - nova_capacity_slot(): a cross-worker file lock so at most one Nova serverless job per test run is in flight, held until the job is terminal. - wait_for_nova_training_job(): replaces the hand-rolled poll loops. It stops the job and skips the test if the job is still waiting for instances after 60 minutes, and stops the job whenever the wait ends with the job still running, so no orphaned job keeps queueing for capacity. --- .../tests/integ/train/nova_capacity.py | 127 ++++++++++++++++++ .../train/test_rlvr_trainer_integration.py | 28 ++-- ...est_sft_trainer_data_mixing_integration.py | 54 ++++---- .../train/test_sft_trainer_integration.py | 27 ++-- 4 files changed, 176 insertions(+), 60 deletions(-) create mode 100644 sagemaker-train/tests/integ/train/nova_capacity.py diff --git a/sagemaker-train/tests/integ/train/nova_capacity.py b/sagemaker-train/tests/integ/train/nova_capacity.py new file mode 100644 index 0000000000..c7a8550d90 --- /dev/null +++ b/sagemaker-train/tests/integ/train/nova_capacity.py @@ -0,0 +1,127 @@ +# 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. +"""Helpers for deep integ tests that fine-tune Nova models on serverless capacity. + +Nova serverless fine-tuning jobs (e.g. ``nova-textgeneration-lite-v2`` SFT/RLVR) run +on a shared, service-owned p5 pool that customers draw from as well, and each job +takes 4 p5 instances. When several of these tests run at once (pytest-xdist runs the +deep suite in parallel) they can occupy most of that pool and leave no room for a +customer job. These helpers keep the tests' footprint to one job at a time: + +* ``nova_capacity_slot`` serializes Nova serverless jobs across all xdist workers of a + test run. Hold it from submission until the job is terminal. +* ``wait_for_nova_training_job`` waits for the job with a bounded wait for capacity: + if the job is still waiting for instances after ``pending_timeout`` it is stopped + and the test is skipped. It also stops the job whenever the wait ends with the job + still running (timeout, error, interrupt), so no orphaned job keeps queueing for + capacity after the test is gone. +""" + +from __future__ import absolute_import + +import logging +import os +import tempfile +import time +from contextlib import contextmanager + +import pytest + +from .. import lock + +logger = logging.getLogger(__name__) + +# The lock file must be shared by every xdist worker of a run. fcntl locks are released +# by the OS when the holding process dies, so a killed run never leaves a stale lock. +NOVA_CAPACITY_LOCK_PATH = os.path.join(tempfile.gettempdir(), "sagemaker_nova_capacity_lock") + +# In CI, Nova serverless jobs that eventually ran waited a few minutes for instances +# (median 3-6 min, max ~100 min while competing with each other). Once the tests are +# serialized, a job still waiting after an hour means the shared pool is full, so we +# give the capacity back instead of queueing behind customers. +NOVA_PENDING_TIMEOUT_SECONDS = 3600 + +TERMINAL_STATUSES = ("Completed", "Failed", "Stopped") +_WAITING_FOR_CAPACITY_STATUSES = ("Starting", "Pending") + + +@contextmanager +def nova_capacity_slot(): + """Allow at most one Nova serverless training job per test run at a time.""" + start = time.time() + logger.info("Waiting for the Nova capacity slot (%s)", NOVA_CAPACITY_LOCK_PATH) + with lock.lock(NOVA_CAPACITY_LOCK_PATH): + logger.info("Acquired the Nova capacity slot after %ds", int(time.time() - start)) + yield + + +def _stop_quietly(training_job): + """Best-effort stop of a training job that is not terminal yet.""" + name = getattr(training_job, "training_job_name", None) + try: + logger.warning("Stopping non-terminal training job %s", name) + training_job.stop() + except Exception as e: # pylint: disable=broad-except + logger.warning("Failed to stop training job %s: %s", name, e) + + +def wait_for_nova_training_job( + training_job, + max_wait_time, + poll_interval=30, + pending_timeout=NOVA_PENDING_TIMEOUT_SECONDS, +): + """Poll a training job until it is terminal, with a bounded wait for capacity. + + Args: + training_job: The ``sagemaker.core.resources.TrainingJob`` returned by + ``train(wait=False)``. + max_wait_time (int): Maximum total seconds to wait for a terminal status. + poll_interval (int): Seconds between status checks. + pending_timeout (int): Seconds after which a job that is still waiting for + instances (secondary status ``Starting``/``Pending``) is stopped and the + test is skipped. + + Returns: + str: The last observed ``training_job_status``. It is terminal unless + ``max_wait_time`` ran out, in which case the job has been stopped. + """ + start = time.time() + status = None + try: + while time.time() - start < max_wait_time: + training_job.refresh() + status = training_job.training_job_status + if status in TERMINAL_STATUSES: + return status + elapsed = time.time() - start + if ( + training_job.secondary_status in _WAITING_FOR_CAPACITY_STATUSES + and elapsed > pending_timeout + ): + pytest.skip( + f"Nova capacity unavailable: training job " + f"{training_job.training_job_name} was still waiting for instances " + f"after {int(elapsed)}s; stopped it to release the shared pool." + ) + time.sleep(poll_interval) + logger.warning( + "Training job %s not terminal after %ds (status: %s)", + training_job.training_job_name, + max_wait_time, + status, + ) + return status + finally: + if status not in TERMINAL_STATUSES: + _stop_quietly(training_job) diff --git a/sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py b/sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py index 6e83b4637f..0719ab1265 100644 --- a/sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py @@ -26,6 +26,8 @@ from sagemaker.ai_registry.evaluator import Evaluator from sagemaker.ai_registry.air_constants import REWARD_FUNCTION +from .nova_capacity import nova_capacity_slot, wait_for_nova_training_job + logger = logging.getLogger(__name__) EVALUATOR_NAME = "test-integ-rlvr-trainer" @@ -184,22 +186,16 @@ def test_rlvr_trainer_nova_workflow(sagemaker_session_us_east_1): rlvr_trainer.hyperparameters.save_steps = 10 rlvr_trainer.hyperparameters.global_batch_size = 32 - training_job = rlvr_trainer.train(wait=False) - logger.info(f"Training job submitted: {training_job.training_job_arn}") - - # Manual wait loop - max_wait_time = 10800 # 3 hour timeout (Nova training takes >1 hour) - poll_interval = 30 - 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 - - time.sleep(poll_interval) + # Nova serverless jobs share a p5 pool with customers: run one at a time and + # give the capacity back if the job cannot get instances (see nova_capacity.py). + with nova_capacity_slot(): + training_job = rlvr_trainer.train(wait=False) + logger.info(f"Training job submitted: {training_job.training_job_arn}") + wait_for_nova_training_job( + training_job, + max_wait_time=10800, # 3 hour timeout (Nova training takes >1 hour) + poll_interval=30, + ) # Verify job completed successfully assert training_job.training_job_status == "Completed" diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_data_mixing_integration.py b/sagemaker-train/tests/integ/train/test_sft_trainer_data_mixing_integration.py index eed716dbec..6b91027902 100644 --- a/sagemaker-train/tests/integ/train/test_sft_trainer_data_mixing_integration.py +++ b/sagemaker-train/tests/integ/train/test_sft_trainer_data_mixing_integration.py @@ -37,6 +37,8 @@ from sagemaker.train.common import TrainingType from sagemaker.train.data_mixing_config import DataMixingConfig +from .nova_capacity import nova_capacity_slot, wait_for_nova_training_job + logging.basicConfig( level=logging.INFO, format="%(levelname)s - %(name)s - %(message)s", @@ -155,34 +157,30 @@ def test_sft_trainer_nova_lite2_with_data_mixing(sagemaker_session_us_east_1, tr overrides={"name": f"sft-nova-datamix-integ-{unique_id}"}, ) - logger.info("Submitting SFT training job with data mixing config...") - try: - training_job = sft_trainer.train(wait=False) - except ValueError as e: - if "Failed to download from S3 Access Point" in str(e): - pytest.skip( - "Skipping: account does not have Forge subscription for data mixing recipes. " - "This is expected for non-subscribed accounts. " - "See https://docs.aws.amazon.com/sagemaker/latest/dg/nova-forge.html#nova-forge-prereq-access" - ) - raise - - assert training_job is not None - logger.info(f"Training job submitted: {training_job.training_job_name}") - - # Manual wait loop — Nova training can take over an hour - max_wait_time = 10800 # 3 hour timeout - poll_interval = 30 # Check every 30 seconds - 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 - - time.sleep(poll_interval) + # Nova serverless jobs share a p5 pool with customers: run one at a time and + # give the capacity back if the job cannot get instances (see nova_capacity.py). + with nova_capacity_slot(): + logger.info("Submitting SFT training job with data mixing config...") + try: + training_job = sft_trainer.train(wait=False) + except ValueError as e: + if "Failed to download from S3 Access Point" in str(e): + pytest.skip( + "Skipping: account does not have Forge subscription for data mixing recipes. " + "This is expected for non-subscribed accounts. " + "See https://docs.aws.amazon.com/sagemaker/latest/dg/nova-forge.html#nova-forge-prereq-access" + ) + raise + + assert training_job is not None + logger.info(f"Training job submitted: {training_job.training_job_name}") + + # Nova training can take over an hour + wait_for_nova_training_job( + training_job, + max_wait_time=10800, # 3 hour timeout + poll_interval=30, + ) # Verify job completed successfully assert ( diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py index 0026c45a69..67a57aff1c 100644 --- a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py @@ -20,6 +20,8 @@ from sagemaker.train.sft_trainer import SFTTrainer from sagemaker.train.common import TrainingType +from .nova_capacity import nova_capacity_slot, wait_for_nova_training_job + @pytest.mark.gpu_intensive def test_sft_trainer_lora_complete_workflow(sagemaker_session, mlflow_resource_arn): @@ -123,22 +125,15 @@ def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1): base_job_name=f"sft-nova-integ-{unique_id}", ) - # Create training job - training_job = sft_trainer_nova.train(wait=False) - - # Manual wait loop - max_wait_time = 10800 # 3 hour timeout (Nova training takes >1 hour) - poll_interval = 30 # Check every 30 seconds - 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 - - time.sleep(poll_interval) + # Nova serverless jobs share a p5 pool with customers: run one at a time and + # give the capacity back if the job cannot get instances (see nova_capacity.py). + with nova_capacity_slot(): + training_job = sft_trainer_nova.train(wait=False) + wait_for_nova_training_job( + training_job, + max_wait_time=10800, # 3 hour timeout (Nova training takes >1 hour) + poll_interval=30, + ) # Verify job completed successfully assert training_job.training_job_status == "Completed"