Skip to content

Commit efc1272

Browse files
mujtaba1747Syed Jafri
andauthored
wip: nova hyperpod integ tests (#5990)
* wip: nova hyperpod integ tests * feat: add polling logic to test * fix: generate bucket name using env var (to pass fortress security scan) --------- Co-authored-by: Syed Jafri <syedjfr@amazon.com>
1 parent 5c93d56 commit efc1272

1 file changed

Lines changed: 142 additions & 0 deletions

File tree

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integration tests for SFT trainer on HyperPod (Nova Micro)"""
14+
from __future__ import absolute_import
15+
16+
import time
17+
import random
18+
import pytest
19+
import boto3
20+
import logging
21+
22+
from sagemaker.train import SFTTrainer
23+
from sagemaker.train.common import TrainingType
24+
from sagemaker.core.training.configs import HyperPodCompute
25+
from sagemaker.core.shapes import StoppingCondition
26+
27+
logger = logging.getLogger(__name__)
28+
29+
S3_BUCKET_PREFIX = "mc-flows-sdk-testing"
30+
S3_PREFIX = "input_data"
31+
32+
33+
@pytest.fixture(scope="module")
34+
def region():
35+
"""Resolve the AWS region for this test module."""
36+
return "us-east-1"
37+
38+
39+
@pytest.fixture(scope="module")
40+
def account_id(region):
41+
"""Resolve the AWS account ID from the caller's credentials."""
42+
return boto3.Session(region_name=region).client("sts").get_caller_identity()["Account"]
43+
44+
45+
@pytest.fixture(scope="module")
46+
def s3_bucket(region, account_id):
47+
"""Construct the S3 bucket name from region and account ID."""
48+
return f"{S3_BUCKET_PREFIX}-{region}-{account_id}"
49+
50+
51+
@pytest.fixture(scope="module")
52+
def s3_client(region):
53+
"""Create an S3 client in the correct region."""
54+
return boto3.client("s3", region_name=region)
55+
56+
57+
@pytest.fixture(scope="module")
58+
def verified_training_dataset(s3_client, s3_bucket):
59+
"""Verify the training dataset exists in S3 and return its path."""
60+
s3_key = f"{S3_PREFIX}/sft-nova/sft_200_samples.jsonl"
61+
s3_path = f"s3://{s3_bucket}/{s3_key}"
62+
try:
63+
bucket_region = s3_client.get_bucket_location(Bucket=s3_bucket)[
64+
"LocationConstraint"
65+
] or "us-east-1"
66+
s3_regional_client = boto3.client("s3", region_name=bucket_region)
67+
s3_regional_client.head_object(Bucket=s3_bucket, Key=s3_key)
68+
except s3_client.exceptions.ClientError as e:
69+
if e.response["Error"]["Code"] in ("404", "NoSuchKey"):
70+
pytest.fail(
71+
f"Training file not found in S3: {s3_path}"
72+
)
73+
else:
74+
raise
75+
return s3_path
76+
77+
78+
@pytest.fixture(scope="module")
79+
def hyperpod_compute():
80+
"""Create HyperPod compute configuration."""
81+
return HyperPodCompute(
82+
cluster_name="pysdk-hp-integ-tests",
83+
namespace="kubeflow",
84+
instance_type="ml.g6.12xlarge",
85+
node_count=1,
86+
)
87+
88+
89+
@pytest.mark.gpu_intensive
90+
@pytest.mark.us_east_1
91+
def test_sft_trainer_nova_micro_hyperpod_lora(
92+
verified_training_dataset, hyperpod_compute, s3_bucket
93+
):
94+
"""Test SFT training workflow with Nova Micro on HyperPod using LORA."""
95+
unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}"
96+
s3_output_path = f"s3://{s3_bucket}/output"
97+
98+
stopping_condition = StoppingCondition(max_runtime_in_seconds=20800)
99+
100+
sft_trainer = SFTTrainer(
101+
model="nova-textgeneration-micro",
102+
training_type=TrainingType.LORA,
103+
compute=hyperpod_compute,
104+
base_job_name=f"hp-micro-sft-integ-{unique_id}",
105+
training_dataset=verified_training_dataset,
106+
s3_output_path=s3_output_path,
107+
stopping_condition=stopping_condition,
108+
overrides={"recipes.training_config.optim_config.lr": 5e-6},
109+
)
110+
111+
sft_trainer.hyperparameters.max_epochs = 1
112+
113+
# Submit — SFTTrainer routes to HyperPod based on compute type
114+
job_name = sft_trainer.train(wait=False)
115+
116+
logger.info(f"HyperPod SFT job submitted: {job_name}")
117+
118+
# Poll for job completion by checking for the manifest in S3
119+
from sagemaker.train.base_trainer import BaseTrainer
120+
121+
max_wait_time = 21600 # 6 hour timeout (HyperPod jobs can take longer)
122+
poll_interval = 60 # Check every 60 seconds
123+
start_time = time.time()
124+
checkpoint_path = None
125+
126+
while time.time() - start_time < max_wait_time:
127+
checkpoint_path = BaseTrainer._resolve_checkpoint_from_manifest(
128+
job_name=job_name,
129+
output_s3_path=s3_output_path,
130+
)
131+
if checkpoint_path:
132+
logger.info(f"Checkpoint resolved: {checkpoint_path}")
133+
break
134+
135+
elapsed = int(time.time() - start_time)
136+
logger.info(f"Waiting for manifest... ({elapsed}s elapsed)")
137+
time.sleep(poll_interval)
138+
139+
assert checkpoint_path is not None, (
140+
f"Job {job_name} did not produce a manifest within {max_wait_time}s"
141+
)
142+
logger.info(f"Training complete. Checkpoint: {checkpoint_path}")

0 commit comments

Comments
 (0)