SageMaker V3 Instance Preferences (Multi-Instance-Type) Example

SageMaker V3 Instance Preferences (Multi-Instance-Type) Example#

Instead of one fixed instance_type, a training job can name an ordered list of up to 5 acceptable instance types. SageMaker launches the job on the first candidate with available capacity, so a busy first choice no longer means resubmitting.

Step 1: Setup Session#

Initialize the SageMaker session, execution role, and a training image.

The image is fixed at submission time while the instance type is not, so resolve it from one of the candidate types you will list in Step 2 and make sure every candidate can run it. Here both candidates are GPU types and share the same GPU image.

from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.core import image_uris
from sagemaker.core.shapes import InstancePreference
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.core.training.configs import Compute, InputData, OutputDataConfig

sagemaker_session = Session()
role = get_execution_role()
region = sagemaker_session.boto_region_name
bucket = sagemaker_session.default_bucket()

training_image = image_uris.retrieve(
    framework="pytorch",
    region=region,
    version="2.0.0",
    py_version="py310",
    instance_type="ml.p5.48xlarge",  # a Step 2 candidate; ml.p4d.24xlarge resolves to the same image
    image_scope="training",
)

Step 2: Train with an ordered list of instance preferences#

The job asks for ml.p5.48xlarge first and falls back to ml.p4d.24xlarge if p5 capacity is unavailable. The top-level instance_count=2 applies to whichever type wins. Leave the classic instance_type field unset — it is mutually exclusive with the preference list.

compute = Compute(
    instance_preferences=[
        InstancePreference(instance_type="ml.p5.48xlarge"),   # priority 1
        InstancePreference(instance_type="ml.p4d.24xlarge"),  # priority 2 (fallback)
    ],
    instance_count=2,  # applies to whichever preference wins
    volume_size_in_gb=500,
)

trainer = ModelTrainer(
    base_job_name="instance-prefs-example",
    training_image=training_image,
    role=role,
    compute=compute,
    output_data_config=OutputDataConfig(
        s3_output_path=f"s3://{bucket}/instance-preferences-example/output"
    ),
)

trainer.train(
    input_data_config=[
        InputData(
            channel_name="train",
            data_source=f"s3://{bucket}/instance-preferences-example/input",
        )
    ],
    wait=False,
)
training_job = trainer._latest_training_job
print(f"Started: {training_job.training_job_name}")

Step 3: Find out which instance type actually ran#

The submitted list is a set of candidates, not the outcome. Once a type is selected, the job reports it in the read-only selected_instance_type / selected_instance_count fields; the top-level instance_type you never set stays empty.

training_job.refresh()
rc = training_job.resource_config
print(f"Status:                  {training_job.training_job_status}")
print(f"Submitted preferences:   {[p.instance_type for p in rc.instance_preferences]}")
print(f"Selected instance type:  {rc.selected_instance_type}")   # None until selected
print(f"Selected instance count: {rc.selected_instance_count}")  # None until selected

Step 4: Per-preference instance counts#

When candidate types have different sizes, give every element its own instance_count and leave the top-level count unset. Counts go one way or the other — never both, never only some elements.

per_preference_compute = Compute(
    instance_preferences=[
        # 2 of the larger type...
        InstancePreference(instance_type="ml.p5.48xlarge", instance_count=2),
        # ...or 4 of the smaller type
        InstancePreference(instance_type="ml.p4d.24xlarge", instance_count=4),
    ],
    volume_size_in_gb=500,
)

Step 5: Per-preference training plans (reserved capacity)#

A preference with a training plan draws from that plan’s reserved capacity, and one without falls back to on-demand — the job starts on whichever has capacity first. Each plan’s instance type must match its preference, one plan per preference, and per-preference training_plan_arns is mutually exclusive with the whole-job training_plan_arn (which instead applies to whichever preference matches its type).

p5_plan_arn = f"arn:aws:sagemaker:{region}:<ACCOUNT_ID>:training-plan/<P5_PLAN_NAME>"
p4d_plan_arn = f"arn:aws:sagemaker:{region}:<ACCOUNT_ID>:training-plan/<P4D_PLAN_NAME>"

plan_backed_compute = Compute(
    instance_preferences=[
        InstancePreference(
            instance_type="ml.p5.48xlarge",
            training_plan_arns=[p5_plan_arn],
        ),
        InstancePreference(
            instance_type="ml.p4d.24xlarge",
            training_plan_arns=[p4d_plan_arn],
        ),
        InstancePreference(instance_type="ml.p4de.24xlarge"),  # on-demand fallback
    ],
    instance_count=2,
)

Notes#

  • Backward compatible: jobs that don’t set instance_preferences behave exactly as before.

  • Not supported with: local mode, training recipes, JumpStart models, AWS Batch training queues, heterogeneous clusters (instance_groups), instance_placement_config, and managed spot training.

  • 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 run on.

  • max_pending_time_in_seconds bounds the total time spent working through the list rather than each preference, and takes effect only when the list includes an accelerated instance type (ml.p, ml.g, ml.trn).

  • Billing is based on the selected instance type and count.

  • Processing jobs support the same feature: see the processing example.

Contents