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

sagemaker/core/lineage/visualizer.py

@@ -0,0 +1,346 @@
+# 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.
+"""This module contains functionality to display lineage data."""
+from __future__ import absolute_import
+
+import logging
+
+from typing import Optional, Any, Iterator
+
+import pandas as pd
+from pandas import DataFrame
+
+from sagemaker.core.lineage._api_types import AssociationSummary
+from sagemaker.core.lineage.association import Association
+
+
+class LineageTableVisualizer(object):
+    """Creates a dataframe containing the lineage assoociations of a SageMaker object."""
+
+    def __init__(self, sagemaker_session):
+        """Init for LineageTableVisualizer.
+
+        Args:
+            sagemaker_session (obj): The sagemaker session used for API requests.
+        """
+        self._session = sagemaker_session
+
+    def show(
+        self,
+        trial_component_name: Optional[str] = None,
+        training_job_name: Optional[str] = None,
+        processing_job_name: Optional[str] = None,
+        pipeline_execution_step: Optional[object] = None,
+        model_package_arn: Optional[str] = None,
+        endpoint_arn: Optional[str] = None,
+        artifact_arn: Optional[str] = None,
+        context_arn: Optional[str] = None,
+        actions_arn: Optional[str] = None,
+    ) -> DataFrame:
+        """Generate a dataframe containing all incoming and outgoing lineage entities.
+
+          Examples:
+          .. code-block:: python
+
+              viz = LineageTableVisualizer(sagemaker_session)
+              df = viz.show(training_job_name=training_job_name)
+              # in a notebook
+              display(df.to_html())
+
+        Args:
+            trial_component_name (str, optional): Name of  a trial component. Defaults to None.
+            training_job_name (str, optional): Name of a training job. Defaults to None.
+            processing_job_name (str, optional): Name of a processing job. Defaults to None.
+            pipeline_execution_step (obj, optional): Pipeline execution step. Defaults to None.
+            model_package_arn (str, optional): Model package arn. Defaults to None.
+            endpoint_arn (str, optional): Endpoint arn. Defaults to None.
+            artifact_arn (str, optional): Artifact arn. Defaults to None.
+            context_arn (str, optional): Context arn. Defaults to None.
+            actions_arn (str, optional): Action arn. Defaults to None.
+
+        Returns:
+            DataFrame: Pandas dataframe containing lineage associations.
+        """
+        start_arn: str = None
+
+        if trial_component_name:
+            start_arn = self._get_start_arn_from_trial_component_name(trial_component_name)
+        elif training_job_name:
+            trial_component_name = training_job_name + "-aws-training-job"
+            start_arn = self._get_start_arn_from_trial_component_name(trial_component_name)
+        elif processing_job_name:
+            trial_component_name = processing_job_name + "-aws-processing-job"
+            start_arn = self._get_start_arn_from_trial_component_name(trial_component_name)
+        elif pipeline_execution_step:
+            start_arn = self._get_start_arn_from_pipeline_execution_step(pipeline_execution_step)
+        elif model_package_arn:
+            start_arn = self._get_start_arn_from_model_package_arn(model_package_arn)
+        elif endpoint_arn:
+            start_arn = self._get_start_arn_from_endpoint_arn(endpoint_arn)
+        elif artifact_arn:
+            start_arn = artifact_arn
+        elif context_arn:
+            start_arn = context_arn
+        elif actions_arn:
+            start_arn = actions_arn
+
+        return self._get_associations_dataframe(start_arn)
+
+    def _get_start_arn_from_pipeline_execution_step(self, pipeline_execution_step: object) -> str:
+        """Given a pipeline exection step retrieve the arn of the lineage entity that represents it.
+
+        Args:
+            pipeline_execution_step (obj): Pipeline execution step.
+
+        Returns:
+            str: The arn of the lineage entity
+        """
+        start_arn: str = None
+
+        if not pipeline_execution_step["Metadata"]:
+            return None
+
+        metadata: Any = pipeline_execution_step["Metadata"]
+        jobs: list = ["TrainingJob", "ProcessingJob", "TransformJob"]
+
+        for job in jobs:
+            if job in metadata and metadata[job]:
+                job_arn = metadata[job]["Arn"]
+                start_arn = self._get_start_arn_from_job_arn(job_arn)
+                break
+
+        if "RegisterModel" in metadata:
+            start_arn = self._get_start_arn_from_model_package_arn(metadata["RegisterModel"]["Arn"])
+
+        return start_arn
+
+    def _get_start_arn_from_job_arn(self, job_arn: str) -> str:
+        """Given a job arn return the lineage entity.
+
+        Args:
+            job_arn (str): Arn of a training, processing, or transform job.
+
+        Returns:
+          str: The arn of the job's lineage entity.
+        """
+        start_arn: str = None
+        response: Any = self._session.sagemaker_client.list_trial_components(SourceArn=job_arn)
+        trial_components: Any = response["TrialComponentSummaries"]
+        if trial_components:
+            start_arn = trial_components[0]["TrialComponentArn"]
+        else:
+            logging.warning("No trial components found for %s", job_arn)
+        return start_arn
+
+    def _get_associations_dataframe(self, arn: str) -> DataFrame:
+        """Create a data frame containing lineage association information.
+
+        Args:
+            arn (str): The arn of the lineage entity of interest.
+
+        Returns:
+            DataFrame: A dataframe with association information.
+        """
+        if arn is None:
+            # no associations
+            return None
+
+        upstream_associations: Iterator[AssociationSummary] = self._get_associations(dest_arn=arn)
+        downstream_associations: Iterator[AssociationSummary] = self._get_associations(src_arn=arn)
+        inputs: list = list(map(self._convert_input_association_to_df_row, upstream_associations))
+        outputs: list = list(
+            map(self._convert_output_association_to_df_row, downstream_associations)
+        )
+        df: DataFrame = pd.DataFrame(
+            inputs + outputs,
+            columns=[
+                "Name/Source",
+                "Direction",
+                "Type",
+                "Association Type",
+                "Lineage Type",
+            ],
+        )
+        return df
+
+    def _get_start_arn_from_trial_component_name(self, tc_name: str) -> str:
+        """Given a trial component name retrieve a start arn.
+
+        Args:
+            tc_name (str): Name of the trial compoonent.
+
+        Returns:
+          str: The arn of the trial component.
+        """
+        response: Any = self._session.sagemaker_client.describe_trial_component(
+            TrialComponentName=tc_name
+        )
+        tc_arn: str = response["TrialComponentArn"]
+        return tc_arn
+
+    def _get_start_arn_from_model_package_arn(self, model_package_arn: str) -> str:
+        """Given a model package arn retrieve the arn lineage entity.
+
+        Args:
+            model_package_arn (str): The arn of a model package.
+
+        Returns:
+            str: The arn of the lineage entity that represents the model package.
+        """
+        response: Any = self._session.sagemaker_client.list_artifacts(SourceUri=model_package_arn)
+        artifacts: Any = response["ArtifactSummaries"]
+        artifact_arn: str = None
+        if artifacts:
+            artifact_arn = artifacts[0]["ArtifactArn"]
+        else:
+            logging.debug("No artifacts found for %s.", model_package_arn)
+        return artifact_arn
+
+    def _get_start_arn_from_endpoint_arn(self, endpoint_arn: str) -> str:
+        """Given an endpoint arn retrieve the arn of the lineage entity.
+
+        Args:
+            endpoint_arn (str): The arn of an endpoint
+
+        Returns:
+            str: The arn of the lineage entity that represents the model package.
+        """
+        response: Any = self._session.sagemaker_client.list_contexts(SourceUri=endpoint_arn)
+        contexts: Any = response["ContextSummaries"]
+        context_arn: str = None
+        if contexts:
+            context_arn = contexts[0]["ContextArn"]
+        else:
+            logging.debug("No contexts found for %s.", endpoint_arn)
+        return context_arn
+
+    def _get_associations(
+        self, src_arn: Optional[str] = None, dest_arn: Optional[str] = None
+    ) -> Iterator[AssociationSummary]:
+        """Given an arn retrieve all associated lineage entities.
+
+        The arn must be one of: experiment, trial, trial component, artifact, action, or context.
+
+        Args:
+            src_arn (str, optional): The arn of the source. Defaults to None.
+            dest_arn (str, optional): The arn of the destination. Defaults to None.
+
+        Returns:
+            array: An array of associations that are either incoming or outgoing from the lineage
+            entity of interest.
+        """
+        if src_arn:
+            associations: Iterator[AssociationSummary] = Association.list(
+                source_arn=src_arn, sagemaker_session=self._session
+            )
+        else:
+            associations: Iterator[AssociationSummary] = Association.list(
+                destination_arn=dest_arn, sagemaker_session=self._session
+            )
+        return associations
+
+    def _convert_input_association_to_df_row(self, association) -> list:
+        """Convert an input association to a data frame row.
+
+        Args:
+            association (obj): ``Association``
+
+        Returns:
+            array: Array of column values for the association data frame.
+        """
+        return self._convert_association_to_df_row(
+            association.source_arn,
+            association.source_name,
+            "Input",
+            association.source_type,
+            association.association_type,
+        )
+
+    def _convert_output_association_to_df_row(self, association) -> list:
+        """Convert an output association to a data frame row.
+
+        Args:
+            association (obj): ``Association``
+
+        Returns:
+            array: Array of column values for the association data frame.
+        """
+        return self._convert_association_to_df_row(
+            association.destination_arn,
+            association.destination_name,
+            "Output",
+            association.destination_type,
+            association.association_type,
+        )
+
+    def _convert_association_to_df_row(
+        self,
+        arn: str,
+        name: str,
+        direction: str,
+        src_dest_type: str,
+        association_type: type,
+    ) -> list:
+        """Convert association data into a data frame row.
+
+        Args:
+            arn (str): The arn of the associated entity.
+            name (str): The name of the associated entity.
+            direction (str): The direction the association is with the entity of interest. Values
+            are 'Input' or 'Output'.
+            src_dest_type (str): The type of the entity that is associated with the entity of
+            interest.
+            association_type ([type]): The type of the association.
+
+        Returns:
+            [type]: [description]
+        """
+        arn_name = arn.split(":")[5]
+        entity_type = arn_name.split("/")[0]
+        name = self._get_friendly_name(name, arn, entity_type)
+        return [name, direction, src_dest_type, association_type, entity_type]
+
+    def _get_friendly_name(self, name: str, arn: str, entity_type: str) -> str:
+        """Get a human readable name from the association.
+
+        Args:
+            name (str): The name of the associated entity
+            arn (str): The arn of the associated entity
+            entity_type (str): The type of the associated entity (artifact, action, etc...)
+
+        Returns:
+            str: The name for the association that will be displayed in the data frame.
+        """
+        if name:
+            return name
+
+        if entity_type == "artifact":
+            artifact = self._session.sagemaker_client.describe_artifact(ArtifactArn=arn)
+            uri = artifact["Source"]["SourceUri"]
+
+            # shorten the uri if the length is more than 40,
+            # e.g s3://flintstone-end-to-end-tests-gamma-us-west-2-069083975568/results/
+            # canary-auto-1608761252626/preprocessed-data/tuning_data/train.txt
+            # become s3://.../preprocessed-data/tuning_data/train.txt
+            if len(uri) > 48:
+                name = uri[:5] + "..." + uri[len(uri) - 40 :]
+
+            # if not then use the full uri
+            if not name:
+                name = uri
+
+        # if still don't have name derive from arn
+        if not name:
+            name = arn.split(":")[5].split("/")[1]
+
+        return name

sagemaker/core/local/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+"""Local mode utilities for SageMaker."""
+from __future__ import absolute_import
+
+from sagemaker.core.local.local_session import LocalSession  # noqa: F401
+
+__all__ = ["LocalSession"]