diff --git a/sagemaker-train/src/sagemaker/train/common_utils/data_mixing_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/data_mixing_utils.py index 2fb52f2995..0f91c9425f 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/data_mixing_utils.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/data_mixing_utils.py @@ -480,16 +480,18 @@ def build_hyperpod_datamix_recipe_from_context( 4. Inject customer_data_percent and nova_data_percentages into data_mixing.sources 5. Validate non-zero categories exist in template's nova_data section 6. Write final YAML to HyperPod CLI recipes directory - 7. Return (relative_recipe_path, image_uri) + 7. Return (recipe_path, image_uri) Args: context: The HyperPodTemplateContext from resolve_hyperpod_datamix_context. validated_config: A DataMixingConfig validated via validate_data_mixing_categories. Returns: - Tuple of (relative_recipe_path, image_uri). relative_recipe_path is relative - to the HyperPod CLI recipes_collection/recipes directory with .yaml extension - removed. image_uri is the container image URI from the context (or None). + Tuple of (recipe_path, image_uri). recipe_path is the absolute filesystem + path (including the .yaml extension) of the generated recipe written under + the HyperPod CLI recipes_collection/recipes directory; it is consumed by the + recipe resolver as a user recipe file. image_uri is the container image URI + from the context (or None). Raises: RuntimeError: If hyperpod_cli is not installed. @@ -652,18 +654,11 @@ def _apply_overrides(recipe: dict, overrides: dict) -> dict: with open(recipe_path, "w") as f: f.write(recipe_output) - relative_path = ( - recipe_path.split(HYPERPOD_RECIPE_PATH, 1)[1] - .lstrip("/").lstrip("\\") - .removesuffix(".yaml") - ) - logger.info( - "Generated HyperPod datamix recipe at '%s' (relative: '%s') from context '%s'.", + "Generated HyperPod datamix recipe at '%s' from context '%s'.", recipe_path, - relative_path, context.recipe_name, ) - return relative_path, context.image_uri + return recipe_path, context.image_uri diff --git a/sagemaker-train/tests/integ/train/test_nova_hyperpod.py b/sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py similarity index 100% rename from sagemaker-train/tests/integ/train/test_nova_hyperpod.py rename to sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py 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 5e7c00b045..9bf1ac1908 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 @@ -36,6 +36,7 @@ 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 @@ -49,8 +50,8 @@ REGION = "us-east-1" DATA_PREFIX = "test-sft-data-mixing-hyperpod-integ" NUM_TRAINING_SAMPLES = 300 -HYPERPOD_CLUSTER_NAME = "riv-rig" -HYPERPOD_INSTANCE_TYPE = "ml.p5.48xlarge" +HYPERPOD_CLUSTER_NAME = "pysdk-hp-integ-tests" +HYPERPOD_INSTANCE_TYPE = "ml.g6.12xlarge" def _generate_training_data() -> str: @@ -125,10 +126,10 @@ def training_resources(sagemaker_session_us_east_1): @pytest.mark.gpu_intensive @pytest.mark.us_east_1 def test_sft_trainer_nova_micro_data_mixing_hyperpod(sagemaker_session_us_east_1, training_resources): - """Test SFT trainer with Nova Lite 2 model and data mixing on HyperPod. + """Test SFT trainer with Nova Micro model and data mixing on HyperPod. This end-to-end test submits a real HyperPod training job with DataMixingConfig - for Nova Lite 2. The SDK resolves the datamix recipe from SageMaker Hub, validates + for Nova Micro. The SDK resolves the datamix recipe from SageMaker Hub, validates categories, and includes the serialized config in the HyperPod override parameters. """ unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}" @@ -184,4 +185,33 @@ def test_sft_trainer_nova_micro_data_mixing_hyperpod(sagemaker_session_us_east_1 assert get_job_result.returncode == 0, ( f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}" ) - logger.info(f"Verified job '{job_name}' exists on the cluster.") + 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" + ) + logger.info(f"Training complete. Checkpoint: {checkpoint_path}") diff --git a/sagemaker-train/tests/unit/train/common_utils/test_data_mixing_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_data_mixing_utils.py index 0ae2044cb3..1e69073f88 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_data_mixing_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_data_mixing_utils.py @@ -13,6 +13,8 @@ """Unit tests for data mixing utility functions.""" from __future__ import absolute_import +import os + import pytest from sagemaker.train.data_mixing_config import DataMixingConfig @@ -823,8 +825,12 @@ def capture_write(path, mode="r", **kwargs): assert parsed["data_mixing"]["sources"]["nova_data"]["en-entertainment"] == 20 assert parsed["data_mixing"]["sources"]["nova_data"]["en-scientific"] == 10 - def test_return_value_is_relative_path_and_image_uri(self): - """Return value should be (relative_path_without_extension, image_uri).""" + def test_return_value_is_absolute_recipe_path_and_image_uri(self): + """Return value should be (absolute_recipe_path_with_extension, image_uri). + + The path must be a loadable filesystem path (absolute, .yaml extension + intact) because the recipe resolver opens it as a user recipe file. + """ from unittest.mock import patch, MagicMock, mock_open from sagemaker.train.common_utils.data_mixing_utils import ( @@ -843,12 +849,14 @@ def test_return_value_is_relative_path_and_image_uri(self): with patch("builtins.open", m_open): result = build_hyperpod_datamix_recipe_from_context(context, config) - relative_path, image_uri = result + recipe_path, image_uri = result - # relative_path should not end with .yaml - assert not relative_path.endswith(".yaml") - # relative_path should contain fine-tuning/nova - assert "fine-tuning/nova" in relative_path + # recipe_path must be an absolute filesystem path + assert os.path.isabs(recipe_path) + # recipe_path must retain the .yaml extension so it can be opened + assert recipe_path.endswith(".yaml") + # recipe_path should live under the HyperPod CLI recipes fine-tuning/nova dir + assert "fine-tuning/nova" in recipe_path # image_uri should match the context assert image_uri == self.IMAGE_URI