diff --git a/sagemaker-serve/tests/integ/test_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_model_customization_deployment.py index fc51705919..77bb18985b 100644 --- a/sagemaker-serve/tests/integ/test_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_model_customization_deployment.py @@ -28,6 +28,11 @@ logger = logging.getLogger(__name__) from sagemaker.core.helper.session_helper import Session +from sagemaker.core.resources import TrainingJob, ModelPackage, InferenceComponent, Endpoint +from sagemaker.core.utils.exceptions import FailedStatusError +from sagemaker.serve import ModelBuilder +from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder +from sagemaker.serve.model_reuse import MODEL_SOURCE_TAG_KEY # This test relies on resources in a specific region AWS_REGION = "us-west-2" @@ -68,7 +73,6 @@ def model_package_arn(): @pytest.fixture def endpoint_name(): """Generate unique endpoint name.""" - import time return f"e2e-{int(time.time())}-{random.randint(100, 10000)}" @@ -80,7 +84,6 @@ def cleanup_endpoints(): for ep_name in endpoints_to_cleanup: try: - from sagemaker.core.resources import Endpoint endpoint = Endpoint.get(endpoint_name=ep_name, region=AWS_REGION) endpoint.delete() except Exception: @@ -92,9 +95,6 @@ class TestModelCustomizationFromTrainingJob: def test_build_from_training_job(self, training_job_name, sagemaker_session): """Test building model from training job.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.serve import ModelBuilder - import time training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, sagemaker_session=sagemaker_session) @@ -112,11 +112,7 @@ def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanu For LORA models, this verifies the two-step deployment: base IC + adapter IC are both created on the same endpoint. """ - from sagemaker.core.resources import TrainingJob, InferenceComponent - from sagemaker.serve import ModelBuilder - import time - from sagemaker.core.utils.exceptions import FailedStatusError training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) @@ -148,7 +144,6 @@ def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanu assert endpoint.endpoint_status == "InService" # Verify model-source tag is present on the endpoint for reuse discovery. - from sagemaker.serve.model_reuse import MODEL_SOURCE_TAG_KEY sm_client = boto3.client("sagemaker", region_name=AWS_REGION) endpoint_tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get("Tags", []) assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in endpoint_tags}, ( @@ -200,8 +195,6 @@ def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanu def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_session): """Test fetching endpoint names for base model.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.serve import ModelBuilder training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, sagemaker_session=sagemaker_session) @@ -209,13 +202,78 @@ def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_ assert isinstance(endpoint_names, set) + def test_deploy_reuse_returns_existing_endpoint(self, training_job_name, endpoint_name, cleanup_endpoints, sagemaker_session): + """deploy(reuse_resources=True) finds and returns an existing tagged endpoint. + + Flow: + 1. First deploy creates an endpoint with the model-source tag. + 2. Second deploy with reuse_resources=True returns the same endpoint. + """ + + training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) + model_builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) + model_builder.accept_eula = True + model_builder.build(model_name=f"test-model-{int(time.time())}-{random.randint(100, 10000)}", region=AWS_REGION) + + try: + endpoint = model_builder.deploy(endpoint_name=endpoint_name) + except (FailedStatusError, ClientError) as e: + msg = str(e) + if "InsufficientInstanceCapacity" in msg or "ResourceLimitExceeded" in msg: + cleanup_endpoints.append(endpoint_name) + pytest.xfail(f"Capacity/quota limit: {e}") + raise + + cleanup_endpoints.append(endpoint_name) + assert endpoint is not None + assert endpoint.endpoint_status == "InService" + + # Second deploy with reuse should find the same endpoint + builder2 = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) + builder2.accept_eula = True + builder2.build(region=AWS_REGION, reuse_resources=True) + endpoint2 = builder2.deploy(reuse_resources=True) + + assert endpoint2 is not None + assert endpoint2.endpoint_arn == endpoint.endpoint_arn, ( + f"Expected reuse to return {endpoint.endpoint_arn}, got {endpoint2.endpoint_arn}" + ) + + def test_build_reuse_skips_model_creation(self, training_job_name, sagemaker_session): + """build(reuse_resources=True) reuses an existing tagged Model.""" + + training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) + # Note: tiny collision risk if parallel CI runs share the same timestamp+random seed + unique_id = f"{int(time.time())}-{random.randint(100, 10000)}" + model_name = f"reuse-build-{unique_id}" + + builder1 = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) + builder1.accept_eula = True + model1 = builder1.build(model_name=model_name, region=AWS_REGION, reuse_resources=False) + assert model1 is not None + assert model1.model_arn is not None + + # Second build with reuse + builder2 = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) + builder2.accept_eula = True + model2 = builder2.build(region=AWS_REGION, reuse_resources=True) + assert model2 is not None + assert model2.model_arn == model1.model_arn, ( + f"Expected build reuse to return {model1.model_arn}, got {model2.model_arn}" + ) + + # Cleanup: delete the model created during this test + sm_client = boto3.client("sagemaker", region_name=AWS_REGION) + try: + sm_client.delete_model(ModelName=model_name) + except Exception: + pass + class TestModelCustomizationFromModelPackage: def test_build_from_model_package(self, model_package_arn, sagemaker_session): """Test building model from model package.""" - from sagemaker.core.resources import ModelPackage - from sagemaker.serve import ModelBuilder model_package = ModelPackage.get(model_package_name=model_package_arn, region=AWS_REGION) model_builder = ModelBuilder(model=model_package, sagemaker_session=sagemaker_session) @@ -227,9 +285,6 @@ def test_build_from_model_package(self, model_package_arn, sagemaker_session): def test_deploy_from_model_package(self, model_package_arn, cleanup_endpoints, sagemaker_session): """Test deploying model from model package.""" - from sagemaker.core.resources import ModelPackage - from sagemaker.serve import ModelBuilder - import time model_package = ModelPackage.get(model_package_name=model_package_arn, region=AWS_REGION) endpoint_name = f"e2e-{int(time.time())}-{random.randint(100, 10000)}" @@ -249,8 +304,6 @@ class TestInstanceTypeAutoDetection: def test_instance_type_from_recipe(self, training_job_name, sagemaker_session): """Test instance type auto-detection from recipe.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.serve import ModelBuilder training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, sagemaker_session=sagemaker_session) @@ -266,8 +319,6 @@ class TestModelCustomizationDetection: def test_is_model_customization_training_job(self, training_job_name, sagemaker_session): """Test detection from training job.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.serve import ModelBuilder training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, sagemaker_session=sagemaker_session) @@ -276,8 +327,6 @@ def test_is_model_customization_training_job(self, training_job_name, sagemaker_ def test_is_model_customization_model_package(self, model_package_arn, sagemaker_session): """Test detection from model package.""" - from sagemaker.core.resources import ModelPackage - from sagemaker.serve import ModelBuilder model_package = ModelPackage.get(model_package_name=model_package_arn, region=AWS_REGION) model_builder = ModelBuilder(model=model_package, sagemaker_session=sagemaker_session) @@ -286,8 +335,6 @@ def test_is_model_customization_model_package(self, model_package_arn, sagemaker def test_fetch_model_package_arn(self, training_job_name, sagemaker_session): """Test fetching model package ARN.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.serve import ModelBuilder training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder(model=training_job, sagemaker_session=sagemaker_session) @@ -303,9 +350,6 @@ class TestTrainerIntegration: def test_sft_trainer_build(self, training_job_name, sagemaker_session): """Test building model from SFTTrainer.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.train.sft_trainer import SFTTrainer - from sagemaker.serve import ModelBuilder training_job = TrainingJob.get( training_job_name=training_job_name, region=AWS_REGION @@ -327,9 +371,6 @@ def test_sft_trainer_build(self, training_job_name, sagemaker_session): def test_dpo_trainer_build(self, training_job_name, sagemaker_session): """Test building model from DPOTrainer.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.train.dpo_trainer import DPOTrainer - from sagemaker.serve import ModelBuilder from unittest.mock import patch training_job = TrainingJob.get( @@ -363,9 +404,6 @@ def test_dpo_trainer_build(self, training_job_name, sagemaker_session): - Improved test assertions to work with new object structures """ -from sagemaker.core.resources import TrainingJob, ModelPackage -from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder - @pytest.mark.serial class TestModelCustomizationDeployment: @@ -374,7 +412,6 @@ class TestModelCustomizationDeployment: @pytest.fixture(scope="class") def setup_config(self, training_job_name): """Setup test configuration.""" - from sagemaker.core.helper.session_helper import get_execution_role return { "training_job_name": training_job_name, "region": AWS_REGION, diff --git a/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py index 09c9e069eb..6655f16f58 100644 --- a/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py @@ -26,13 +26,13 @@ import pytest import random from sagemaker.serve import ModelBuilder +from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder from sagemaker.serve.model_reuse import MODEL_SOURCE_TAG_KEY -from sagemaker.core.resources import TrainingJob +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.resources import TrainingJob, Endpoint logger = logging.getLogger(__name__) -from sagemaker.core.helper.session_helper import Session - # This test relies on resources in a specific region AWS_REGION = "us-east-1" os.environ.setdefault("AWS_DEFAULT_REGION", AWS_REGION) @@ -172,7 +172,6 @@ def cleanup_endpoints(): for ep_name in endpoints_to_cleanup: try: - from sagemaker.core.resources import Endpoint endpoint = Endpoint.get(endpoint_name=ep_name, region=AWS_REGION) endpoint.delete() except Exception: @@ -203,7 +202,13 @@ def test_build_from_training_job(self, training_job_name, sagemaker_session): assert model_builder.instance_type is not None def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanup_endpoints, sagemaker_session): - """Test deploying a Nova model from a training job and invoking it.""" + """Test deploying a Nova model from a training job, invoking it, and reusing it. + + For Nova models, this verifies: + 1. Endpoint is created InService with model-source tag. + 2. Endpoint is invokable. + 3. deploy(reuse_resources=True) from a new ModelBuilder returns the same endpoint. + """ training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) model_builder = ModelBuilder( model=training_job, @@ -251,6 +256,34 @@ def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanu assert response_body is not None, f"Empty response from invoke on {endpoint_name}" assert isinstance(response_body, dict) + # Verify reuse: a new ModelBuilder with reuse_resources=True should + # find and return the same endpoint without creating a new one. + builder2 = ModelBuilder( + model=training_job, + instance_type=NOVA_INSTANCE_TYPE, + sagemaker_session=sagemaker_session, + ) + builder2.accept_eula = True + builder2.build(region=AWS_REGION, reuse_resources=True) + + # build(reuse_resources=True) should reuse the existing Model (no new one created) + sm_client = boto3.client("sagemaker", region_name=AWS_REGION) + model_name_1 = model_builder.built_model.model_name + models_with_prefix = sm_client.list_models( + NameContains=model_name_1, MaxResults=10 + ).get("Models", []) + assert len(models_with_prefix) == 1, ( + f"Expected 1 model with name '{model_name_1}', got {len(models_with_prefix)}. " + f"build(reuse_resources=True) should not create a new Model." + ) + + endpoint2 = builder2.deploy(reuse_resources=True) + + assert endpoint2 is not None + assert endpoint2.endpoint_arn == endpoint.endpoint_arn, ( + f"Expected reuse to return {endpoint.endpoint_arn}, got {endpoint2.endpoint_arn}" + ) + def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_session): """Test fetching endpoint names for base model.""" training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) @@ -259,6 +292,54 @@ def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_ assert isinstance(endpoint_names, set) + def test_build_reuse_skips_model_creation(self, training_job_name, sagemaker_session): + """build(reuse_resources=True) reuses an existing tagged Model. + + First build creates a Model, second build with reuse finds it. + """ + training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) + # Note: tiny collision risk if parallel CI runs share the same timestamp+random seed + unique_id = f"{int(time.time())}-{random.randint(100, 10000)}" + model_name = f"nova-reuse-build-{unique_id}" + + # First build + builder1 = ModelBuilder( + model=training_job, + instance_type=NOVA_INSTANCE_TYPE, + sagemaker_session=sagemaker_session, + ) + builder1.accept_eula = True + model1 = builder1.build( + model_name=model_name, + region=AWS_REGION, + reuse_resources=False, + ) + assert model1 is not None + assert model1.model_arn is not None + + # Second build with reuse + builder2 = ModelBuilder( + model=training_job, + instance_type=NOVA_INSTANCE_TYPE, + sagemaker_session=sagemaker_session, + ) + builder2.accept_eula = True + model2 = builder2.build( + region=AWS_REGION, + reuse_resources=True, + ) + assert model2 is not None + assert model2.model_arn == model1.model_arn, ( + f"Expected build reuse to return {model1.model_arn}, got {model2.model_arn}" + ) + + # Cleanup: delete the model created during this test + sm_client = boto3.client("sagemaker", region_name=AWS_REGION) + try: + sm_client.delete_model(ModelName=model_name) + except Exception: + pass + @pytest.mark.us_east_1 class TestModelCustomizationFromModelPackage: @@ -333,7 +414,6 @@ def test_is_model_customization_training_job(self, training_job_name, sagemaker_ def test_is_model_customization_model_package(self, model_package_arn, sagemaker_session): """Test detection from a Nova model package.""" - from sagemaker.core.resources import ModelPackage model_package = ModelPackage.get(model_package_name=model_package_arn, region=AWS_REGION) model_builder = ModelBuilder(model=model_package, sagemaker_session=sagemaker_session) @@ -360,7 +440,6 @@ class TestTrainerIntegration: def test_sft_trainer_build(self, training_job_name, sagemaker_session): """Test building a model from a Nova SFTTrainer object.""" - from sagemaker.train.sft_trainer import SFTTrainer training_job = TrainingJob.get( training_job_name=training_job_name, region=AWS_REGION @@ -387,7 +466,6 @@ def test_sft_trainer_build(self, training_job_name, sagemaker_session): def test_rlvr_trainer_build(self, training_job_name, sagemaker_session): """Test building a model from a Nova RLVRTrainer object.""" - from sagemaker.train.rlvr_trainer import RLVRTrainer training_job = TrainingJob.get( training_job_name=training_job_name, region=AWS_REGION @@ -420,7 +498,6 @@ class TestNovaBedrockDeployment: @pytest.fixture(scope="class") def role_arn(self): """Execution role ARN with Bedrock permissions.""" - from sagemaker.core.helper.session_helper import get_execution_role return get_execution_role() @pytest.fixture(scope="class") @@ -440,8 +517,6 @@ def deployed_nova_model(self, training_job_name, role_arn, bedrock_client): """Deploy a Nova model to Bedrock from its TrainingJob and yield the deployment details, cleaning up the custom model and deployment after. """ - from sagemaker.core.resources import TrainingJob - from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder unique = f"{int(time.time())}-{random.randint(1000, 9999)}" custom_model_name = f"nova-integ-{unique}" @@ -499,14 +574,33 @@ def test_nova_bedrock_deployment_active(self, deployed_nova_model, bedrock_clien ) assert deployment.get("status") == "Active" - def test_nova_bedrock_custom_model_tagged_for_reuse(self, deployed_nova_model, bedrock_client): - """The Nova custom model should carry the model-source tag that powers reuse.""" + def test_nova_bedrock_custom_model_tagged_for_reuse(self, deployed_nova_model, training_job_name, role_arn, bedrock_client): + """The Nova custom model should carry the model-source tag and be discoverable via reuse.""" + model_arn = deployed_nova_model["model_arn"] tags = bedrock_client.list_tags_for_resource(resourceARN=model_arn).get("tags", []) assert MODEL_SOURCE_TAG_KEY in {t["key"] for t in tags}, ( f"Custom model {model_arn} missing model-source tag for reuse" ) + # Verify reuse: a second deploy with reuse_resources=True should find the + # existing model instead of creating a new one. + training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) + builder2 = BedrockModelBuilder(model=training_job) + + unique = f"{int(time.time())}-{random.randint(1000, 9999)}" + response2 = builder2.deploy( + custom_model_name=f"nova-reuse-integ-{unique}", + deployment_name=f"nova-reuse-integ-{unique}-dep", + role_arn=role_arn, + reuse_resources=True, + ) + + reused_model_arn = response2.get("modelArn") or response2.get("importedModelArn") + assert reused_model_arn == model_arn, ( + f"Expected reuse to return {model_arn}, got {reused_model_arn}" + ) + @pytest.mark.slow def test_nova_bedrock_invoke(self, deployed_nova_model, bedrock_runtime): """Invoke the deployed Nova model on Bedrock end-to-end.""" @@ -531,3 +625,33 @@ def test_nova_bedrock_invoke(self, deployed_nova_model, bedrock_runtime): assert isinstance(result, dict) text = result["output"]["message"]["content"][0]["text"] assert isinstance(text, str) and len(text) > 0 + + @pytest.mark.slow + def test_nova_bedrock_reuse_returns_existing_model( + self, deployed_nova_model, training_job_name, role_arn, bedrock_client + ): + """deploy(reuse_resources=True) finds the existing custom model instead of creating new. + + Uses the model created by the deployed_nova_model fixture (already Active + and tagged with model-source). A second deploy with reuse should return + the same model ARN. + """ + + existing_model_arn = deployed_nova_model["model_arn"] + + training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) + builder2 = BedrockModelBuilder(model=training_job) + + unique = f"{int(time.time())}-{random.randint(1000, 9999)}" + response2 = builder2.deploy( + custom_model_name=f"nova-reuse-integ-{unique}", + deployment_name=f"nova-reuse-integ-{unique}-dep", + role_arn=role_arn, + reuse_resources=True, + ) + + reused_model_arn = response2.get("modelArn") or response2.get("importedModelArn") + + assert reused_model_arn == existing_model_arn, ( + f"Expected reuse to return {existing_model_arn}, got {reused_model_arn}" + ) diff --git a/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py b/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py index 880900f4b8..556e646a3a 100644 --- a/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_mtrl_trainer_integration.py @@ -247,3 +247,34 @@ def test_evaluate_comparison(self, attached_trainer, config): f"[{config['env_name']}] Comparison eval failed with status: {status}, " f"reason: {execution.status.failure_reason}" ) + + +@pytest.mark.gpu_intensive +class TestMTRLShowMetrics: + """Integration tests for show_metrics() on completed MTRL jobs.""" + + def test_show_metrics_on_completed_job(self, config): + """show_metrics() and stream_logs() on a completed MTRL job run without error.""" + job = MultiTurnRLTrainer.attach(job_name=config["existing_job_name"]) + assert job.job_status == "Completed" + + trainer = MultiTurnRLTrainer( + model=config["base_model"], + agent_env=config["agent_core_arn"], + training_dataset=config["dataset"], + output_model_package_group=config["model_package_group"], + mlflow_app_arn=config["mlflow_resource_arn"], + s3_output_path=config["s3_output_path"], + role=config["role"], + accept_eula=True, + ) + trainer._latest_job = job + + result = trainer.show_metrics() + logger.info( + f"[{config['env_name']}] show_metrics() returned: {type(result).__name__}" + ) + + # stream_logs() should exit quickly for a completed job + trainer.stream_logs(poll=2) + logger.info(f"[{config['env_name']}] stream_logs() completed without error") 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 78e301b5a3..b9262cb6e5 100644 --- a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py @@ -23,8 +23,8 @@ from sagemaker.train.common import TrainingType @pytest.mark.gpu_intensive -def test_sft_trainer_lora_complete_workflow(sagemaker_session): - """Test complete SFT training workflow with LORA.""" +def test_sft_trainer_lora_complete_workflow(sagemaker_session, mlflow_resource_arn): + """Test complete SFT training workflow with LORA, including show_metrics via MLflow.""" unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}" sft_trainer = SFTTrainer( @@ -33,6 +33,7 @@ def test_sft_trainer_lora_complete_workflow(sagemaker_session): model_package_group="arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models", training_dataset="s3://mc-flows-sdk-testing/input_data/sft/sample_data_256_final.jsonl", s3_output_path="s3://mc-flows-sdk-testing/output/", + mlflow_resource_arn=mlflow_resource_arn, accept_eula=True, base_job_name=f"sft-lora-integ-{unique_id}", ) @@ -59,6 +60,14 @@ def test_sft_trainer_lora_complete_workflow(sagemaker_session): assert hasattr(training_job, 'output_model_package_arn') assert training_job.output_model_package_arn is not None + # Verify show_metrics() works via MLflow path for OSS models + result = sft_trainer.show_metrics() + # OSS MLflow path renders inline; may return None or a DataFrame + print(f"show_metrics() returned: {type(result).__name__}") + + # Verify stream_logs() exits without error on a completed job + sft_trainer.stream_logs(poll=2) + @pytest.mark.gpu_intensive def test_sft_trainer_with_validation_dataset(sagemaker_session): @@ -99,7 +108,7 @@ def test_sft_trainer_with_validation_dataset(sagemaker_session): @pytest.mark.gpu_intensive @pytest.mark.us_east_1 def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1): - """Test SFT trainer with Nova model.""" + """Test SFT trainer with Nova model, including show_metrics() after completion.""" # sagemaker_session_us_east_1 fixture is defined in conftest.py (us-east-1 region) unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}" @@ -137,6 +146,34 @@ def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1): assert hasattr(training_job, 'output_model_package_arn') assert training_job.output_model_package_arn is not None + # Verify show_metrics() returns valid training metrics + # Use non-interactive backend so plt.show() doesn't require a display in CI + import matplotlib + matplotlib.use("Agg") + + df = sft_trainer_nova.show_metrics() + assert df is not None, "show_metrics() returned None" + assert not df.empty, "show_metrics() returned empty DataFrame" + assert "global_step" in df.columns, ( + f"Expected 'global_step' column, got: {list(df.columns)}" + ) + assert len(df) > 0 + + # Verify metric filter works + df_filtered = sft_trainer_nova.show_metrics(metrics=["training_loss"]) + assert not df_filtered.empty + assert set(df_filtered.columns) == {"global_step", "training_loss"} + + # Verify step range filter works + min_step = int(df["global_step"].min()) + max_step = int(df["global_step"].max()) + if max_step > min_step: + mid = (min_step + max_step) // 2 + df_range = sft_trainer_nova.show_metrics(starting_step=mid, ending_step=max_step) + assert not df_range.empty + assert df_range["global_step"].min() >= mid + assert df_range["global_step"].max() <= max_step + def test_sft_trainer_lora_invalid_instance_type_raises(sagemaker_session): """An unsupported instance type must raise before a job is submitted. diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py b/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py index 28020feb9d..5c1baccf0e 100644 --- a/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py +++ b/sagemaker-train/tests/integ/train/test_sft_trainer_serverful_smtj.py @@ -163,6 +163,44 @@ def test_sft_trainer_serverful_smtj(sagemaker_session_us_east_1, training_resour ) logger.info(f"Training job completed successfully: {training_job.training_job_name}") + # Verify show_metrics() returns valid training metrics after completion + # Use non-interactive backend so plt.show() doesn't require a display in CI + import matplotlib + matplotlib.use("Agg") + + df = sft_trainer.show_metrics() + assert df is not None, "show_metrics() returned None" + assert not df.empty, "show_metrics() returned empty DataFrame" + assert "global_step" in df.columns, ( + f"Expected 'global_step' column, got: {list(df.columns)}" + ) + assert len(df) > 0 + logger.info( + f"show_metrics() returned {len(df)} rows, columns: {list(df.columns)}" + ) + + # Verify metric filter + df_filtered = sft_trainer.show_metrics(metrics=["training_loss"]) + assert not df_filtered.empty + assert set(df_filtered.columns) == {"global_step", "training_loss"} + + # Verify step range filter + min_step = int(df["global_step"].min()) + max_step = int(df["global_step"].max()) + if max_step > min_step: + mid = (min_step + max_step) // 2 + df_range = sft_trainer.show_metrics(starting_step=mid, ending_step=max_step) + assert not df_range.empty + assert df_range["global_step"].min() >= mid + assert df_range["global_step"].max() <= max_step + logger.info( + f"Step range [{mid}, {max_step}] returned {len(df_range)}/{len(df)} rows" + ) + + # Verify stream_logs() exits without error on a completed job + sft_trainer.stream_logs(poll=2) + logger.info("stream_logs() completed without error") + @pytest.mark.us_east_1 def test_sft_trainer_serverful_smtj_invalid_instance_type_raises( diff --git a/sagemaker-train/tests/integ/train/test_stream_logs_trainer.py b/sagemaker-train/tests/integ/train/test_stream_logs_trainer.py index b6a952bf49..940cc989c7 100644 --- a/sagemaker-train/tests/integ/train/test_stream_logs_trainer.py +++ b/sagemaker-train/tests/integ/train/test_stream_logs_trainer.py @@ -15,6 +15,7 @@ import logging import time +from datetime import datetime, timezone import boto3 import pytest @@ -60,13 +61,17 @@ def test_stream_logs_with_start_time(self, sagemaker_session): """AgentRFTJob.stream_logs() respects start_time parameter.""" job = AgentRFTJob.get(MTRL_JOB_NAME, session=sagemaker_session.boto_session) - future_ms = int((time.time() + 86400) * 1000) + # Use a timestamp from when the job was running (extracted from job name) + + job_start = datetime(2026, 7, 29, 12, 9, 59, tzinfo=timezone.utc) + start_time_ms = int(job_start.timestamp() * 1000) + start = time.time() - job.stream_logs(poll=2, start_time=future_ms) + job.stream_logs(poll=2, start_time=start_time_ms) elapsed = time.time() - start assert elapsed < 30 - print(f"✓ stream_logs(start_time=future) completed in {elapsed:.1f}s") + print(f"✓ stream_logs(start_time=job_start) completed in {elapsed:.1f}s")