From a77c9281511ddf88379567b51ec9839d09da235c Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Mon, 6 Jul 2026 14:38:27 -0700 Subject: [PATCH] fix: always apply evaluator identity keywords and allow explicit domain_id Evaluator.create() gated the identity SearchKeywords (@subtype/@evaluatortype/@contenttype) behind `method is not None`. Since `method` is only resolved for RewardFunction, RewardPrompt evaluators were registered with no identity keywords and were therefore invisible in SageMaker Studio (which filters hub content by these keywords). The asset still existed in the hub, so the failure was silent. Always attach the identity keywords regardless of evaluator type or whether a method was resolved; only `method` remains conditional. Also add an optional `domain_id` parameter to Evaluator.create() and DataSet.create(). The @domain keyword is already auto-injected inside Studio, but the domain cannot be inferred outside Studio (laptop/CI), leaving assets un-scoped and hidden. Callers can now pass domain_id explicitly; auto-detection is skipped when it is provided. Adds regression tests covering RewardFunction and RewardPrompt identity keywords and explicit domain_id for both Evaluator and DataSet. Ref: P467494019 --- .../src/sagemaker/ai_registry/dataset.py | 22 ++++--- .../src/sagemaker/ai_registry/evaluator.py | 39 +++++++----- .../ai_registry/test_dataset_domain_id.py | 44 +++++++++++++ .../tests/unit/ai_registry/test_evaluator.py | 62 ++++++++++++++++++- .../ai_registry/test_evaluator_domain_id.py | 40 +++++++++++- 5 files changed, 184 insertions(+), 23 deletions(-) diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset.py b/sagemaker-train/src/sagemaker/ai_registry/dataset.py index eb1bd3505a..86d8f907e6 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset.py +++ b/sagemaker-train/src/sagemaker/ai_registry/dataset.py @@ -230,13 +230,14 @@ def create( description: str = "", tags: Optional[List[Tuple[str, str]]] = None, role: Optional[str] = None, + domain_id: Optional[str] = None, sagemaker_session: Optional[Session] = None, ) -> "DataSet": """Create a new DataSet Hub AIR entity. - + Creates a new version if entity already exists. This is the primary entry point for users. Uploads to S3 internally if local file input is provided. - + Args: name: Name of the dataset source: S3 URI or local file path for the dataset @@ -245,20 +246,27 @@ def create( description: Description of the dataset tags: Optional list of (key, value) tag tuples role: Optional IAM role ARN. If not provided, uses default execution role. + domain_id: Optional SageMaker Studio domain ID used to tag the dataset so it is + visible in Studio. If not provided, it is auto-detected from the Studio + environment; supply it explicitly when creating datasets outside Studio + (e.g. from a laptop or CI) so they still appear in the target domain. sagemaker_session: Optional SageMaker session. If not provided, uses default session. - + Returns: DataSet: The created dataset instance - + Raises: ValueError: If validation fails or required parameters are missing """ # Get or create session for domain ID extraction if sagemaker_session is None: sagemaker_session = Session() - - # Extract domain ID if available (only works in Studio environments) - domain_id = _get_current_domain_id(sagemaker_session) + + # Use the caller-provided domain ID when given; otherwise auto-detect it (only + # works in Studio environments). This keeps datasets discoverable in Studio even + # when created from environments where the domain cannot be inferred. + if domain_id is None: + domain_id = _get_current_domain_id(sagemaker_session) # Validate dataset file cls._validate_dataset_file(source) diff --git a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py index 5d0b47939f..fff522ac9a 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py +++ b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py @@ -211,19 +211,25 @@ def create( source: Optional[str] = None, wait: bool = True, role: Optional[str] = None, + domain_id: Optional[str] = None, sagemaker_session: Optional[Session] = None, ) -> "Evaluator": """Create a new Evaluator entity in the AI Registry. - + Args: name: Name of the evaluator type: Type of evaluator (RewardFunction or RewardPrompt) source: Lambda ARN, S3 URI, or local file path depending on evaluator type wait: Whether to wait for the evaluator to be available + role: Optional IAM role ARN. If not provided, uses default execution role. + domain_id: Optional SageMaker Studio domain ID used to tag the evaluator so it + is visible in Studio. If not provided, it is auto-detected from the Studio + environment; supply it explicitly when creating evaluators outside Studio + (e.g. from a laptop or CI) so they still appear in the target domain. Returns: Evaluator: Newly created Evaluator instance - + Raises: ValueError: If source is required but not provided, or if type is unsupported """ @@ -231,8 +237,11 @@ def create( if sagemaker_session is None: sagemaker_session = Session() - # Extract domain ID if available (only works in Studio environments) - domain_id = _get_current_domain_id(sagemaker_session) + # Use the caller-provided domain ID when given; otherwise auto-detect it (only + # works in Studio environments). This keeps evaluators discoverable in Studio even + # when created from environments where the domain cannot be inferred. + if domain_id is None: + domain_id = _get_current_domain_id(sagemaker_session) sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session=sagemaker_session) role = TrainDefaults.get_role(role=role, sagemaker_session=sagemaker_session) @@ -257,17 +266,19 @@ def create( if source and source.startswith(LAMBDA_ARN_PREFIX): content_type = EVALUATOR_BYOLAMBDA - # Prepare tags for SearchKeywords + # Prepare tags for SearchKeywords. + # The identity keywords (@subtype/@evaluatortype/@contenttype) describe what the + # asset IS and must ALWAYS be applied so the evaluator is discoverable and visible + # in Studio, regardless of evaluator type or whether a method was resolved. Only + # `method` is conditional (reward prompts do not have a method). + tags = [ + ("@" + "subtype", EVALUATOR_HUB_CONTENT_SUBTYPE.lower()), + ("@" + DOC_KEY_SUB_TYPE.lower(), type.lower()), + ("@" + "contenttype", content_type.lower()), + ] if method is not None: - tags = [ - (TAG_KEY_METHOD, method.value), - ("@" + "subtype", EVALUATOR_HUB_CONTENT_SUBTYPE.lower()), - ("@" + DOC_KEY_SUB_TYPE.lower(), type.lower()), - ("@" + "contenttype", content_type.lower()) - ] - else: - tags = [] - + tags.insert(0, (TAG_KEY_METHOD, method.value)) + # Add domain-id to SearchKeywords if available if domain_id: tags.append((TAG_KEY_DOMAIN_ID, domain_id)) diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py b/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py index d1e73ee665..da2d8ceca7 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py @@ -189,3 +189,47 @@ def test_domain_id_added_without_customization_technique( # Verify domain-id is still in tags assert any(tag[0] == '@domain' and tag[1] == mock_domain_id for tag in tags) + + @patch('sagemaker.core.helper.session_helper.Session') + @patch('sagemaker.ai_registry.dataset._get_current_domain_id') + @patch('sagemaker.ai_registry.dataset.AIRHub') + @patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session') + @patch('sagemaker.train.defaults.TrainDefaults.get_role') + def test_explicit_domain_id_used_without_auto_detection( + self, mock_get_role, mock_get_session, mock_air_hub, mock_get_domain_id, mock_session, sample_dataset_file + ): + """An explicit domain_id is tagged and auto-detection is not invoked. + + Covers the non-Studio case (P467494019) where the domain cannot be inferred and + must be supplied by the caller. + """ + mock_session_instance = Mock() + mock_session.return_value = mock_session_instance + mock_get_session.return_value = mock_session_instance + mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role" + + mock_air_hub.upload_to_s3 = Mock() + mock_air_hub.import_hub_content = Mock() + mock_air_hub.describe_hub_content = Mock(return_value={ + 'HubContentName': 'test-dataset', + 'HubContentArn': 'arn:aws:sagemaker:us-west-2:123:hub-content/test', + 'HubContentVersion': '1.0.0', + 'HubContentStatus': 'Available', + 'CreationTime': '2024-01-01', + 'LastModifiedTime': '2024-01-01', + 'HubContentDocument': '{"DatasetS3Bucket": "bucket", "DatasetS3Prefix": "prefix"}' + }) + + with patch('sagemaker.ai_registry.dataset.DataSet.wait'): + DataSet.create( + name="test-dataset", + source=sample_dataset_file, + customization_technique=CustomizationTechnique.SFT, + domain_id="d-explicit123", + ) + + # Auto-detection must be skipped when domain_id is explicitly provided. + mock_get_domain_id.assert_not_called() + + tags = mock_air_hub.import_hub_content.call_args[1]['tags'] + assert any(tag[0] == '@domain' and tag[1] == 'd-explicit123' for tag in tags) diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py index 3c4b233473..eeeedfce58 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py @@ -18,10 +18,17 @@ from sagemaker.ai_registry.evaluator import Evaluator, EvaluatorMethod from sagemaker.ai_registry.air_constants import ( RESPONSE_KEY_HUB_CONTENT_VERSION, RESPONSE_KEY_HUB_CONTENT_ARN, - RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME, REWARD_FUNCTION + RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME, + REWARD_FUNCTION, REWARD_PROMPT ) +def _keywords_from_import_call(mock_air_hub): + """Extract the SearchKeyword strings passed to import_hub_content.""" + tags = mock_air_hub.import_hub_content.call_args.kwargs["tags"] + return [f"{tag[0]}:{tag[1]}" for tag in tags] + + class TestEvaluator: @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_create_with_lambda_arn(self, mock_air_hub): @@ -81,6 +88,59 @@ def test_create_with_byoc(self, mock_air_hub, mock_boto3): call_kwargs = mock_lambda_client.create_function.call_args[1] assert call_kwargs["Handler"] == "lambda_function.lambda_handler" + @patch('sagemaker.ai_registry.evaluator.AIRHub') + def test_create_reward_function_always_sets_identity_keywords(self, mock_air_hub): + """RewardFunction evaluators must always carry the full identity keyword set.""" + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_CREATION_TIME: "2024-01-01", + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01" + } + + Evaluator.create( + name="test-evaluator", + source="arn:aws:lambda:us-west-2:123456789012:function:test", + type=REWARD_FUNCTION, + wait=False + ) + + keywords = _keywords_from_import_call(mock_air_hub) + assert "@subtype:aws/evaluator" in keywords + assert "@evaluatortype:rewardfunction" in keywords + assert "@contenttype:byolambda" in keywords + assert "method:lambda" in keywords + + @patch('sagemaker.ai_registry.evaluator.AIRHub') + def test_create_reward_prompt_always_sets_identity_keywords(self, mock_air_hub): + """RewardPrompt evaluators have no method but must still be Studio-visible. + + Regression test for P467494019: identity keywords (@subtype/@evaluatortype/ + @contenttype) were previously gated on ``method is not None``, so reward prompts + registered with no keywords and were invisible in Studio. + """ + mock_air_hub.upload_to_s3.return_value = "s3://bucket/path/prompt.txt" + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_CREATION_TIME: "2024-01-01", + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01" + } + + Evaluator.create( + name="test-prompt-evaluator", + source="s3://bucket/path/prompt.txt", + type=REWARD_PROMPT, + wait=False + ) + + keywords = _keywords_from_import_call(mock_air_hub) + assert "@subtype:aws/evaluator" in keywords + assert "@evaluatortype:rewardprompt" in keywords + assert "@contenttype:byocode" in keywords + @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_get_all(self, mock_air_hub): mock_air_hub.list_hub_content.return_value = { diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py index 3507eb4fe7..f6cdaf44af 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py @@ -100,6 +100,44 @@ def test_domain_id_not_added_when_unavailable( # Get the tags argument call_args = mock_air_hub.import_hub_content.call_args tags = call_args[1]['tags'] - + # Verify domain-id is NOT in tags assert not any(tag[0] == '@domain' for tag in tags) + + @patch('sagemaker.core.helper.session_helper.Session') + @patch('sagemaker.ai_registry.evaluator._get_current_domain_id') + @patch('sagemaker.ai_registry.evaluator.AIRHub') + def test_explicit_domain_id_used_without_auto_detection( + self, mock_air_hub, mock_get_domain_id, mock_session + ): + """An explicit domain_id is tagged and auto-detection is not invoked. + + Covers the non-Studio case (P467494019) where the domain cannot be inferred and + must be supplied by the caller. + """ + mock_session.return_value = Mock() + mock_air_hub.import_hub_content = Mock() + mock_air_hub.describe_hub_content = Mock(return_value={ + 'HubContentName': 'test-evaluator', + 'HubContentArn': 'arn:aws:sagemaker:us-west-2:123:hub-content/test', + 'HubContentVersion': '1.0.0', + 'HubContentStatus': 'Available', + 'CreationTime': '2024-01-01', + 'LastModifiedTime': '2024-01-01', + 'HubContentDocument': '{"SubType": "AWS/Evaluator", "JsonContent": "{}"}' + }) + + with patch('sagemaker.ai_registry.evaluator.Evaluator.wait'): + with patch('sagemaker.ai_registry.evaluator.Evaluator._handle_reward_function', return_value=(EvaluatorMethod.LAMBDA, 'arn:aws:lambda:...')): + Evaluator.create( + name="test-evaluator", + type="RewardFunction", + source="arn:aws:lambda:us-west-2:123:function:test", + domain_id="d-explicit123", + ) + + # Auto-detection must be skipped when domain_id is explicitly provided. + mock_get_domain_id.assert_not_called() + + tags = mock_air_hub.import_hub_content.call_args[1]['tags'] + assert any(tag[0] == '@domain' and tag[1] == 'd-explicit123' for tag in tags)