test(train): Serialize Nova serverless integ jobs

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)

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"

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 (

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"