sagemaker-core 1.0.4__py3-none-any.whl → 1.0.6__py3-none-any.whl

sagemaker_core/tools/data_extractor.py

@@ -5,6 +5,8 @@ from pydantic import BaseModel
 
 from sagemaker_core.tools.constants import (
     ADDITIONAL_OPERATION_FILE_PATH,
+    FEATURE_STORE_SERVICE_JSON_FILE_PATH,
+    METRICS_SERVICE_JSON_FILE_PATH,
     SERVICE_JSON_FILE_PATH,
     RUNTIME_SERVICE_JSON_FILE_PATH,
 )
@@ -13,6 +15,8 @@ from sagemaker_core.tools.constants import (
 class ServiceJsonData(BaseModel):
     sagemaker: dict
     sagemaker_runtime: dict
+    sagemaker_feature_store: dict
+    sagemaker_metrics: dict
 
 
 @lru_cache(maxsize=1)
@@ -21,15 +25,26 @@ def load_service_jsons() -> ServiceJsonData:
         service_json = json.load(file)
     with open(RUNTIME_SERVICE_JSON_FILE_PATH, "r") as file:
         runtime_service_json = json.load(file)
-    return ServiceJsonData(sagemaker=service_json, sagemaker_runtime=runtime_service_json)
+    with open(FEATURE_STORE_SERVICE_JSON_FILE_PATH, "r") as file:
+        feature_store_service_json = json.load(file)
+    with open(METRICS_SERVICE_JSON_FILE_PATH, "r") as file:
+        metrics_service_json = json.load(file)
+    return ServiceJsonData(
+        sagemaker=service_json,
+        sagemaker_runtime=runtime_service_json,
+        sagemaker_feature_store=feature_store_service_json,
+        sagemaker_metrics=metrics_service_json,
+    )
 
 
 @lru_cache(maxsize=1)
 def load_combined_shapes_data() -> dict:
     service_json_data = load_service_jsons()
     return {
-        **service_json_data.sagemaker["shapes"],
         **service_json_data.sagemaker_runtime["shapes"],
+        **service_json_data.sagemaker_feature_store["shapes"],
+        **service_json_data.sagemaker_metrics["shapes"],
+        **service_json_data.sagemaker["shapes"],
     }
 
 
@@ -37,8 +52,10 @@ def load_combined_shapes_data() -> dict:
 def load_combined_operations_data() -> dict:
     service_json_data = load_service_jsons()
     return {
-        **service_json_data.sagemaker["operations"],
         **service_json_data.sagemaker_runtime["operations"],
+        **service_json_data.sagemaker_feature_store["operations"],
+        **service_json_data.sagemaker_metrics["operations"],
+        **service_json_data.sagemaker["operations"],
     }
 
 

sagemaker_core/tools/resources_codegen.py

@@ -64,9 +64,6 @@ from sagemaker_core.tools.templates import (
     UPDATE_METHOD_TEMPLATE,
     POPULATE_DEFAULTS_DECORATOR_TEMPLATE,
     CREATE_METHOD_TEMPLATE_WITHOUT_DEFAULTS,
-    INVOKE_METHOD_TEMPLATE,
-    INVOKE_ASYNC_METHOD_TEMPLATE,
-    INVOKE_WITH_RESPONSE_STREAM_METHOD_TEMPLATE,
     IMPORT_METHOD_TEMPLATE,
     FAILED_STATUS_ERROR_TEMPLATE,
     GET_NAME_METHOD_TEMPLATE,
@@ -188,8 +185,8 @@ class ResourcesCodeGen:
             "from rich.style import Style",
             "from sagemaker_core.main.code_injection.codec import transform",
             "from sagemaker_core.main.code_injection.constants import Color",
-            "from sagemaker_core.main.utils import SageMakerClient, SageMakerRuntimeClient, ResourceIterator, Unassigned, get_textual_rich_logger, "
-            "snake_to_pascal, pascal_to_snake, is_not_primitive, is_not_str_dict, is_snake_case, is_primitive_list, serialize",
+            "from sagemaker_core.main.utils import SageMakerClient, ResourceIterator, Unassigned, get_textual_rich_logger, "
+            "snake_to_pascal, pascal_to_snake, is_not_primitive, is_not_str_dict, is_primitive_list, serialize",
             "from sagemaker_core.main.intelligent_defaults_helper import load_default_configs_for_resource_name, get_config_value",
             "from sagemaker_core.main.shapes import *",
             "from sagemaker_core.main.exceptions import *",
@@ -434,30 +431,6 @@ class ResourcesCodeGen:
             ):
                 resource_class += add_indent(wait_for_delete_method, 4)
 
-            if invoke_method := self._evaluate_method(
-                resource_name,
-                "invoke",
-                object_methods,
-                resource_attributes=resource_attributes,
-            ):
-                resource_class += add_indent(invoke_method, 4)
-
-            if invoke_async_method := self._evaluate_method(
-                resource_name,
-                "invoke_async",
-                object_methods,
-                resource_attributes=resource_attributes,
-            ):
-                resource_class += add_indent(invoke_async_method, 4)
-
-            if invoke_with_response_stream_method := self._evaluate_method(
-                resource_name,
-                "invoke_with_response_stream",
-                object_methods,
-                resource_attributes=resource_attributes,
-            ):
-                resource_class += add_indent(invoke_with_response_stream_method, 4)
-
             if import_method := self._evaluate_method(resource_name, "import", class_methods):
                 resource_class += add_indent(import_method, 4)
 
@@ -1172,174 +1145,6 @@ class ResourcesCodeGen:
         # Return the formatted method
         return formatted_method
 
-    def generate_invoke_method(self, resource_name: str, **kwargs) -> str:
-        """
-        Auto-generate the INVOKE ASYNC method for a resource.
-
-        Args:
-            resource_name (str): The resource name.
-
-        Returns:
-            str: The formatted Update Method template.
-
-        """
-        # Get the operation and shape for the 'create' method
-        operation_name = "Invoke" + resource_name
-        operation_metadata = self.operations[operation_name]
-        operation_input_shape_name = operation_metadata["input"]["shape"]
-
-        # Generate the arguments for the 'create' method
-        invoke_args = self._generate_method_args(
-            operation_input_shape_name, kwargs["resource_attributes"]
-        )
-
-        operation_input_args = self._generate_operation_input_necessary_args(
-            operation_metadata, kwargs["resource_attributes"]
-        )
-
-        # Convert the resource name to snake case
-        resource_lower = convert_to_snake_case(resource_name)
-
-        # Convert the operation name to snake case
-        operation = convert_to_snake_case(operation_name)
-
-        # generate docstring
-        docstring = self._generate_docstring(
-            title=f"Invoke a {resource_name} resource",
-            operation_name=operation_name,
-            resource_name=resource_name,
-            operation_input_shape_name=operation_input_shape_name,
-            include_session_region=False,
-            include_return_resource_docstring=False,
-            return_string=f"Returns:\n" f"    The Invoke response.\n",
-            exclude_resource_attrs=kwargs["resource_attributes"],
-        )
-        # Format the method using the CREATE_METHOD_TEMPLATE
-        formatted_method = INVOKE_METHOD_TEMPLATE.format(
-            docstring=docstring,
-            service_name="sagemaker-runtime",
-            invoke_args=invoke_args,
-            resource_name=resource_name,
-            resource_lower=resource_lower,
-            operation_input_args=operation_input_args,
-            operation=operation,
-        )
-
-        # Return the formatted method
-        return formatted_method
-
-    def generate_invoke_async_method(self, resource_name: str, **kwargs) -> str:
-        """
-        Auto-generate the INVOKE method for a resource.
-
-        Args:
-            resource_name (str): The resource name.
-
-        Returns:
-            str: The formatted Update Method template.
-
-        """
-        # Get the operation and shape for the 'create' method
-        operation_name = "Invoke" + resource_name + "Async"
-        operation_metadata = self.operations[operation_name]
-        operation_input_shape_name = operation_metadata["input"]["shape"]
-
-        # Generate the arguments for the 'create' method
-        invoke_args = self._generate_method_args(
-            operation_input_shape_name, kwargs["resource_attributes"]
-        )
-
-        operation_input_args = self._generate_operation_input_necessary_args(
-            operation_metadata, kwargs["resource_attributes"]
-        )
-
-        # Convert the resource name to snake case
-        resource_lower = convert_to_snake_case(resource_name)
-
-        # Convert the operation name to snake case
-        operation = convert_to_snake_case(operation_name)
-
-        # generate docstring
-        docstring = self._generate_docstring(
-            title=f"Invoke Async a {resource_name} resource",
-            operation_name=operation_name,
-            resource_name=resource_name,
-            operation_input_shape_name=operation_input_shape_name,
-            include_session_region=False,
-            include_return_resource_docstring=False,
-            return_string=f"Returns:\n" f"    The Invoke response.\n",
-            exclude_resource_attrs=kwargs["resource_attributes"],
-        )
-        # Format the method using the CREATE_METHOD_TEMPLATE
-        formatted_method = INVOKE_ASYNC_METHOD_TEMPLATE.format(
-            docstring=docstring,
-            service_name="sagemaker-runtime",
-            create_args=invoke_args,
-            resource_name=resource_name,
-            resource_lower=resource_lower,
-            operation_input_args=operation_input_args,
-            operation=operation,
-        )
-
-        # Return the formatted method
-        return formatted_method
-
-    def generate_invoke_with_response_stream_method(self, resource_name: str, **kwargs) -> str:
-        """
-        Auto-generate the INVOKE with response stream method for a resource.
-
-        Args:
-            resource_name (str): The resource name.
-
-        Returns:
-            str: The formatted Update Method template.
-
-        """
-        # Get the operation and shape for the 'create' method
-        operation_name = "Invoke" + resource_name + "WithResponseStream"
-        operation_metadata = self.operations[operation_name]
-        operation_input_shape_name = operation_metadata["input"]["shape"]
-
-        # Generate the arguments for the 'create' method
-        invoke_args = self._generate_method_args(
-            operation_input_shape_name, kwargs["resource_attributes"]
-        )
-
-        operation_input_args = self._generate_operation_input_necessary_args(
-            operation_metadata, kwargs["resource_attributes"]
-        )
-
-        # Convert the resource name to snake case
-        resource_lower = convert_to_snake_case(resource_name)
-
-        # Convert the operation name to snake case
-        operation = convert_to_snake_case(operation_name)
-
-        # generate docstring
-        docstring = self._generate_docstring(
-            title=f"Invoke with response stream a {resource_name} resource",
-            operation_name=operation_name,
-            resource_name=resource_name,
-            operation_input_shape_name=operation_input_shape_name,
-            include_session_region=False,
-            include_return_resource_docstring=False,
-            return_string=f"Returns:\n" f"    The Invoke response.\n",
-            exclude_resource_attrs=kwargs["resource_attributes"],
-        )
-        # Format the method using the CREATE_METHOD_TEMPLATE
-        formatted_method = INVOKE_WITH_RESPONSE_STREAM_METHOD_TEMPLATE.format(
-            docstring=docstring,
-            service_name="sagemaker-runtime",
-            create_args=invoke_args,
-            resource_name=resource_name,
-            resource_lower=resource_lower,
-            operation_input_args=operation_input_args,
-            operation=operation,
-        )
-
-        # Return the formatted method
-        return formatted_method
-
     def generate_get_method(self, resource_name: str) -> str:
         """
         Auto-generate the GET method (describe API) for a resource.

sagemaker_core/tools/resources_extractor.py

@@ -97,7 +97,7 @@ class ResourcesExtractor:
             if resource_name not in self.resource_methods:
                 self.resource_methods[resource_name] = dict()
             for operation_name, operation in resource_operations.items():
-                self.actions_under_resource.update(operation_name)
+                self.actions_under_resource.add(operation_name)
                 method = Method(**operation)
                 method.get_docstring_title(self.operations[operation_name])
                 self.resource_methods[resource_name][operation["method_name"]] = method

sagemaker_core/tools/templates.py

@@ -197,77 +197,6 @@ def update(
     return self
 """
 
-INVOKE_METHOD_TEMPLATE = """
-@Base.add_validate_call
-def invoke(self, 
-{invoke_args}
-) -> Optional[object]:
-{docstring}
-    logger.info(f"Invoking {resource_lower} resource.")
-    client = SageMakerRuntimeClient(service_name="{service_name}").client
-    operation_input_args = {{
-{operation_input_args}
-    }}
-    logger.debug(f"Input request: {{operation_input_args}}")
-    # serialize the input request
-    operation_input_args = serialize(operation_input_args)
-    logger.debug(f"Serialized input request: {{operation_input_args}}")
-
-    # create the resource
-    response = client.{operation}(**operation_input_args)
-    logger.debug(f"Response: {{response}}")
-
-    return response
-"""
-
-INVOKE_ASYNC_METHOD_TEMPLATE = """
-@Base.add_validate_call
-def invoke_async(self, 
-{create_args}
-) -> Optional[object]:
-{docstring}
-    logger.info(f"Invoking {resource_lower} resource Async.")
-    client = SageMakerRuntimeClient(service_name="{service_name}").client
-    
-    operation_input_args = {{
-{operation_input_args}
-    }}
-    logger.debug(f"Input request: {{operation_input_args}}")
-    # serialize the input request
-    operation_input_args = serialize(operation_input_args)
-    logger.debug(f"Serialized input request: {{operation_input_args}}")
-
-    # create the resource
-    response = client.{operation}(**operation_input_args)
-    logger.debug(f"Response: {{response}}")
-
-    return response
-"""
-
-INVOKE_WITH_RESPONSE_STREAM_METHOD_TEMPLATE = """
-@Base.add_validate_call
-def invoke_with_response_stream(self, 
-{create_args}
-) -> Optional[object]:
-{docstring}
-    logger.info(f"Invoking {resource_lower} resource with Response Stream.")
-    client = SageMakerRuntimeClient(service_name="{service_name}").client
-
-    operation_input_args = {{
-{operation_input_args}
-    }}
-    logger.debug(f"Input request: {{operation_input_args}}")
-    # serialize the input request
-    operation_input_args = serialize(operation_input_args)
-    logger.debug(f"Serialized input request: {{operation_input_args}}")
-
-    # create the resource
-    response = client.{operation}(**operation_input_args)
-    logger.debug(f"Response: {{response}}")
-
-    return response
-"""
-
 POPULATE_DEFAULTS_DECORATOR_TEMPLATE = """
 def populate_inputs_decorator(create_func):
     @functools.wraps(create_func)
@@ -639,7 +568,7 @@ class Base(BaseModel):
     
     @classmethod
     def get_sagemaker_client(cls, session = None, region_name = None, service_name = 'sagemaker'):
-        return SageMakerClient(session=session, region_name=region_name, service_name=service_name).client
+        return SageMakerClient(session=session, region_name=region_name).get_client(service_name=service_name)
     
     @staticmethod
     def get_updated_kwargs_with_configured_attributes(

{sagemaker_core-1.0.4.dist-info → sagemaker_core-1.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sagemaker-core
-Version: 1.0.4
+Version: 1.0.6
 Summary: An python package for sagemaker core functionalities
 Author-email: AWS <sagemaker-interests@amazon.com>
 Project-URL: Repository, https://github.com/aws/sagemaker-core.git

{sagemaker_core-1.0.4.dist-info → sagemaker_core-1.0.6.dist-info}/RECORD

@@ -6,29 +6,29 @@ sagemaker_core/main/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 sagemaker_core/main/config_schema.py,sha256=IDvL1YuMOwVg4nb0D4ysU51tc1UjlQASyz1QZoncs7I,56907
 sagemaker_core/main/exceptions.py,sha256=87DUlrmHxaWoiYNlpNY9ixxFMPRk_dIGPsA2e_xdVwQ,5602
 sagemaker_core/main/intelligent_defaults_helper.py,sha256=5SDM6UavZtp-k5LhqRL7GRIDgzFB5UsC_p7YuiSPK9A,8334
-sagemaker_core/main/resources.py,sha256=3ivOSUiWqH8VWmOwxz1Oa45aet7CTsc8mZott-ge7yc,1299296
-sagemaker_core/main/shapes.py,sha256=mxsgFQYYMCWrZm_GF3w_wwgfTon0AgheZoXKu17ux1c,683409
+sagemaker_core/main/resources.py,sha256=S4wuM705rkah3cFO-kfWAXPgYq7SI3PG2Qh6ZHtfTfk,1314805
+sagemaker_core/main/shapes.py,sha256=kuefQBAgOogQynQikMq-6RxNezaVXE3ChwJedvVMAFw,693437
 sagemaker_core/main/user_agent.py,sha256=4sZybDXkzRoZnOnVDQ8p8zFTfiRJdsH7amDWInVQ4xU,2708
-sagemaker_core/main/utils.py,sha256=I7WsC-DdbnAZg_PtqcLPywXlqjAwj7ZYbeHhv93Z1O4,18840
+sagemaker_core/main/utils.py,sha256=lXkJyiCow5uf32l0EmkimB0RKVk2BS7OM2fYoLsOfD4,18346
 sagemaker_core/main/code_injection/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sagemaker_core/main/code_injection/base.py,sha256=11_Jif0nOzfbLGlXaacKf-wcizzfS64U0OSZGoVffFU,1733
 sagemaker_core/main/code_injection/codec.py,sha256=2DjmeD2uND307UqDefvVEpE0rZ8yfFU3Bi3TvQCQveI,7658
 sagemaker_core/main/code_injection/constants.py,sha256=2ICExGge8vAWx7lSTW0JGh-bH1korkvpOpDu5M63eI4,980
-sagemaker_core/main/code_injection/shape_dag.py,sha256=mQR562skpJVs3VbNS-dtGvF94JXVcRBZS5b0cpKH3x0,648111
+sagemaker_core/main/code_injection/shape_dag.py,sha256=iIkfDlMpNdzeLY958ZB3VTD55BJC1sid3qry-YOG-rM,655019
 sagemaker_core/resources/__init__.py,sha256=EAYTFMN-nPjnPjjBbhIUeaL67FLKNPd7qbcbl9VIrws,31
 sagemaker_core/shapes/__init__.py,sha256=RnbIu9eTxKt-DNsOFJabrWIgrrtS9_SdAozP9JBl_ic,28
 sagemaker_core/tools/__init__.py,sha256=xX79JImxCVzrWMnjgntLCve2G5I-R4pRar5s20kT9Rs,56
 sagemaker_core/tools/codegen.py,sha256=mKWVi2pWnPxyIoWUEPYjEc9Gw7D9bCOrHqa00yzIZ1o,2005
-sagemaker_core/tools/constants.py,sha256=lJ-KqRtv0F4mtZBqUNCX5RbOhXmPUsIslLPVVFJunDA,3080
-sagemaker_core/tools/data_extractor.py,sha256=NpP22uCaWgvn9uOHsFEkgtn3h2nOIKDvJTLihN01f1A,1368
+sagemaker_core/tools/constants.py,sha256=8oM0nHuzXWnOJTIdXsu50f9RsEz4rgmmZKWhLkrTP-s,3309
+sagemaker_core/tools/data_extractor.py,sha256=pNfmTA0NUA96IgfLrla7a36Qjc1NljbwgZYaOhouKqQ,2113
 sagemaker_core/tools/method.py,sha256=4Hmu4UWpiBgUTZljYdW1KIKDduDxf_nfhCyuWgLVMWI,717
-sagemaker_core/tools/resources_codegen.py,sha256=m7B2vrkfokLmzND9tCkty7EhVMk1FLevTd2LmJd_GKg,89974
-sagemaker_core/tools/resources_extractor.py,sha256=B9Twlf2vb7AkZUKLnxSdQZ_rNLPz0ROk9kuq-6uB1As,14526
+sagemaker_core/tools/resources_codegen.py,sha256=ShsWjABBkSpPlp3gByzpe36M6hfYfBofODlUynqVXTY,82405
+sagemaker_core/tools/resources_extractor.py,sha256=hN61ehZbPnhFW-2FIVDi7NsEz4rLvGr-WoglHQGfrug,14523
 sagemaker_core/tools/shapes_codegen.py,sha256=_ve959bwH8usZ6dPlpXxi2on9t0hLpcmhRWnaWHCWMQ,11745
 sagemaker_core/tools/shapes_extractor.py,sha256=4KjgDmhlPM4G1f1NeYbORKlXs1s7Q_sm_NK31S_ROQ0,11950
-sagemaker_core/tools/templates.py,sha256=IvL9ykLvOXDNCMnIZn30o-6Cmx4fw-BIJz3nl9LD4dc,24614
-sagemaker_core-1.0.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sagemaker_core-1.0.4.dist-info/METADATA,sha256=62vOPMIkSDkfR_ekA4qjnxJ96kC4hbzlhAtd2w52Qk0,4877
-sagemaker_core-1.0.4.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-sagemaker_core-1.0.4.dist-info/top_level.txt,sha256=R3GAZZ1zC5JxqdE_0x2Lu_WYi2Xfke7VsiP3L5zngfA,15
-sagemaker_core-1.0.4.dist-info/RECORD,,
+sagemaker_core/tools/templates.py,sha256=nze_A01EpegYUwoR_gRv2qBNKNFruBY8L3RiIX5lz3M,22458
+sagemaker_core-1.0.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sagemaker_core-1.0.6.dist-info/METADATA,sha256=yoth9VivBfunXlVEWY3wT575kCh-QM4zZn44LjSn3UQ,4877
+sagemaker_core-1.0.6.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+sagemaker_core-1.0.6.dist-info/top_level.txt,sha256=R3GAZZ1zC5JxqdE_0x2Lu_WYi2Xfke7VsiP3L5zngfA,15
+sagemaker_core-1.0.6.dist-info/RECORD,,

{sagemaker_core-1.0.4.dist-info → sagemaker_core-1.0.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any