feat(core): add InstancePreferences for multi-instance-type training and processing

[papriwal]

Caution

DO NOT MERGE

Description

Adds support for InstancePreferences, an ordered list of candidate instance types for training and processing jobs. When set, the service tries each candidate in list order and launches on the first type with available capacity, then reports the resolved choice back via SelectedInstanceType and SelectedInstanceCount.

This makes capacity-constrained jobs easier to schedule: instead of picking one instance type and retrying on InsufficientCapacityError, you supply a ranked list once.

Training

from sagemaker.core.shapes import InstancePreference

compute = Compute(
    instance_preferences=[
        InstancePreference(instance_type="ml.p5.48xlarge", instance_count=2),
        InstancePreference(instance_type="ml.p4d.24xlarge", instance_count=4),
    ],
    volume_size_in_gb=100,
)

instance_preferences is mutually exclusive with instance_type, instance_groups and instance_placement_config. Each entry may carry its own instance_count, or you can leave them unset and use the uniform ResourceConfig.InstanceCount for whichever type wins. training_plan_arns is training-only and capped at one entry per preference.

Processing

ProcessingClusterConfig gains the same three members. InstanceType and InstanceCount are no longer required on that shape, since a preferences-only config supplies them per candidate — VolumeSizeInGB remains the only required member.

What's in this change

Area	Files	Notes
sagemaker-core src	9	training/processing configs and utils, Spark processor plumbing, plus regenerated shapes.py, shape_dag.py and the bundled service model
sagemaker-core tests	5	unit coverage for compute configs, processing and the service model; a new processing integ test
sagemaker-train	4	model_trainer.py validation, defaults.py handling, unit and integ tests
Docs	2	docs/training/index.rst, docs/ml_ops/index.rst
Examples	2	one training and one processing notebook

Generated artifacts were produced by sagemaker-core's own codegen from the bundled sample model, not hand-edited. New shapes: InstancePreference, InstancePreferenceList, ProcessingInstancePreference, ProcessingInstancePreferenceList, TrainingPlanArnList.

Unsupported combinations are rejected with a clear error rather than failing at the service: local mode, training recipes, and JumpStart training all reject instance_preferences up front.

Notes for reviewers

Requires botocore ≥ 1.43.90. These shapes first ship in that release. On an older botocore the client rejects the new members with a ParamValidationError; only callers using this feature are affected. The dependency floor is bumped as part of the release, not in this PR.

One unrelated model line moves. The bundled service-2.json refresh also brings MaxPendingTimeInSeconds.min from 7200 to 1800, matching botocore 1.43.90+. It shows up in the diff but is not part of this feature — the checked-in model was simply behind on that field.

IntPipeVar on five members is not reproducible by codegen. shapes.py widens instance_count on InstancePreference, ProcessingInstancePreference and ProcessingClusterConfig, plus selected_instance_count on ResourceConfig and ProcessingClusterConfig, to IntPipeVar so pipeline variables can be passed. PIPE_VAR_OVERRIDES in tools/constants.py is keyed by shape name and currently only covers ResourceConfig, so re-running codegen narrows these back to Optional[int]. Happy to add the missing override entries here or in a follow-up — worth deciding, since the periodic botocore sync regenerates this file. The two selected_instance_count fields are Describe-response only, so a pipeline variable arguably shouldn't apply to them at all.

Testing

• Unit tests for training and processing config construction, serialization, and the mutually-exclusive validation paths
• A service-model test asserting the new shapes and the relaxed required list
• Integration tests for training and processing that assert the full round trip, including that the service echoes SelectedInstanceType / SelectedInstanceCount; they skip unless the required role, image and output-path environment variables are set
• Both notebooks run end to end

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @papriwal's task in 3m 18s —— View job

---

Code review: feat(core): add InstancePreferences for multi-instance-type training and processing

• Read the full PR diff
• Verify validation logic (training + processing count/exclusivity rules)
• Check imports (Unassigned, List/Dict/Union), serialization of output-only fields
• Cross-check Spark image-resolution branch and defaults handling
• Post review

---

Overall this is a clean, well-structured, and unusually well-tested change. The mutually-exclusive/count-mode validation is mirrored consistently across the training and processing planes, the output-only selected_* fields are correctly kept out of create requests, and the unsupported-combination guards (local mode / recipes / JumpStart) fail fast with clear messages. I verified the two things most likely to bite:

• Unassigned is imported in both modules/utils.py and training/utils.py, so the new validators won't NameError.
• Output-only selected_instance_type / selected_instance_count are stripped before the API call. The modules-path _to_resource_config keeps None values in filtered_dict (unlike the training path, which drops them with and v is not None), but serialize()/_serialize_shape drop None/Unassigned, so the create request never sends them. No functional bug — just an asymmetry between the two _to_resource_config implementations that's a little surprising to read.

No correctness bugs found. A few minor, non-blocking observations:

1. Codegen drift on the IntPipeVar widening (most actionable)

You already flagged this in the PR description, and I agree it's the one thing worth resolving in-PR rather than deferring. shapes.py widens instance_count (on InstancePreference, ProcessingInstancePreference, ProcessingClusterConfig) and selected_instance_count to IntPipeVar, but PIPE_VAR_OVERRIDES in tools/constants.py only keys ResourceConfig. The next periodic botocore sync will regenerate shapes.py and silently narrow these back to Optional[int], breaking pipeline-variable callers with no test catching it. Recommend adding the override entries here (and, per your own note, dropping the two Describe-only selected_instance_count fields from the widening since a pipeline variable can't meaningfully apply to a response-only field).

2. Docs/notebooks describe client-side rejections that aren't enforced client-side

• docs/training/index.rst lists instance_preferences as mutually exclusive with managed spot training, but validate_instance_preferences only checks instance_type / instance_groups / instance_placement_config — not enable_managed_spot_training. Similarly the training docs say AWS Batch queues reject it, with no client guard.
• The processing notebook says "Not supported with: local mode and FrameworkProcessor", but Processor has no local-mode guard, whereas the training path added explicit ModelTrainer guards for local/recipe/JumpStart.

None of these are bugs (the service is the source of truth and will reject), but the enforcement is asymmetric with the training side and the docs read as if the SDK validates them. Either add the matching client-side checks or soften the wording to "rejected by the service."

3. Public-API asymmetry: typed objects vs. raw dicts

Training takes List[InstancePreference] (typed pydantic models) while processing takes List[Dict[str, Union[str,int]]] with PascalCase keys ({"InstanceType": ..., "InstanceCount": ...}). This mirrors the existing Processor style, so it's a defensible choice, but it's an inconsistency users will notice and there's no shared validation helper between the two validate/_validate_processing_... functions (near-duplicated logic). Worth a comment in the docstrings pointing users at the dict shape (the processing docstring does this — good).

Minor

• test_pyspark_degenerate_preferences_reach_base_validation passes instance_preferences=[{}]; the resulting error is the count-mode message ("every element must set its own InstanceCount") rather than a clear "missing InstanceType" one. A dedicated "each preference requires InstanceType" check would give a friendlier message, though the server also enforces it.
• The unrelated MaxPendingTimeInSeconds.min 7200 → 1800 change and the botocore ≥ 1.43.90 floor are well-called-out in the description; just make sure the dependency floor bump actually lands with the release as noted, since older botocore raises ParamValidationError on these new members.

Nice work — the test coverage (round-trip, per-preference vs. uniform count, Spark image agreement/divergence, service-model pinning, defaults handling) is thorough and the docstrings explaining the count modes are excellent.
· branch feature-smtj-instance-preferences-latest

[codecov]

Codecov Report

❌ Patch coverage is 95.93496% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.75%. Comparing base (10ada9b) to head (6fc2324).
⚠️ Report is 30 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6246      +/-   ##
==========================================
+ Coverage   70.68%   70.75%   +0.07%     
==========================================
  Files         553      553              
  Lines       67891    68005     +114     
==========================================
+ Hits        47990    48120     +130     
+ Misses      19901    19885      -16     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jam-jee]

not sure if we need to update service-2.json manually? It would be auto-updated based on boto spec?

[deeppcs]

This was just updated manually to enable local testing before Trebuchet release. This should be a no-op now that these fields are already released in the public spec?