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..a44d0db7e7 100644 --- a/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py +++ b/sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py @@ -10,14 +10,145 @@ # 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.train import ModelTrainer -from sagemaker.train.configs import Compute +from sagemaker.train.configs import Compute, InputData + +logger = logging.getLogger(__name__) + +# A trainable classical-ML model keeps these tests fast/cheap on CPU. +TRAINABLE_MODEL_ID = "catboost-regression-model" +HUB_NAME_PREFIX = "sdk-integ-train-hub" +ALIASED_REFERENCE_NAME = "sdk-integ-aliased-catboost" +PRIVATE_MODEL_NAME = "sdk-integ-private-catboost" + + +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.list_hub_contents(HubName=hub_name, HubContentType=content_type) + if any( + c["HubContentName"] == name and c.get("HubContentStatus") == "Available" + for c in resp.get("HubContentSummaries", []) + ): + 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) + + +def _default_training_dataset(region, model_id): + """Resolve the model's default training dataset S3 URI from JS metadata.""" + from sagemaker.core.jumpstart.accessors import JumpStartModelsAccessor + + specs = JumpStartModelsAccessor.get_model_specs(region=region, model_id=model_id, version="*") + if not getattr(specs, "training_supported", False): + return None + key = getattr(specs, "default_training_dataset_key", None) + if not key: + return None + return f"s3://jumpstart-cache-prod-{region}/{key}" + + +@pytest.fixture(scope="module") +def sagemaker_session(): + from sagemaker.core.helper.session_helper import 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: + pytest.skip(f"Cannot create private hub (missing permissions?): {e}") + + for _ in range(30): + if sm.describe_hub(HubName=hub_name)["HubStatus"] == "InService": + break + time.sleep(2) + else: + pytest.skip(f"Hub {hub_name} did not reach InService") + + yield hub_name + _delete_hub(sm, hub_name) @pytest.mark.parametrize( @@ -42,7 +173,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 +185,125 @@ 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: + pytest.skip(f"Cannot create hub content reference: {e}") + if not _wait_for_content(sm, private_hub, TRAINABLE_MODEL_ID, "ModelReference"): + pytest.skip(f"ModelReference {TRAINABLE_MODEL_ID} not available in {private_hub}") + + dataset = _default_training_dataset(region, TRAINABLE_MODEL_ID) + if dataset is None: + pytest.skip(f"{TRAINABLE_MODEL_ID} is not trainable / has no default dataset") + + 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, + ) + model_trainer.train(input_data_config=[InputData(channel_name="training", data_source=dataset)]) + + +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: + pytest.skip(f"Cannot create aliased hub content reference: {e}") + if not _wait_for_content(sm, private_hub, ALIASED_REFERENCE_NAME, "ModelReference"): + pytest.skip(f"Aliased reference {ALIASED_REFERENCE_NAME} not available") + + dataset = _default_training_dataset(region, TRAINABLE_MODEL_ID) + if dataset is None: + pytest.skip(f"{TRAINABLE_MODEL_ID} is not trainable / has no default dataset") + + 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, + ) + model_trainer.train(input_data_config=[InputData(channel_name="training", data_source=dataset)]) + + +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) + region = _region(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: + pytest.skip(f"Cannot read public model document: {e}") + + 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: + pytest.skip(f"import_hub_content for a private Model not permitted/supported: {e}") + if not _wait_for_content(sm, private_hub, PRIVATE_MODEL_NAME, "Model"): + pytest.skip(f"Private Model {PRIVATE_MODEL_NAME} not available in {private_hub}") + + dataset = _default_training_dataset(region, TRAINABLE_MODEL_ID) + if dataset is None: + pytest.skip(f"{TRAINABLE_MODEL_ID} is not trainable / has no default dataset") + + 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, + ) + model_trainer.train(input_data_config=[InputData(channel_name="training", data_source=dataset)])