chore: create JumpStartInternalMetadata class

Author: evakravi

src/sagemaker/jumpstart/estimator.py

@@ -35,6 +35,7 @@
 from sagemaker.jumpstart.factory.estimator import get_deploy_kwargs, get_fit_kwargs, get_init_kwargs
 from sagemaker.jumpstart.factory.model import get_default_predictor
 from sagemaker.jumpstart.session_utils import get_model_id_version_from_training_job
+from sagemaker.jumpstart.types import JumpStartInternalMetadata
 from sagemaker.jumpstart.utils import (
     validate_model_id_and_get_type,
     resolve_model_sagemaker_config_field,
@@ -534,14 +535,16 @@ def _validate_model_id_and_get_type_hook():
             if not self.model_type and not hub_arn:
                 raise ValueError(INVALID_MODEL_ID_ERROR_MSG.format(model_id=model_id))
 
-        self._specs = verify_model_region_and_return_specs(
-            region=self.region,
-            model_id=self.model_id,
-            version=self.model_version,
-            hub_arn=self.hub_arn,
-            model_type=self.model_type,
-            scope=JumpStartScriptScope.TRAINING,
-            sagemaker_session=self.sagemaker_session,
+        self._internal = JumpStartInternalMetadata(
+            specs=verify_model_region_and_return_specs(
+                region=self.region,
+                model_id=self.model_id,
+                version=self.model_version,
+                hub_arn=self.hub_arn,
+                model_type=self.model_type,
+                scope=JumpStartScriptScope.TRAINING,
+                sagemaker_session=self.sagemaker_session,
+            )
         )
 
         estimator_init_kwargs = get_init_kwargs(

src/sagemaker/jumpstart/model.py

@@ -38,7 +38,7 @@
     get_register_kwargs,
 )
 from sagemaker.jumpstart.session_utils import get_model_id_version_from_endpoint
-from sagemaker.jumpstart.types import JumpStartSerializablePayload
+from sagemaker.jumpstart.types import JumpStartInternalMetadata, JumpStartSerializablePayload
 from sagemaker.jumpstart.utils import (
     validate_model_id_and_get_type,
     verify_model_region_and_return_specs,
@@ -312,14 +312,16 @@ def _validate_model_id_and_type():
             if not self.model_type and not hub_arn:
                 raise ValueError(INVALID_MODEL_ID_ERROR_MSG.format(model_id=model_id))
 
-        self._specs = verify_model_region_and_return_specs(
-            region=self.region,
-            model_id=self.model_id,
-            version=self.model_version,
-            hub_arn=self.hub_arn,
-            model_type=self.model_type,
-            scope=JumpStartScriptScope.INFERENCE,
-            sagemaker_session=self.sagemaker_session,
+        self._internal = JumpStartInternalMetadata(
+            specs=verify_model_region_and_return_specs(
+                region=self.region,
+                model_id=self.model_id,
+                version=self.model_version,
+                hub_arn=self.hub_arn,
+                model_type=self.model_type,
+                scope=JumpStartScriptScope.INFERENCE,
+                sagemaker_session=self.sagemaker_session,
+            )
         )
 
         self._model_data_is_set = model_data is not None

src/sagemaker/jumpstart/types.py

@@ -2482,3 +2482,17 @@ def __init__(
         self.source_uri = source_uri
         self.model_card = model_card
         self.accept_eula = accept_eula
+
+
+class JumpStartInternalMetadata(JumpStartDataHolderType):
+    """Data class for storing internal/private fields for JumpStart models."""
+
+    slots = ["specs"]
+
+    def __init__(self, specs: JumpStartModelSpecs):
+        """Initializes a JumpStartInternalMetadata object.
+
+        Args:
+            specs (JumpStartModelSpecs): specs for model.
+        """
+        self.specs = specs