diff --git a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py index f8fad16a77..7cf9c55f9e 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py +++ b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py @@ -98,6 +98,7 @@ def import_hub_content( hub_content_document: str, hub_content_version: str = AIR_HUB_CONTENT_DEFAULT_VERSION, tags: Optional[tuple] = None, + description: Optional[str] = None, session: Optional[Session] = None, ): """Import hub content into the AI Registry hub. @@ -108,6 +109,7 @@ def import_hub_content( document_schema_version: Schema version of the document hub_content_document: JSON document content tags: Optional tuple of tags + description: Optional description of the hub content session: Boto3 session Returns: @@ -127,6 +129,8 @@ def import_hub_content( } if tags: request["HubContentSearchKeywords"] = [f"{tag[0]}:{tag[1]}" for tag in tags] + if description: + request["HubContentDescription"] = description return client.import_hub_content(**request) @classmethod diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset.py b/sagemaker-train/src/sagemaker/ai_registry/dataset.py index df79a73265..b3279db433 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset.py +++ b/sagemaker-train/src/sagemaker/ai_registry/dataset.py @@ -15,6 +15,7 @@ from __future__ import annotations import json +import logging import os import tempfile from datetime import datetime @@ -55,6 +56,8 @@ from sagemaker.core.helper.session_helper import Session from sagemaker.train.defaults import TrainDefaults +logger = logging.getLogger(__name__) + class DataSet(AIRHubEntity): """Dataset entity for AI Registry.""" @@ -358,6 +361,7 @@ def create( document_schema_version=DATASET_DOCUMENT_SCHEMA_VERSION, hub_content_document=document_str, tags=tags, + description=description, session=sagemaker_session ) @@ -519,7 +523,7 @@ def create_version( self, source: str, customization_technique: Optional[CustomizationTechnique] = None - ) -> bool: + ) -> Optional["DataSet"]: """Create a new version of this dataset. Args: @@ -527,7 +531,7 @@ def create_version( customization_technique: Customization technique to use. If None, uses existing technique. Returns: - True if version created successfully, False otherwise + DataSet: The newly created version, or None if creation failed """ try: # Get current dataset metadata @@ -545,19 +549,27 @@ def create_version( technique = customization_technique or (CustomizationTechnique(existing_technique) if existing_technique else None) # Create new version - DataSet.create( + new_dataset = DataSet.create( name=self.name, source=source, customization_technique=technique, tags=[ (TAG_KEY_CUSTOMIZATION_TECHNIQUE, technique.value), (TAG_KEY_METHOD, keywords.get(TAG_KEY_METHOD, "")) - ] if technique else [(TAG_KEY_METHOD, keywords.get(TAG_KEY_METHOD, ""))] + ] if technique else [(TAG_KEY_METHOD, keywords.get(TAG_KEY_METHOD, ""))], + sagemaker_session=self.sagemaker_session, + ) + logger.info( + "Created new version %s for dataset %s, arn: %s", + new_dataset.version, self.name, new_dataset.arn ) - return True + return new_dataset except Exception as e: - print(f"Failed to create new version for dataset {self.name} with exception : {e}") - return False + logger.error( + "Failed to create new version for dataset %s with exception : %s", + self.name, e + ) + return None @staticmethod def _parse_keywords(search_keywords: List[str]) -> dict: diff --git a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py index fff522ac9a..fbddd361bb 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py +++ b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py @@ -15,6 +15,7 @@ import io import json +import logging import os import zipfile from collections.abc import Sequence @@ -40,6 +41,7 @@ RESPONSE_KEY_LAST_MODIFIED_TIME, RESPONSE_KEY_FUNCTION_ARN, RESPONSE_KEY_HUB_CONTENT_NAME, RESPONSE_KEY_HUB_CONTENT_STATUS, RESPONSE_KEY_HUB_CONTENT_DOCUMENT, RESPONSE_KEY_HUB_CONTENT_SEARCH_KEYWORDS, + RESPONSE_KEY_HUB_CONTENT_DESCRIPTION, DOC_KEY_JSON_CONTENT, DOC_KEY_REFERENCE, DOC_KEY_SUB_TYPE, REWARD_FUNCTION, REWARD_PROMPT, ) @@ -52,6 +54,9 @@ from sagemaker.train.common_utils.finetune_utils import _get_current_domain_id from sagemaker.train.defaults import TrainDefaults +logger = logging.getLogger(__name__) + + class EvaluatorMethod(Enum): """Enum for Evaluator method types.""" BYOC = "byoc" @@ -103,6 +108,7 @@ def __init__( status: Optional[HubContentStatus] = None, created_time: Optional[datetime] = None, updated_time: Optional[datetime] = None, + description: Optional[str] = None, sagemaker_session: Optional[Session] = None ) -> None: """Initialize Evaluator entity. @@ -117,9 +123,10 @@ def __init__( status: Current status of the evaluator created_time: Creation timestamp updated_time: Last update timestamp + description: Description of the evaluator sagemaker_session: Optional SageMaker session. """ - super().__init__(name, version, arn, status, created_time, updated_time,sagemaker_session) + super().__init__(name, version, arn, status, created_time, updated_time, description, sagemaker_session) self.method = method self.type = type self.reference = reference @@ -160,6 +167,7 @@ def refresh(self): self.reference = json_content.get(DOC_KEY_REFERENCE, "") self.type = json_content.get(DOC_KEY_SUB_TYPE, "") self.status = response[RESPONSE_KEY_HUB_CONTENT_STATUS] + self.description = response.get(RESPONSE_KEY_HUB_CONTENT_DESCRIPTION, "") method_str = keywords.get(TAG_KEY_METHOD) self.method = EvaluatorMethod(method_str) if method_str else None self.created = response.get(RESPONSE_KEY_CREATION_TIME) @@ -198,6 +206,7 @@ def get(cls, name: str, sagemaker_session=None) -> "Evaluator": method=EvaluatorMethod(method_str) if method_str else None, reference=reference, status=response[RESPONSE_KEY_HUB_CONTENT_STATUS], + description=response.get(RESPONSE_KEY_HUB_CONTENT_DESCRIPTION, ""), created_time=response.get(RESPONSE_KEY_CREATION_TIME), updated_time=response.get(RESPONSE_KEY_LAST_MODIFIED_TIME), ) @@ -212,6 +221,7 @@ def create( wait: bool = True, role: Optional[str] = None, domain_id: Optional[str] = None, + description: Optional[str] = None, sagemaker_session: Optional[Session] = None, ) -> "Evaluator": """Create a new Evaluator entity in the AI Registry. @@ -226,6 +236,7 @@ def create( 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. + description: Optional description of the evaluator. Returns: Evaluator: Newly created Evaluator instance @@ -294,6 +305,7 @@ def create( document_schema_version=EVALUATOR_DOCUMENT_SCHEMA_VERSION, hub_content_document=hub_content_document, tags=tags, + description=description, session=sagemaker_session, ) @@ -310,6 +322,7 @@ def create( created_time=describe_response[RESPONSE_KEY_CREATION_TIME], updated_time=describe_response[RESPONSE_KEY_LAST_MODIFIED_TIME], reference=reference, + description=description, sagemaker_session=sagemaker_session ) @@ -495,21 +508,29 @@ def get_versions(self) -> List["Evaluator"]: return evaluators @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="Evaluator.create_version") - def create_version(self, source: str) -> bool: + def create_version(self, source: str) -> "Evaluator": """Create a new version of this evaluator. Args: source: Lambda ARN or local file path for the function Returns: - bool: True if version created successfully, False otherwise + Evaluator: The newly created version. + + Raises: + RuntimeError: If version creation fails. """ try: - Evaluator.create( + new_evaluator = Evaluator.create( name=self.name, type=self.type, source=source, + sagemaker_session=self.sagemaker_session, + ) + logger.info( + "Created new version %s for evaluator %s, arn: %s", + new_evaluator.version, self.name, new_evaluator.arn ) - return True + return new_evaluator except Exception as e: raise RuntimeError(f"[PySDK Error] Failed to create new version: {str(e)}") diff --git a/sagemaker-train/tests/integ/ai_registry/test_dataset.py b/sagemaker-train/tests/integ/ai_registry/test_dataset.py index 46f5ce987a..a3f6aeeb4d 100644 --- a/sagemaker-train/tests/integ/ai_registry/test_dataset.py +++ b/sagemaker-train/tests/integ/ai_registry/test_dataset.py @@ -184,9 +184,10 @@ def test_dataset_wait(self, unique_name, sample_jsonl_file, cleanup_list): def test_create_dataset_version(self, unique_name, sample_jsonl_file, cleanup_list): """Test creating new dataset version.""" dataset = DataSet.create(name=unique_name, source=sample_jsonl_file, wait=False) - result = dataset.create_version(sample_jsonl_file) + new_version = dataset.create_version(sample_jsonl_file) cleanup_list.append(dataset) - assert result is True + assert isinstance(new_version, DataSet) + assert new_version.name == dataset.name def test_dataset_validation_invalid_extension(self, unique_name): """Test dataset validation with invalid file extension.""" diff --git a/sagemaker-train/tests/integ/ai_registry/test_evaluator.py b/sagemaker-train/tests/integ/ai_registry/test_evaluator.py index 51497c6cfd..c63e32a41e 100644 --- a/sagemaker-train/tests/integ/ai_registry/test_evaluator.py +++ b/sagemaker-train/tests/integ/ai_registry/test_evaluator.py @@ -196,8 +196,9 @@ def test_create_evaluator_version(self, unique_name, sample_prompt_file, cleanup Evaluator.delete_by_name(name=unique_name) evaluator = Evaluator.create(name=unique_name, type=REWARD_PROMPT, source=sample_prompt_file, wait=False) # cleanup_list.append(evaluator) - result = evaluator.create_version(source=sample_prompt_file) - assert result is True + new_version = evaluator.create_version(source=sample_prompt_file) + assert isinstance(new_version, Evaluator) + assert new_version.name == evaluator.name Evaluator.delete_by_name(name=unique_name) def test_create_reward_prompt_without_source_fails(self, unique_name): diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset.py b/sagemaker-train/tests/unit/ai_registry/test_dataset.py index 28641567a6..24f1b57bbd 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset.py @@ -374,13 +374,25 @@ def test_create_version_success(self, mock_air_hub, mock_create): "HubContentDocument": "{}", "HubContentSearchKeywords": ["customization_technique:sft", "method:generated"] } - mock_create.return_value = Mock() - - dataset = DataSet("test", "arn", "1.0.0", "s3://bucket/prefix", HubContentStatus.AVAILABLE, "desc", CustomizationTechnique.SFT) + session = Mock() + session.sagemaker_config = {"SchemaVersion": "1.0"} + new_dataset = Mock(spec=DataSet) + new_dataset.arn = "test-arn-v2" + new_dataset.version = "2.0.0" + mock_create.return_value = new_dataset + + dataset = DataSet( + "test", "arn", "1.0.0", "s3://bucket/prefix", HubContentStatus.AVAILABLE, "desc", + CustomizationTechnique.SFT, sagemaker_session=session, + ) result = dataset.create_version("s3://bucket/new-data") - - assert result is True + + assert result is new_dataset mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert create_kwargs["name"] == "test" + assert create_kwargs["source"] == "s3://bucket/new-data" + assert create_kwargs["sagemaker_session"] is session @patch('sagemaker.ai_registry.dataset.AIRHub') def test_create_version_failure(self, mock_air_hub): @@ -389,7 +401,45 @@ def test_create_version_failure(self, mock_air_hub): dataset = DataSet("test", "arn", "1.0.0", "s3://bucket/prefix", HubContentStatus.AVAILABLE, "desc", CustomizationTechnique.SFT) result = dataset.create_version("s3://bucket/new-data") - assert result is False + assert result is None + + + @patch('sagemaker.train.defaults.resolve_and_validate_role', return_value="arn:aws:iam::123456789012:role/SageMakerRole") + @patch('boto3.client') + @patch('sagemaker.core.helper.session_helper.Session') + @patch('sagemaker.train.common_utils.finetune_utils._get_current_domain_id') + @patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_file') + @patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_format') + @patch('sagemaker.ai_registry.dataset.AIRHub') + def test_create_threads_description_to_import_hub_content( + self, mock_air_hub, mock_validate_format, mock_validate_file, mock_get_domain_id, + mock_session, mock_boto_client, mock_resolve_role + ): + """description passed to DataSet.create is forwarded to AIRHub.import_hub_content.""" + mock_get_domain_id.return_value = None + mock_session_instance = Mock() + mock_session_instance.get_caller_identity_arn.return_value = "arn:aws:iam::123456789012:role/SageMakerRole" + mock_session.return_value = mock_session_instance + mock_sts_client = Mock() + mock_sts_client.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto_client.return_value = mock_sts_client + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + "HubContentArn": "test-arn", + "HubContentVersion": "1.0.0", + "CreationTime": "2024-01-01", + "LastModifiedTime": "2024-01-01", + } + + DataSet.create( + name="test-dataset", + source="s3://bucket/data.jsonl", + description="my description", + wait=False, + ) + + mock_air_hub.import_hub_content.assert_called_once() + assert mock_air_hub.import_hub_content.call_args.kwargs["description"] == "my description" class TestDataSetCreateWithContentMetadata: diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py index eeeedfce58..0cf3341e28 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py @@ -19,6 +19,8 @@ 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, + RESPONSE_KEY_HUB_CONTENT_NAME, RESPONSE_KEY_HUB_CONTENT_STATUS, + RESPONSE_KEY_HUB_CONTENT_DOCUMENT, RESPONSE_KEY_HUB_CONTENT_SEARCH_KEYWORDS, REWARD_FUNCTION, REWARD_PROMPT ) @@ -30,6 +32,50 @@ def _keywords_from_import_call(mock_air_hub): class TestEvaluator: + @patch('sagemaker.ai_registry.evaluator.AIRHub') + def test_create_threads_description_to_import_hub_content(self, mock_air_hub): + """description passed to Evaluator.create is forwarded to AIRHub.import_hub_content.""" + 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, + description="my evaluator description", + wait=False, + ) + + mock_air_hub.import_hub_content.assert_called_once() + assert mock_air_hub.import_hub_content.call_args.kwargs["description"] == "my evaluator description" + + @patch('sagemaker.ai_registry.evaluator.AIRHub') + def test_get_reads_back_description(self, mock_air_hub): + """Evaluator.get() populates description from the HubContentDescription response field.""" + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_NAME: "test-evaluator", + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", + RESPONSE_KEY_HUB_CONTENT_STATUS: "Available", + RESPONSE_KEY_HUB_CONTENT_DOCUMENT: json.dumps({ + "SubType": "AWS/Evaluator", + "JsonContent": json.dumps({"Reference": "ref"}), + }), + "HubContentDescription": "stored description", + RESPONSE_KEY_HUB_CONTENT_SEARCH_KEYWORDS: [], + RESPONSE_KEY_CREATION_TIME: "2024-01-01", + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01", + } + + evaluator = Evaluator.get("test-evaluator") + + assert evaluator.description == "stored description" + @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_create_with_lambda_arn(self, mock_air_hub): mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} @@ -192,13 +238,25 @@ def test_get_versions(self, mock_air_hub): @patch('sagemaker.ai_registry.evaluator.Evaluator.create') def test_create_version_success(self, mock_create): - mock_create.return_value = MagicMock() - - evaluator = Evaluator("test", "1.0.0", "arn", "AWS/Evaluator", method=EvaluatorMethod.LAMBDA, reference="lambda-arn") + session = MagicMock() + session.sagemaker_config = {"SchemaVersion": "1.0"} + new_evaluator = MagicMock() + new_evaluator.arn = "test-arn-v2" + new_evaluator.version = "2.0.0" + mock_create.return_value = new_evaluator + + evaluator = Evaluator( + "test", "1.0.0", "arn", "AWS/Evaluator", method=EvaluatorMethod.LAMBDA, + reference="lambda-arn", sagemaker_session=session, + ) result = evaluator.create_version("arn:aws:lambda:us-west-2:123456789012:function:new") - - assert result is True + + assert result is new_evaluator mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert create_kwargs["name"] == "test" + assert create_kwargs["source"] == "arn:aws:lambda:us-west-2:123456789012:function:new" + assert create_kwargs["sagemaker_session"] is session @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_create_version_failure(self, mock_air_hub):