fix(train): refresh the entity's own version in DataSet and Evaluator

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:

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:

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"

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"