sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

{sagemaker_core → sagemaker/core}/tools/shapes_extractor.py

@@ -16,17 +16,13 @@ import pprint
 from functools import lru_cache
 from typing import Optional, Any
 
-from sagemaker_core.tools.constants import (
-    BASIC_JSON_TYPES_TO_PYTHON_TYPES,
-    SHAPE_DAG_FILE_PATH,
-    SHAPES_WITH_JSON_FIELD_ALIAS,
-)
-from sagemaker_core.main.utils import (
+from sagemaker.core.tools.constants import BASIC_JSON_TYPES_TO_PYTHON_TYPES, SHAPE_DAG_FILE_PATH
+from sagemaker.core.utils.utils import (
     reformat_file_with_black,
     convert_to_snake_case,
     snake_to_pascal,
 )
-from sagemaker_core.tools.data_extractor import load_combined_shapes_data
+from sagemaker.core.tools.data_extractor import load_combined_shapes_data
 
 
 class ShapesExtractor:
@@ -103,11 +99,6 @@ class ShapesExtractor:
                 _dag[shape] = {"type": "structure", "members": []}
                 for member, member_attrs in shape_data["members"].items():
                     shape_node_member = {"name": member, "shape": member_attrs["shape"]}
-                    # Add alias if field name is json, to address the Bug: https://github.com/aws/sagemaker-python-sdk/issues/4944
-                    if shape in SHAPES_WITH_JSON_FIELD_ALIAS and member == "Json":
-                        shape_node_member["name"] = "JsonFormat"
-                        shape_node_member["alias"] = "json"
-
                     member_shape_dict = _all_shapes[member_attrs["shape"]]
                     shape_node_member["type"] = member_shape_dict["type"]
                     _dag[shape]["members"].append(shape_node_member)
@@ -126,21 +117,19 @@ class ShapesExtractor:
                 _dag[shape]["value_type"] = _all_shapes[_map_value_shape]["type"]
         return _dag
 
-    def _evaluate_list_type(self, member_shape, add_shapes_prefix=True):
+    def _evaluate_list_type(self, member_shape):
         list_shape_name = member_shape["member"]["shape"]
         list_shape_member = self.combined_shapes[list_shape_name]
         list_shape_type = list_shape_member["type"]
         if list_shape_type == "list":
-            member_type = f"List[{self._evaluate_list_type(list_shape_member, add_shapes_prefix)}]"
+            member_type = f"List[{self._evaluate_list_type(list_shape_member)}]"
         elif list_shape_type == "map":
-            member_type = f"List[{self._evaluate_map_type(list_shape_member, add_shapes_prefix)}]"
+            member_type = f"List[{self._evaluate_map_type(list_shape_member)}]"
         elif list_shape_type == "structure":
             # handling an edge case of nested structure
             if list_shape_name == "SearchExpression":
                 member_type = f"List['{list_shape_name}']"
             else:
-                if add_shapes_prefix:
-                    list_shape_name = f"shapes.{list_shape_name}"
                 member_type = f"List[{list_shape_name}]"
         elif list_shape_type in BASIC_JSON_TYPES_TO_PYTHON_TYPES.keys():
             member_type = f"List[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[list_shape_type]}]"
@@ -150,7 +139,7 @@ class ShapesExtractor:
             )
         return member_type
 
-    def _evaluate_map_type(self, member_shape, add_shapes_prefix=True):
+    def _evaluate_map_type(self, member_shape):
         map_key_shape_name = member_shape["key"]["shape"]
         map_value_shape_name = member_shape["value"]["shape"]
         map_key_shape = self.combined_shapes[map_key_shape_name]
@@ -163,8 +152,6 @@ class ShapesExtractor:
                 "Unhandled map shape key type encountered, needs extra logic to handle this"
             )
         if map_value_shape_type == "structure":
-            if add_shapes_prefix:
-                map_value_shape_name = f"shapes.{map_value_shape_name}"
             member_type = (
                 f"Dict[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[map_key_shape_type]}, "
                 f"{map_value_shape_name}]"
@@ -172,12 +159,12 @@ class ShapesExtractor:
         elif map_value_shape_type == "list":
             member_type = (
                 f"Dict[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[map_key_shape_type]}, "
-                f"{self._evaluate_list_type(map_value_shape, add_shapes_prefix)}]"
+                f"{self._evaluate_list_type(map_value_shape)}]"
             )
         elif map_value_shape_type == "map":
             member_type = (
                 f"Dict[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[map_key_shape_type]}, "
-                f"{self._evaluate_map_type(map_value_shape, add_shapes_prefix)}]"
+                f"{self._evaluate_map_type(map_value_shape)}]"
             )
         else:
             member_type = (
@@ -187,13 +174,9 @@ class ShapesExtractor:
         return member_type
 
     def generate_data_shape_members_and_string_body(
-        self,
-        shape,
-        resource_plan: Optional[Any] = None,
-        required_override=(),
-        add_shapes_prefix=True,
+        self, shape, resource_plan: Optional[Any] = None, required_override=()
     ):
-        shape_members = self.generate_shape_members(shape, required_override, add_shapes_prefix)
+        shape_members = self.generate_shape_members(shape, required_override)
         resource_names = None
         if resource_plan is not None:
             resource_names = [row["resource_name"] for _, row in resource_plan.iterrows()]
@@ -207,31 +190,29 @@ class ShapesExtractor:
                 and snake_to_pascal(attr[: -len("_name")]) in resource_names
             ):
                 if value.startswith("Optional"):
-                    init_data_body += f"{attr}: Optional[Union[str, object]] = Unassigned()\n"
+                    init_data_body += (
+                        f"{attr}: Optional[Union[StrPipeVar, object]] = Unassigned()\n"
+                    )
                 else:
-                    init_data_body += f"{attr}: Union[str, object]\n"
+                    init_data_body += f"{attr}: Union[StrPipeVar, object]\n"
             elif attr == "lambda":
                 init_data_body += f"# {attr}: {value}\n"
             else:
                 init_data_body += f"{attr}: {value}\n"
         return shape_members, init_data_body
 
-    def generate_data_shape_string_body(
-        self, shape, resource_plan, required_override=(), add_shapes_prefix=True
-    ):
+    def generate_data_shape_string_body(self, shape, resource_plan, required_override=()):
         return self.generate_data_shape_members_and_string_body(
-            shape, resource_plan, required_override, add_shapes_prefix
+            shape, resource_plan, required_override
         )[1]
 
-    def generate_data_shape_members(
-        self, shape, resource_plan, required_override=(), add_shapes_prefix=True
-    ):
+    def generate_data_shape_members(self, shape, resource_plan, required_override=()):
         return self.generate_data_shape_members_and_string_body(
-            shape, resource_plan, required_override, add_shapes_prefix
+            shape, resource_plan, required_override
         )[0]
 
     @lru_cache
-    def generate_shape_members(self, shape, required_override=(), add_shapes_prefix=True):
+    def generate_shape_members(self, shape, required_override=()):
         shape_dict = self.combined_shapes[shape]
         members = shape_dict["members"]
         required_args = list(required_override) or shape_dict.get("required", [])
@@ -239,48 +220,29 @@ class ShapesExtractor:
         # bring the required members in front
         ordered_members = {key: members[key] for key in required_args if key in members}
         ordered_members.update(members)
-        field_aliases = {}
-
         for member_name, member_attrs in ordered_members.items():
             member_shape_name = member_attrs["shape"]
             if self.combined_shapes[member_shape_name]:
                 member_shape = self.combined_shapes[member_shape_name]
                 member_shape_type = member_shape["type"]
                 if member_shape_type == "structure":
-                    if add_shapes_prefix:
-                        member_shape_name = f"shapes.{member_shape_name}"
-                    member_type = f"{member_shape_name}"
+                    member_type = member_shape_name
                 elif member_shape_type == "list":
-                    member_type = self._evaluate_list_type(member_shape, add_shapes_prefix)
+                    member_type = self._evaluate_list_type(member_shape)
                 elif member_shape_type == "map":
-                    member_type = self._evaluate_map_type(member_shape, add_shapes_prefix)
+                    member_type = self._evaluate_map_type(member_shape)
                 else:
                     # Shape is a simple type like string
                     member_type = BASIC_JSON_TYPES_TO_PYTHON_TYPES[member_shape_type]
             else:
                 raise Exception("The Shape definition mush exist. The Json Data might be corrupt")
-
-            is_required = member_name in required_args
-            # Add alias if field name is json, to address the Bug: https://github.com/aws/sagemaker-python-sdk/issues/4944
-            if shape in SHAPES_WITH_JSON_FIELD_ALIAS and member_name == "Json":
-                updated_member_name_snake_case = "json_format"
-                field_aliases[updated_member_name_snake_case] = "json"
-                init_data_body[f"{updated_member_name_snake_case}"] = (
-                    (
-                        f"{member_type} = Field(alias='{field_aliases[updated_member_name_snake_case]}')"
-                    )
-                    if is_required
-                    else f"Optional[{member_type}] = Field(default=Unassigned(), alias='json')"
-                )
+            member_name_snake_case = convert_to_snake_case(member_name)
+            if member_name in required_args:
+                init_data_body[f"{member_name_snake_case}"] = f"{member_type}"
             else:
-                member_name_snake_case = convert_to_snake_case(member_name)
-                if is_required:
-                    init_data_body[f"{member_name_snake_case}"] = f"{member_type}"
-                else:
-                    init_data_body[f"{member_name_snake_case}"] = (
-                        f"Optional[{member_type}] = Unassigned()"
-                    )
-
+                init_data_body[f"{member_name_snake_case}"] = (
+                    f"Optional[{member_type}] = Unassigned()"
+                )
         return init_data_body
 
     @lru_cache

{sagemaker_core → sagemaker/core}/tools/templates.py

@@ -40,7 +40,7 @@ def create(
     cls,
 {create_args}
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> Optional["{resource_name}"]:
 {docstring}
     logger.info("Creating {resource_lower} resource.")
@@ -71,7 +71,7 @@ def create(
     cls,
 {create_args}
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> Optional["{resource_name}"]:
 {docstring}
     logger.info("Creating {resource_lower} resource.")
@@ -102,7 +102,7 @@ def load(
     cls,
 {import_args}
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> Optional["{resource_name}"]:
 {docstring}
     logger.info(f"Importing {resource_lower} resource.")
@@ -214,7 +214,7 @@ def get(
     cls,
 {describe_args}
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> Optional["{resource_name}"]:
 {docstring}
     operation_input_args = {{
@@ -454,7 +454,7 @@ if "delete_failed" in current_status.lower() or "deletefailed" in current_status
 
 DELETED_STATUS_CHECK = """
 if current_status.lower() == "deleted":
-    print("Resource was deleted.")
+    logger.info("Resource was deleted.")
     return
 """
 
@@ -504,7 +504,7 @@ def get_all(
     cls,
 {get_all_args}
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> ResourceIterator["{resource}"]:
 {docstring}
     client = Base.get_sagemaker_client(session=session, region_name=region, service_name="{service_name}")
@@ -512,7 +512,6 @@ def get_all(
     operation_input_args = {{
 {operation_input_args}
     }}
-{extract_name_mapping}
 {custom_key_mapping}
     # serialize the input request
     operation_input_args = serialize(operation_input_args)
@@ -529,7 +528,7 @@ GET_ALL_METHOD_NO_ARGS_TEMPLATE = '''
 def get_all(
     cls,
     session: Optional[Session] = None,
-    region: Optional[str] = None,
+    region: Optional[StrPipeVar] = None,
 ) -> ResourceIterator["{resource}"]:
     """
     Get all {resource} resources.
@@ -543,7 +542,6 @@ def get_all(
 
     """
     client = Base.get_sagemaker_client(session=session, region_name=region, service_name="{service_name}")
-{extract_name_mapping}
 {custom_key_mapping}
     return ResourceIterator(
 {resource_iterator_args}
@@ -602,7 +600,8 @@ DESERIALIZE_INPUT_AND_RESPONSE_TO_CLS_TEMPLATE = """
 
 RESOURCE_BASE_CLASS_TEMPLATE = """
 class Base(BaseModel):
-    model_config = ConfigDict(protected_namespaces=(), validate_assignment=True, extra="forbid")
+    model_config = ConfigDict(protected_namespaces=(), validate_assignment=True, extra="forbid", arbitrary_types_allowed=True)
+    config_manager: ClassVar[SageMakerConfig] = SageMakerConfig()
     
     @classmethod
     def get_sagemaker_client(cls, session = None, region_name = None, service_name = 'sagemaker'):
@@ -615,17 +614,17 @@ class Base(BaseModel):
         try:
             for configurable_attribute in config_schema_for_resource:
                 if kwargs.get(configurable_attribute) is None:
-                    resource_defaults = load_default_configs_for_resource_name(
+                    resource_defaults = Base.config_manager.load_default_configs_for_resource_name(
                         resource_name=resource_name
                     )
-                    global_defaults = load_default_configs_for_resource_name(
+                    global_defaults = Base.config_manager.load_default_configs_for_resource_name(
                         resource_name="GlobalDefaults"
                     )
-                    if config_value := get_config_value(
+                    if config_value := Base.config_manager.get_resolved_config_value(
                         configurable_attribute, resource_defaults, global_defaults
                     ):
                         resource_name = snake_to_pascal(configurable_attribute)
-                        class_object = getattr(shapes, resource_name, None) or globals().get(resource_name)
+                        class_object = globals()[resource_name]
                         kwargs[configurable_attribute] = class_object(**config_value)
         except BaseException as e:
             logger.debug("Could not load Default Configs. Continuing.", exc_info=True)
@@ -671,9 +670,7 @@ class Base(BaseModel):
     @staticmethod
     def _get_chained_attribute(item_value: Any):
         resource_name = type(item_value).__name__
-        class_object = globals().get(resource_name) or getattr(shapes, resource_name, None)
-        if class_object is None:
-            return item_value
+        class_object = globals()[resource_name]
         return class_object(**Base.populate_chained_attributes(
             resource_name=resource_name,
             operation_input_args=vars(item_value)
@@ -714,3 +711,28 @@ Raises:
             error_message = e.response['Error']['Message']
             error_code = e.response['Error']['Code']
         ```"""
+
+SERIALIZE_INPUT_ENDPOINT_TEMPLATE = """
+    use_serializer = False
+    if ((self.serializer is not None and self.deserializer is None) or
+    (self.serializer is None and self.deserializer is not None)):
+        raise ValueError("Both serializer and deserializer must be provided together, or neither should be provided")
+    if self.serializer is not None and self.deserializer is not None:
+        use_serializer = True
+    if use_serializer:
+        body = self.serializer.serialize(body)
+    operation_input_args = {{
+{operation_input_args}
+    }}
+    # serialize the input request
+    operation_input_args = serialize(operation_input_args)
+    logger.debug(f"Serialized input request: {{operation_input_args}}")"""
+
+DESERIALIZE_RESPONSE_ENDPOINT_TEMPLATE = """
+    transformed_response = transform(response, 'InvokeEndpointOutput')
+    # Deserialize the body if a deserializer is provided
+    if use_serializer:
+        body_content = transformed_response["body"]
+        deserialized_body = self.deserializer.deserialize(body_content, transformed_response["content_type"])
+        transformed_response["body"] = deserialized_body
+    return {return_type_conversion}(**transformed_response)"""

sagemaker/core/training/__init__.py

@@ -0,0 +1,14 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Training configuration and utilities."""
+from __future__ import absolute_import