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)