Skip to content

fix(train): Enforce Nova Restricted storage on all MTRL MPG branches - #6186

Open
SifeiLi wants to merge 1 commit into
aws:masterfrom
SifeiLi:fix-nova-rmpg-mpg-resolution
Open

fix(train): Enforce Nova Restricted storage on all MTRL MPG branches#6186
SifeiLi wants to merge 1 commit into
aws:masterfrom
SifeiLi:fix-nova-rmpg-mpg-resolution

Conversation

@SifeiLi

@SifeiLi SifeiLi commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

Restricted managed storage was only applied when MultiTurnRLTrainer auto-created a fresh model package group. The explicit output_model_package_group branch and the continued-customization branch (group derived from the source ModelPackage) performed no Nova check.

Centralize the requirement in _nova_managed_configuration() and apply it on every resolution path for both the output and intermediate checkpoint groups. Existing groups that do not satisfy the requirement now fail fast at trainer construction with an actionable error, since groups cannot be converted after creation. Non-Nova (OSS) models are unaffected.

Adds unit tests for the explicit, object, derived, and pre-existing default-named group branches, plus OSS pass-through.

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.

Tested below cases

Integration Tests — Happy Path (real SageMaker, us-east-1, wheel build)

# Case Verification Method Result
I1 Nova + explicit Standard group Trainer construction raises ValueError (bug repro — old SDK proceeded silently) ✅ PASS
I2 Nova + no group (auto-create) DescribeModelPackageGroup on both auto-created groups returns ManagedStorageType=Restricted (output + intermediate checkpoint) ✅ PASS
I3 Nova + explicit Restricted group Accepted, resolves to group ARN ✅ PASS

Integration Tests — Deny Cases (real SageMaker, us-east-1, wheel build)

# Case Why It Matters Result
D1 Explicit Standard output group (name string) Original repro path ✅ PASS
D2 Pre-existing default-named Standard group (nova-…-mtrl-mpg) Residue left by the old buggy SDK in customer accounts is not silently reused ✅ PASS
D3 Explicit Standard intermediate checkpoint group Second resolver enforced, not just output ✅ PASS
D4 ModelPackageGroup object referencing Standard group Caller-constructed objects cannot bypass validation (authoritative re-fetch) ✅ PASS
D5 Group derived from source ModelPackage (continued customization) The exact branch from the issue — Nova chains no longer inherit Standard groups ✅ PASS

Notes: All rejections produce the actionable message "ModelPackageGroup 'X' uses 'Standard' managed storage, but model 'Y' requires 'Restricted'. Existing groups cannot be converted…". Integration runs used the built wheel
(sagemaker_train-1.20.0) in a clean venv; construction exercised the full real path (live hub metadata for nova-textgeneration-lite-v2, real MPG describe/create, S3 validation). All test fixtures deleted post-run. Non-Nova (OSS)
models: no validation, no extra API calls — behavior unchanged.

Restricted managed storage was only applied when MultiTurnRLTrainer
auto-created a fresh model package group. The explicit
output_model_package_group branch and the continued-customization
branch (group derived from the source ModelPackage) performed no Nova
check, so a Nova run could silently land its output or intermediate
checkpoints in a Standard-storage group, or split lineage across
groups depending on the branch taken.

Centralize the requirement in _nova_managed_configuration() and apply
it on every resolution path for both the output and intermediate
checkpoint groups. Existing groups that do not satisfy the requirement
now fail fast at trainer construction with an actionable error, since
groups cannot be converted after creation. Non-Nova (OSS) models are
unaffected.

Adds unit tests for the explicit, object, derived, and pre-existing
default-named group branches, plus OSS pass-through.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant