diff --git a/docs/ml_ops/index.rst b/docs/ml_ops/index.rst index 31923939c7..5fd4d6cbbe 100644 --- a/docs/ml_ops/index.rst +++ b/docs/ml_ops/index.rst @@ -317,6 +317,26 @@ Run data preprocessing with ``ScriptProcessor`` (sklearn) or ``FrameworkProcesso :doc:`SKLearn example <../v3-examples/ml-ops-examples/v3-processing-job-sklearn>` ยท :doc:`PyTorch example <../v3-examples/ml-ops-examples/v3-processing-job-pytorch/v3-pytorch-processing-example>` +**Instance Preferences:** pass an ordered list of candidate instance types and the platform runs the job on the first type with available capacity. + +.. code-block:: python + + from sagemaker.core.processing import Processor + + processor = Processor( + role=role, image_uri=processing_image, volume_size_in_gb=100, + instance_preferences=[ + {"InstanceType": "ml.m5.4xlarge", "InstanceCount": 2}, + {"InstanceType": "ml.m5.2xlarge", "InstanceCount": 4}, + ], + ) + + processor.run(job_name="instance-prefs-processing") + +Up to 5 candidates are allowed, each instance type at most once, and exactly one is selected; the list is mutually exclusive with ``instance_type``. Counts use exactly one of two modes โ a top-level ``instance_count`` shared by whichever candidate wins, or an ``InstanceCount`` on every candidate โ and mixed, partial, or omitted counts are rejected. Selection is based on capacity, not on workload fit, so list only types the job can genuinely run on. The winner is reported as ``SelectedInstanceType`` / ``SelectedInstanceCount`` on the job's ``ClusterConfig``, and billing is for that type and count. Supported on ``Processor``, ``ScriptProcessor``, ``PySparkProcessor``, and ``SparkJarProcessor``; training plans are training-only and do not apply to processing. + +:doc:`Instance Preferences example <../v3-examples/ml-ops-examples/v3-processing-instance-preferences>` + Batch Transform Jobs @@ -728,5 +748,6 @@ Explore comprehensive MLOps examples: ../v3-examples/ml-ops-examples/v3-model-registry-example/v3-model-registry-example ../v3-examples/ml-ops-examples/v3-processing-job-pytorch/v3-pytorch-processing-example ../v3-examples/ml-ops-examples/v3-processing-job-sklearn + ../v3-examples/ml-ops-examples/v3-processing-instance-preferences ../v3-examples/ml-ops-examples/v3-emr-serverless-step-example ../v3-examples/ml-ops-examples/v3-mlflow-train-inference-e2e-example diff --git a/docs/training/index.rst b/docs/training/index.rst index a69e7bf6f5..858afbe321 100644 --- a/docs/training/index.rst +++ b/docs/training/index.rst @@ -388,6 +388,80 @@ Key points: +Instance Preferences +-------------------- + + +Provide an ordered list of candidate instance types and the platform launches the job on the first type with available capacity, instead of failing when one scarce type is unavailable. + +**Ordered Candidates:** + +.. code-block:: python + + from sagemaker.train.model_trainer import ModelTrainer + from sagemaker.core.training.configs import Compute, SourceCode + from sagemaker.core.shapes import InstancePreference + + compute = Compute( + instance_preferences=[ + InstancePreference(instance_type="ml.p5.48xlarge"), + InstancePreference(instance_type="ml.p4d.24xlarge"), + InstancePreference(instance_type="ml.g5.12xlarge"), + ], + instance_count=2, + ) + + model_trainer = ModelTrainer( + training_image=training_image, + source_code=SourceCode(source_dir="./source", entry_script="train.py"), + compute=compute, + base_job_name="instance-preferences-training", + ) + + model_trainer.train() + +**Per-Candidate Instance Counts and Training Plans:** + +.. code-block:: python + + compute = Compute( + instance_preferences=[ + InstancePreference( + instance_type="ml.p5.48xlarge", + instance_count=2, + training_plan_arns=[p5_plan_arn], + ), + InstancePreference(instance_type="ml.p4d.24xlarge", instance_count=4), + ], + ) + +**Identifying the Instance Type That Ran:** + +.. code-block:: python + + training_job = model_trainer._latest_training_job + training_job.refresh() + + resource_config = training_job.resource_config + print(resource_config.selected_instance_type) # None until a candidate is selected + print(resource_config.selected_instance_count) + +Key points: + +- Up to 5 candidates per job, each instance type at most once, and exactly one type is selected +- ``instance_preferences`` is mutually exclusive with ``instance_type``, ``instance_groups``, ``instance_placement_config``, and managed spot training +- Not supported in local mode, with training recipes or JumpStart models, or for jobs submitted through AWS Batch training queues; jobs that leave ``instance_preferences`` unset are unaffected +- Counts use exactly one of two modes: a top-level ``instance_count`` shared by whichever candidate wins, or an ``instance_count`` on every candidate to size each type differently. Mixed, partial, and omitted counts are rejected +- A candidate with ``training_plan_arns`` draws from that plan's reserved capacity (one plan, whose instance type must match the candidate's); one without a plan uses on-demand. Per-candidate plans are mutually exclusive with the job-level ``training_plan_arn``, which instead applies to whichever candidate matches its type +- Selection is based on capacity, not on workload fit: cross-type differences such as architecture or GPU memory are not validated, so list only types your job can genuinely run on +- While no candidate has capacity the job stays pending and keeps retrying the list. ``max_pending_time_in_seconds`` bounds the total time spent working through the list rather than each candidate, and takes effect only when the list includes an accelerated instance type (``ml.p``, ``ml.g``, ``ml.trn``) +- The winner is reported as ``selected_instance_type`` / ``selected_instance_count`` on describe; the top-level instance type is not returned +- Job-level settings such as volume size, volume KMS key, and keep-alive period apply to the selected type, as does billing + +:doc:`Full example notebook <../v3-examples/training-examples/instance-preferences-example>` + + + AWS Batch Training Queues ------------------------- @@ -438,6 +512,7 @@ Key points: - Batch manages capacity allocation and job scheduling automatically - Resources (Service Environments, Job Queues) can be created via console or programmatically - Supports FIFO and priority-based scheduling +- Queued jobs request a single ``instance_type``; ``instance_preferences`` is not supported :doc:`Full example notebook <../v3-examples/training-examples/aws_batch/sm-training-queues_getting_started_with_model_trainer>` @@ -555,4 +630,5 @@ Training Examples Hyperparameter Training <../v3-examples/training-examples/hyperparameter-training-example> Training with JumpStart Models <../v3-examples/training-examples/jumpstart-training-example> Custom Distributed Training <../v3-examples/training-examples/custom-distributed-training-example> + Instance Preferences <../v3-examples/training-examples/instance-preferences-example> AWS Batch for Training <../v3-examples/training-examples/aws_batch/sm-training-queues_getting_started_with_model_trainer> diff --git a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json index 6b551c5fd6..b3e24cdd40 100644 --- a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json +++ b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json @@ -28787,6 +28787,33 @@ }, "documentation":"
Defines an instance group for heterogeneous cluster training. When requesting a training job using the CreateTrainingJob API, you can configure multiple instance groups .
" }, + "InstancePreference":{ + "type":"structure", + "required":["InstanceType"], + "members":{ + "InstanceType":{ + "shape":"TrainingInstanceType", + "documentation":"The ML compute instance type for this candidate.
" + }, + "InstanceCount":{ + "shape":"TrainingInstanceCount", + "documentation":"The number of ML compute instances to use for this candidate. Optional; mutually exclusive with the uniform ResourceConfig.InstanceCount (either every preference sets its own count, or the uniform count is used for whichever type wins).
The training plan(s) to target for this candidate type. Optional; list is capped at 1. Training-only; mutually exclusive with the whole-job ResourceConfig.TrainingPlanArn.
Defines a single candidate instance type in an ordered InstancePreferences list. When a training job specifies InstancePreferences, the platform tries each candidate type in list order and launches the job on the first type with available capacity.
An ordered list of candidate instance types for a training job (array position is the priority). Currently limited to a minimum of 1 and a maximum of 5 entries; the limit is enforced by the service.
", + "max":5, + "min":1 + }, "InstanceGroupHealthCheckConfiguration":{ "type":"structure", "required":[ @@ -35023,7 +35050,7 @@ "documentation":"Maximum job scheduler pending time in seconds.
", "box":true, "max":2419200, - "min":7200 + "min":1800 }, "MaxPercentageOfInputDatasetLabeled":{ "type":"integer", @@ -40147,8 +40174,6 @@ "ProcessingClusterConfig":{ "type":"structure", "required":[ - "InstanceCount", - "InstanceType", "VolumeSizeInGB" ], "members":{ @@ -40167,6 +40192,18 @@ "VolumeKmsKeyId":{ "shape":"KmsKeyId", "documentation":"The Amazon Web Services Key Management Service (Amazon Web Services KMS) key that Amazon SageMaker uses to encrypt data on the storage volume attached to the ML compute instance(s) that run the processing job.
Certain Nitro-based instances include local storage, dependent on the instance type. Local storage volumes are encrypted using a hardware module on the instance. You can't request a VolumeKmsKeyId when using an instance type with local storage.
For a list of instance types that support local instance storage, see Instance Store Volumes.
For more information about local instance storage encryption, see SSD Instance Store Volumes.
An ordered list of candidate instance types (maximum 5). When set, the platform tries each candidate in list order and launches the job on the first type with available capacity. Mutually exclusive with InstanceType and InstanceCount.
The instance type the job was launched on when InstancePreferences is used. Output-only (returned in DescribeProcessingJob); ignored on CreateProcessingJob.
The resolved number of instances the job was launched with when InstancePreferences is used. Output-only (returned in DescribeProcessingJob); ignored on CreateProcessingJob.
Configuration for the cluster used to run a processing job.
" @@ -40237,6 +40274,28 @@ "max":100, "min":1 }, + "ProcessingInstancePreference":{ + "type":"structure", + "required":["InstanceType"], + "members":{ + "InstanceType":{ + "shape":"ProcessingInstanceType", + "documentation":"The ML compute instance type for this candidate.
" + }, + "InstanceCount":{ + "shape":"ProcessingInstanceCount", + "documentation":"The number of ML compute instances to use for this candidate. Optional; mutually exclusive with the uniform ProcessingClusterConfig.InstanceCount (either every preference sets its own count, or the uniform count is used for whichever type wins).
Defines a single candidate instance type in an ordered InstancePreferences list for a processing job. InstanceType is required; InstanceCount is optional per element. Per-type training plans are training-only and do not apply to processing.
An ordered list of candidate instance types for a processing job (array position is the priority). Currently limited to a minimum of 1 and a maximum of 5 entries; the limit is enforced by the service.
", + "max":5, + "min":1 + }, "ProcessingInstanceType":{ "type":"string", "enum":[ @@ -42998,6 +43057,19 @@ "InstancePlacementConfig":{ "shape":"InstancePlacementConfig", "documentation":"Configuration for how training job instances are placed and allocated within UltraServers. Only applicable for UltraServer capacity.
" + }, + "InstancePreferences":{ + "shape":"InstancePreferenceList", + "documentation":"An ordered list of candidate instance types (maximum 5). When set, the platform tries each candidate in list order and launches the job on the first type with available capacity. Mutually exclusive with InstanceType, InstanceGroups, and InstancePlacementConfig.
The instance type the job was launched on when InstancePreferences is used. Output-only (returned in DescribeTrainingJob); ignored on CreateTrainingJob.
The resolved number of instances the job was launched with when InstancePreferences is used. Output-only (returned in DescribeTrainingJob); ignored on CreateTrainingJob.
Describes the resources, including machine learning (ML) compute instances and ML storage volumes, to use for model training.
" @@ -46934,6 +47006,13 @@ "type":"list", "member":{"shape":"TrainingPlanArn"} }, + "TrainingPlanArnList":{ + "type":"list", + "member":{"shape":"TrainingPlanArn"}, + "documentation":"The training plan(s) targeted by a single instance preference. Currently limited to exactly 1 training plan per preference; the limit is enforced by the service.
", + "max":1, + "min":1 + }, "TrainingPlanDurationHours":{ "type":"long", "box":true, diff --git a/sagemaker-core/src/sagemaker/core/modules/configs.py b/sagemaker-core/src/sagemaker/core/modules/configs.py index a23c5e14c4..865018f50c 100644 --- a/sagemaker-core/src/sagemaker/core/modules/configs.py +++ b/sagemaker-core/src/sagemaker/core/modules/configs.py @@ -74,7 +74,7 @@ "MetricDefinition", ] -from sagemaker.core.modules.utils import convert_unassigned_to_none +from sagemaker.core.modules.utils import convert_unassigned_to_none, validate_instance_preferences class BaseConfig(BaseModel): @@ -147,6 +147,11 @@ class Compute(shapes.ResourceConfig): subsequent training jobs. instance_groups (Optional[List[InstanceGroup]]): A list of instance groups for heterogeneous clusters to be used in the training job. + instance_preferences (Optional[List[InstancePreference]]): + An ordered list of candidate instance types (maximum 5). When set, the platform tries + each candidate in list order and launches the job on the first type with available + capacity. Mutually exclusive with ``instance_type``, ``instance_groups``, and + ``instance_placement_config``. enable_managed_spot_training (Optional[bool]): To train models using managed spot training, choose True. Managed spot training provides a fully managed and scalable infrastructure for training machine learning @@ -159,8 +164,10 @@ class Compute(shapes.ResourceConfig): @model_validator(mode="after") def _model_validator(self) -> "Compute": - """Convert Unassigned values to None.""" - return convert_unassigned_to_none(self) + """Convert Unassigned values to None and validate instance_preferences.""" + converted = convert_unassigned_to_none(self) + validate_instance_preferences(converted) + return converted def _to_resource_config(self) -> shapes.ResourceConfig: """Convert to a sagemaker_core.shapes.ResourceConfig object.""" @@ -169,6 +176,11 @@ def _to_resource_config(self) -> shapes.ResourceConfig: filtered_dict = { k: v for k, v in compute_config_dict.items() if k in resource_config_fields } + # Preserve the nested InstancePreference model objects instead of the + # dumped dicts, so pydantic does not re-validate their optional scalar + # fields (e.g. an unset per-preference instance_count) as Unassigned(). + if self.instance_preferences: + filtered_dict["instance_preferences"] = self.instance_preferences return shapes.ResourceConfig(**filtered_dict) diff --git a/sagemaker-core/src/sagemaker/core/modules/utils.py b/sagemaker-core/src/sagemaker/core/modules/utils.py index 94dc2dff22..d50df1d48a 100644 --- a/sagemaker-core/src/sagemaker/core/modules/utils.py +++ b/sagemaker-core/src/sagemaker/core/modules/utils.py @@ -192,3 +192,73 @@ def _run_clone_command_silent(repo_url, dest_dir): logger.error(f"Failed to clone repository: {repo_url}") logger.error(f"Error output:\n{e}") raise + + +def validate_instance_preferences(compute) -> None: + """Client-side validation for Compute.instance_preferences (server remains + the source of truth). + + - instance_preferences is mutually exclusive with the classic + single-cluster fields instance_type / instance_groups / + instance_placement_config. + - 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 + instance_count on EVERY element with the top-level unset. Both-set, + partial, and neither are rejected. + - Training plans: the top-level (whole-job) training_plan_arn is mutually + exclusive with per-preference training_plan_arns. + + List-size limits (max preferences, max plans per preference) are + deliberately NOT enforced client-side: they are server-side configurable, + so raising them must not require a new SDK release. + + No-op when instance_preferences is not set. + """ + preferences = getattr(compute, "instance_preferences", None) + if not preferences or isinstance(preferences, Unassigned): + return + + def _value(obj, field): + value = getattr(obj, field, None) + if value is None or isinstance(value, Unassigned): + return None + return value + + for field in ("instance_type", "instance_groups", "instance_placement_config"): + if _value(compute, field) is not None: + raise ValueError( + f"instance_preferences is mutually exclusive with {field}; " + "specify either a single fixed cluster or instance_preferences, not both." + ) + + instance_types = [_value(p, "instance_type") for p in preferences] + duplicates = sorted( + {t for t in instance_types if t is not None and instance_types.count(t) > 1} + ) + if duplicates: + raise ValueError( + "instance_preferences must not contain duplicate instance types: " f"{duplicates}." + ) + + per_preference_counts = [_value(p, "instance_count") is not None for p in preferences] + if _value(compute, "instance_count") is not None: + if any(per_preference_counts): + raise ValueError( + "The top-level instance_count and per-preference instance_count are " + "mutually exclusive; set the top-level instance_count (applies to " + "whichever preference wins) or an instance_count on every element of " + "instance_preferences, not both." + ) + elif not all(per_preference_counts): + raise ValueError( + "When the top-level instance_count is not set, every element of " + "instance_preferences must set its own instance_count." + ) + + per_preference_plans = [_value(p, "training_plan_arns") for p in preferences] + if _value(compute, "training_plan_arn") is not None and any(per_preference_plans): + raise ValueError( + "The top-level (whole-job) training_plan_arn and per-preference " + "training_plan_arns are mutually exclusive; set one or the other, not both." + ) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 34385619c8..9a65dba45f 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -87,6 +87,57 @@ logger = logging.getLogger(__name__) +def _validate_processing_instance_preferences( + instance_type=None, + instance_count=None, + instance_preferences=None, +): + """Client-side validation for Processor.instance_preferences (the service + remains the source of truth). + + - instance_preferences is mutually exclusive with instance_type (a single + fixed cluster). The top-level instance_count is NOT exclusive: it is the + shared count for whichever preference wins. + - Instance types must not repeat across preferences. + - Count mode: exactly one of the top-level instance_count with no + per-preference InstanceCount, or an InstanceCount on EVERY element with + the top-level unset. Both-set, partial, and neither are rejected. + + No-op when instance_preferences is not set. + """ + if not instance_preferences: + return + + if instance_type is not None: + raise ValueError( + "instance_preferences is mutually exclusive with instance_type; " + "specify either a single instance_type (+instance_count) or " + "instance_preferences, not both." + ) + types = [preference.get("InstanceType") for preference in instance_preferences] + duplicates = sorted({t for t in types if t is not None and types.count(t) > 1}) + if duplicates: + raise ValueError( + f"instance_preferences must not contain duplicate instance types: {duplicates}." + ) + per_pref_counts = [ + preference.get("InstanceCount") is not None for preference in instance_preferences + ] + if instance_count is not None: + if any(per_pref_counts): + raise ValueError( + "The top-level instance_count and per-preference InstanceCount " + "are mutually exclusive; set the top-level instance_count " + "(applies to whichever preference wins) or an InstanceCount on " + "every element of instance_preferences, not both." + ) + elif not all(per_pref_counts): + raise ValueError( + "When the top-level instance_count is not set, every element of " + "instance_preferences must set its own InstanceCount." + ) + + class Processor(object): """Handles Amazon SageMaker Processing tasks.""" @@ -108,6 +159,7 @@ def __init__( env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, tags: Optional[Tags] = None, network_config: Optional[NetworkConfig] = None, + instance_preferences: Optional[List[Dict[str, Union[str, int]]]] = None, ): """Initializes a ``Processor`` instance. @@ -151,10 +203,25 @@ def __init__( A :class:`~sagemaker.network.NetworkConfig` object that configures network isolation, encryption of inter-container traffic, security group IDs, and subnets. + instance_preferences (list[dict]): An ordered list of candidate instance types + (maximum 5). Each element is a dict of the form + ``{"InstanceType": "ml.m5.4xlarge", "InstanceCount": 2}``. When set, the + platform tries each candidate in list order and launches the job on the first + type with available capacity. Mutually exclusive with ``instance_type``. + The top-level ``instance_count`` is the shared count for whichever preference + wins; alternatively set ``InstanceCount`` on EVERY element (per-preference + mode) and leave ``instance_count`` unset โ never both, never partial + (default: None). """ self.image_uri = image_uri self.instance_count = instance_count self.instance_type = instance_type + self.instance_preferences = instance_preferences + _validate_processing_instance_preferences( + instance_type=instance_type, + instance_count=instance_count, + instance_preferences=instance_preferences, + ) self.entrypoint = entrypoint self.volume_size_in_gb = volume_size_in_gb self.max_runtime_in_seconds = max_runtime_in_seconds @@ -648,13 +715,23 @@ def _get_process_args(self, inputs, outputs, experiment_config): process_request_args["output_config"]["KmsKeyId"] = self.output_kms_key process_request_args["experiment_config"] = experiment_config process_request_args["job_name"] = self._current_job_name - process_request_args["resources"] = { - "ClusterConfig": { + if self.instance_preferences: + cluster_config = { + "InstancePreferences": self.instance_preferences, + "VolumeSizeInGB": self.volume_size_in_gb, + } + # Uniform-count mode: the shared top-level count applies to + # whichever preference wins. In per-preference mode the counts + # live on each element and the top-level key is omitted. + if self.instance_count is not None: + cluster_config["InstanceCount"] = self.instance_count + else: + cluster_config = { "InstanceType": self.instance_type, "InstanceCount": self.instance_count, "VolumeSizeInGB": self.volume_size_in_gb, } - } + process_request_args["resources"] = {"ClusterConfig": cluster_config} if self.volume_kms_key is not None: process_request_args["resources"]["ClusterConfig"][ "VolumeKmsKeyId" @@ -703,6 +780,7 @@ def __init__( env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, tags: Optional[Tags] = None, network_config: Optional[NetworkConfig] = None, + instance_preferences: Optional[List[Dict[str, Union[str, int]]]] = None, ): """Initializes a ``ScriptProcessor`` instance. @@ -747,6 +825,9 @@ def __init__( A :class:`~sagemaker.network.NetworkConfig` object that configures network isolation, encryption of inter-container traffic, security group IDs, and subnets. + instance_preferences (list[dict]): Ordered instance-type candidates for the + processing job (mutually exclusive with ``instance_type``); each element + is ``{"InstanceType": str, "InstanceCount": Optional[int]}``. """ self._CODE_CONTAINER_BASE_PATH = "/opt/ml/processing/input/" self._CODE_CONTAINER_INPUT_NAME = "code" @@ -774,6 +855,7 @@ def __init__( env=env, tags=format_tags(tags), network_config=network_config, + instance_preferences=instance_preferences, ) @_telemetry_emitter(feature=Feature.PROCESSING, func_name="ScriptProcessor.run") diff --git a/sagemaker-core/src/sagemaker/core/shapes/shapes.py b/sagemaker-core/src/sagemaker/core/shapes/shapes.py index e0e79a5751..f301ff246f 100644 --- a/sagemaker-core/src/sagemaker/core/shapes/shapes.py +++ b/sagemaker-core/src/sagemaker/core/shapes/shapes.py @@ -1655,6 +1655,23 @@ class InstanceGroup(Base): instance_group_name: StrPipeVar +class InstancePreference(Base): + """ + InstancePreference + Defines a single candidate instance type in an ordered InstancePreferences list. When a training job specifies InstancePreferences, the platform tries each candidate type in list order and launches the job on the first type with available capacity. + + Attributes + ---------------------- + instance_type: The ML compute instance type for this candidate. + instance_count: The number of ML compute instances to use for this candidate. Optional; mutually exclusive with the uniform ResourceConfig.InstanceCount (either every preference sets its own count, or the uniform count is used for whichever type wins). + training_plan_arns: The training plan(s) to target for this candidate type. Optional; list is capped at 1. Training-only; mutually exclusive with the whole-job ResourceConfig.TrainingPlanArn. + """ + + instance_type: StrPipeVar + instance_count: Optional[IntPipeVar] = Unassigned() + training_plan_arns: Optional[List[StrPipeVar]] = Unassigned() + + class PlacementSpecification(Base): """ PlacementSpecification @@ -1700,6 +1717,9 @@ class ResourceConfig(Base): instance_groups: The configuration of a heterogeneous cluster in JSON format. training_plan_arn: The Amazon Resource Name (ARN); of the training plan to use for this resource configuration. instance_placement_config: Configuration for how training job instances are placed and allocated within UltraServers. Only applicable for UltraServer capacity. + instance_preferences: An ordered list of candidate instance types (max 5). When set, the platform tries each candidate in list order and launches the job on the first type with available capacity. Mutually exclusive with instance_type, instance_groups, and instance_placement_config. + selected_instance_type: The instance type the job was launched on when instance_preferences is used. Output-only (returned in DescribeTrainingJob); ignored on CreateTrainingJob. + selected_instance_count: The resolved number of instances the job was launched with when instance_preferences is used. Output-only (returned in DescribeTrainingJob); ignored on CreateTrainingJob. """ instance_type: Optional[StrPipeVar] = Unassigned() @@ -1710,6 +1730,9 @@ class ResourceConfig(Base): instance_groups: Optional[List[InstanceGroup]] = Unassigned() training_plan_arn: Optional[StrPipeVar] = Unassigned() instance_placement_config: Optional[InstancePlacementConfig] = Unassigned() + instance_preferences: Optional[List[InstancePreference]] = Unassigned() + selected_instance_type: Optional[StrPipeVar] = Unassigned() + selected_instance_count: Optional[IntPipeVar] = Unassigned() class StoppingCondition(Base): @@ -9319,6 +9342,21 @@ class ProcessingOutputConfig(Base): kms_key_id: Optional[StrPipeVar] = Unassigned() +class ProcessingInstancePreference(Base): + """ + ProcessingInstancePreference + Defines a single candidate instance type in an ordered InstancePreferences list for a processing job. When a processing job specifies InstancePreferences, the platform tries each candidate type in list order and launches the job on the first type with available capacity. Per-type training plans are training-only and do not apply to processing. + + Attributes + ---------------------- + instance_type: The ML compute instance type for this candidate. + instance_count: The number of ML compute instances to use for this candidate. Optional; mutually exclusive with the uniform ProcessingClusterConfig.InstanceCount (either every preference sets its own count, or the uniform count is used for whichever type wins). + """ + + instance_type: StrPipeVar + instance_count: Optional[IntPipeVar] = Unassigned() + + class ProcessingClusterConfig(Base): """ ProcessingClusterConfig @@ -9330,12 +9368,18 @@ class ProcessingClusterConfig(Base): instance_type: The ML compute instance type for the processing job. volume_size_in_gb: The size of the ML storage volume in gigabytes that you want to provision. You must specify sufficient ML storage for your scenario. Certain Nitro-based instances include local storage with a fixed total size, dependent on the instance type. When using these instances for processing, Amazon SageMaker mounts the local instance storage instead of Amazon EBS gp2 storage. You can't request a VolumeSizeInGB greater than the total size of the local instance storage. For a list of instance types that support local instance storage, including the total size per instance type, see Instance Store Volumes. volume_kms_key_id: The Amazon Web Services Key Management Service (Amazon Web Services KMS) key that Amazon SageMaker uses to encrypt data on the storage volume attached to the ML compute instance(s) that run the processing job. Certain Nitro-based instances include local storage, dependent on the instance type. Local storage volumes are encrypted using a hardware module on the instance. You can't request a VolumeKmsKeyId when using an instance type with local storage. For a list of instance types that support local instance storage, see Instance Store Volumes. For more information about local instance storage encryption, see SSD Instance Store Volumes. + instance_preferences: An ordered list of candidate instance types (max 5). When set, the platform tries each candidate in list order and launches the job on the first type with available capacity. Mutually exclusive with instance_type/instance_count. + selected_instance_type: The instance type the job was launched on when instance_preferences is used. Output-only (returned in DescribeProcessingJob); ignored on CreateProcessingJob. + selected_instance_count: The resolved number of instances the job was launched with when instance_preferences is used. Output-only (returned in DescribeProcessingJob); ignored on CreateProcessingJob. """ - instance_count: int - instance_type: StrPipeVar + instance_count: Optional[IntPipeVar] = Unassigned() + instance_type: Optional[StrPipeVar] = Unassigned() volume_size_in_gb: int volume_kms_key_id: Optional[StrPipeVar] = Unassigned() + instance_preferences: Optional[List[ProcessingInstancePreference]] = Unassigned() + selected_instance_type: Optional[StrPipeVar] = Unassigned() + selected_instance_count: Optional[IntPipeVar] = Unassigned() class ProcessingResources(Base): diff --git a/sagemaker-core/src/sagemaker/core/spark/processing.py b/sagemaker-core/src/sagemaker/core/spark/processing.py index 82cdef954c..6f959ea4bb 100644 --- a/sagemaker-core/src/sagemaker/core/spark/processing.py +++ b/sagemaker-core/src/sagemaker/core/spark/processing.py @@ -17,6 +17,7 @@ post-processing, feature engineering, data validation, and model evaluation on SageMaker using Spark and PySpark. """ + from __future__ import absolute_import import json @@ -96,6 +97,7 @@ def __init__( env=None, tags=None, network_config=None, + instance_preferences=None, ): """Initialize a ``_SparkProcessorBase`` instance. @@ -150,9 +152,45 @@ def __init__( session = sagemaker_session or Session() region = session.boto_region_name - self.image_uri = self._retrieve_image_uri( - image_uri, framework_version, py_version, container_version, region, instance_type - ) + # One image per job, winner unknown at create: auto-resolve only when + # every candidate type yields the same image. + if image_uri is None and instance_type is None and instance_preferences: + candidate_uris = {} + for pref in instance_preferences: + candidate_type = pref.get("InstanceType") + if not candidate_type: + continue + try: + candidate_uris[candidate_type] = self._retrieve_image_uri( + None, + framework_version, + py_version, + container_version, + region, + candidate_type, + ) + except ValueError as e: + # e.g. GPU candidates: Spark images have no gpu variant. + raise ValueError( + f"Cannot auto-resolve a Spark image for instance_preferences " + f"candidate {candidate_type} ({e}); pass image_uri explicitly." + ) from e + if len(set(candidate_uris.values())) > 1: + raise ValueError( + "instance_preferences candidates resolve to different container " + f"images ({candidate_uris}); pass image_uri explicitly." + ) + if candidate_uris: + self.image_uri = next(iter(candidate_uris.values())) + else: + # No typed candidates: base Processor validation owns the reject. + self.image_uri = self._retrieve_image_uri( + image_uri, framework_version, py_version, container_version, region, None + ) + else: + self.image_uri = self._retrieve_image_uri( + image_uri, framework_version, py_version, container_version, region, instance_type + ) env = env or {} command = [_SparkProcessorBase._default_command] @@ -172,6 +210,7 @@ def __init__( env=env, tags=format_tags(tags), network_config=network_config, + instance_preferences=instance_preferences, ) def get_run_args( @@ -707,6 +746,7 @@ def __init__( env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, tags: Optional[Tags] = None, network_config: Optional[NetworkConfig] = None, + instance_preferences: Optional[List[Dict[str, Union[str, int]]]] = None, ): """Initialize an ``PySparkProcessor`` instance. @@ -775,6 +815,7 @@ def __init__( env=env, tags=format_tags(tags), network_config=network_config, + instance_preferences=instance_preferences, ) def get_run_args( @@ -984,6 +1025,7 @@ def __init__( env: Optional[Dict[str, Union[str, PipelineVariable]]] = None, tags: Optional[Tags] = None, network_config: Optional[NetworkConfig] = None, + instance_preferences: Optional[List[Dict[str, Union[str, int]]]] = None, ): """Initialize a ``SparkJarProcessor`` instance. @@ -1051,6 +1093,7 @@ def __init__( env=env, tags=format_tags(tags), network_config=network_config, + instance_preferences=instance_preferences, ) def get_run_args( diff --git a/sagemaker-core/src/sagemaker/core/training/configs.py b/sagemaker-core/src/sagemaker/core/training/configs.py index 0b03e18ff1..6ba49005a9 100644 --- a/sagemaker-core/src/sagemaker/core/training/configs.py +++ b/sagemaker-core/src/sagemaker/core/training/configs.py @@ -49,7 +49,7 @@ DatasetSource, ) -from sagemaker.core.training.utils import convert_unassigned_to_none +from sagemaker.core.training.utils import convert_unassigned_to_none, validate_instance_preferences __all__ = [ "BaseConfig", @@ -175,6 +175,11 @@ class Compute(shapes.ResourceConfig): A list of instance groups for heterogeneous clusters to be used in the training job. training_plan_arn (Optional[StrPipeVar]): The Amazon Resource Name (ARN) of the training plan to use for this resource configuration. + instance_preferences (Optional[List[InstancePreference]]): + An ordered list of candidate instance types (maximum 5). When set, the platform tries + each candidate in list order and launches the job on the first type with available + capacity. Mutually exclusive with ``instance_type``, ``instance_groups``, and + ``instance_placement_config``. enable_managed_spot_training (Optional[BoolPipeVar]): To train models using managed spot training, choose True. Managed spot training provides a fully managed and scalable infrastructure for training machine learning @@ -187,8 +192,10 @@ class Compute(shapes.ResourceConfig): @model_validator(mode="after") def _model_validator(self) -> "Compute": - """Convert Unassigned values to None.""" - return convert_unassigned_to_none(self) + """Convert Unassigned values to None and validate instance_preferences.""" + converted = convert_unassigned_to_none(self) + validate_instance_preferences(converted) + return converted def _to_resource_config(self) -> shapes.ResourceConfig: """Convert to a sagemaker.core.shapes.ResourceConfig object.""" @@ -201,6 +208,11 @@ def _to_resource_config(self) -> shapes.ResourceConfig: } if not filtered_dict: return None + # Preserve the nested InstancePreference model objects instead of the + # dumped dicts, so pydantic does not re-validate their optional scalar + # fields (e.g. an unset per-preference instance_count) as Unassigned(). + if self.instance_preferences: + filtered_dict["instance_preferences"] = self.instance_preferences return shapes.ResourceConfig(**filtered_dict) diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 0d03d2fcfb..12734b1d2a 100644 --- a/sagemaker-core/src/sagemaker/core/training/utils.py +++ b/sagemaker-core/src/sagemaker/core/training/utils.py @@ -256,3 +256,73 @@ def resolve_nova_checkpoint_uri( "Could not resolve the Nova checkpoint URI from any known output layout. " + " ".join(errors) ) + + +def validate_instance_preferences(compute) -> None: + """Client-side validation for Compute.instance_preferences (server remains + the source of truth). + + - instance_preferences is mutually exclusive with the classic + single-cluster fields instance_type / instance_groups / + instance_placement_config. + - 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 + instance_count on EVERY element with the top-level unset. Both-set, + partial, and neither are rejected. + - Training plans: the top-level (whole-job) training_plan_arn is mutually + exclusive with per-preference training_plan_arns. + + List-size limits (max preferences, max plans per preference) are + deliberately NOT enforced client-side: they are server-side configurable, + so raising them must not require a new SDK release. + + No-op when instance_preferences is not set. + """ + preferences = getattr(compute, "instance_preferences", None) + if not preferences or isinstance(preferences, Unassigned): + return + + def _value(obj, field): + value = getattr(obj, field, None) + if value is None or isinstance(value, Unassigned): + return None + return value + + for field in ("instance_type", "instance_groups", "instance_placement_config"): + if _value(compute, field) is not None: + raise ValueError( + f"instance_preferences is mutually exclusive with {field}; " + "specify either a single fixed cluster or instance_preferences, not both." + ) + + instance_types = [_value(p, "instance_type") for p in preferences] + duplicates = sorted( + {t for t in instance_types if t is not None and instance_types.count(t) > 1} + ) + if duplicates: + raise ValueError( + "instance_preferences must not contain duplicate instance types: " f"{duplicates}." + ) + + per_preference_counts = [_value(p, "instance_count") is not None for p in preferences] + if _value(compute, "instance_count") is not None: + if any(per_preference_counts): + raise ValueError( + "The top-level instance_count and per-preference instance_count are " + "mutually exclusive; set the top-level instance_count (applies to " + "whichever preference wins) or an instance_count on every element of " + "instance_preferences, not both." + ) + elif not all(per_preference_counts): + raise ValueError( + "When the top-level instance_count is not set, every element of " + "instance_preferences must set its own instance_count." + ) + + per_preference_plans = [_value(p, "training_plan_arns") for p in preferences] + if _value(compute, "training_plan_arn") is not None and any(per_preference_plans): + raise ValueError( + "The top-level (whole-job) training_plan_arn and per-preference " + "training_plan_arns are mutually exclusive; set one or the other, not both." + ) diff --git a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py index 439fef722d..2fa99017c3 100644 --- a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py +++ b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py @@ -10163,6 +10163,19 @@ "member_type": "structure", "type": "list", }, + "InstancePreference": { + "members": [ + {"name": "InstanceType", "shape": "TrainingInstanceType", "type": "string"}, + {"name": "InstanceCount", "shape": "TrainingInstanceCount", "type": "integer"}, + {"name": "TrainingPlanArns", "shape": "TrainingPlanArnList", "type": "list"}, + ], + "type": "structure", + }, + "InstancePreferenceList": { + "member_shape": "InstancePreference", + "member_type": "structure", + "type": "list", + }, "InstanceRequirementsEniConfiguration": { "members": [ {"name": "CustomerEni", "shape": "String", "type": "string"}, @@ -14787,6 +14800,17 @@ {"name": "InstanceType", "shape": "ProcessingInstanceType", "type": "string"}, {"name": "VolumeSizeInGB", "shape": "ProcessingVolumeSizeInGB", "type": "integer"}, {"name": "VolumeKmsKeyId", "shape": "KmsKeyId", "type": "string"}, + { + "name": "InstancePreferences", + "shape": "ProcessingInstancePreferenceList", + "type": "list", + }, + {"name": "SelectedInstanceType", "shape": "ProcessingInstanceType", "type": "string"}, + { + "name": "SelectedInstanceCount", + "shape": "ProcessingInstanceCount", + "type": "integer", + }, ], "type": "structure", }, @@ -14815,6 +14839,18 @@ "member_type": "structure", "type": "list", }, + "ProcessingInstancePreference": { + "members": [ + {"name": "InstanceType", "shape": "ProcessingInstanceType", "type": "string"}, + {"name": "InstanceCount", "shape": "ProcessingInstanceCount", "type": "integer"}, + ], + "type": "structure", + }, + "ProcessingInstancePreferenceList": { + "member_shape": "ProcessingInstancePreference", + "member_type": "structure", + "type": "list", + }, "ProcessingJob": { "members": [ {"name": "ProcessingInputs", "shape": "ProcessingInputs", "type": "list"}, @@ -15820,6 +15856,9 @@ "shape": "InstancePlacementConfig", "type": "structure", }, + {"name": "InstancePreferences", "shape": "InstancePreferenceList", "type": "list"}, + {"name": "SelectedInstanceType", "shape": "TrainingInstanceType", "type": "string"}, + {"name": "SelectedInstanceCount", "shape": "TrainingInstanceCount", "type": "integer"}, ], "type": "structure", }, @@ -17204,6 +17243,11 @@ ], "type": "structure", }, + "TrainingPlanArnList": { + "member_shape": "TrainingPlanArn", + "member_type": "string", + "type": "list", + }, "TrainingPlanArns": { "member_shape": "TrainingPlanArn", "member_type": "string", diff --git a/sagemaker-core/tests/integ/processing/__init__.py b/sagemaker-core/tests/integ/processing/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py b/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py new file mode 100644 index 0000000000..d309079367 --- /dev/null +++ b/sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py @@ -0,0 +1,133 @@ +# 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. +"""End-to-end integration test for Instance Preferences on processing jobs. + +Launches a REAL processing job whose ``ClusterConfig`` carries an ordered +``InstancePreferences`` list (no ``InstanceType``) through the ``Processor`` +path, then asserts the Describe contract: + +- the create request is accepted with ``InstancePreferences`` + the uniform + ``InstanceCount`` only; +- Describe echoes ``InstancePreferences`` and does not return the top-level + ``InstanceType`` the customer never set; +- once the job leaves the pre-instance states, ``SelectedInstanceType`` / + ``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. +""" + +from __future__ import absolute_import + +import os +import time +import uuid + +import pytest + +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") + +PREFERENCE_TYPES = [ + t.strip() + for t in os.environ.get( + "PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.xlarge,ml.m5.large" + ).split(",") + if t.strip() +] +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) + + +def test_processor_instance_preferences_e2e(): + """Create a real processing job with InstancePreferences and verify the winner.""" + from sagemaker.core.helper.session_helper import Session + from sagemaker.core.processing import Processor + + client = _sagemaker_client() + session = Session(sagemaker_client=client) + + processor = Processor( + role=ROLE_ARN, + image_uri=IMAGE_URI, + instance_count=1, + instance_preferences=[{"InstanceType": t} for t in PREFERENCE_TYPES], + volume_size_in_gb=30, + max_runtime_in_seconds=1800, + sagemaker_session=session, + ) + + 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: + described = client.describe_processing_job(ProcessingJobName=job_name) + status = described["ProcessingJobStatus"] + 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 diff --git a/sagemaker-core/tests/unit/test_compute_configs.py b/sagemaker-core/tests/unit/test_compute_configs.py index 7564b99447..5f1e6a1ebf 100644 --- a/sagemaker-core/tests/unit/test_compute_configs.py +++ b/sagemaker-core/tests/unit/test_compute_configs.py @@ -72,3 +72,253 @@ def test_compute_is_not_hyperpod_instance(self): """Compute is not an instance of HyperPodCompute.""" compute = Compute(instance_type="ml.p5.48xlarge") assert not isinstance(compute, HyperPodCompute) + + +class TestComputeInstancePreferences: + """Instance Preferences (multi-instance-type) support on the Compute config.""" + + def test_training_compute_instance_preferences_round_trip(self): + from sagemaker.core.shapes.shapes import InstancePreference + + prefs = [ + InstancePreference(instance_type="ml.p5.48xlarge"), + InstancePreference(instance_type="ml.p4d.24xlarge"), + ] + compute = Compute(instance_preferences=prefs, instance_count=2) + rc = compute._to_resource_config() + assert [p.instance_type for p in rc.instance_preferences] == [ + "ml.p5.48xlarge", + "ml.p4d.24xlarge", + ] + assert rc.instance_count == 2 + + def test_training_compute_per_preference_count(self): + """A per-preference (unset uniform) count must round-trip without error.""" + from sagemaker.core.shapes.shapes import InstancePreference + from sagemaker.core.utils.utils import Unassigned + + prefs = [ + InstancePreference(instance_type="ml.p5.48xlarge", instance_count=2), + InstancePreference(instance_type="ml.p4d.24xlarge", instance_count=4), + ] + compute = Compute(instance_preferences=prefs) + rc = compute._to_resource_config() + assert rc.instance_preferences[0].instance_count == 2 + assert rc.instance_preferences[1].instance_count == 4 + + def test_training_compute_per_preference_training_plan(self): + from sagemaker.core.shapes.shapes import InstancePreference + + prefs = [ + InstancePreference( + instance_type="ml.p5.48xlarge", + training_plan_arns=[ + "arn:aws:sagemaker:us-west-2:111122223333:training-plan/p5-plan" + ], + ), + InstancePreference(instance_type="ml.p4d.24xlarge"), + ] + rc = Compute(instance_preferences=prefs, instance_count=1)._to_resource_config() + assert rc.instance_preferences[0].training_plan_arns == [ + "arn:aws:sagemaker:us-west-2:111122223333:training-plan/p5-plan" + ] + + def test_selected_fields_not_sent_on_create(self): + """selected_instance_type/count are output-only and must not be populated on create.""" + from sagemaker.core.shapes.shapes import InstancePreference + from sagemaker.core.utils.utils import Unassigned + + rc = Compute( + instance_preferences=[InstancePreference(instance_type="ml.p5.48xlarge")], + instance_count=1, + )._to_resource_config() + # training Compute filters out None/Unassigned values -> stays Unassigned (not sent) + assert isinstance(rc.selected_instance_type, Unassigned) + assert isinstance(rc.selected_instance_count, Unassigned) + + def test_single_type_still_works(self): + """Classic single-type path is unchanged when instance_preferences is not set.""" + from sagemaker.core.utils.utils import Unassigned + + rc = Compute(instance_type="ml.m5.xlarge", instance_count=1)._to_resource_config() + assert rc.instance_type == "ml.m5.xlarge" + assert rc.instance_count == 1 + # instance_preferences is filtered out on create -> stays Unassigned (not sent) + assert isinstance(rc.instance_preferences, Unassigned) + + +class TestComputeInstancePreferencesClientValidation: + """Client-side V1/V4 validation on both Compute classes (server remains + the source of truth).""" + + @pytest.fixture(params=["training", "modules"]) + def compute_cls(self, request): + if request.param == "training": + from sagemaker.core.training.configs import Compute as ComputeCls + else: + from sagemaker.core.modules.configs import Compute as ComputeCls + return ComputeCls + + def test_instance_type_rejected_with_preferences(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="mutually exclusive with instance_type"): + compute_cls( + instance_type="ml.m5.xlarge", + instance_count=1, + instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")], + ) + + def test_uniform_count_rejected_with_per_preference_counts(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="top-level instance_count and per-preference"): + compute_cls( + instance_count=1, + instance_preferences=[ + InstancePreference(instance_type="ml.m5.xlarge", instance_count=2), + InstancePreference(instance_type="ml.m4.xlarge"), + ], + ) + + def test_partial_per_preference_counts_rejected(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="every element"): + compute_cls( + instance_preferences=[ + InstancePreference(instance_type="ml.m5.xlarge", instance_count=2), + InstancePreference(instance_type="ml.m4.xlarge"), + ], + ) + + def test_no_count_at_all_rejected(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="every element"): + compute_cls( + instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")], + ) + + def test_duplicate_instance_types_rejected(self, compute_cls): + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="duplicate instance types"): + compute_cls( + instance_count=1, + instance_preferences=[ + InstancePreference(instance_type="ml.m5.xlarge"), + InstancePreference(instance_type="ml.m5.xlarge"), + ], + ) + + def test_whole_job_plan_rejected_with_per_preference_plans(self, compute_cls): + """V8: whole-job training_plan_arn XOR per-preference training_plan_arns.""" + from sagemaker.core.shapes.shapes import InstancePreference + + with pytest.raises(ValueError, match="training_plan_arn and per-preference"): + compute_cls( + instance_count=1, + training_plan_arn=( + "arn:aws:sagemaker:us-west-2:111122223333:training-plan/whole-job" + ), + instance_preferences=[ + InstancePreference( + instance_type="ml.p5.48xlarge", + training_plan_arns=[ + "arn:aws:sagemaker:us-west-2:111122223333:training-plan/p5" + ], + ), + InstancePreference(instance_type="ml.p4d.24xlarge"), + ], + ) + + def test_whole_job_plan_allowed_without_per_preference_plans(self, compute_cls): + """A whole-job training_plan_arn with NO per-preference plans is valid.""" + from sagemaker.core.shapes.shapes import InstancePreference + + compute_cls( + instance_count=1, + 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"), + ], + ) + + +class TestModulesComputeInstancePreferences: + """Instance Preferences support on the modules.configs Compute class.""" + + def test_modules_compute_instance_preferences_round_trip(self): + from sagemaker.core.modules.configs import Compute as ModulesCompute + from sagemaker.core.shapes.shapes import InstancePreference + + prefs = [ + InstancePreference(instance_type="ml.p5.48xlarge"), + InstancePreference(instance_type="ml.p4d.24xlarge"), + ] + 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", + ] + + 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() + assert rc.instance_type == "ml.m5.xlarge" + + +class TestProcessingClusterConfigInstancePreferences: + """InstancePreferences on the ProcessingClusterConfig shape.""" + + def test_processing_cluster_config_accepts_instance_preferences(self): + from sagemaker.core.shapes.shapes import ( + ProcessingClusterConfig, + ProcessingInstancePreference, + ) + from sagemaker.core.utils.utils import Unassigned + + pcc = ProcessingClusterConfig( + instance_preferences=[ + ProcessingInstancePreference(instance_type="ml.m5.4xlarge"), + ProcessingInstancePreference(instance_type="ml.m5.2xlarge"), + ], + volume_size_in_gb=100, + ) + assert [p.instance_type for p in pcc.instance_preferences] == [ + "ml.m5.4xlarge", + "ml.m5.2xlarge", + ] + # instance_type / instance_count are now optional (mutually exclusive with prefs) + assert isinstance(pcc.instance_type, Unassigned) + assert isinstance(pcc.instance_count, Unassigned) + + def test_processing_instance_preference_has_no_training_plan_arns(self): + """Processing preferences are a separate shape without training plans (training-only).""" + import pydantic + import pytest + from sagemaker.core.shapes.shapes import ProcessingInstancePreference + + with pytest.raises(pydantic.ValidationError): + ProcessingInstancePreference( + instance_type="ml.m5.4xlarge", + training_plan_arns=["arn:aws:sagemaker:us-west-2:111122223333:training-plan/p"], + ) + + def test_processing_cluster_config_single_type_still_works(self): + from sagemaker.core.shapes.shapes import ProcessingClusterConfig + + pcc = ProcessingClusterConfig( + instance_type="ml.m5.xlarge", instance_count=1, volume_size_in_gb=30 + ) + assert pcc.instance_type == "ml.m5.xlarge" + assert pcc.instance_count == 1 diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index ca811424d1..559ac4179c 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -2035,3 +2035,260 @@ def test_run_with_s3_source_dir_prefix(self, mock_session): wait=False, ) assert processor.latest_job == mock_job + + +class TestProcessorInstancePreferences: + """Instance Preferences (multi-instance-type) support on Processor.""" + + def test_uniform_count_mode_emits_preferences_and_shared_count(self, mock_session): + """Uniform mode: top-level instance_count applies to whichever type wins.""" + preferences = [ + {"InstanceType": "ml.m5.4xlarge"}, + {"InstanceType": "ml.m5.2xlarge"}, + ] + processor = Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_count=2, + instance_preferences=preferences, + volume_size_in_gb=100, + sagemaker_session=mock_session, + ) + processor._current_job_name = "job-name" + args = processor._get_process_args([], [], None) + cluster_config = args["resources"]["ClusterConfig"] + assert cluster_config["InstancePreferences"] == preferences + assert cluster_config["InstanceCount"] == 2 + assert cluster_config["VolumeSizeInGB"] == 100 + # the classic single-type key is not emitted in instance-preferences mode + assert "InstanceType" not in cluster_config + + def test_per_preference_count_mode_omits_top_level_count(self, mock_session): + """Per-preference mode: every element carries its count; top-level unset.""" + preferences = [ + {"InstanceType": "ml.m5.4xlarge", "InstanceCount": 2}, + {"InstanceType": "ml.m5.2xlarge", "InstanceCount": 4}, + ] + processor = Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_preferences=preferences, + volume_size_in_gb=100, + sagemaker_session=mock_session, + ) + processor._current_job_name = "job-name" + args = processor._get_process_args([], [], None) + cluster_config = args["resources"]["ClusterConfig"] + assert cluster_config["InstancePreferences"] == preferences + assert "InstanceCount" not in cluster_config + assert "InstanceType" not in cluster_config + + def test_get_process_args_classic_single_type_unchanged(self, mock_session): + processor = Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_type="ml.m5.xlarge", + instance_count=2, + sagemaker_session=mock_session, + ) + processor._current_job_name = "job-name" + args = processor._get_process_args([], [], None) + cluster_config = args["resources"]["ClusterConfig"] + assert cluster_config["InstanceType"] == "ml.m5.xlarge" + assert cluster_config["InstanceCount"] == 2 + assert "InstancePreferences" not in cluster_config + + def test_instance_preferences_mutually_exclusive_with_instance_type(self, mock_session): + with pytest.raises(ValueError, match="mutually exclusive with instance_type"): + Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_type="ml.m5.xlarge", + instance_count=1, + instance_preferences=[{"InstanceType": "ml.m5.4xlarge", "InstanceCount": 1}], + sagemaker_session=mock_session, + ) + + def test_uniform_count_rejected_with_per_preference_counts(self, mock_session): + """V4: reject when both the uniform and any per-preference count are set.""" + with pytest.raises(ValueError, match="top-level instance_count and per-preference"): + Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_count=2, + instance_preferences=[ + {"InstanceType": "ml.m5.4xlarge", "InstanceCount": 2}, + {"InstanceType": "ml.m5.2xlarge"}, + ], + sagemaker_session=mock_session, + ) + + def test_partial_per_preference_counts_rejected(self, mock_session): + """V4: without a uniform count, EVERY element must set InstanceCount.""" + with pytest.raises(ValueError, match="every element"): + Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_preferences=[ + {"InstanceType": "ml.m5.4xlarge", "InstanceCount": 2}, + {"InstanceType": "ml.m5.2xlarge"}, + ], + sagemaker_session=mock_session, + ) + + def test_no_count_at_all_rejected(self, mock_session): + """V4: neither a uniform count nor per-preference counts is invalid.""" + with pytest.raises(ValueError, match="every element"): + Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_preferences=[ + {"InstanceType": "ml.m5.4xlarge"}, + {"InstanceType": "ml.m5.2xlarge"}, + ], + sagemaker_session=mock_session, + ) + + def test_duplicate_instance_types_rejected(self, mock_session): + with pytest.raises(ValueError, match="duplicate instance types"): + Processor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_count=1, + instance_preferences=[ + {"InstanceType": "ml.m5.4xlarge"}, + {"InstanceType": "ml.m5.4xlarge"}, + ], + sagemaker_session=mock_session, + ) + + +class TestScriptAndSparkProcessorInstancePreferences: + """instance_preferences plumbs through the ScriptProcessor/Spark subclass chain.""" + + _PREFS = [ + {"InstanceType": "ml.m5.xlarge", "InstanceCount": 1}, + {"InstanceType": "ml.m4.xlarge", "InstanceCount": 2}, + ] + + def test_script_processor_forwards_instance_preferences(self, mock_session): + from sagemaker.core.processing import ScriptProcessor + + processor = ScriptProcessor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + command=["python3"], + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + assert processor.instance_preferences == self._PREFS + assert processor.instance_type is None + + def test_script_processor_validation_applies(self, mock_session): + from sagemaker.core.processing import ScriptProcessor + + with pytest.raises(ValueError, match="mutually exclusive with instance_type"): + ScriptProcessor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + command=["python3"], + instance_type="ml.m5.xlarge", + instance_count=1, + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + + def test_pyspark_processor_forwards_instance_preferences(self, mock_session): + from sagemaker.core.spark.processing import PySparkProcessor + + mock_session.boto_region_name = "us-west-2" + processor = PySparkProcessor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + assert processor.instance_preferences == self._PREFS + assert processor.instance_type is None + + def test_pyspark_image_resolution_requires_candidate_agreement(self, mock_session): + from sagemaker.core.spark import processing as spark_processing + + mock_session.boto_region_name = "us-west-2" + # Agreement: every candidate resolves to the same image -> used. + with patch.object( + spark_processing.image_uris, "retrieve", return_value="resolved-uri" + ) as mock_retrieve: + processor = spark_processing.PySparkProcessor( + role="arn:aws:iam::123456789012:role/role", + framework_version="3.5", + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + assert processor.image_uri == "resolved-uri" + called_types = {c.kwargs["instance_type"] for c in mock_retrieve.call_args_list} + assert called_types == {"ml.m5.xlarge", "ml.m4.xlarge"} + + def test_pyspark_image_resolution_rejects_divergent_candidates(self, mock_session): + from sagemaker.core.spark import processing as spark_processing + + mock_session.boto_region_name = "us-west-2" + # Divergence: candidates resolve to different images -> explicit image_uri required. + with patch.object( + spark_processing.image_uris, "retrieve", side_effect=["cpu-uri", "gpu-uri"] + ): + with pytest.raises(ValueError, match="pass image_uri explicitly"): + spark_processing.PySparkProcessor( + role="arn:aws:iam::123456789012:role/role", + framework_version="3.5", + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + + def test_pyspark_unresolvable_candidate_names_it(self, mock_session): + from sagemaker.core.spark import processing as spark_processing + + mock_session.boto_region_name = "us-west-2" + # Real behavior: Spark has no GPU image, so retrieve raises for GPU + # candidates; the error must name the candidate and the remedy. + with patch.object( + spark_processing.image_uris, + "retrieve", + side_effect=["cpu-uri", ValueError("Unsupported processor: gpu")], + ): + with pytest.raises( + ValueError, match=r"candidate ml\.m4\.xlarge.*pass image_uri explicitly" + ): + spark_processing.PySparkProcessor( + role="arn:aws:iam::123456789012:role/role", + framework_version="3.5", + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + + def test_pyspark_degenerate_preferences_reach_base_validation(self, mock_session): + from sagemaker.core.spark import processing as spark_processing + + mock_session.boto_region_name = "us-west-2" + # No InstanceType on any element: must NOT crash in image resolution; + # the base Processor validation owns the reject. + with patch.object(spark_processing.image_uris, "retrieve", return_value="uri"): + with pytest.raises(ValueError): + spark_processing.PySparkProcessor( + role="arn:aws:iam::123456789012:role/role", + framework_version="3.5", + instance_preferences=[{}], + sagemaker_session=mock_session, + ) + + def test_sparkjar_processor_forwards_instance_preferences(self, mock_session): + from sagemaker.core.spark.processing import SparkJarProcessor + + mock_session.boto_region_name = "us-west-2" + processor = SparkJarProcessor( + role="arn:aws:iam::123456789012:role/role", + image_uri="image-uri", + instance_preferences=self._PREFS, + sagemaker_session=mock_session, + ) + assert processor.instance_preferences == self._PREFS diff --git a/sagemaker-core/tests/unit/test_service_model_instance_preferences.py b/sagemaker-core/tests/unit/test_service_model_instance_preferences.py new file mode 100644 index 0000000000..10c05fb1ad --- /dev/null +++ b/sagemaker-core/tests/unit/test_service_model_instance_preferences.py @@ -0,0 +1,126 @@ +# 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. +"""Pin the Instance Preferences contract in the bundled SageMaker service model. + +The bundled ``sample/sagemaker/2017-07-24/service-2.json`` is both the codegen +input for ``shapes.py`` and the model the runtime botocore loader injects, so +these tests guard the Instance Preferences API surface against accidental +regeneration/edit drift: + +- training and processing preferences are DISTINCT shapes + (``InstancePreference`` vs ``ProcessingInstancePreference``), matching the + service model where the two planes use different member types; +- per-preference training plans are training-only and capped at 1 + (``TrainingPlanArnList``); +- the output-only ``SelectedInstanceType``/``SelectedInstanceCount`` exist on + both ``ResourceConfig`` and ``ProcessingClusterConfig``; +- ``ProcessingClusterConfig`` no longer requires ``InstanceType``/ + ``InstanceCount`` (mutually exclusive with ``InstancePreferences``, + enforced server-side). +""" +from __future__ import absolute_import + +import json +import os + +import pytest + +SERVICE_JSON = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "sample", + "sagemaker", + "2017-07-24", + "service-2.json", +) + + +@pytest.fixture(scope="module") +def shapes(): + with open(SERVICE_JSON) as f: + return json.load(f)["shapes"] + + +class TestTrainingInstancePreferenceModel: + def test_instance_preference_shape(self, shapes): + shape = shapes["InstancePreference"] + assert shape["type"] == "structure" + assert shape["required"] == ["InstanceType"] + members = shape["members"] + assert members["InstanceType"]["shape"] == "TrainingInstanceType" + assert members["InstanceCount"]["shape"] == "TrainingInstanceCount" + assert members["TrainingPlanArns"]["shape"] == "TrainingPlanArnList" + + def test_instance_preference_list_capped_at_5(self, shapes): + lst = shapes["InstancePreferenceList"] + assert lst["type"] == "list" + assert lst["member"]["shape"] == "InstancePreference" + assert lst["min"] == 1 + assert lst["max"] == 5 + + def test_training_plan_arn_list_capped_at_1(self, shapes): + lst = shapes["TrainingPlanArnList"] + assert lst["type"] == "list" + assert lst["member"]["shape"] == "TrainingPlanArn" + assert lst["min"] == 1 + assert lst["max"] == 1 + + def test_resource_config_members(self, shapes): + members = shapes["ResourceConfig"]["members"] + assert members["InstancePreferences"]["shape"] == "InstancePreferenceList" + assert members["SelectedInstanceType"]["shape"] == "TrainingInstanceType" + assert members["SelectedInstanceCount"]["shape"] == "TrainingInstanceCount" + + +class TestProcessingInstancePreferenceModel: + def test_processing_instance_preference_shape(self, shapes): + shape = shapes["ProcessingInstancePreference"] + assert shape["type"] == "structure" + assert shape["required"] == ["InstanceType"] + members = shape["members"] + assert members["InstanceType"]["shape"] == "ProcessingInstanceType" + assert members["InstanceCount"]["shape"] == "ProcessingInstanceCount" + # per-type training plans are training-only + assert "TrainingPlanArns" not in members + + def test_processing_instance_preference_list_capped_at_5(self, shapes): + lst = shapes["ProcessingInstancePreferenceList"] + assert lst["type"] == "list" + assert lst["member"]["shape"] == "ProcessingInstancePreference" + assert lst["min"] == 1 + assert lst["max"] == 5 + + def test_processing_cluster_config_members(self, shapes): + pcc = shapes["ProcessingClusterConfig"] + members = pcc["members"] + assert members["InstancePreferences"]["shape"] == "ProcessingInstancePreferenceList" + assert members["SelectedInstanceType"]["shape"] == "ProcessingInstanceType" + assert members["SelectedInstanceCount"]["shape"] == "ProcessingInstanceCount" + + def test_processing_cluster_config_type_count_not_required(self, shapes): + # InstanceType/InstanceCount are mutually exclusive with + # InstancePreferences; the exclusivity is enforced server-side, so the + # client model must not hard-require them. + required = shapes["ProcessingClusterConfig"]["required"] + assert "InstanceType" not in required + assert "InstanceCount" not in required + assert "VolumeSizeInGB" in required + + +class TestPreferenceShapesAreDistinct: + def test_training_and_processing_preferences_do_not_share_shapes(self, shapes): + assert ( + shapes["ResourceConfig"]["members"]["InstancePreferences"]["shape"] + != shapes["ProcessingClusterConfig"]["members"]["InstancePreferences"]["shape"] + ) diff --git a/sagemaker-train/src/sagemaker/train/defaults.py b/sagemaker-train/src/sagemaker/train/defaults.py index dbd5cdfb0d..c52030f33a 100644 --- a/sagemaker-train/src/sagemaker/train/defaults.py +++ b/sagemaker-train/src/sagemaker/train/defaults.py @@ -139,7 +139,10 @@ def get_compute(compute: Optional[Compute] = None) -> Compute: volume_size_in_gb=DEFAULT_VOLUME_SIZE, ) logger.info(f"Compute not provided. Using default:\n{compute}") - if not compute.instance_groups: + if not compute.instance_groups and not compute.instance_preferences: + # When instance_preferences is set, the top-level instance_type / + # instance_count must stay unset (mutually exclusive with the + # preference list; the uniform count, when used, is customer-set). if compute.instance_type is None: compute.instance_type = DEFAULT_INSTANCE_TYPE logger.info(f"Instance type not provided. Using default:\n{DEFAULT_INSTANCE_TYPE}") diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 241b6ac3be..9980266b00 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -846,6 +846,13 @@ def train( ) else: + if self.compute is not None and getattr( + self.compute, "instance_preferences", None + ): + raise ValueError( + "Local mode training does not support 'instance_preferences'. " + "Set a single 'instance_type' on Compute for local mode." + ) local_container = _LocalContainer( training_job_name=training_request["training_job_name"], instance_type=training_request["resource_config"].instance_type, @@ -1327,6 +1334,11 @@ def from_recipe( SourceModelPackageArn and ModelPackageGroupArn. If also specified in recipe, direct param wins on conflict. """ + if getattr(compute, "instance_preferences", None): + raise ValueError( + "Training recipes do not support ``instance_preferences``. " + "Set a single ``instance_type`` in Compute when using training recipes." + ) if compute.instance_type is None: raise ValueError("Must set ``instance_type`` in Compute when using training recipes.") device_type = _determine_device_type(compute.instance_type) @@ -1569,6 +1581,11 @@ def from_jumpstart_config( f"Training is not supported for the model ID: {jumpstart_config.model_id}.\n" "Please check that the model ID is available for training." ) + if compute and getattr(compute, "instance_preferences", None): + raise ValueError( + "JumpStart training does not support ``instance_preferences``. " + "Set a single ``instance_type`` in Compute for JumpStart models." + ) if compute and document.SupportedTrainingInstanceTypes: if compute.instance_type not in document.SupportedTrainingInstanceTypes: raise ValueError( diff --git a/sagemaker-train/tests/integ/train/test_instance_preferences.py b/sagemaker-train/tests/integ/train/test_instance_preferences.py new file mode 100644 index 0000000000..c276809952 --- /dev/null +++ b/sagemaker-train/tests/integ/train/test_instance_preferences.py @@ -0,0 +1,150 @@ +# 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. +"""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. +""" + +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") + +PREFERENCE_TYPES = [ + t.strip() + for t in os.environ.get( + "INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.xlarge,ml.m4.xlarge" + ).split(",") + if t.strip() +] +WAIT_TIMEOUT_SECONDS = 30 * 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 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]}" + + 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 + + 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), + ) + + 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) + 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 diff --git a/sagemaker-train/tests/unit/train/test_defaults.py b/sagemaker-train/tests/unit/train/test_defaults.py index 1c02b72e01..e6416af697 100644 --- a/sagemaker-train/tests/unit/train/test_defaults.py +++ b/sagemaker-train/tests/unit/train/test_defaults.py @@ -252,6 +252,50 @@ def test_default_compute_has_correct_values(self): assert result.instance_count == 1 assert result.volume_size_in_gb == 30 + def test_no_default_instance_type_with_instance_preferences(self): + """instance_preferences must not get a default top-level instance_type/count. + + The top-level instance_type is mutually exclusive with + instance_preferences (server-side V1 validation); injecting the + default would make every instance-preferences job fail with + 'InstanceType cannot be specified with InstancePreferences'. + """ + from sagemaker.core.shapes.shapes import InstancePreference + + compute = Compute( + instance_preferences=[ + InstancePreference(instance_type="ml.m5.xlarge"), + InstancePreference(instance_type="ml.m4.xlarge"), + ], + instance_count=1, + ) + result = TrainDefaults.get_compute(compute=compute) + + assert result.instance_type is None + assert [p.instance_type for p in result.instance_preferences] == [ + "ml.m5.xlarge", + "ml.m4.xlarge", + ] + # the customer-set uniform count is preserved untouched + assert result.instance_count == 1 + # volume size default still applies (whole-job knob, not exclusive) + assert result.volume_size_in_gb == DEFAULT_VOLUME_SIZE + + def test_no_default_instance_count_with_per_preference_counts(self): + """Per-preference count mode must not get the default uniform count.""" + from sagemaker.core.shapes.shapes import InstancePreference + + compute = Compute( + instance_preferences=[ + InstancePreference(instance_type="ml.m5.xlarge", instance_count=2), + InstancePreference(instance_type="ml.m4.xlarge", instance_count=4), + ], + ) + result = TrainDefaults.get_compute(compute=compute) + + assert result.instance_type is None + assert result.instance_count is None + class TestTrainDefaultsGetStoppingCondition: """Test TrainDefaults.get_stopping_condition method.""" diff --git a/v3-examples/ml-ops-examples/v3-processing-instance-preferences.ipynb b/v3-examples/ml-ops-examples/v3-processing-instance-preferences.ipynb new file mode 100644 index 0000000000..ec5ae9462d --- /dev/null +++ b/v3-examples/ml-ops-examples/v3-processing-instance-preferences.ipynb @@ -0,0 +1,147 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SageMaker V3 Processing with Instance Preferences (Multi-Instance-Type)\n", + "\n", + "A processing job can name an **ordered list of up to 5 candidate instance types** via the\n", + "`Processor`'s `instance_preferences` parameter. SageMaker launches the job on the first\n", + "candidate with available capacity and reports the choice as `SelectedInstanceType` /\n", + "`SelectedInstanceCount` on the job's `ClusterConfig`.\n", + "\n", + "Training jobs support the same feature \u2014 see the\n", + "[training example](../training-examples/instance-preferences-example.ipynb).\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Setup Session" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sagemaker.core.helper.session_helper import Session, get_execution_role\n", + "from sagemaker.core import image_uris\n", + "\n", + "sagemaker_session = Session()\n", + "role = get_execution_role()\n", + "region = sagemaker_session.boto_region_name\n", + "\n", + "processing_image = image_uris.retrieve(\n", + " framework=\"sklearn\",\n", + " region=region,\n", + " version=\"1.2-1\",\n", + " instance_type=\"ml.m5.xlarge\",\n", + " image_scope=\"training\",\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Run a processing job with an ordered list of instance preferences" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sagemaker.core.processing import Processor\n", + "\n", + "processor = Processor(\n", + " role=role,\n", + " image_uri=processing_image,\n", + " instance_preferences=[\n", + " {\"InstanceType\": \"ml.m5.4xlarge\"},\n", + " {\"InstanceType\": \"ml.m5.2xlarge\"},\n", + " ],\n", + " instance_count=2, # applies to whichever preference wins\n", + " volume_size_in_gb=100,\n", + ")\n", + "\n", + "processor.run(wait=False, logs=False, job_name=\"instance-prefs-processing-example\")\n", + "\n", + "describe = processor.sagemaker_session.sagemaker_client.describe_processing_job(\n", + " ProcessingJobName=\"instance-prefs-processing-example\"\n", + ")\n", + "cluster_config = describe[\"ProcessingResources\"][\"ClusterConfig\"]\n", + "print(f\"Submitted preferences: {cluster_config.get('InstancePreferences')}\")\n", + "print(f\"Selected instance type: {cluster_config.get('SelectedInstanceType')}\")\n", + "print(f\"Selected instance count: {cluster_config.get('SelectedInstanceCount')}\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Processing with per-preference instance counts\n", + "\n", + "Give **every** preference its own `InstanceCount` instead of the shared top-level\n", + "`instance_count` when the candidate types differ in size. Leave `instance_count` unset here.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "per_preference_processor = Processor(\n", + " role=role,\n", + " image_uri=processing_image,\n", + " instance_preferences=[\n", + " # 2 of the larger type...\n", + " {\"InstanceType\": \"ml.m5.4xlarge\", \"InstanceCount\": 2},\n", + " # ...or 4 of the smaller type\n", + " {\"InstanceType\": \"ml.m5.2xlarge\", \"InstanceCount\": 4},\n", + " ],\n", + " volume_size_in_gb=100,\n", + ")\n", + "\n", + "per_preference_processor.run(\n", + " wait=False, logs=False, job_name=\"instance-prefs-processing-per-pref-example\"\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Notes\n", + "\n", + "- **Backward compatible**: jobs that don't set `instance_preferences` behave exactly as\n", + " before.\n", + "- **Not supported with**: local mode and `FrameworkProcessor`. Training plans are\n", + " training-only and do not apply to processing.\n", + "- Each instance type may appear only once, and exactly one type is selected per job.\n", + "- Selection is based on capacity, not on workload fit: cross-type differences such as\n", + " architecture or GPU memory are not validated, so list only types your job can run on.\n", + "- Billing is based on the **selected** instance type and count.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3-examples/training-examples/instance-preferences-example.ipynb b/v3-examples/training-examples/instance-preferences-example.ipynb new file mode 100644 index 0000000000..fd5d67c5b7 --- /dev/null +++ b/v3-examples/training-examples/instance-preferences-example.ipynb @@ -0,0 +1,225 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SageMaker V3 Instance Preferences (Multi-Instance-Type) Example\n", + "\n", + "Instead of one fixed `instance_type`, a training job can name an **ordered list of up to 5\n", + "acceptable instance types**. SageMaker launches the job on the first candidate with\n", + "available capacity, so a busy first choice no longer means resubmitting.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Setup Session\n", + "\n", + "Initialize the SageMaker session, execution role, and a training image.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sagemaker.core.helper.session_helper import Session, get_execution_role\n", + "from sagemaker.core import image_uris\n", + "from sagemaker.core.shapes import InstancePreference\n", + "from sagemaker.train.model_trainer import ModelTrainer\n", + "from sagemaker.core.training.configs import Compute, InputData, OutputDataConfig\n", + "\n", + "sagemaker_session = Session()\n", + "role = get_execution_role()\n", + "region = sagemaker_session.boto_region_name\n", + "bucket = sagemaker_session.default_bucket()\n", + "\n", + "training_image = image_uris.retrieve(\n", + " framework=\"pytorch\",\n", + " region=region,\n", + " version=\"2.0.0\",\n", + " py_version=\"py310\",\n", + " instance_type=\"ml.m5.xlarge\",\n", + " image_scope=\"training\",\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Train with an ordered list of instance preferences\n", + "\n", + "The job asks for `ml.p5.48xlarge` first and falls back to `ml.p4d.24xlarge` if p5\n", + "capacity is unavailable. The top-level `instance_count=2` applies to whichever type wins.\n", + "Leave the classic `instance_type` field unset \u2014 it is mutually exclusive with the\n", + "preference list.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "compute = Compute(\n", + " instance_preferences=[\n", + " InstancePreference(instance_type=\"ml.p5.48xlarge\"), # priority 1\n", + " InstancePreference(instance_type=\"ml.p4d.24xlarge\"), # priority 2 (fallback)\n", + " ],\n", + " instance_count=2, # applies to whichever preference wins\n", + " volume_size_in_gb=500,\n", + ")\n", + "\n", + "trainer = ModelTrainer(\n", + " base_job_name=\"instance-prefs-example\",\n", + " training_image=training_image,\n", + " role=role,\n", + " compute=compute,\n", + " output_data_config=OutputDataConfig(\n", + " s3_output_path=f\"s3://{bucket}/instance-preferences-example/output\"\n", + " ),\n", + ")\n", + "\n", + "trainer.train(\n", + " input_data_config=[\n", + " InputData(\n", + " channel_name=\"train\",\n", + " data_source=f\"s3://{bucket}/instance-preferences-example/input\",\n", + " )\n", + " ],\n", + " wait=False,\n", + ")\n", + "training_job = trainer._latest_training_job\n", + "print(f\"Started: {training_job.training_job_name}\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Find out which instance type actually ran\n", + "\n", + "The submitted list is a set of *candidates*, not the outcome. Once a type is selected, the\n", + "job reports it in the read-only `selected_instance_type` / `selected_instance_count` fields;\n", + "the top-level `instance_type` you never set stays empty.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "training_job.refresh()\n", + "rc = training_job.resource_config\n", + "print(f\"Status: {training_job.training_job_status}\")\n", + "print(f\"Submitted preferences: {[p.instance_type for p in rc.instance_preferences]}\")\n", + "print(f\"Selected instance type: {rc.selected_instance_type}\") # None until selected\n", + "print(f\"Selected instance count: {rc.selected_instance_count}\") # None until selected\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Per-preference instance counts\n", + "\n", + "When candidate types have different sizes, give **every** element its own `instance_count`\n", + "and leave the top-level count unset. Counts go one way or the other \u2014 never both, never only\n", + "some elements.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "per_preference_compute = Compute(\n", + " instance_preferences=[\n", + " # 2 of the larger type...\n", + " InstancePreference(instance_type=\"ml.p5.48xlarge\", instance_count=2),\n", + " # ...or 4 of the smaller type\n", + " InstancePreference(instance_type=\"ml.p4d.24xlarge\", instance_count=4),\n", + " ],\n", + " volume_size_in_gb=500,\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Per-preference training plans (reserved capacity)\n", + "\n", + "A preference with a training plan draws from that plan's reserved capacity, and one without\n", + "falls back to on-demand \u2014 the job starts on whichever has capacity first. Each plan's\n", + "instance type must match its preference, one plan per preference, and per-preference\n", + "`training_plan_arns` is mutually exclusive with the whole-job `training_plan_arn` (which\n", + "instead applies to whichever preference matches its type).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "p5_plan_arn = f\"arn:aws:sagemaker:{region}: