diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/document.py b/sagemaker-core/src/sagemaker/core/jumpstart/document.py index d9feb40984..35fdfa0994 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/document.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/document.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """This module contains utilites for JumpStart model metadata.""" + from __future__ import absolute_import import json @@ -47,26 +48,54 @@ def get_hub_content_and_document( logger.debug("No sagemaker session provided. Using default session.") hub_name = jumpstart_config.hub_name if jumpstart_config.hub_name else SAGEMAKER_PUBLIC_HUB - hub_content_type = "Model" if hub_name == SAGEMAKER_PUBLIC_HUB else "ModelReference" region = sagemaker_session.boto_region_name - try: - hub_content = HubContent.get( - hub_name=hub_name, - hub_content_name=jumpstart_config.model_id, - hub_content_version=jumpstart_config.model_version, - hub_content_type=hub_content_type, - session=sagemaker_session.boto_session, - region=region, - ) - except ClientError as e: - if e.response["Error"]["Code"] == "ResourceNotFound": - logger.error( - f"Hub content {jumpstart_config.model_id} not found in {hub_name}.\n" - "Please check that the Model ID is availble in the specified hub." + # The hub content may be filed under an alias that differs from the public + # model_id, so honor hub_content_name when provided. + hub_content_name = ( + jumpstart_config.hub_content_name + if getattr(jumpstart_config, "hub_content_name", None) + else jumpstart_config.model_id + ) + + # A private hub can contain either a ModelReference (a pointer to a public + # JumpStart model) or a privately-owned Model authored directly into the + # hub. We cannot tell which from the name alone, so probe: try + # ModelReference first, then fall back to Model. The public hub only holds + # Models. This mirrors ModelBuilder's resolution in accessors.py. + if hub_name == SAGEMAKER_PUBLIC_HUB: + content_types_to_try = ["Model"] + else: + content_types_to_try = ["ModelReference", "Model"] + + hub_content = None + last_error: Optional[ClientError] = None + for content_type in content_types_to_try: + try: + hub_content = HubContent.get( + hub_name=hub_name, + hub_content_name=hub_content_name, + hub_content_version=jumpstart_config.model_version, + hub_content_type=content_type, + session=sagemaker_session.boto_session, + region=region, ) - raise e + break + except ClientError as e: + if e.response["Error"]["Code"] == "ResourceNotFound": + last_error = e + continue + raise e + + if hub_content is None: + logger.error( + f"Hub content {hub_content_name} not found in {hub_name} as any of " + f"{content_types_to_try}.\n" + "Please check that the Model ID (or hub_content_name) is available " + "in the specified hub." + ) + raise last_error logger.info( f"hub_content_name: {hub_content.hub_content_name}, " diff --git a/sagemaker-core/tests/unit/jumpstart/test_document.py b/sagemaker-core/tests/unit/jumpstart/test_document.py index 08c8b6d2c0..653db1290b 100644 --- a/sagemaker-core/tests/unit/jumpstart/test_document.py +++ b/sagemaker-core/tests/unit/jumpstart/test_document.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Test for JumpStart Document.""" + from __future__ import absolute_import import json @@ -81,3 +82,145 @@ def test_get_hub_content_document_failure(jumpstart_session): get_hub_content_and_document( jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session ) + + +# --------------------------------------------------------------------------- +# Tests for private-hub content-type probing + hub_content_name alias support. +# +# A private hub can contain either a ModelReference (a pointer to a public +# model) or a privately-owned Model. get_hub_content_and_document() must not +# guess from the hub name; it probes ModelReference first, then falls back to +# Model. The public hub only holds Models. It also honors hub_content_name when +# the content is filed under an alias differing from model_id. +# +# Note: distinct model_id / hub_name values are used per test to avoid the +# module-level lru_cache on get_hub_content_and_document returning a stale +# result across tests. +# --------------------------------------------------------------------------- + + +def _hub_content(hub_name, name, content_type, doc): + return HubContent( + hub_name=hub_name, + hub_content_name=name, + hub_content_version="1.0.0", + hub_content_type=content_type, + hub_content_document=json.dumps(doc), + ) + + +def _not_found(): + return ClientError( + error_response={"Error": {"Code": "ResourceNotFound"}}, + operation_name="DescribeHubContent", + ) + + +def _load_doc(): + cur_dir = os.path.dirname(os.path.abspath(__file__)) + with open(os.path.join(cur_dir, "hub_content_document.json"), "r") as f: + return json.load(f) + + +def test_public_hub_uses_model_type_only(jumpstart_session): + """Public hub: resolve as Model, and never probe ModelReference.""" + doc = _load_doc() + jumpstart_config = JumpStartConfig(model_id="probe-public-model") + + with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get: + mock_get.return_value = _hub_content( + "SageMakerPublicHub", "probe-public-model", "Model", doc + ) + hub_content, _ = get_hub_content_and_document( + jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session + ) + + assert hub_content.hub_content_type == "Model" + # Public hub must be looked up exactly once, as Model. + assert mock_get.call_count == 1 + assert mock_get.call_args.kwargs["hub_content_type"] == "Model" + + +def test_private_hub_resolves_model_reference_first(jumpstart_session): + """Private hub holding a ModelReference: first probe (ModelReference) hits.""" + doc = _load_doc() + jumpstart_config = JumpStartConfig(model_id="probe-ref-model", hub_name="my-private-hub-ref") + + with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get: + mock_get.return_value = _hub_content( + "my-private-hub-ref", "probe-ref-model", "ModelReference", doc + ) + hub_content, _ = get_hub_content_and_document( + jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session + ) + + assert hub_content.hub_content_type == "ModelReference" + # ModelReference is tried first and succeeds -> single call. + assert mock_get.call_count == 1 + assert mock_get.call_args.kwargs["hub_content_type"] == "ModelReference" + + +def test_private_hub_falls_back_to_model(jumpstart_session): + """Private hub holding a privately-owned Model: ModelReference misses, then + the Model fallback resolves it (the core of the fix).""" + doc = _load_doc() + jumpstart_config = JumpStartConfig( + model_id="probe-private-model", hub_name="my-private-hub-model" + ) + + with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get: + mock_get.side_effect = [ + _not_found(), # ModelReference lookup misses + _hub_content( # Model fallback resolves + "my-private-hub-model", "probe-private-model", "Model", doc + ), + ] + hub_content, _ = get_hub_content_and_document( + jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session + ) + + assert hub_content.hub_content_type == "Model" + # Two probes: ModelReference (miss) then Model (hit). + assert mock_get.call_count == 2 + assert [c.kwargs["hub_content_type"] for c in mock_get.call_args_list] == [ + "ModelReference", + "Model", + ] + + +def test_private_hub_honors_hub_content_name_alias(jumpstart_session): + """When hub_content_name is set (alias differs from model_id), the lookup + must use the alias, not the model_id.""" + doc = _load_doc() + jumpstart_config = JumpStartConfig( + model_id="probe-alias-public-id", + hub_name="my-private-hub-alias", + hub_content_name="the-alias-name", + ) + + with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get: + mock_get.return_value = _hub_content( + "my-private-hub-alias", "the-alias-name", "ModelReference", doc + ) + get_hub_content_and_document( + jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session + ) + + # Lookup used the alias, not the model_id. + assert mock_get.call_args.kwargs["hub_content_name"] == "the-alias-name" + + +def test_private_hub_not_found_as_either_type_raises(jumpstart_session): + """Private hub where neither ModelReference nor Model exists: raise.""" + jumpstart_config = JumpStartConfig( + model_id="probe-missing-model", hub_name="my-private-hub-missing" + ) + + with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get: + mock_get.side_effect = [_not_found(), _not_found()] + with pytest.raises(ClientError): + get_hub_content_and_document( + jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session + ) + # Both content types were attempted before giving up. + assert mock_get.call_count == 2 diff --git a/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py b/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py index 298ea85e3e..18c54ed9c5 100644 --- a/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py +++ b/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py @@ -10,15 +10,222 @@ # 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. -"""This module contains the Integ Tests for JumpStart Training.""" +"""This module contains the Integ Tests for JumpStart Training. + +Coverage: + * Public JumpStart models (model_id only). + * Private-hub ModelReference (a pointer to a public model), including an + aliased reference whose hub content name differs from the public model_id. + * A privately-owned Model authored directly into a private hub. + +The private-hub / private-model tests each create their own temporary hub, +populate it, run training, and tear the hub down. They skip gracefully if the +environment lacks permissions to create hubs or import content. +""" + from __future__ import absolute_import +import time +import uuid +import logging + import pytest +from botocore.exceptions import ClientError from sagemaker.core.jumpstart import JumpStartConfig +from sagemaker.core.helper.session_helper import Session from sagemaker.train import ModelTrainer from sagemaker.train.configs import Compute +logger = logging.getLogger(__name__) + +# A trainable classical-ML model keeps these tests fast/cheap on CPU. +TRAINABLE_MODEL_ID = "catboost-regression-model" +# Gated, trainable model reused from the v2 private-hub parity tests. Exercises +# the accept_eula / ModelAccessConfig path for a gated ModelReference. Gated +# models resolve to GPU and require real EULA acceptance, so the test that uses +# it runs a real training job and is marked slow_test + gpu_intensive (scheduled +# CI, not PR checks). The instance type is intentionally left to the SDK: +# from_jumpstart_config validates a supplied instance_type against the model's +# SupportedTrainingInstanceTypes and raises if it is not in the list, so +# resolving the model's own default is safer than hardcoding one here. +GATED_TRAINABLE_MODEL_ID = "meta-textgeneration-llama-3-2-1b" +HUB_NAME_PREFIX = "sdk-integ-train-hub" +ALIASED_REFERENCE_NAME = "sdk-integ-aliased-catboost" +PRIVATE_MODEL_NAME = "sdk-integ-private-catboost" + +# Only these error codes represent "this restricted account is not allowed to +# set up the fixture" and warrant a graceful skip. Any other ClientError is a +# real failure (e.g. a service-side regression in create_hub_content_reference +# or import_hub_content) and must fail loudly so the test does not silently +# stop guarding the fix while CI stays green. +_SKIPPABLE_SETUP_ERROR_CODES = frozenset( + { + "AccessDeniedException", + "AccessForbiddenException", + "UnauthorizedOperation", + } +) + + +def _skip_if_unauthorized(e, message): + """Skip only on an expected authorization error; re-raise everything else.""" + if e.response.get("Error", {}).get("Code") in _SKIPPABLE_SETUP_ERROR_CODES: + pytest.skip(f"{message}: {e}") + raise + + +def _assert_reference_channels(model_trainer): + """Assert the SDK resolved a ModelReference into hub-aware training channels. + + Resolving a ModelReference must produce a container image and attach a + HubAccessConfig(hub_content_arn=...) to the model channel (defaults.py, + hub_content_type == "ModelReference" branch). Asserting this on the + SDK-resolved channels — rather than passing an explicit training channel to + train() — is what actually guards the fix; a non-gated model would train + fine even if this plumbing regressed. + """ + assert model_trainer.training_image + model_channels = [ + c for c in model_trainer.input_data_config if getattr(c, "channel_name", None) == "model" + ] + assert len(model_channels) == 1 + hub_access_config = model_channels[0].data_source.s3_data_source.hub_access_config + assert hub_access_config is not None + assert hub_access_config.hub_content_arn + + +def _assert_gated_reference_channels(model_trainer): + """Assert accept_eula flowed into the model channel's ModelAccessConfig. + + Verified against defaults.py get_model_artifact_input: the resolved "model" + channel always carries + data_source.s3_data_source.model_access_config = ModelAccessConfig( + accept_eula=jumpstart_config.accept_eula). Using a gated model_id is what + makes accept_eula=True meaningful (a gated model is unusable without it); the + assertion itself is the same plumbing every JumpStart model uses. + + Asserted before .train() so the gated ModelAccessConfig plumbing is pinned + even if the training job itself later fails for an unrelated capacity/quota + reason. + """ + assert model_trainer.training_image + model_channels = [ + c for c in model_trainer.input_data_config if getattr(c, "channel_name", None) == "model" + ] + assert len(model_channels) == 1 + model_access_config = model_channels[0].data_source.s3_data_source.model_access_config + assert model_access_config is not None, "gated reference resolved without a ModelAccessConfig" + assert model_access_config.accept_eula is True + + +def _assert_owned_model_channels(model_trainer): + """Assert a privately-owned Model resolved into direct (non-brokered) channels. + + An owned Model (not a reference) must resolve to a model channel with a real + S3 artifact URI and NO HubAccessConfig — the inverse of the reference case. + """ + assert model_trainer.training_image + model_channels = [ + c for c in model_trainer.input_data_config if getattr(c, "channel_name", None) == "model" + ] + assert len(model_channels) == 1 + s3_source = model_channels[0].data_source.s3_data_source + assert s3_source.s3_uri + assert s3_source.hub_access_config is None + + +def _sm_client(sagemaker_session): + return sagemaker_session.boto_session.client("sagemaker") + + +def _region(sagemaker_session): + return sagemaker_session.boto_region_name + + +def _execution_role(sagemaker_session): + """Resolve a SageMaker execution role from the running environment.""" + return sagemaker_session.get_caller_identity_arn() + + +def _public_model_arn(region, model_id): + return f"arn:aws:sagemaker:{region}:aws:hub-content/" f"SageMakerPublicHub/Model/{model_id}" + + +def _wait_for_content(sm, hub_name, name, content_type, timeout=300, poll=10): + deadline = time.time() + timeout + while time.time() < deadline: + try: + resp = sm.describe_hub_content( + HubName=hub_name, + HubContentName=name, + HubContentType=content_type, + ) + if resp.get("HubContentStatus") == "Available": + return True + except ClientError: + pass + time.sleep(poll) + return False + + +def _delete_hub(sm, hub_name): + for content_type in ("ModelReference", "Model"): + try: + resp = sm.list_hub_contents(HubName=hub_name, HubContentType=content_type) + except ClientError: + continue + for c in resp.get("HubContentSummaries", []): + try: + if content_type == "ModelReference": + sm.delete_hub_content_reference( + HubName=hub_name, + HubContentType=content_type, + HubContentName=c["HubContentName"], + ) + else: + sm.delete_hub_content( + HubName=hub_name, + HubContentType=content_type, + HubContentName=c["HubContentName"], + HubContentVersion=c["HubContentVersion"], + ) + except ClientError as e: + logger.warning("Failed to delete hub content %s: %s", c, e) + try: + sm.delete_hub(HubName=hub_name) + except ClientError as e: + logger.warning("Failed to delete hub %s: %s", hub_name, e) + + +@pytest.fixture(scope="module") +def sagemaker_session(): + return Session() + + +@pytest.fixture(scope="module") +def private_hub(sagemaker_session): + """Create a temporary private hub; tear it (and its contents) down after.""" + sm = _sm_client(sagemaker_session) + hub_name = f"{HUB_NAME_PREFIX}-{uuid.uuid4().hex[:8]}" + try: + sm.create_hub( + HubName=hub_name, + HubDescription="SDK integ test JumpStart training private hub", + ) + except ClientError as e: + _skip_if_unauthorized(e, "Cannot create private hub (missing permissions?)") + + for _ in range(30): + if sm.describe_hub(HubName=hub_name)["HubStatus"] == "InService": + break + time.sleep(2) + else: + pytest.fail(f"Hub {hub_name} did not reach InService") + + yield hub_name + _delete_hub(sm, hub_name) + @pytest.mark.parametrize( "test_case", @@ -42,7 +249,7 @@ ], ) def test_jumpstart_train(test_case): - """Test JumpStart training.""" + """Test JumpStart training from a public model_id.""" jumpstart = JumpStartConfig( model_id=test_case["model_id"], accept_eula=test_case.get("accept_eula", False), @@ -54,3 +261,192 @@ def test_jumpstart_train(test_case): compute=test_case.get("compute"), ) model_trainer.train() + + +def test_jumpstart_train_from_private_hub_reference(private_hub, sagemaker_session): + """Train from a ModelReference (pointer to a public model) in a private hub.""" + sm = _sm_client(sagemaker_session) + region = _region(sagemaker_session) + + try: + sm.create_hub_content_reference( + HubName=private_hub, + SageMakerPublicHubContentArn=_public_model_arn(region, TRAINABLE_MODEL_ID), + ) + except ClientError as e: + _skip_if_unauthorized(e, "Cannot create hub content reference") + if not _wait_for_content(sm, private_hub, TRAINABLE_MODEL_ID, "ModelReference"): + pytest.fail( + f"ModelReference {TRAINABLE_MODEL_ID} did not become Available in {private_hub}" + ) + + jumpstart = JumpStartConfig(model_id=TRAINABLE_MODEL_ID, hub_name=private_hub, accept_eula=True) + model_trainer = ModelTrainer.from_jumpstart_config( + jumpstart, + role=_execution_role(sagemaker_session), + base_job_name="sdk-integ-train-ref", + compute=Compute(instance_type="ml.m5.xlarge"), + sagemaker_session=sagemaker_session, + ) + + # Assert the fix's plumbing on the SDK-resolved channels before training: + # resolving a ModelReference must attach a HubAccessConfig(hub_content_arn=...) + # to the model channel. This is the assertion that would catch the plumbing + # regressing; a non-gated model would otherwise train fine even if it broke. + _assert_reference_channels(model_trainer) + + # Train on the SDK-resolved channels (no explicit training channel), so the + # hub-aware channel construction under test is actually exercised. + model_trainer.train() + + +@pytest.mark.slow_test +@pytest.mark.gpu_intensive +def test_jumpstart_train_from_gated_reference(private_hub, sagemaker_session): + """Train from a GATED ModelReference in a private hub, verifying the + accept_eula / ModelAccessConfig path. + + Gated models resolve to a GPU instance and require real EULA acceptance, so + this runs a real training job and is marked gpu_intensive (submits a real job + that consumes training capacity; scheduled CI, not PR checks) as well as + slow_test. The resolved channels are asserted before training so the fix's + ModelAccessConfig/HubAccessConfig plumbing is pinned even if the job itself + later fails for an unrelated capacity/quota reason.""" + sm = _sm_client(sagemaker_session) + region = _region(sagemaker_session) + + try: + sm.create_hub_content_reference( + HubName=private_hub, + SageMakerPublicHubContentArn=_public_model_arn(region, GATED_TRAINABLE_MODEL_ID), + ) + except ClientError as e: + _skip_if_unauthorized(e, "Cannot create gated hub content reference") + if not _wait_for_content(sm, private_hub, GATED_TRAINABLE_MODEL_ID, "ModelReference"): + pytest.fail( + f"Gated reference {GATED_TRAINABLE_MODEL_ID} did not become Available in {private_hub}" + ) + + jumpstart = JumpStartConfig( + model_id=GATED_TRAINABLE_MODEL_ID, + hub_name=private_hub, + accept_eula=True, + ) + model_trainer = ModelTrainer.from_jumpstart_config( + jumpstart, + role=_execution_role(sagemaker_session), + base_job_name="sdk-integ-train-gated-ref", + # No compute: let from_jumpstart_config resolve the gated model's own + # default (GPU) instance type. Passing one risks a ValueError if it is + # not in the model's SupportedTrainingInstanceTypes. + sagemaker_session=sagemaker_session, + ) + + # Pin the fix's plumbing on the SDK-resolved channels before training: a gated + # ModelReference must resolve with accept_eula flowed into a ModelAccessConfig, + # plus the HubAccessConfig every reference gets. + _assert_reference_channels(model_trainer) + _assert_gated_reference_channels(model_trainer) + + # Train on the SDK-resolved channels (no explicit training channel), so the + # hub-aware, gated channel construction under test is actually exercised + # end-to-end against a real training job. + model_trainer.train() + + +def test_jumpstart_train_from_aliased_reference(private_hub, sagemaker_session): + """Train from a ModelReference filed under an alias that differs from the + public model_id (exercises hub_content_name resolution).""" + sm = _sm_client(sagemaker_session) + region = _region(sagemaker_session) + + try: + sm.create_hub_content_reference( + HubName=private_hub, + SageMakerPublicHubContentArn=_public_model_arn(region, TRAINABLE_MODEL_ID), + HubContentName=ALIASED_REFERENCE_NAME, + ) + except ClientError as e: + _skip_if_unauthorized(e, "Cannot create aliased hub content reference") + if not _wait_for_content(sm, private_hub, ALIASED_REFERENCE_NAME, "ModelReference"): + pytest.fail(f"Aliased reference {ALIASED_REFERENCE_NAME} did not become Available") + + jumpstart = JumpStartConfig( + model_id=TRAINABLE_MODEL_ID, + hub_name=private_hub, + hub_content_name=ALIASED_REFERENCE_NAME, + accept_eula=True, + ) + model_trainer = ModelTrainer.from_jumpstart_config( + jumpstart, + role=_execution_role(sagemaker_session), + base_job_name="sdk-integ-train-alias", + compute=Compute(instance_type="ml.m5.xlarge"), + sagemaker_session=sagemaker_session, + ) + + # The alias must have been threaded through resolution (not the model_id). + assert model_trainer._jumpstart_config.hub_content_name == ALIASED_REFERENCE_NAME + # ...and resolving it as a ModelReference must attach the hub-aware channels. + _assert_reference_channels(model_trainer) + + # Train on the SDK-resolved channels (no explicit training channel). + model_trainer.train() + + +def test_jumpstart_train_from_private_owned_model(private_hub, sagemaker_session): + """Train from a privately-owned Model authored directly into a private hub + (content-type Model, not a ModelReference). Exercises the document.py + fallback-to-Model resolution probe.""" + sm = _sm_client(sagemaker_session) + + # Author a private Model by importing a trainable public model's document + # into the private hub as content-type Model. + try: + public = sm.describe_hub_content( + HubName="SageMakerPublicHub", + HubContentType="Model", + HubContentName=TRAINABLE_MODEL_ID, + ) + except ClientError as e: + _skip_if_unauthorized(e, "Cannot read public model document") + + try: + sm.import_hub_content( + HubName=private_hub, + HubContentName=PRIVATE_MODEL_NAME, + HubContentType="Model", + HubContentDocument=public["HubContentDocument"], + DocumentSchemaVersion=public.get("DocumentSchemaVersion", "2.0.0"), + HubContentDisplayName=public.get("HubContentDisplayName", PRIVATE_MODEL_NAME), + HubContentDescription="Privately owned model for integ test", + HubContentMarkdown=public.get("HubContentMarkdown", ""), + HubContentSearchKeywords=public.get("HubContentSearchKeywords", []), + ) + except ClientError as e: + # Only skip if the account is simply not allowed to author a private + # Model. Any other failure is a real regression in the owned-Model path + # (the core case this fix enables) and must fail loudly. + _skip_if_unauthorized(e, "import_hub_content for a private Model not permitted") + if not _wait_for_content(sm, private_hub, PRIVATE_MODEL_NAME, "Model"): + pytest.fail(f"Private Model {PRIVATE_MODEL_NAME} did not become Available in {private_hub}") + + jumpstart = JumpStartConfig( + model_id=TRAINABLE_MODEL_ID, + hub_name=private_hub, + hub_content_name=PRIVATE_MODEL_NAME, + accept_eula=True, + ) + model_trainer = ModelTrainer.from_jumpstart_config( + jumpstart, + role=_execution_role(sagemaker_session), + base_job_name="sdk-integ-train-private", + compute=Compute(instance_type="ml.m5.xlarge"), + sagemaker_session=sagemaker_session, + ) + + # Owned Model (fallback-to-Model probe): direct S3 artifact, no HubAccessConfig. + _assert_owned_model_channels(model_trainer) + + # Train on the SDK-resolved channels (no explicit training channel). + model_trainer.train() diff --git a/sagemaker-train/tests/unit/train/test_defaults.py b/sagemaker-train/tests/unit/train/test_defaults.py index 1c02b72e01..236b612aca 100644 --- a/sagemaker-train/tests/unit/train/test_defaults.py +++ b/sagemaker-train/tests/unit/train/test_defaults.py @@ -26,6 +26,7 @@ ) from sagemaker.train.configs import Compute, StoppingCondition from sagemaker.core.shapes import InstanceGroup +from sagemaker.core.jumpstart.configs import JumpStartConfig class TestDefaultConstants: @@ -678,3 +679,146 @@ def test_sets_default_volume_size_when_instance_groups_and_no_document_volume( sagemaker_session=mock_session, ) assert result.volume_size_in_gb == DEFAULT_VOLUME_SIZE + + +# Gated model reused from the v2 private-hub parity tests (llama-3.2-1b). +# A gated ModelReference is the case the reviewer asked to cover: it must +# both propagate accept_eula into ModelAccessConfig and attach a +# HubAccessConfig (because it resolves as a ModelReference). +GATED_MODEL_ID = "meta-textgeneration-llama-3-2-1b" + + +class TestJumpStartTrainDefaultsGatedModelReferenceEula: + """EULA / ModelAccessConfig handling for a gated ModelReference in a private hub. + + These are the training-side analogue of the v2 gated private-hub test. They + are fully mocked at the resolver seam (get_hub_content_and_document) so they + are fast and credential-free, and they assert the two things that must work + for a gated reference: + 1. accept_eula flows into ModelAccessConfig.accept_eula on the S3 source. + 2. A HubAccessConfig (brokered artifact access) is attached because the + content resolves as a ModelReference. + """ + + def _gated_reference_hub_content(self): + """A hub_content mock standing in for a gated ModelReference.""" + hub_content = MagicMock() + hub_content.hub_content_type = "ModelReference" + hub_content.hub_content_name = GATED_MODEL_ID + hub_content.hub_content_arn = ( + "arn:aws:sagemaker:us-west-2:123456789012:hub-content/" + f"my-private-hub/ModelReference/{GATED_MODEL_ID}" + ) + return hub_content + + def _training_components_model(self): + """A minimal training-components model with a resolvable artifact URI.""" + tcm = MagicMock() + tcm.TrainingArtifactUri = "s3://jumpstart-cache-prod-us-west-2/artifacts/model.tar.gz" + tcm.TrainingArtifactCompressionType = "None" + tcm.DefaultTrainingDatasetUri = "s3://jumpstart-cache-prod-us-west-2/datasets/train/" + return tcm + + @patch("sagemaker.train.defaults.JumpStartTrainDefaults._get_training_variant") + @patch("sagemaker.train.defaults.JumpStartTrainDefaults._get_training_components_model") + @patch("sagemaker.train.defaults.get_hub_content_and_document") + @patch("sagemaker.train.defaults.TrainDefaults.get_sagemaker_session") + def test_model_artifact_input_gated_reference_sets_accept_eula_and_hub_access( + self, mock_get_session, mock_get_hub_content, mock_get_tcm, mock_get_variant + ): + """Gated ModelReference -> model channel carries accept_eula=True + HubAccessConfig.""" + mock_get_session.return_value = MagicMock() + hub_content = self._gated_reference_hub_content() + mock_get_hub_content.return_value = (hub_content, MagicMock()) + mock_get_tcm.return_value = self._training_components_model() + # No instance-type variant; fall back to the base TrainingArtifactUri. + mock_get_variant.return_value = None + + jumpstart_config = JumpStartConfig( + model_id=GATED_MODEL_ID, + hub_name="my-private-hub", + accept_eula=True, + ) + + result = JumpStartTrainDefaults.get_model_artifact_input( + jumpstart_config=jumpstart_config, + compute=Compute(instance_type="ml.g5.2xlarge", instance_count=1), + input_data_config=None, + environment={}, + sagemaker_session=mock_get_session.return_value, + ) + + model_channels = [c for c in result if c.channel_name == "model"] + assert len(model_channels) == 1 + s3_source = model_channels[0].data_source.s3_data_source + # 1. accept_eula propagated into ModelAccessConfig. + assert s3_source.model_access_config is not None + assert s3_source.model_access_config.accept_eula is True + # 2. HubAccessConfig attached because the content is a ModelReference. + assert s3_source.hub_access_config is not None + assert s3_source.hub_access_config.hub_content_arn == hub_content.hub_content_arn + + @patch("sagemaker.train.defaults.JumpStartTrainDefaults._get_training_components_model") + @patch("sagemaker.train.defaults.get_hub_content_and_document") + @patch("sagemaker.train.defaults.TrainDefaults.get_sagemaker_session") + def test_training_dataset_input_gated_reference_sets_accept_eula_and_hub_access( + self, mock_get_session, mock_get_hub_content, mock_get_tcm + ): + """Gated ModelReference -> default training channel also carries the EULA + slip.""" + mock_get_session.return_value = MagicMock() + hub_content = self._gated_reference_hub_content() + mock_get_hub_content.return_value = (hub_content, MagicMock()) + mock_get_tcm.return_value = self._training_components_model() + + jumpstart_config = JumpStartConfig( + model_id=GATED_MODEL_ID, + hub_name="my-private-hub", + accept_eula=True, + ) + + result = JumpStartTrainDefaults.get_training_dataset_input( + jumpstart_config=jumpstart_config, + input_data_config=None, + sagemaker_session=mock_get_session.return_value, + ) + + train_channels = [c for c in result if c.channel_name in ("training", "train")] + assert len(train_channels) == 1 + s3_source = train_channels[0].data_source + assert s3_source.model_access_config is not None + assert s3_source.model_access_config.accept_eula is True + assert s3_source.hub_access_config is not None + assert s3_source.hub_access_config.hub_content_arn == hub_content.hub_content_arn + + @patch("sagemaker.train.defaults.JumpStartTrainDefaults._get_training_variant") + @patch("sagemaker.train.defaults.JumpStartTrainDefaults._get_training_components_model") + @patch("sagemaker.train.defaults.get_hub_content_and_document") + @patch("sagemaker.train.defaults.TrainDefaults.get_sagemaker_session") + def test_model_artifact_input_gated_reference_defaults_accept_eula_false( + self, mock_get_session, mock_get_hub_content, mock_get_tcm, mock_get_variant + ): + """When accept_eula is left at its default, ModelAccessConfig.accept_eula is False.""" + mock_get_session.return_value = MagicMock() + hub_content = self._gated_reference_hub_content() + mock_get_hub_content.return_value = (hub_content, MagicMock()) + mock_get_tcm.return_value = self._training_components_model() + mock_get_variant.return_value = None + + # accept_eula not set -> defaults to False on JumpStartConfig. + jumpstart_config = JumpStartConfig( + model_id=GATED_MODEL_ID, + hub_name="my-private-hub", + ) + + result = JumpStartTrainDefaults.get_model_artifact_input( + jumpstart_config=jumpstart_config, + compute=Compute(instance_type="ml.g5.2xlarge", instance_count=1), + input_data_config=None, + environment={}, + sagemaker_session=mock_get_session.return_value, + ) + + model_channels = [c for c in result if c.channel_name == "model"] + assert len(model_channels) == 1 + s3_source = model_channels[0].data_source.s3_data_source + assert s3_source.model_access_config.accept_eula is False