Update service-2.json with latest public botocore service model

sagemaker-core/src/sagemaker/core/tools/codegen.py

@@ -44,8 +44,15 @@ def generate_code(
         service_json=service_json_data.sagemaker
     )
 
-    shapes_code_gen.generate_shapes()
-    reformat_file_with_black(".")
+    # Collect paths of all generated files, then reformat only those.
+    generated_files = []
+
+    generated_files.append(shapes_code_gen.generate_shapes())
+    generated_files.append(resources_code_gen.generate_resources())
+    generated_files.append(resources_code_gen.generate_config_schema())
+
+    for filepath in generated_files:
+        reformat_file_with_black(filepath)
 
 
 """

sagemaker-core/src/sagemaker/core/tools/constants.py

@@ -56,6 +56,7 @@
     "str": "string",
     "StrPipeVar": "string",
     "int": "integer",
+    "IntPipeVar": "integer",
     "bool": "boolean",
     "float": "double",
     "datetime.datetime": "timestamp",
@@ -108,6 +109,7 @@
 RESOURCES_CODEGEN_FILE_NAME = "resources.py"
 
 SHAPES_CODEGEN_FILE_NAME = "shapes.py"
+SHAPES_CODEGEN_OUTPUT_DIR = os.getcwd() + "/src/sagemaker/core/shapes"
 
 CONFIG_SCHEMA_FILE_NAME = "config_schema.py"
 
@@ -117,4 +119,18 @@
 # E.g. DescribeInferenceComponent returns empty ComputeResourceRequirements for adapter ICs.
 REQUIRED_TO_OPTIONAL_OVERRIDES = {
     "InferenceComponentComputeResourceRequirements": ["MinMemoryRequiredInMb"],
+    # ModelPackageName is not applicable to versioned model packages (group-based).
+    # ModelPackageSecurityConfig.KmsKeyId is absent when no KMS key is configured.
+    "DescribeModelPackageOutput": ["ModelPackageName"],
+    "ModelPackageSecurityConfig": ["KmsKeyId"],
+}
+
+# Members where the generated primitive type should be replaced with a PipelineVariable
+# Key: shape name, Value: dict of member name -> replacement type.
+PIPE_VAR_OVERRIDES = {
+    "ResourceConfig": {
+        "InstanceCount": "IntPipeVar",
+        "VolumeSizeInGB": "IntPipeVar",
+        "KeepAlivePeriodInSeconds": "IntPipeVar",
+    },
 }

sagemaker-core/src/sagemaker/core/tools/resources_codegen.py

@@ -247,13 +247,16 @@ def generate_resources(
         self,
         output_folder: str = GENERATED_CLASSES_LOCATION,
         file_name: str = RESOURCES_CODEGEN_FILE_NAME,
-    ) -> None:
+    ) -> str:
         """
         Generate the resources file.
 
         Args:
             output_folder (str, optional): The output folder path. Defaults to "GENERATED_CLASSES_LOCATION".
             file_name (str, optional): The output file name. Defaults to "RESOURCES_CODEGEN_FILE_NAME".
+
+        Returns:
+            str: The path to the generated output file.
         """
         # Check if the output folder exists, if not, create it
         os.makedirs(output_folder, exist_ok=True)
@@ -304,6 +307,8 @@ def generate_resources(
                 if resource_class:
                     file.write(f"{resource_class}\n\n")
 
+        return output_file
+
     def _evaluate_method(
         self, resource_name: str, method_name: str, methods: list, **kwargs
     ) -> str:
@@ -1934,13 +1939,15 @@ def generate_get_all_method(self, resource_name: str) -> str:
         )
         return formatted_method
 
-    def generate_config_schema(self):
+    def generate_config_schema(self) -> str:
         """
         Generates the Config Schema that is used by json Schema to validate config jsons .
         This function creates a python file with a variable that is consumed in the scripts to further fetch configs.
 
         Input for generating the Schema is the service JSON that is already loaded in the class
 
+        Returns:
+            str: The path to the generated output file.
         """
         resource_properties = {}
 
@@ -2001,6 +2008,8 @@ def generate_config_schema(self):
                 f"SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {json.dumps(combined_config_schema, indent=4)}"
             )
 
+        return output
+
     def _cleanup_class_attributes_types(self, class_attributes: dict) -> dict:
         """
         Helper function that creates a direct mapping of attribute to type without default parameters assigned and without Optionals
@@ -2030,7 +2039,7 @@ def _get_dict_with_default_configurable_attributes(self, class_attributes: dict)
             Dict with attributes that can be configurable
 
         """
-        PYTHON_TYPES = ["StrPipeVar", "datetime.datetime", "bool", "int", "float"]
+        PYTHON_TYPES = ["StrPipeVar", "IntPipeVar", "datetime.datetime", "bool", "int", "float"]
         default_attributes = {}
         for key, value in class_attributes.items():
             if value in PYTHON_TYPES or value.startswith("List"):

sagemaker-core/src/sagemaker/core/tools/resources_extractor.py

@@ -162,12 +162,18 @@ def _extract_resources_plan(self):
             [key[len("Import") :] for key in self.actions if key.startswith("Import")]
         )
 
+        # Resource names that collide with critical imports (e.g. boto3.session.Session)
+        EXCLUDED_RESOURCES = {"Session"}
+
         self.resources.update(
-            self.create_resources
-            | self.add_resources
-            | self.start_resources
-            | self.register_resources
-            | self.import_resources
+            (
+                self.create_resources
+                | self.add_resources
+                | self.start_resources
+                | self.register_resources
+                | self.import_resources
+            )
+            - EXCLUDED_RESOURCES
         )
 
         self._filter_actions_for_resources(self.resources)

sagemaker-core/src/sagemaker/core/tools/shapes_codegen.py

@@ -22,6 +22,7 @@
     LICENCES_STRING,
     GENERATED_CLASSES_LOCATION,
     SHAPES_CODEGEN_FILE_NAME,
+    SHAPES_CODEGEN_OUTPUT_DIR,
 )
 from sagemaker.core.tools.shapes_extractor import ShapesExtractor
 from sagemaker.core.utils.utils import (
@@ -209,7 +210,7 @@ def generate_imports(self):
         imports += "from pydantic import BaseModel, ConfigDict\n"
         imports += "from typing import List, Dict, Optional, Any, Union\n"
         imports += "from sagemaker.core.utils.utils import Unassigned\n"
-        imports += "from sagemaker.core.helper.pipeline_variable import StrPipeVar\n"
+        imports += "from sagemaker.core.helper.pipeline_variable import StrPipeVar, IntPipeVar\n"
         imports += "\n"
         imports += "# Suppress Pydantic warnings about field names shadowing parent attributes\n"
         imports += "warnings.filterwarnings('ignore', message='.*shadows an attribute.*')\n"
@@ -252,13 +253,14 @@ def _filter_input_output_shapes(self, shape):
 
     def generate_shapes(
         self,
-        output_folder=GENERATED_CLASSES_LOCATION,
+        output_folder=SHAPES_CODEGEN_OUTPUT_DIR,
         file_name=SHAPES_CODEGEN_FILE_NAME,
-    ) -> None:
+    ) -> str:
         """
         Generates the shape classes and writes them to the specified output folder.
 
         :param output_folder: The path to the output folder.
+        :return: The path to the generated output file.
         """
         # Check if the output folder exists, if not, create it
         os.makedirs(output_folder, exist_ok=True)
@@ -294,3 +296,5 @@ def generate_shapes(
                         # Generate and write data class for shape
                         shape_class = self.generate_data_class_for_shape(shape)
                         file.write(shape_class)
+
+        return output_file

sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py

@@ -19,6 +19,7 @@
 from sagemaker.core.tools.constants import (
     BASIC_JSON_TYPES_TO_PYTHON_TYPES,
     REQUIRED_TO_OPTIONAL_OVERRIDES,
+    PIPE_VAR_OVERRIDES,
     SHAPE_DAG_FILE_PATH,
 )
 from sagemaker.core.utils.utils import (
@@ -222,8 +223,7 @@ def generate_shape_members(self, shape, required_override=()):
         required_args = list(required_override) or shape_dict.get("required", [])
         # Remove members that are known to be optional despite the service model
         required_args = [
-            r for r in required_args
-            if r not in REQUIRED_TO_OPTIONAL_OVERRIDES.get(shape, [])
+            r for r in required_args if r not in REQUIRED_TO_OPTIONAL_OVERRIDES.get(shape, [])
         ]
         init_data_body = {}
         # bring the required members in front
@@ -242,7 +242,9 @@ def generate_shape_members(self, shape, required_override=()):
                     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]
+                    member_type = PIPE_VAR_OVERRIDES.get(shape, {}).get(
+                        member_name, BASIC_JSON_TYPES_TO_PYTHON_TYPES[member_shape_type]
+                    )
             else:
                 raise Exception("The Shape definition mush exist. The Json Data might be corrupt")
             member_name_snake_case = convert_to_snake_case(member_name)