Skip to content

Commit 69f054e

Browse files
authored
fix: always apply evaluator identity keywords and allow explicit domain_id (#5989)
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
1 parent efc1272 commit 69f054e

5 files changed

Lines changed: 184 additions & 23 deletions

File tree

‎sagemaker-train/src/sagemaker/ai_registry/dataset.py‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -230,13 +230,14 @@ def create(
230230
description: str = "",
231231
tags: Optional[List[Tuple[str, str]]] = None,
232232
role: Optional[str] = None,
233+
domain_id: Optional[str] = None,
233234
sagemaker_session: Optional[Session] = None,
234235
) -> "DataSet":
235236
"""Create a new DataSet Hub AIR entity.
236-
237+
237238
Creates a new version if entity already exists. This is the primary entry point
238239
for users. Uploads to S3 internally if local file input is provided.
239-
240+
240241
Args:
241242
name: Name of the dataset
242243
source: S3 URI or local file path for the dataset
@@ -245,20 +246,27 @@ def create(
245246
description: Description of the dataset
246247
tags: Optional list of (key, value) tag tuples
247248
role: Optional IAM role ARN. If not provided, uses default execution role.
249+
domain_id: Optional SageMaker Studio domain ID used to tag the dataset so it is
250+
visible in Studio. If not provided, it is auto-detected from the Studio
251+
environment; supply it explicitly when creating datasets outside Studio
252+
(e.g. from a laptop or CI) so they still appear in the target domain.
248253
sagemaker_session: Optional SageMaker session. If not provided, uses default session.
249-
254+
250255
Returns:
251256
DataSet: The created dataset instance
252-
257+
253258
Raises:
254259
ValueError: If validation fails or required parameters are missing
255260
"""
256261
# Get or create session for domain ID extraction
257262
if sagemaker_session is None:
258263
sagemaker_session = Session()
259-
260-
# Extract domain ID if available (only works in Studio environments)
261-
domain_id = _get_current_domain_id(sagemaker_session)
264+
265+
# Use the caller-provided domain ID when given; otherwise auto-detect it (only
266+
# works in Studio environments). This keeps datasets discoverable in Studio even
267+
# when created from environments where the domain cannot be inferred.
268+
if domain_id is None:
269+
domain_id = _get_current_domain_id(sagemaker_session)
262270

263271
# Validate dataset file
264272
cls._validate_dataset_file(source)

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

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -211,28 +211,37 @@ def create(
211211
source: Optional[str] = None,
212212
wait: bool = True,
213213
role: Optional[str] = None,
214+
domain_id: Optional[str] = None,
214215
sagemaker_session: Optional[Session] = None,
215216
) -> "Evaluator":
216217
"""Create a new Evaluator entity in the AI Registry.
217-
218+
218219
Args:
219220
name: Name of the evaluator
220221
type: Type of evaluator (RewardFunction or RewardPrompt)
221222
source: Lambda ARN, S3 URI, or local file path depending on evaluator type
222223
wait: Whether to wait for the evaluator to be available
224+
role: Optional IAM role ARN. If not provided, uses default execution role.
225+
domain_id: Optional SageMaker Studio domain ID used to tag the evaluator so it
226+
is visible in Studio. If not provided, it is auto-detected from the Studio
227+
environment; supply it explicitly when creating evaluators outside Studio
228+
(e.g. from a laptop or CI) so they still appear in the target domain.
223229
224230
Returns:
225231
Evaluator: Newly created Evaluator instance
226-
232+
227233
Raises:
228234
ValueError: If source is required but not provided, or if type is unsupported
229235
"""
230236
# Get or create session for domain ID extraction
231237
if sagemaker_session is None:
232238
sagemaker_session = Session()
233239

234-
# Extract domain ID if available (only works in Studio environments)
235-
domain_id = _get_current_domain_id(sagemaker_session)
240+
# Use the caller-provided domain ID when given; otherwise auto-detect it (only
241+
# works in Studio environments). This keeps evaluators discoverable in Studio even
242+
# when created from environments where the domain cannot be inferred.
243+
if domain_id is None:
244+
domain_id = _get_current_domain_id(sagemaker_session)
236245
sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session=sagemaker_session)
237246
role = TrainDefaults.get_role(role=role, sagemaker_session=sagemaker_session)
238247

@@ -257,17 +266,19 @@ def create(
257266
if source and source.startswith(LAMBDA_ARN_PREFIX):
258267
content_type = EVALUATOR_BYOLAMBDA
259268

260-
# Prepare tags for SearchKeywords
269+
# Prepare tags for SearchKeywords.
270+
# The identity keywords (@subtype/@evaluatortype/@contenttype) describe what the
271+
# asset IS and must ALWAYS be applied so the evaluator is discoverable and visible
272+
# in Studio, regardless of evaluator type or whether a method was resolved. Only
273+
# `method` is conditional (reward prompts do not have a method).
274+
tags = [
275+
("@" + "subtype", EVALUATOR_HUB_CONTENT_SUBTYPE.lower()),
276+
("@" + DOC_KEY_SUB_TYPE.lower(), type.lower()),
277+
("@" + "contenttype", content_type.lower()),
278+
]
261279
if method is not None:
262-
tags = [
263-
(TAG_KEY_METHOD, method.value),
264-
("@" + "subtype", EVALUATOR_HUB_CONTENT_SUBTYPE.lower()),
265-
("@" + DOC_KEY_SUB_TYPE.lower(), type.lower()),
266-
("@" + "contenttype", content_type.lower())
267-
]
268-
else:
269-
tags = []
270-
280+
tags.insert(0, (TAG_KEY_METHOD, method.value))
281+
271282
# Add domain-id to SearchKeywords if available
272283
if domain_id:
273284
tags.append((TAG_KEY_DOMAIN_ID, domain_id))

‎sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,3 +189,47 @@ def test_domain_id_added_without_customization_technique(
189189

190190
# Verify domain-id is still in tags
191191
assert any(tag[0] == '@domain' and tag[1] == mock_domain_id for tag in tags)
192+
193+
@patch('sagemaker.core.helper.session_helper.Session')
194+
@patch('sagemaker.ai_registry.dataset._get_current_domain_id')
195+
@patch('sagemaker.ai_registry.dataset.AIRHub')
196+
@patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session')
197+
@patch('sagemaker.train.defaults.TrainDefaults.get_role')
198+
def test_explicit_domain_id_used_without_auto_detection(
199+
self, mock_get_role, mock_get_session, mock_air_hub, mock_get_domain_id, mock_session, sample_dataset_file
200+
):
201+
"""An explicit domain_id is tagged and auto-detection is not invoked.
202+
203+
Covers the non-Studio case (P467494019) where the domain cannot be inferred and
204+
must be supplied by the caller.
205+
"""
206+
mock_session_instance = Mock()
207+
mock_session.return_value = mock_session_instance
208+
mock_get_session.return_value = mock_session_instance
209+
mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role"
210+
211+
mock_air_hub.upload_to_s3 = Mock()
212+
mock_air_hub.import_hub_content = Mock()
213+
mock_air_hub.describe_hub_content = Mock(return_value={
214+
'HubContentName': 'test-dataset',
215+
'HubContentArn': 'arn:aws:sagemaker:us-west-2:123:hub-content/test',
216+
'HubContentVersion': '1.0.0',
217+
'HubContentStatus': 'Available',
218+
'CreationTime': '2024-01-01',
219+
'LastModifiedTime': '2024-01-01',
220+
'HubContentDocument': '{"DatasetS3Bucket": "bucket", "DatasetS3Prefix": "prefix"}'
221+
})
222+
223+
with patch('sagemaker.ai_registry.dataset.DataSet.wait'):
224+
DataSet.create(
225+
name="test-dataset",
226+
source=sample_dataset_file,
227+
customization_technique=CustomizationTechnique.SFT,
228+
domain_id="d-explicit123",
229+
)
230+
231+
# Auto-detection must be skipped when domain_id is explicitly provided.
232+
mock_get_domain_id.assert_not_called()
233+
234+
tags = mock_air_hub.import_hub_content.call_args[1]['tags']
235+
assert any(tag[0] == '@domain' and tag[1] == 'd-explicit123' for tag in tags)

‎sagemaker-train/tests/unit/ai_registry/test_evaluator.py‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,17 @@
1818
from sagemaker.ai_registry.evaluator import Evaluator, EvaluatorMethod
1919
from sagemaker.ai_registry.air_constants import (
2020
RESPONSE_KEY_HUB_CONTENT_VERSION, RESPONSE_KEY_HUB_CONTENT_ARN,
21-
RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME, REWARD_FUNCTION
21+
RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME,
22+
REWARD_FUNCTION, REWARD_PROMPT
2223
)
2324

2425

26+
def _keywords_from_import_call(mock_air_hub):
27+
"""Extract the SearchKeyword strings passed to import_hub_content."""
28+
tags = mock_air_hub.import_hub_content.call_args.kwargs["tags"]
29+
return [f"{tag[0]}:{tag[1]}" for tag in tags]
30+
31+
2532
class TestEvaluator:
2633
@patch('sagemaker.ai_registry.evaluator.AIRHub')
2734
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):
8188
call_kwargs = mock_lambda_client.create_function.call_args[1]
8289
assert call_kwargs["Handler"] == "lambda_function.lambda_handler"
8390

91+
@patch('sagemaker.ai_registry.evaluator.AIRHub')
92+
def test_create_reward_function_always_sets_identity_keywords(self, mock_air_hub):
93+
"""RewardFunction evaluators must always carry the full identity keyword set."""
94+
mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"}
95+
mock_air_hub.describe_hub_content.return_value = {
96+
RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0",
97+
RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn",
98+
RESPONSE_KEY_CREATION_TIME: "2024-01-01",
99+
RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01"
100+
}
101+
102+
Evaluator.create(
103+
name="test-evaluator",
104+
source="arn:aws:lambda:us-west-2:123456789012:function:test",
105+
type=REWARD_FUNCTION,
106+
wait=False
107+
)
108+
109+
keywords = _keywords_from_import_call(mock_air_hub)
110+
assert "@subtype:aws/evaluator" in keywords
111+
assert "@evaluatortype:rewardfunction" in keywords
112+
assert "@contenttype:byolambda" in keywords
113+
assert "method:lambda" in keywords
114+
115+
@patch('sagemaker.ai_registry.evaluator.AIRHub')
116+
def test_create_reward_prompt_always_sets_identity_keywords(self, mock_air_hub):
117+
"""RewardPrompt evaluators have no method but must still be Studio-visible.
118+
119+
Regression test for P467494019: identity keywords (@subtype/@evaluatortype/
120+
@contenttype) were previously gated on ``method is not None``, so reward prompts
121+
registered with no keywords and were invisible in Studio.
122+
"""
123+
mock_air_hub.upload_to_s3.return_value = "s3://bucket/path/prompt.txt"
124+
mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"}
125+
mock_air_hub.describe_hub_content.return_value = {
126+
RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0",
127+
RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn",
128+
RESPONSE_KEY_CREATION_TIME: "2024-01-01",
129+
RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01"
130+
}
131+
132+
Evaluator.create(
133+
name="test-prompt-evaluator",
134+
source="s3://bucket/path/prompt.txt",
135+
type=REWARD_PROMPT,
136+
wait=False
137+
)
138+
139+
keywords = _keywords_from_import_call(mock_air_hub)
140+
assert "@subtype:aws/evaluator" in keywords
141+
assert "@evaluatortype:rewardprompt" in keywords
142+
assert "@contenttype:byocode" in keywords
143+
84144
@patch('sagemaker.ai_registry.evaluator.AIRHub')
85145
def test_get_all(self, mock_air_hub):
86146
mock_air_hub.list_hub_content.return_value = {

‎sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,44 @@ def test_domain_id_not_added_when_unavailable(
100100
# Get the tags argument
101101
call_args = mock_air_hub.import_hub_content.call_args
102102
tags = call_args[1]['tags']
103-
103+
104104
# Verify domain-id is NOT in tags
105105
assert not any(tag[0] == '@domain' for tag in tags)
106+
107+
@patch('sagemaker.core.helper.session_helper.Session')
108+
@patch('sagemaker.ai_registry.evaluator._get_current_domain_id')
109+
@patch('sagemaker.ai_registry.evaluator.AIRHub')
110+
def test_explicit_domain_id_used_without_auto_detection(
111+
self, mock_air_hub, mock_get_domain_id, mock_session
112+
):
113+
"""An explicit domain_id is tagged and auto-detection is not invoked.
114+
115+
Covers the non-Studio case (P467494019) where the domain cannot be inferred and
116+
must be supplied by the caller.
117+
"""
118+
mock_session.return_value = Mock()
119+
mock_air_hub.import_hub_content = Mock()
120+
mock_air_hub.describe_hub_content = Mock(return_value={
121+
'HubContentName': 'test-evaluator',
122+
'HubContentArn': 'arn:aws:sagemaker:us-west-2:123:hub-content/test',
123+
'HubContentVersion': '1.0.0',
124+
'HubContentStatus': 'Available',
125+
'CreationTime': '2024-01-01',
126+
'LastModifiedTime': '2024-01-01',
127+
'HubContentDocument': '{"SubType": "AWS/Evaluator", "JsonContent": "{}"}'
128+
})
129+
130+
with patch('sagemaker.ai_registry.evaluator.Evaluator.wait'):
131+
with patch('sagemaker.ai_registry.evaluator.Evaluator._handle_reward_function', return_value=(EvaluatorMethod.LAMBDA, 'arn:aws:lambda:...')):
132+
Evaluator.create(
133+
name="test-evaluator",
134+
type="RewardFunction",
135+
source="arn:aws:lambda:us-west-2:123:function:test",
136+
domain_id="d-explicit123",
137+
)
138+
139+
# Auto-detection must be skipped when domain_id is explicitly provided.
140+
mock_get_domain_id.assert_not_called()
141+
142+
tags = mock_air_hub.import_hub_content.call_args[1]['tags']
143+
assert any(tag[0] == '@domain' and tag[1] == 'd-explicit123' for tag in tags)

0 commit comments

Comments
 (0)