sagemaker-core 0.1.3__py3-none-any.whl

sagemaker_core/tools/resources_extractor.py

@@ -0,0 +1,373 @@
+# 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.
+"""A class for extracting resource information from a service JSON."""
+import logging
+from typing import Optional
+
+import pandas as pd
+
+from sagemaker_core.tools.constants import CLASS_METHODS, OBJECT_METHODS
+from sagemaker_core.tools.data_extractor import (
+    load_additional_operations_data,
+    load_combined_operations_data,
+    load_combined_shapes_data,
+)
+from sagemaker_core.tools.method import Method
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+"""
+This class is used to extract the resources and its actions from the service-2.json file.
+"""
+
+
+class ResourcesExtractor:
+    """
+    A class for extracting resource information from a service JSON.
+
+    Args:
+        service_json (dict): The Botocore service.json containing the shape definitions.
+
+    Attributes:
+        service_json (dict): The service JSON containing operations and shapes.
+        operations (dict): The operations defined in the service JSON.
+        shapes (dict): The shapes defined in the service JSON.
+        resource_actions (dict): A dictionary mapping resources to their associated actions.
+        actions_under_resource (set): A set of actions that are performed on resources.
+        create_resources (set): A set of resources that can be created.
+        add_resources (set): A set of resources that can be added.
+        start_resources (set): A set of resources that can be started.
+        register_resources (set): A set of resources that can be registered.
+        import_resources (set): A set of resources that can be imported.
+        resources (set): A set of all resources.
+        df (DataFrame): A DataFrame containing resource information.
+
+    Methods:
+        _filter_actions_for_resources(resources): Filters actions based on the given resources.
+        _extract_resources_plan(): Extracts the resource plan from the service JSON.
+        _get_status_chain_and_states(shape_name, status_chain): Recursively extracts the status chain and states for a given shape.
+        _extract_resource_plan_as_dataframe(): Builds a DataFrame containing resource information.
+        get_resource_plan(): Returns the resource plan DataFrame.
+    """
+
+    RESOURCE_TO_ADDITIONAL_METHODS = {
+        "Cluster": ["DescribeClusterNode", "ListClusterNodes"],
+    }
+
+    def __init__(
+        self,
+        combined_shapes: Optional[dict] = None,
+        combined_operations: Optional[dict] = None,
+    ):
+        """
+        Initializes a ResourceExtractor object.
+
+        Args:
+            service_json (dict): The service JSON containing operations and shapes.
+        """
+        self.operations = combined_operations or load_combined_operations_data()
+        self.shapes = combined_shapes or load_combined_shapes_data()
+        self.additional_operations = load_additional_operations_data()
+        # contains information about additional methods only now.
+        # TODO: replace resource_actions with resource_methods to include all methods
+        self.resource_methods = {}
+        self.resource_actions = {}
+        self.actions_under_resource = set()
+
+        self._extract_resources_plan()
+
+    def _filter_additional_operations(self):
+        """
+        Extracts information from additional operations defined in additional_operations.json
+
+        Returns:
+            None
+        """
+        for resource_name, resource_operations in self.additional_operations.items():
+            self.resources.add(resource_name)
+            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)
+                method = Method(**operation)
+                method.get_docstring_title(self.operations[operation_name])
+                self.resource_methods[resource_name][operation["method_name"]] = method
+                self.actions.remove(operation_name)
+
+    def _filter_actions_for_resources(self, resources):
+        """
+        Filters actions based on the given resources.
+
+        Args:
+            resources (set): A set of resources.
+
+        Returns:
+            None
+        """
+        for resource in sorted(resources, key=len, reverse=True):
+            filtered_actions = set(
+                [
+                    a
+                    for a in self.actions
+                    if a.endswith(resource)
+                    or (a.startswith("List") and a.endswith(resource + "s"))
+                    or a.startswith("Invoke" + resource)
+                ]
+            )
+            self.actions_under_resource.update(filtered_actions)
+            self.resource_actions[resource] = filtered_actions
+
+            self.actions = self.actions - filtered_actions
+
+    def _extract_resources_plan(self):
+        """
+        Extracts the resource plan from the service JSON.
+
+        Returns:
+            None
+        """
+        self.actions = set(self.operations.keys())
+        log.info(f"Total actions - {len(self.actions)}")
+
+        # Filter out additional operations and resources first
+        self.resources = set()
+        self._filter_additional_operations()
+
+        self.create_resources = set(
+            [key[len("Create") :] for key in self.actions if key.startswith("Create")]
+        )
+
+        self.add_resources = set(
+            [key[len("Add") :] for key in self.actions if key.startswith("Add")]
+        )
+
+        self.start_resources = set(
+            [key[len("Start") :] for key in self.actions if key.startswith("Start")]
+        )
+
+        self.register_resources = set(
+            [key[len("Register") :] for key in self.actions if key.startswith("Register")]
+        )
+
+        self.import_resources = set(
+            [key[len("Import") :] for key in self.actions if key.startswith("Import")]
+        )
+
+        self.resources.update(
+            self.create_resources
+            | self.add_resources
+            | self.start_resources
+            | self.register_resources
+            | self.import_resources
+        )
+
+        self._filter_actions_for_resources(self.resources)
+
+        log.info(f"Total resource - {len(self.resources)}")
+
+        log.info(f"Total actions_under_resource - {len(self.actions_under_resource)}")
+
+        log.info(f"Unsupported actions: - {len(self.actions)}")
+
+        self._extract_resource_plan_as_dataframe()
+
+    def get_status_chain_and_states(self, resource_name):
+        """
+        Extract the status chain and states for a given resource.
+
+        Args:
+            resource_name (str): The name of the resource
+
+        Returns:
+            status_chain (list): The status chain for the resource.
+            resource_states (list): The states associated with the resource.
+        """
+        resource_operation = self.operations["Describe" + resource_name]
+        resource_operation_output_shape_name = resource_operation["output"]["shape"]
+        output_members_data = self.shapes[resource_operation_output_shape_name]["members"]
+        if len(output_members_data) == 1:
+            single_member_name = next(iter(output_members_data))
+            single_member_shape_name = output_members_data[single_member_name]["shape"]
+            status_chain = []
+            status_chain.append(
+                {"name": single_member_name, "shape_name": single_member_shape_name}
+            )
+            resource_status_chain, resource_states = self._get_status_chain_and_states(
+                single_member_shape_name, status_chain
+            )
+        else:
+            resource_status_chain, resource_states = self._get_status_chain_and_states(
+                resource_operation_output_shape_name
+            )
+
+        return resource_status_chain, resource_states
+
+    def _get_status_chain_and_states(self, shape_name, status_chain: list = None):
+        """
+        Recursively extracts the status chain and states for a given shape.
+
+        Args:
+            shape_name (str): The name of the shape.
+            status_chain (list): The current status chain.
+
+        Returns:
+            status_chain (list): The status chain for the shape.
+            resource_states (list): The states associated with the shape.
+        """
+        if status_chain is None:
+            status_chain = []
+
+        member_data = self.shapes[shape_name]["members"]
+        status_name = next((member for member in member_data if "status" in member.lower()), None)
+        if status_name is None:
+            return [], []
+
+        status_shape_name = member_data[status_name]["shape"]
+
+        status_chain.append({"name": status_name, "shape_name": status_shape_name})
+
+        if "enum" in self.shapes[status_shape_name]:
+            resource_states = self.shapes[status_shape_name]["enum"]
+            return status_chain, resource_states
+        else:
+            status_chain, resource_states = self._get_status_chain_and_states(
+                status_shape_name, status_chain
+            )
+            return status_chain, resource_states
+
+    def _extract_resource_plan_as_dataframe(self):
+        """
+        Builds a DataFrame containing resource information.
+
+        Returns:
+            None
+        """
+        self.df = pd.DataFrame(
+            columns=[
+                "resource_name",
+                "type",
+                "class_methods",
+                "object_methods",
+                "chain_resource_name",
+                "additional_methods",
+                "raw_actions",
+                "resource_status_chain",
+                "resource_states",
+            ]
+        )
+
+        for resource, actions in sorted(self.resource_actions.items()):
+            class_methods = set()
+            object_methods = set()
+            additional_methods = set()
+            chain_resource_names = set()
+            resource_status_chain = set()
+            resource_states = set()
+
+            for action in actions:
+                action_low = action.lower()
+                resource_low = resource.lower()
+
+                if action_low.split(resource_low)[0] == "describe":
+                    class_methods.add("get")
+                    object_methods.add("refresh")
+
+                    output_shape_name = self.operations[action]["output"]["shape"]
+                    output_members_data = self.shapes[output_shape_name]["members"]
+
+                    resource_status_chain, resource_states = self.get_status_chain_and_states(
+                        resource
+                    )
+
+                    if resource_low.endswith("job") or resource_low.endswith("jobv2"):
+                        object_methods.add("wait")
+                    elif resource_states and resource_low != "action":
+                        object_methods.add("wait_for_status")
+
+                    if "Deleting" in resource_states or "DELETING" in resource_states:
+                        object_methods.add("wait_for_delete")
+
+                    continue
+
+                if action_low.split(resource_low)[0] == "create":
+                    shape_name = self.operations[action]["input"]["shape"]
+                    input = self.shapes[shape_name]
+                    for member in input["members"]:
+                        if member.endswith("Name") or member.endswith("Names"):
+                            chain_resource_name = member[: -len("Name")]
+
+                            if (
+                                chain_resource_name != resource
+                                and chain_resource_name in self.resources
+                            ):
+                                chain_resource_names.add(chain_resource_name)
+                action_split = action_low.split(resource_low)
+                if action_split[0] == "invoke":
+                    if not action_split[1]:
+                        invoke_method = "invoke"
+                    elif action_split[1] == "async":
+                        invoke_method = "invoke_async"
+                    else:
+                        invoke_method = "invoke_with_response_stream"
+                    object_methods.add(invoke_method)
+                elif action_split[0] in CLASS_METHODS:
+                    if action_low.split(resource_low)[0] == "list":
+                        class_methods.add("get_all")
+                    else:
+                        class_methods.add(action_low.split(resource_low)[0])
+                elif action_split[0] in OBJECT_METHODS:
+                    object_methods.add(action_split[0])
+                else:
+                    additional_methods.add(action)
+
+            if resource in self.RESOURCE_TO_ADDITIONAL_METHODS:
+                additional_methods.update(self.RESOURCE_TO_ADDITIONAL_METHODS[resource])
+
+            new_row = pd.DataFrame(
+                {
+                    "resource_name": [resource],
+                    "type": ["resource"],
+                    "class_methods": [list(sorted(class_methods))],
+                    "object_methods": [list(sorted(object_methods))],
+                    "chain_resource_name": [list(sorted(chain_resource_names))],
+                    "additional_methods": [list(sorted(additional_methods))],
+                    "raw_actions": [list(sorted(actions))],
+                    "resource_status_chain": [list(resource_status_chain)],
+                    "resource_states": [list(resource_states)],
+                }
+            )
+
+            self.df = pd.concat([self.df, new_row], ignore_index=True)
+
+        self.df.to_csv("resource_plan.csv", index=False)
+
+    def get_resource_plan(self):
+        """
+        Returns the resource plan DataFrame.
+
+        Returns:
+            df (DataFrame): The resource plan DataFrame.
+        """
+        return self.df
+
+    def get_resource_methods(self):
+        """
+        Returns the resource methods dict.
+
+        Returns:
+            resource_methods (dict): The resource methods dict.
+        """
+        return self.resource_methods
+
+
+resource_extractor = ResourcesExtractor()

sagemaker_core/tools/shapes_codegen.py

@@ -0,0 +1,284 @@
+# 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.
+"""A class for generating class structure from Service Model JSON.
+
+To run the script be sure to set the PYTHONPATH
+export PYTHONPATH=<sagemaker-code-gen repo directory>:$PYTHONPATH
+"""
+import os
+
+from sagemaker_core.code_injection.codec import pascal_to_snake
+from sagemaker_core.tools.constants import (
+    LICENCES_STRING,
+    GENERATED_CLASSES_LOCATION,
+    SHAPES_CODEGEN_FILE_NAME,
+)
+from sagemaker_core.tools.shapes_extractor import ShapesExtractor
+from sagemaker_core.util.util import add_indent, convert_to_snake_case, remove_html_tags
+from sagemaker_core.tools.templates import SHAPE_CLASS_TEMPLATE, SHAPE_BASE_CLASS_TEMPLATE
+from sagemaker_core.tools.data_extractor import (
+    load_combined_shapes_data,
+    load_combined_operations_data,
+)
+from .resources_extractor import ResourcesExtractor
+
+
+class ShapesCodeGen:
+    """
+    Generates shape classes based on an input Botocore service.json.
+
+    Args:
+        service_json (dict): The Botocore service.json containing the shape definitions.
+
+    Attributes:
+        service_json (dict): The Botocore service.json containing the shape definitions.
+        shapes_extractor (ShapesExtractor): An instance of the ShapesExtractor class.
+        shape_dag (dict): Shape DAG generated from service.json
+
+    Methods:
+        build_graph(): Builds a directed acyclic graph (DAG) representing the dependencies between shapes.
+        topological_sort(): Performs a topological sort on the DAG to determine the order in which shapes should be generated.
+        generate_data_class_for_shape(shape): Generates a data class for a given shape.
+        _generate_doc_string_for_shape(shape): Generates the docstring for a given shape.
+        generate_imports(): Generates the import statements for the generated shape classes.
+        generate_base_class(): Generates the base class for the shape classes.
+        _filter_input_output_shapes(shape): Filters out shapes that are used as input or output for operations.
+        generate_shapes(output_folder): Generates the shape classes and writes them to the specified output folder.
+    """
+
+    def __init__(self):
+        self.combined_shapes = load_combined_shapes_data()
+        self.combined_operations = load_combined_operations_data()
+        self.shapes_extractor = ShapesExtractor()
+        self.shape_dag = self.shapes_extractor.get_shapes_dag()
+        self.resources_extractor = ResourcesExtractor()
+        self.resources_plan = self.resources_extractor.get_resource_plan()
+        self.resource_methods = self.resources_extractor.get_resource_methods()
+
+    def build_graph(self):
+        """
+        Builds a directed acyclic graph (DAG) representing the dependencies between shapes.
+
+        Steps:
+        1. Loop over the Service Json shapes.
+            1.1. If dependency(members) found, add association of node -> dependency.
+                1.1.1. Sometimes members are not shape themselves, but have associated links to actual shapes.
+                    In that case add link to node -> dependency (actual)
+                        CreateExperimentRequest -> [ExperimentEntityName, ExperimentDescription, TagList]
+            1.2. else leaf node found (no dependent members), add association of node -> None.
+
+        :return: A dict which defines the structure of the DAG in the format:
+            {key : [dependencies]}
+            Example input:
+                {'CreateExperimentRequest': ['ExperimentEntityName', 'ExperimentEntityName',
+                    'ExperimentDescription', 'TagList'],
+                'CreateExperimentResponse': ['ExperimentArn'],
+                'DeleteExperimentRequest': ['ExperimentEntityName'],
+                'DeleteExperimentResponse': ['ExperimentArn']}
+        """
+        graph = {}
+
+        for node, attributes in self.combined_shapes.items():
+            if "members" in attributes:
+                for member, member_attributes in attributes["members"].items():
+                    # add shapes and not shape attribute
+                    # i.e. ExperimentEntityName taken over ExperimentName
+                    if member_attributes["shape"] in self.combined_shapes.keys():
+                        node_deps = graph.get(node, [])
+                        # evaluate the member shape and then append to node deps
+                        member_shape = self.combined_shapes[member_attributes["shape"]]
+                        if member_shape["type"] == "list":
+                            node_deps.append(member_shape["member"]["shape"])
+                        elif member_shape["type"] == "map":
+                            node_deps.append(member_shape["key"]["shape"])
+                            node_deps.append(member_shape["value"]["shape"])
+                        else:
+                            node_deps.append(member_attributes["shape"])
+                        graph[node] = node_deps
+            else:
+                graph[node] = None
+        return graph
+
+    def topological_sort(self):
+        """
+        Performs a topological sort on the DAG to determine the order in which shapes should be generated.
+
+        :return: A list of shape names in the order of topological sort.
+        """
+        graph = self.build_graph()
+        visited = set()
+        stack = []
+
+        def dfs(node):
+            visited.add(node)
+            # unless leaf node is reached do dfs
+            if graph.get(node) is not None:
+                for neighbor in graph.get(node, []):
+                    if neighbor not in visited:
+                        dfs(neighbor)
+            stack.append(node)
+
+        for node in graph:
+            if node not in visited:
+                dfs(node)
+
+        return stack
+
+    def generate_data_class_for_shape(self, shape):
+        """
+        Generates a data class for a given shape.
+
+        :param shape: The name of the shape.
+        :return: The generated data class as a string.
+        """
+        class_name = shape
+        init_data = self.shapes_extractor.generate_data_shape_string_body(
+            shape, self.resources_plan
+        )
+        try:
+            data_class_members = add_indent(init_data, 4)
+        except Exception:
+            print("DEBUG HELP\n", init_data)
+            raise
+        return SHAPE_CLASS_TEMPLATE.format(
+            class_name=class_name + "(Base)",
+            data_class_members=data_class_members,
+            docstring=self._generate_doc_string_for_shape(shape),
+            class_name_snake=pascal_to_snake(class_name),
+        )
+
+    def _generate_doc_string_for_shape(self, shape):
+        """
+        Generates the docstring for a given shape.
+
+        :param shape: The name of the shape.
+        :return: The generated docstring as a string.
+        """
+        shape_dict = self.combined_shapes[shape]
+
+        docstring = f" {shape}"
+        if "documentation" in shape_dict:
+            docstring += f"\n \t  {shape_dict['documentation']}"
+
+        docstring += "\n\n \t Attributes"
+        docstring += "\n\t----------------------"
+
+        if "members" in shape_dict:
+            for member, member_attributes in shape_dict["members"].items():
+                docstring += f"\n \t{convert_to_snake_case(member)}"
+                if "documentation" in member_attributes:
+                    docstring += f": \t {member_attributes['documentation']}"
+
+        return remove_html_tags(docstring)
+
+    def generate_license(self):
+        """
+        Generates the license string.
+
+        Returns:
+            str: The license string.
+        """
+        return LICENCES_STRING
+
+    def generate_imports(self):
+        """
+        Generates the import statements for the generated shape classes.
+
+        :return: The generated import statements as a string.
+        """
+        imports = "import datetime\n"
+        imports += "\n"
+        imports += "from pydantic import BaseModel, ConfigDict\n"
+        imports += "from typing import List, Dict, Optional, Any, Union\n"
+        imports += "from sagemaker_core.generated.utils import Unassigned"
+        imports += "\n"
+        return imports
+
+    def generate_base_class(self):
+        """
+        Generates the base class for the shape classes.
+
+        :return: The generated base class as a string.
+        """
+        # more customizations would be added later
+        return SHAPE_BASE_CLASS_TEMPLATE.format(
+            class_name="Base(BaseModel)",
+        )
+
+    def _filter_input_output_shapes(self, shape):
+        """
+        Filters out shapes that are used as input or output for operations.
+
+        :param shape: The name of the shape.
+        :return: True if the shape should be generated, False otherwise.
+        """
+        operation_input_output_shapes = set()
+        for operation, attrs in self.combined_operations.items():
+            if attrs.get("input"):
+                operation_input_output_shapes.add(attrs["input"]["shape"])
+            if attrs.get("output"):
+                operation_input_output_shapes.add(attrs["output"]["shape"])
+
+        required_output_shapes = set()
+        for resource_name in self.resource_methods:
+            for method in self.resource_methods[resource_name].values():
+                required_output_shapes.add(method.return_type)
+
+        if shape in operation_input_output_shapes and shape not in required_output_shapes:
+            return False
+        return True
+
+    def generate_shapes(
+        self,
+        output_folder=GENERATED_CLASSES_LOCATION,
+        file_name=SHAPES_CODEGEN_FILE_NAME,
+    ) -> None:
+        """
+        Generates the shape classes and writes them to the specified output folder.
+
+        :param output_folder: The path to the output folder.
+        """
+        # Check if the output folder exists, if not, create it
+        os.makedirs(output_folder, exist_ok=True)
+
+        # Create the full path for the output file
+        output_file = os.path.join(output_folder, file_name)
+
+        # Open the output file
+        with open(output_file, "w") as file:
+            # Generate and write the license to the file
+            license = self.generate_license()
+            file.write(license)
+
+            # Generate and write the imports to the file
+            imports = self.generate_imports()
+            file.write(imports)
+
+            # Generate and write Base Class
+            base_class = self.generate_base_class()
+            file.write(base_class)
+            file.write("\n\n")
+
+            # Iterate through shapes in topological order and generate classes
+            topological_order = self.topological_sort()
+            for shape in topological_order:
+
+                # Extract the necessary data for the shape
+                if self._filter_input_output_shapes(shape):
+                    shape_dict = self.combined_shapes[shape]
+                    shape_type = shape_dict["type"]
+                    if shape_type == "structure":
+
+                        # Generate and write data class for shape
+                        shape_class = self.generate_data_class_for_shape(shape)
+                        file.write(shape_class)