diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset.py b/sagemaker-train/src/sagemaker/ai_registry/dataset.py index 2184be558c..d3ef453a35 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset.py +++ b/sagemaker-train/src/sagemaker/ai_registry/dataset.py @@ -127,7 +127,10 @@ def refresh(self): return self response = AIRHub.describe_hub_content( - DATASET_HUB_CONTENT_TYPE, self.name, session=self.sagemaker_session + DATASET_HUB_CONTENT_TYPE, + self.name, + hub_content_version=self.version, + session=self.sagemaker_session, ) doc = json.loads(response[RESPONSE_KEY_HUB_CONTENT_DOCUMENT]) try: diff --git a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py index 6d197ea117..d4a228a2b7 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py +++ b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py @@ -167,7 +167,10 @@ def refresh(self): return self response = AIRHub.describe_hub_content( - EVALUATOR_HUB_CONTENT_TYPE, self.name, session=self.sagemaker_session + EVALUATOR_HUB_CONTENT_TYPE, + self.name, + hub_content_version=self.version, + session=self.sagemaker_session, ) doc = json.loads(response[RESPONSE_KEY_HUB_CONTENT_DOCUMENT]) try: diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset.py b/sagemaker-train/tests/unit/ai_registry/test_dataset.py index eadb686b2c..50d6353d24 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset.py @@ -596,3 +596,23 @@ def test_create_passes_content_metadata_to_document( "arn:aws:sagemaker:us-west-2:123:feature-group/fg2", ] assert document["ContentMetadata"]["ExtractionMethod"] == "FeatureStoreDatasetBuilder" + + +class TestDataSetRefreshVersion: + @patch("sagemaker.ai_registry.air_hub_entity.AIRHub") + @patch("sagemaker.ai_registry.dataset.AIRHub") + def test_refresh_describes_own_version(self, mock_air_hub, mock_entity_air_hub): + mock_air_hub.describe_hub_content.return_value = { + "HubContentName": "test-dataset", + "HubContentArn": "test-arn", + "HubContentVersion": "1.0.0", + "HubContentStatus": "Available", + "HubContentDocument": json.dumps({"DatasetS3Bucket": "b", "DatasetS3Prefix": "p"}), + "HubContentSearchKeywords": [], + } + dataset = DataSet(name="test-dataset", arn="test-arn", version="1.0.0", status="Available") + + dataset.refresh() + + assert mock_air_hub.describe_hub_content.call_args.kwargs["hub_content_version"] == "1.0.0" + assert dataset.version == "1.0.0" diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py index e302d78616..1879472490 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py @@ -259,3 +259,25 @@ def test_create_version_failure(self, mock_air_hub): with pytest.raises(RuntimeError, match="Failed to create new version: Error"): evaluator.create_version("arn:aws:lambda:us-west-2:123456789012:function:new") + + +class TestEvaluatorRefreshVersion: + @patch("sagemaker.ai_registry.air_hub_entity.AIRHub") + @patch("sagemaker.ai_registry.evaluator.AIRHub") + def test_refresh_describes_own_version(self, mock_air_hub, mock_entity_air_hub): + mock_air_hub.describe_hub_content.return_value = { + "HubContentName": "test-eval", + "HubContentArn": "test-arn", + "HubContentVersion": "1.0.0", + "HubContentStatus": "Available", + "HubContentDocument": json.dumps( + {"SubType": "AWS/Evaluator", "JsonContent": json.dumps({"Reference": "ref"})} + ), + "HubContentSearchKeywords": ["method:lambda"], + } + evaluator = Evaluator(name="test-eval", version="1.0.0", arn="test-arn") + + evaluator.refresh() + + assert mock_air_hub.describe_hub_content.call_args.kwargs["hub_content_version"] == "1.0.0" + assert evaluator.version == "1.0.0"