fix: always apply evaluator identity keywords and allow explicit domain_id

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)

sagemaker-train/src/sagemaker/ai_registry/evaluator.py

@@ -211,28 +211,37 @@ 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
         """
         # 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 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))

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)

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 = {

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)