diff --git a/sagemaker-core/src/sagemaker/core/modules/utils.py b/sagemaker-core/src/sagemaker/core/modules/utils.py index d50df1d48a..9f88da497c 100644 --- a/sagemaker-core/src/sagemaker/core/modules/utils.py +++ b/sagemaker-core/src/sagemaker/core/modules/utils.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Utils module.""" + from __future__ import absolute_import import os @@ -200,7 +201,7 @@ def validate_instance_preferences(compute) -> None: - instance_preferences is mutually exclusive with the classic single-cluster fields instance_type / instance_groups / - instance_placement_config. + instance_placement_config, and with managed spot training. - Instance types must not repeat across preferences. - Count mode: exactly one of the top-level instance_count (applies to whichever preference wins) with no per-preference instance_count, or an @@ -231,6 +232,12 @@ def _value(obj, field): f"instance_preferences is mutually exclusive with {field}; " "specify either a single fixed cluster or instance_preferences, not both." ) + # Spot is a bool: only an explicit True conflicts (False/None is the default). + if _value(compute, "enable_managed_spot_training") is True: + raise ValueError( + "instance_preferences is mutually exclusive with managed spot training " + "(enable_managed_spot_training=True)." + ) instance_types = [_value(p, "instance_type") for p in preferences] duplicates = sorted( diff --git a/sagemaker-core/src/sagemaker/core/tools/constants.py b/sagemaker-core/src/sagemaker/core/tools/constants.py index 6d80d1aef9..0768664920 100644 --- a/sagemaker-core/src/sagemaker/core/tools/constants.py +++ b/sagemaker-core/src/sagemaker/core/tools/constants.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Constants used in the code_generator modules.""" + import os CLASS_METHODS = set(["create", "add", "register", "import", "list", "get"]) @@ -134,5 +135,16 @@ "InstanceCount": "IntPipeVar", "VolumeSizeInGB": "IntPipeVar", "KeepAlivePeriodInSeconds": "IntPipeVar", + "SelectedInstanceCount": "IntPipeVar", + }, + "InstancePreference": { + "InstanceCount": "IntPipeVar", + }, + "ProcessingClusterConfig": { + "InstanceCount": "IntPipeVar", + "SelectedInstanceCount": "IntPipeVar", + }, + "ProcessingInstancePreference": { + "InstanceCount": "IntPipeVar", }, } diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 12734b1d2a..916574f0c2 100644 --- a/sagemaker-core/src/sagemaker/core/training/utils.py +++ b/sagemaker-core/src/sagemaker/core/training/utils.py @@ -264,7 +264,7 @@ def validate_instance_preferences(compute) -> None: - instance_preferences is mutually exclusive with the classic single-cluster fields instance_type / instance_groups / - instance_placement_config. + instance_placement_config, and with managed spot training. - Instance types must not repeat across preferences. - Count mode: exactly one of the top-level instance_count (applies to whichever preference wins) with no per-preference instance_count, or an @@ -295,6 +295,12 @@ def _value(obj, field): f"instance_preferences is mutually exclusive with {field}; " "specify either a single fixed cluster or instance_preferences, not both." ) + # Spot is a bool: only an explicit True conflicts (False/None is the default). + if _value(compute, "enable_managed_spot_training") is True: + raise ValueError( + "instance_preferences is mutually exclusive with managed spot training " + "(enable_managed_spot_training=True)." + ) instance_types = [_value(p, "instance_type") for p in preferences] duplicates = sorted( diff --git a/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py b/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py index d309079367..6abf57ee47 100644 --- a/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py +++ b/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py @@ -24,14 +24,11 @@ ``SelectedInstanceCount`` report the resolved winner, which must be one of the submitted preferences. -Configuration (environment variables; SKIPPED when unset): - -- ``PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN`` - execution role ARN -- ``PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI`` - processing image -- ``PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` - optional comma-separated - candidate types (default: ml.m5.xlarge, ml.m5.large) -- ``SAGEMAKER_ENDPOINT`` - optional endpoint - override. +Runs in the standard integration-test account: the execution role is the +suite's ``SageMakerRole`` and the image is resolved through ``image_uris``, the +same way the other ``sagemaker-core`` integration tests obtain theirs. +``PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` (comma-separated) +overrides the candidate types for accounts where the defaults lack quota. """ from __future__ import absolute_import @@ -40,11 +37,12 @@ import time import uuid -import pytest +import boto3 -ROLE_ARN = os.environ.get("PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN") -IMAGE_URI = os.environ.get("PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI") -ENDPOINT = os.environ.get("SAGEMAKER_ENDPOINT") +from sagemaker.core import image_uris + +ROLE = "SageMakerRole" +REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-west-2")) PREFERENCE_TYPES = [ t.strip() @@ -55,23 +53,27 @@ ] WAIT_TIMEOUT_SECONDS = 30 * 60 POLL_SECONDS = 30 - -pytestmark = pytest.mark.skipif( - not (ROLE_ARN and IMAGE_URI), - reason=( - "Processing Instance Preferences integ test requires " - "PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN and " - "PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI" - ), -) - - -def _sagemaker_client(): - """Build a boto3 SageMaker client, honoring SAGEMAKER_ENDPOINT.""" - import boto3 - - region = os.environ.get("AWS_REGION", "us-west-2") - return boto3.client("sagemaker", region_name=region, endpoint_url=ENDPOINT) +TERMINAL = ("Completed", "Failed", "Stopped") + + +def _stop_quietly(client, job_name): + """The job has done its part once a winner is selected; do not leave it + running to max_runtime on shared quota.""" + try: + if ( + client.describe_processing_job(ProcessingJobName=job_name)["ProcessingJobStatus"] + in TERMINAL + ): + return + client.stop_processing_job(ProcessingJobName=job_name) + except Exception: # pylint: disable=broad-except + pass + + +def _processing_image(): + return image_uris.retrieve( + "sklearn", REGION, version="1.2-1", py_version="py3", instance_type=PREFERENCE_TYPES[0] + ) def test_processor_instance_preferences_e2e(): @@ -79,12 +81,12 @@ def test_processor_instance_preferences_e2e(): from sagemaker.core.helper.session_helper import Session from sagemaker.core.processing import Processor - client = _sagemaker_client() + client = boto3.client("sagemaker", region_name=REGION) session = Session(sagemaker_client=client) processor = Processor( - role=ROLE_ARN, - image_uri=IMAGE_URI, + role=ROLE, + image_uri=_processing_image(), instance_count=1, instance_preferences=[{"InstanceType": t} for t in PREFERENCE_TYPES], volume_size_in_gb=30, @@ -95,39 +97,43 @@ def test_processor_instance_preferences_e2e(): job_name = f"instance-prefs-proc-integ-{uuid.uuid4().hex[:8]}" processor.run(wait=False, logs=False, job_name=job_name) - # --- Create accepted; Describe echoes the request contract ------------- - described = client.describe_processing_job(ProcessingJobName=job_name) - cluster_config = described["ProcessingResources"]["ClusterConfig"] - assert [p["InstanceType"] for p in cluster_config.get("InstancePreferences", [])] == ( - PREFERENCE_TYPES - ), f"Describe did not echo InstancePreferences: {cluster_config}" - # The customer never set the top-level InstanceType; it must not come - # back populated on Describe. - assert "InstanceType" not in cluster_config, cluster_config - - # --- Wait for a terminal-or-resolved state ------------------------------ - deadline = time.time() + WAIT_TIMEOUT_SECONDS - status = described["ProcessingJobStatus"] - while time.time() < deadline: + try: + + # --- Create accepted; Describe echoes the request contract ------------- described = client.describe_processing_job(ProcessingJobName=job_name) + cluster_config = described["ProcessingResources"]["ClusterConfig"] + assert [p["InstanceType"] for p in cluster_config.get("InstancePreferences", [])] == ( + PREFERENCE_TYPES + ), f"Describe did not echo InstancePreferences: {cluster_config}" + # The customer never set the top-level InstanceType; it must not come + # back populated on Describe. + assert "InstanceType" not in cluster_config, cluster_config + + # --- Wait for a terminal-or-resolved state ------------------------------ + deadline = time.time() + WAIT_TIMEOUT_SECONDS status = described["ProcessingJobStatus"] + while time.time() < deadline: + described = client.describe_processing_job(ProcessingJobName=job_name) + status = described["ProcessingJobStatus"] + cluster_config = described["ProcessingResources"]["ClusterConfig"] + if status in TERMINAL: + break + if status == "InProgress" and cluster_config.get("SelectedInstanceType"): + break + time.sleep(POLL_SECONDS) + + failure_reason = described.get("FailureReason", "") + + # --- Full contract: resolved winner is surfaced and is a submitted pref - cluster_config = described["ProcessingResources"]["ClusterConfig"] - if status in ("Completed", "Failed", "Stopped"): - break - if status == "InProgress" and cluster_config.get("SelectedInstanceType"): - break - time.sleep(POLL_SECONDS) - - failure_reason = described.get("FailureReason", "") - - # --- Full contract: resolved winner is surfaced and is a submitted pref - - cluster_config = described["ProcessingResources"]["ClusterConfig"] - selected_type = cluster_config.get("SelectedInstanceType") - selected_count = cluster_config.get("SelectedInstanceCount") - assert selected_type in PREFERENCE_TYPES, ( - f"SelectedInstanceType={selected_type!r} not among submitted preferences " - f"{PREFERENCE_TYPES} (job={job_name}, status={status}, " - f"failure={failure_reason!r})" - ) - assert selected_count == 1 - assert "InstanceType" not in cluster_config, cluster_config + selected_type = cluster_config.get("SelectedInstanceType") + selected_count = cluster_config.get("SelectedInstanceCount") + assert selected_type in PREFERENCE_TYPES, ( + f"SelectedInstanceType={selected_type!r} not among submitted preferences " + f"{PREFERENCE_TYPES} (job={job_name}, status={status}, " + f"failure={failure_reason!r})" + ) + assert selected_count == 1 + assert "InstanceType" not in cluster_config, cluster_config + finally: + _stop_quietly(client, job_name) diff --git a/sagemaker-core/tests/unit/test_compute_configs.py b/sagemaker-core/tests/unit/test_compute_configs.py index 5f1e6a1ebf..bf3b6c2981 100644 --- a/sagemaker-core/tests/unit/test_compute_configs.py +++ b/sagemaker-core/tests/unit/test_compute_configs.py @@ -1,4 +1,5 @@ """Unit tests for Compute and HyperPodCompute config classes.""" + import pytest from sagemaker.core.training.configs import Compute, HyperPodCompute @@ -169,6 +170,26 @@ def test_instance_type_rejected_with_preferences(self, compute_cls): instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")], ) + def test_managed_spot_rejected_with_preferences(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="mutually exclusive with managed spot training"): + compute_cls( + instance_count=1, + enable_managed_spot_training=True, + instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")], + ) + + def test_managed_spot_false_allowed_with_preferences(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + compute = compute_cls( + instance_count=1, + enable_managed_spot_training=False, + instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")], + ) + assert compute.enable_managed_spot_training is False + def test_uniform_count_rejected_with_per_preference_counts(self, compute_cls): from sagemaker.core.shapes.shapes import InstancePreference @@ -239,9 +260,7 @@ def test_whole_job_plan_allowed_without_per_preference_plans(self, compute_cls): compute_cls( instance_count=1, - training_plan_arn=( - "arn:aws:sagemaker:us-west-2:111122223333:training-plan/whole-job" - ), + training_plan_arn=("arn:aws:sagemaker:us-west-2:111122223333:training-plan/whole-job"), instance_preferences=[ InstancePreference(instance_type="ml.p5.48xlarge"), InstancePreference(instance_type="ml.p4d.24xlarge"), @@ -260,9 +279,7 @@ def test_modules_compute_instance_preferences_round_trip(self): InstancePreference(instance_type="ml.p5.48xlarge"), InstancePreference(instance_type="ml.p4d.24xlarge"), ] - rc = ModulesCompute( - instance_preferences=prefs, instance_count=2 - )._to_resource_config() + rc = ModulesCompute(instance_preferences=prefs, instance_count=2)._to_resource_config() assert [p.instance_type for p in rc.instance_preferences] == [ "ml.p5.48xlarge", "ml.p4d.24xlarge", @@ -271,9 +288,7 @@ def test_modules_compute_instance_preferences_round_trip(self): def test_modules_compute_single_type_still_works(self): from sagemaker.core.modules.configs import Compute as ModulesCompute - rc = ModulesCompute( - instance_type="ml.m5.xlarge", instance_count=1 - )._to_resource_config() + rc = ModulesCompute(instance_type="ml.m5.xlarge", instance_count=1)._to_resource_config() assert rc.instance_type == "ml.m5.xlarge" diff --git a/sagemaker-core/tests/unit/tools/test_shapes_extractor.py b/sagemaker-core/tests/unit/tools/test_shapes_extractor.py index 7881142c18..54fdcbf2e6 100644 --- a/sagemaker-core/tests/unit/tools/test_shapes_extractor.py +++ b/sagemaker-core/tests/unit/tools/test_shapes_extractor.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Unit tests for sagemaker.core.tools.shapes_extractor module.""" + from __future__ import absolute_import import pytest @@ -343,3 +344,42 @@ def test_get_required_members_none(self, mock_reformat): result = extractor.get_required_members("TestShape") assert len(result) == 0 + + +class TestInstancePreferencesPipeVarOverrides: + """The IntPipeVar annotations on instance-preferences count members come from + PIPE_VAR_OVERRIDES, not the service model. If an override is dropped, codegen + silently narrows the member back to int and pipeline variables stop being + accepted -- so assert the generated type directly.""" + + @pytest.fixture + def extractor(self, tmp_path): + # Constructing the extractor regenerates shape_dag.py; point that write + # at a temp file so the unit test leaves the checked-in file alone. + with ( + patch("sagemaker.core.tools.shapes_extractor.reformat_file_with_black"), + patch( + "sagemaker.core.tools.shapes_extractor.SHAPE_DAG_FILE_PATH", + str(tmp_path / "shape_dag.py"), + ), + ): + return ShapesExtractor() + + @pytest.mark.parametrize( + "shape, member", + [ + ("ResourceConfig", "instance_count"), + ("ResourceConfig", "selected_instance_count"), + ("InstancePreference", "instance_count"), + ("ProcessingClusterConfig", "instance_count"), + ("ProcessingClusterConfig", "selected_instance_count"), + ("ProcessingInstancePreference", "instance_count"), + ], + ) + def test_count_members_generate_as_int_pipe_var(self, extractor, shape, member): + members = extractor.generate_shape_members(shape) + assert member in members, f"{shape}.{member} missing from generated members" + assert "IntPipeVar" in members[member], ( + f"{shape}.{member} generated as {members[member]!r}; " + "expected IntPipeVar via PIPE_VAR_OVERRIDES" + ) diff --git a/sagemaker-train/tests/integ/train/shallow/test_instance_preferences.py b/sagemaker-train/tests/integ/train/shallow/test_instance_preferences.py new file mode 100644 index 0000000000..8bd2f65620 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_instance_preferences.py @@ -0,0 +1,166 @@ +# 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. +"""Shallow submission tests for ``Compute.instance_preferences``. + +Each accepted case submits a real ``CreateTrainingJob`` carrying an ordered +``InstancePreferences`` list (no top-level ``InstanceType``), asserts the +service returned an ARN and echoed the list on Describe, then stops the job. +Rejected cases assert the request is refused with the documented reason. + +Which instance type wins, and how the job runs on it, is training behaviour and +belongs in the deep suites. +""" + +from __future__ import absolute_import + +import os + +import pytest +from sagemaker.core import shapes +from sagemaker.core.shapes import InstancePreference +from sagemaker.core.training.configs import Compute, SourceCode +from sagemaker.train.model_trainer import ModelTrainer + +from .harness import ( + MAX_RUNTIME_IN_SECONDS, + assert_rejected, + assert_submitted, + cpu_image, + submitted, + unique_name, +) + +DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "data") +PARAM_SCRIPT_SOURCE_DIR = os.path.join(DATA_DIR, "params_script") + +# Two small CPU types. INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES (comma-separated) +# overrides them for accounts where these lack quota. +PREFERENCE_TYPES = [ + t.strip() + for t in os.environ.get( + "INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.large,ml.m5.xlarge" + ).split(",") + if t.strip() +] + + +def _trainer(sagemaker_session, name, compute): + return ModelTrainer( + sagemaker_session=sagemaker_session, + training_image=cpu_image(sagemaker_session), + source_code=SourceCode( + source_dir=PARAM_SCRIPT_SOURCE_DIR, + requirements="requirements.txt", + entry_script="train.py", + ), + compute=compute, + stopping_condition=shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS), + base_job_name=name, + ) + + +def _preferences(counts=None): + if counts is None: + return [InstancePreference(instance_type=t) for t in PREFERENCE_TYPES] + return [ + InstancePreference(instance_type=t, instance_count=c) + for t, c in zip(PREFERENCE_TYPES, counts) + ] + + +def _echoed_types(job): + job.refresh() + prefs = job.resource_config.instance_preferences + return [p.instance_type for p in prefs] if prefs else [] + + +class TestInstancePreferencesAccepted: + def test_uniform_count_is_accepted_and_echoed(self, sagemaker_session): + name = unique_name("shallow-ip-uniform") + trainer = _trainer( + sagemaker_session, + name, + Compute(instance_preferences=_preferences(), instance_count=1), + ) + + with submitted(trainer) as job: + assert_submitted(job) + assert _echoed_types(job) == PREFERENCE_TYPES + # The customer never set a top-level type; the service must not invent one. + assert not job.resource_config.instance_type + + def test_per_preference_counts_are_accepted_and_echoed(self, sagemaker_session): + name = unique_name("shallow-ip-per-pref") + trainer = _trainer( + sagemaker_session, + name, + Compute(instance_preferences=_preferences(counts=[1, 1])), + ) + + with submitted(trainer) as job: + assert_submitted(job) + assert _echoed_types(job) == PREFERENCE_TYPES + + +class TestInstancePreferencesRejected: + """Client-side rules must fail fast: a payload that reached the service with + both a fixed type and a list would be a regression in the SDK, not a + behaviour to leave for the backend to catch.""" + + def test_instance_type_with_preferences_is_rejected(self, sagemaker_session): + with pytest.raises(ValueError, match="mutually exclusive with instance_type"): + Compute( + instance_type=PREFERENCE_TYPES[0], + instance_count=1, + instance_preferences=_preferences(), + ) + + def test_managed_spot_with_preferences_is_rejected(self, sagemaker_session): + with pytest.raises(ValueError, match="mutually exclusive with managed spot training"): + Compute( + instance_count=1, + enable_managed_spot_training=True, + instance_preferences=_preferences(), + ) + + def test_duplicate_types_are_rejected(self, sagemaker_session): + with pytest.raises(ValueError, match="duplicate instance types"): + Compute( + instance_count=1, + instance_preferences=[ + InstancePreference(instance_type=PREFERENCE_TYPES[0]), + InstancePreference(instance_type=PREFERENCE_TYPES[0]), + ], + ) + + def test_server_rejects_more_than_five_preferences(self, sagemaker_session): + """The list cap is deliberately not enforced client-side (it is a + server-side tunable), so this is the one rule that must be asserted + against the service.""" + six = [ + "ml.m5.large", + "ml.m5.xlarge", + "ml.m5.2xlarge", + "ml.m4.xlarge", + "ml.c5.xlarge", + "ml.c5.2xlarge", + ] + trainer = _trainer( + sagemaker_session, + unique_name("shallow-ip-six"), + Compute( + instance_preferences=[InstancePreference(instance_type=t) for t in six], + instance_count=1, + ), + ) + assert_rejected(trainer, ("InstancePreferences", "Member must have length")) diff --git a/sagemaker-train/tests/integ/train/test_instance_preferences.py b/sagemaker-train/tests/integ/train/test_instance_preferences.py index c276809952..80b60e8a0e 100644 --- a/sagemaker-train/tests/integ/train/test_instance_preferences.py +++ b/sagemaker-train/tests/integ/train/test_instance_preferences.py @@ -10,141 +10,128 @@ # 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. -"""End-to-end integration test for Instance Preferences (multi-instance-type). - -Launches a REAL training job whose ``ResourceConfig`` carries an ordered -``InstancePreferences`` list (no top-level ``InstanceType``) through the -``ModelTrainer`` + ``Compute`` path, then asserts the Describe contract: - -- the create request is accepted with ``instance_preferences`` only; -- once the job leaves PENDING, ``selected_instance_type`` / - ``selected_instance_count`` report the resolved winner, which must be one - of the submitted preferences; -- the top-level ``instance_type`` the customer never set is not echoed back - populated. - -Configuration (environment variables; the test is SKIPPED when unset so the -suite stays green on hosts without the test-account setup): - -- ``INSTANCE_PREFERENCES_TEST_ROLE_ARN`` - SageMaker execution role ARN -- ``INSTANCE_PREFERENCES_TEST_IMAGE_URI`` - training image the account can pull -- ``INSTANCE_PREFERENCES_TEST_S3_OUTPUT`` - s3:// output path -- ``INSTANCE_PREFERENCES_TEST_S3_INPUT`` - optional s3:// input channel -- ``INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` - optional comma-separated - candidate types (default: ml.m5.xlarge, ml.m4.xlarge) -- ``SAGEMAKER_ENDPOINT`` - optional endpoint override; - honored by the sagemaker-core client loader. +"""End-to-end training test for ``Compute.instance_preferences``. + +Submits a real training job with an ordered list of CPU instance types and +no top-level ``instance_type``, then follows it until the service has picked +a winner and the job has run to completion on it. The shallow suite covers +acceptance and rejection; this test covers the part only the service can +prove: which type was selected and that training actually ran on it. """ from __future__ import absolute_import import os import time -import uuid - -import pytest -from sagemaker.core.utils.utils import Unassigned - -ROLE_ARN = os.environ.get("INSTANCE_PREFERENCES_TEST_ROLE_ARN") -IMAGE_URI = os.environ.get("INSTANCE_PREFERENCES_TEST_IMAGE_URI") -S3_OUTPUT = os.environ.get("INSTANCE_PREFERENCES_TEST_S3_OUTPUT") -S3_INPUT = os.environ.get("INSTANCE_PREFERENCES_TEST_S3_INPUT") +from sagemaker.core.shapes import InstancePreference, StoppingCondition +from sagemaker.train.configs import Compute, SourceCode +from sagemaker.train.model_trainer import ModelTrainer +DATA_DIR = os.path.join(os.path.dirname(__file__), "../..", "data") +PARAM_SCRIPT_SOURCE_CODE = SourceCode( + source_dir=f"{DATA_DIR}/params_script", + requirements="requirements.txt", + entry_script="train.py", +) +HYPERPARAMETERS = { + "integer": 1, + "boolean": True, + "float": 3.14, + "string": "Hello World", + "list": [1, 2, 3], + "dict": { + "string": "value", + "integer": 3, + "float": 3.14, + "list": [1, 2, 3], + "dict": {"key": "value"}, + "boolean": True, + }, +} +DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310" + +# CPU-only candidates. INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES (comma-separated) +# overrides them for accounts where these lack quota. PREFERENCE_TYPES = [ t.strip() for t in os.environ.get( - "INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.xlarge,ml.m4.xlarge" + "INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.large,ml.m5.xlarge" ).split(",") if t.strip() ] -WAIT_TIMEOUT_SECONDS = 30 * 60 +TERMINAL = ("Completed", "Failed", "Stopped") +MAX_RUNTIME_SECONDS = 1800 +WAIT_TIMEOUT_SECONDS = 40 * 60 POLL_SECONDS = 30 -pytestmark = pytest.mark.skipif( - not (ROLE_ARN and IMAGE_URI and S3_OUTPUT), - reason=( - "Instance Preferences integ test requires INSTANCE_PREFERENCES_TEST_ROLE_ARN, " - "INSTANCE_PREFERENCES_TEST_IMAGE_URI and INSTANCE_PREFERENCES_TEST_S3_OUTPUT" - ), -) - - -def _get_value(field): - """None for Unassigned/None, else the raw value.""" - if field is None or isinstance(field, Unassigned): - return None - return field +def _wait(client, job_name, until): + deadline = time.time() + WAIT_TIMEOUT_SECONDS + while True: + described = client.describe_training_job(TrainingJobName=job_name) + if until(described) or time.time() >= deadline: + return described + time.sleep(POLL_SECONDS) -def test_model_trainer_instance_preferences_e2e(): - """Create a real training job with instance_preferences and verify the winner.""" - from sagemaker.core.resources import TrainingJob - from sagemaker.core.shapes import shapes - from sagemaker.train.model_trainer import ModelTrainer - from sagemaker.train.configs import Compute, InputData, OutputDataConfig - job_prefix = f"instance-prefs-integ-{uuid.uuid4().hex[:8]}" +def _stop_quietly(client, job_name): + try: + if client.describe_training_job(TrainingJobName=job_name)["TrainingJobStatus"] in TERMINAL: + return + client.stop_training_job(TrainingJobName=job_name) + except Exception: # pylint: disable=broad-except + pass - compute = Compute( - instance_preferences=[shapes.InstancePreference(instance_type=t) for t in PREFERENCE_TYPES], - instance_count=1, - volume_size_in_gb=50, - ) - # The Compute config must forward the preference list and leave the - # top-level instance type unset. - resource_config = compute._to_resource_config() - assert [p.instance_type for p in resource_config.instance_preferences] == PREFERENCE_TYPES - assert _get_value(resource_config.instance_type) is None +def test_instance_preferences_select_a_winner_and_complete(sagemaker_session): trainer = ModelTrainer( - base_job_name=job_prefix, - training_image=IMAGE_URI, - role=ROLE_ARN, - compute=compute, - output_data_config=OutputDataConfig(s3_output_path=S3_OUTPUT), + sagemaker_session=sagemaker_session, + training_image=DEFAULT_CPU_IMAGE, + hyperparameters=HYPERPARAMETERS, + source_code=PARAM_SCRIPT_SOURCE_CODE, + compute=Compute( + instance_preferences=[InstancePreference(instance_type=t) for t in PREFERENCE_TYPES], + instance_count=1, + ), + stopping_condition=StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_SECONDS), + base_job_name="instance-prefs-e2e", ) + client = sagemaker_session.sagemaker_client - input_data_config = None - if S3_INPUT: - input_data_config = [InputData(channel_name="train", data_source=S3_INPUT)] - - trainer.train(input_data_config=input_data_config, wait=False) + trainer.train(wait=False, logs=False) job_name = trainer._latest_training_job.training_job_name - # --- Create accepted; Describe echoes the request contract ------------- - described = TrainingJob.get(training_job_name=job_name) - assert described.training_job_name == job_name - echoed = described.resource_config - # Top-level instance type was never set; must not come back populated. - assert _get_value(echoed.instance_type) is None - assert [p.instance_type for p in _get_value(echoed.instance_preferences) or []] == ( - PREFERENCE_TYPES - ), f"Describe did not echo instance_preferences (job={job_name})" - - # Wait until terminal or winner visible (Selected* propagation lags the - # secondary-status transitions; gate on the winner, not on secondary). - deadline = time.time() + WAIT_TIMEOUT_SECONDS - status = described.training_job_status - while time.time() < deadline: - described.refresh() - status = described.training_job_status - if status in ("Completed", "Failed", "Stopped"): - break - if _get_value(described.resource_config.selected_instance_type) is not None: - break - time.sleep(POLL_SECONDS) - - failure_reason = _get_value(described.failure_reason) or "" - - # --- Full contract: resolved winner is surfaced and is a submitted pref - - final_rc = described.resource_config - selected_type = _get_value(final_rc.selected_instance_type) - selected_count = _get_value(final_rc.selected_instance_count) - assert selected_type in PREFERENCE_TYPES, ( - f"selected_instance_type={selected_type!r} not among submitted " - f"preferences {PREFERENCE_TYPES} (job={job_name}, status={status}, " - f"failure={failure_reason!r})" - ) - assert selected_count == 1 - assert _get_value(final_rc.instance_type) is None + try: + described = client.describe_training_job(TrainingJobName=job_name) + resource_config = described["ResourceConfig"] + assert [p["InstanceType"] for p in resource_config["InstancePreferences"]] == ( + PREFERENCE_TYPES + ), resource_config + assert "InstanceType" not in resource_config, resource_config + + described = _wait( + client, + job_name, + lambda d: d["TrainingJobStatus"] in TERMINAL + or d["ResourceConfig"].get("SelectedInstanceType"), + ) + resource_config = described["ResourceConfig"] + selected_type = resource_config.get("SelectedInstanceType") + assert selected_type in PREFERENCE_TYPES, ( + f"SelectedInstanceType={selected_type!r} not among {PREFERENCE_TYPES} " + f"(job={job_name}, status={described['TrainingJobStatus']}, " + f"failure={described.get('FailureReason', '')!r})" + ) + assert resource_config.get("SelectedInstanceCount") == 1 + assert "InstanceType" not in resource_config, resource_config + + described = _wait(client, job_name, lambda d: d["TrainingJobStatus"] in TERMINAL) + assert described["TrainingJobStatus"] == "Completed", ( + f"job={job_name} ended {described['TrainingJobStatus']} on {selected_type}: " + f"{described.get('FailureReason', '')!r}" + ) + # The winner must not change once training has run on it. + assert described["ResourceConfig"]["SelectedInstanceType"] == selected_type + finally: + _stop_quietly(client, job_name)