From a07a2e724558e0704ffb5019b4dc05f1f35dd550 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshayshah0803@gmail.com>
Date: Fri, 9 Oct 2026 21:00:38 -0500
Subject: [PATCH] fix(train): refresh the entity's own version in DataSet and
 Evaluator

DataSet.refresh() and Evaluator.refresh() called describe_hub_content without
a version, unlike the base AIRHubEntity.refresh(), so DescribeHubContent
returned the latest version and overwrote the object with it. Calling
refresh() or wait() on an older version (e.g. one from get_versions())
silently turned it into the latest version, and wait() polled the latest
version's status instead of its own.

Pass self.version to describe_hub_content in both overrides.
---
 .../src/sagemaker/ai_registry/dataset.py      |  5 ++++-
 .../src/sagemaker/ai_registry/evaluator.py    |  5 ++++-
 .../tests/unit/ai_registry/test_dataset.py    | 20 +++++++++++++++++
 .../tests/unit/ai_registry/test_evaluator.py  | 22 +++++++++++++++++++
 4 files changed, 50 insertions(+), 2 deletions(-)

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"