Skip to content

Commit f054786

Browse files
mujtaba1747Syed Jafri
andauthored
datamixing recipe path fix (#6073)
* cleanup: use a single logger in base_trainer * fix: update recipe resolver unit tests * Release 3.16.0 (2026-07-15): Bump VERSION files and internal dependency pins to 3.16.0 / 2.16.0 / 1.16.0. Update CHANGELOGs across root and submodules (core, train, serve, mlops). * fix: use absolute path in data mixing recipe path construction --------- Co-authored-by: Syed Jafri <syedjfr@amazon.com>
1 parent 4356ca8 commit f054786

4 files changed

Lines changed: 58 additions & 25 deletions

File tree

‎sagemaker-train/src/sagemaker/train/common_utils/data_mixing_utils.py‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -480,16 +480,18 @@ def build_hyperpod_datamix_recipe_from_context(
480480
4. Inject customer_data_percent and nova_data_percentages into data_mixing.sources
481481
5. Validate non-zero categories exist in template's nova_data section
482482
6. Write final YAML to HyperPod CLI recipes directory
483-
7. Return (relative_recipe_path, image_uri)
483+
7. Return (recipe_path, image_uri)
484484
485485
Args:
486486
context: The HyperPodTemplateContext from resolve_hyperpod_datamix_context.
487487
validated_config: A DataMixingConfig validated via validate_data_mixing_categories.
488488
489489
Returns:
490-
Tuple of (relative_recipe_path, image_uri). relative_recipe_path is relative
491-
to the HyperPod CLI recipes_collection/recipes directory with .yaml extension
492-
removed. image_uri is the container image URI from the context (or None).
490+
Tuple of (recipe_path, image_uri). recipe_path is the absolute filesystem
491+
path (including the .yaml extension) of the generated recipe written under
492+
the HyperPod CLI recipes_collection/recipes directory; it is consumed by the
493+
recipe resolver as a user recipe file. image_uri is the container image URI
494+
from the context (or None).
493495
494496
Raises:
495497
RuntimeError: If hyperpod_cli is not installed.
@@ -652,18 +654,11 @@ def _apply_overrides(recipe: dict, overrides: dict) -> dict:
652654
with open(recipe_path, "w") as f:
653655
f.write(recipe_output)
654656

655-
relative_path = (
656-
recipe_path.split(HYPERPOD_RECIPE_PATH, 1)[1]
657-
.lstrip("/").lstrip("\\")
658-
.removesuffix(".yaml")
659-
)
660-
661657
logger.info(
662-
"Generated HyperPod datamix recipe at '%s' (relative: '%s') from context '%s'.",
658+
"Generated HyperPod datamix recipe at '%s' from context '%s'.",
663659
recipe_path,
664-
relative_path,
665660
context.recipe_name,
666661
)
667662

668-
return relative_path, context.image_uri
663+
return recipe_path, context.image_uri
669664

sagemaker-train/tests/integ/train/test_nova_hyperpod.py renamed to sagemaker-train/tests/integ/train/test_nova_sft_hyperpod.py

File renamed without changes.

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

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import pytest
3737
from sagemaker.train.sft_trainer import SFTTrainer
3838
from sagemaker.train.common import TrainingType
39+
from sagemaker.train.base_trainer import BaseTrainer
3940
from sagemaker.train.data_mixing_config import DataMixingConfig
4041
from sagemaker.core.training.configs import HyperPodCompute
4142

@@ -49,8 +50,8 @@
4950
REGION = "us-east-1"
5051
DATA_PREFIX = "test-sft-data-mixing-hyperpod-integ"
5152
NUM_TRAINING_SAMPLES = 300
52-
HYPERPOD_CLUSTER_NAME = "riv-rig"
53-
HYPERPOD_INSTANCE_TYPE = "ml.p5.48xlarge"
53+
HYPERPOD_CLUSTER_NAME = "pysdk-hp-integ-tests"
54+
HYPERPOD_INSTANCE_TYPE = "ml.g6.12xlarge"
5455

5556

5657
def _generate_training_data() -> str:
@@ -125,10 +126,10 @@ def training_resources(sagemaker_session_us_east_1):
125126
@pytest.mark.gpu_intensive
126127
@pytest.mark.us_east_1
127128
def test_sft_trainer_nova_micro_data_mixing_hyperpod(sagemaker_session_us_east_1, training_resources):
128-
"""Test SFT trainer with Nova Lite 2 model and data mixing on HyperPod.
129+
"""Test SFT trainer with Nova Micro model and data mixing on HyperPod.
129130
130131
This end-to-end test submits a real HyperPod training job with DataMixingConfig
131-
for Nova Lite 2. The SDK resolves the datamix recipe from SageMaker Hub, validates
132+
for Nova Micro. The SDK resolves the datamix recipe from SageMaker Hub, validates
132133
categories, and includes the serialized config in the HyperPod override parameters.
133134
"""
134135
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
184185
assert get_job_result.returncode == 0, (
185186
f"hyperpod get-job failed for '{job_name}': {get_job_result.stderr}"
186187
)
187-
logger.info(f"Verified job '{job_name}' exists on the cluster.")
188+
logger.info(f"Verified job '{job_name}' exists on the cluster using hp-cli.")
189+
190+
# Poll for job completion by checking for the manifest in S3.
191+
# The manifest is written under {output_s3_path}/{job_name}/manifest.json
192+
# once training finishes, so its presence confirms end-to-end completion.
193+
194+
output_s3_path = training_resources["s3_output_path"]
195+
max_wait_time = 21600 # 6 hour timeout (HyperPod jobs can take longer)
196+
poll_interval = 60 # Check every 60 seconds
197+
start_time = time.time()
198+
checkpoint_path = None
199+
200+
while time.time() - start_time < max_wait_time:
201+
checkpoint_path = BaseTrainer._resolve_checkpoint_from_manifest(
202+
job_name=job_name,
203+
output_s3_path=output_s3_path,
204+
sagemaker_session=sagemaker_session_us_east_1,
205+
)
206+
if checkpoint_path:
207+
logger.info(f"Checkpoint resolved: {checkpoint_path}")
208+
break
209+
210+
elapsed = int(time.time() - start_time)
211+
logger.info(f"Waiting for manifest... ({elapsed}s elapsed)")
212+
time.sleep(poll_interval)
213+
214+
assert checkpoint_path is not None, (
215+
f"Job {job_name} did not produce a manifest within {max_wait_time}s"
216+
)
217+
logger.info(f"Training complete. Checkpoint: {checkpoint_path}")

‎sagemaker-train/tests/unit/train/common_utils/test_data_mixing_utils.py‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
"""Unit tests for data mixing utility functions."""
1414
from __future__ import absolute_import
1515

16+
import os
17+
1618
import pytest
1719

1820
from sagemaker.train.data_mixing_config import DataMixingConfig
@@ -823,8 +825,12 @@ def capture_write(path, mode="r", **kwargs):
823825
assert parsed["data_mixing"]["sources"]["nova_data"]["en-entertainment"] == 20
824826
assert parsed["data_mixing"]["sources"]["nova_data"]["en-scientific"] == 10
825827

826-
def test_return_value_is_relative_path_and_image_uri(self):
827-
"""Return value should be (relative_path_without_extension, image_uri)."""
828+
def test_return_value_is_absolute_recipe_path_and_image_uri(self):
829+
"""Return value should be (absolute_recipe_path_with_extension, image_uri).
830+
831+
The path must be a loadable filesystem path (absolute, .yaml extension
832+
intact) because the recipe resolver opens it as a user recipe file.
833+
"""
828834
from unittest.mock import patch, MagicMock, mock_open
829835

830836
from sagemaker.train.common_utils.data_mixing_utils import (
@@ -843,12 +849,14 @@ def test_return_value_is_relative_path_and_image_uri(self):
843849
with patch("builtins.open", m_open):
844850
result = build_hyperpod_datamix_recipe_from_context(context, config)
845851

846-
relative_path, image_uri = result
852+
recipe_path, image_uri = result
847853

848-
# relative_path should not end with .yaml
849-
assert not relative_path.endswith(".yaml")
850-
# relative_path should contain fine-tuning/nova
851-
assert "fine-tuning/nova" in relative_path
854+
# recipe_path must be an absolute filesystem path
855+
assert os.path.isabs(recipe_path)
856+
# recipe_path must retain the .yaml extension so it can be opened
857+
assert recipe_path.endswith(".yaml")
858+
# recipe_path should live under the HyperPod CLI recipes fine-tuning/nova dir
859+
assert "fine-tuning/nova" in recipe_path
852860
# image_uri should match the context
853861
assert image_uri == self.IMAGE_URI
854862

0 commit comments

Comments
 (0)