change: implement local JsonGet function

Co-authored-by: Namrata Madan <[email protected]>

Author: 2 people

src/sagemaker/local/entities.py

@@ -625,7 +625,7 @@ def describe(self):
 
 
 class _LocalPipeline(object):
-    """Placeholder docstring"""
+    """Class representing a local SageMaker Pipeline"""
 
     _executions = {}
 
@@ -645,7 +645,7 @@ def __init__(
         self.last_modified_time = now_time
 
     def describe(self):
-        """Placeholder docstring"""
+        """Describe Pipeline"""
         response = {
             "PipelineArn": self.pipeline.name,
             "PipelineDefinition": self.pipeline.definition(),
@@ -659,7 +659,7 @@ def describe(self):
         return response
 
     def start(self, **kwargs):
-        """Placeholder docstring"""
+        """Start a pipeline execution. Returns a _LocalPipelineExecution object."""
         from sagemaker.local.pipeline import LocalPipelineExecutor
 
         execution_id = str(uuid4())
@@ -670,7 +670,7 @@ def start(self, **kwargs):
 
 
 class _LocalPipelineExecution(object):
-    """Placeholder docstring"""
+    """Class representing a local SageMaker pipeline execution."""
 
     def __init__(
         self,
@@ -693,7 +693,7 @@ def __init__(
         self.blockout_steps = {}
 
     def describe(self):
-        """Placeholder docstring"""
+        """Describe Pipeline Execution."""
         response = {
             "CreationTime": self.creation_time,
             "LastModifiedTime": self.creation_time,
@@ -708,8 +708,14 @@ def describe(self):
         return filtered_response
 
     def list_steps(self):
-        """Placeholder docstring"""
-        # TODO
+        """List pipeline execution steps."""
+        return {
+            "PipelineExecutionSteps": [
+                step.to_list_steps_response()
+                for step in self.step_execution.values()
+                if step.status is not None
+            ]
+        }
 
     def update_execution_success(self):
         """Mark execution as succeeded."""
@@ -730,8 +736,8 @@ def update_step_failure(self, step_name, failure_message):
         self.step_execution.get(step_name).update_step_failure(failure_message)
 
     def mark_step_executing(self, step_name):
-        """Update step's status to EXECUTING"""
-        self.step_execution.get(step_name).status = _LocalExecutionStatus.EXECUTING.value
+        """Update pipelines step's status to EXECUTING and start_time to now."""
+        self.step_execution.get(step_name).mark_step_executing()
 
     def _initialize_step_execution(self, steps):
         """Initialize step_execution dict."""
@@ -751,7 +757,9 @@ def _initialize_step_execution(self, steps):
                     "Step type {} is not supported in local mode.".format(step.step_type.value)
                 )
                 raise ClientError(error_msg, "start_pipeline_execution")
-            self.step_execution[step.name] = _LocalPipelineStepExecution(step.name, step.step_type)
+            self.step_execution[step.name] = _LocalPipelineExecutionStep(
+                step.name, step.step_type, step.description, step.display_name
+            )
             if step.step_type == StepTypeEnum.CONDITION:
                 self._initialize_step_execution(step.if_steps + step.else_steps)
 
@@ -790,44 +798,105 @@ def _construct_validation_exception_message(exception_msg):
         return {"Error": {"Code": "ValidationException", "Message": exception_msg}}
 
 
-class _LocalPipelineStepExecution(object):
-    """Placeholder docstring"""
+class _LocalPipelineExecutionStep(object):
+    """Class representing a local pipeline execution step."""
 
     def __init__(
         self,
-        step_name,
+        name,
         step_type,
-        last_modified_time=None,
+        description,
+        display_name=None,
+        start_time=None,
+        end_time=None,
         status=None,
         properties=None,
         failure_reason=None,
     ):
-        self.step_name = step_name
-        self.step_type = step_type
-        self.status = status or _LocalExecutionStatus.STARTING.value
+        from sagemaker.workflow.steps import StepTypeEnum
+
+        self.name = name
+        self.type = step_type
+        self.description = description
+        self.display_name = display_name
+        self.status = status
         self.failure_reason = failure_reason
         self.properties = properties or {}
-        self.creation_time = datetime.datetime.now()
-        self.last_modified_time = last_modified_time or self.creation_time
+        self.start_time = start_time
+        self.end_time = end_time
+        self._step_type_to_output_format_map = {
+            StepTypeEnum.TRAINING: self._construct_training_metadata,
+            StepTypeEnum.PROCESSING: self._construct_processing_metadata,
+            StepTypeEnum.TRANSFORM: self._construct_transform_metadata,
+            StepTypeEnum.CONDITION: self._construct_condition_metadata,
+            StepTypeEnum.FAIL: self._construct_fail_metadata,
+        }
 
     def update_step_properties(self, properties):
         """Update pipeline step execution output properties."""
-        logger.info("Successfully completed step %s.", self.step_name)
+        logger.info("Successfully completed step %s.", self.name)
         self.properties = deepcopy(properties)
         self.status = _LocalExecutionStatus.SUCCEEDED.value
+        self.end_time = datetime.datetime.now()
 
     def update_step_failure(self, failure_message):
         """Update pipeline step execution failure status and message."""
         logger.error(failure_message)
         self.failure_reason = failure_message
         self.status = _LocalExecutionStatus.FAILED.value
-        raise StepExecutionException(self.step_name, failure_message)
+        self.end_time = datetime.datetime.now()
+        raise StepExecutionException(self.name, failure_message)
+
+    def mark_step_executing(self):
+        """Update pipelines step's status to EXECUTING and start_time to now"""
+        self.status = _LocalExecutionStatus.EXECUTING.value
+        self.start_time = datetime.datetime.now()
+
+    def to_list_steps_response(self):
+        """Convert to response dict for list_steps calls."""
+        response = {
+            "EndTime": self.end_time,
+            "FailureReason": self.failure_reason,
+            "Metadata": self._construct_metadata(),
+            "StartTime": self.start_time,
+            "StepDescription": self.description,
+            "StepDisplayName": self.display_name,
+            "StepName": self.name,
+            "StepStatus": self.status,
+        }
+        filtered_response = {k: v for k, v in response.items() if v is not None}
+        return filtered_response
+
+    def _construct_metadata(self):
+        """Constructs the metadata shape for the list_steps_response."""
+        if self.properties:
+            return self._step_type_to_output_format_map[self.type]()
+        return None
+
+    def _construct_training_metadata(self):
+        """Construct training job metadata response."""
+        return {"TrainingJob": {"Arn": self.properties.TrainingJobArn}}
+
+    def _construct_processing_metadata(self):
+        """Construct processing job metadata response."""
+        return {"ProcessingJob": {"Arn": self.properties.ProcessingJobArn}}
+
+    def _construct_transform_metadata(self):
+        """Construct transform job metadata response."""
+        return {"TransformJob": {"Arn": self.properties.TransformJobArn}}
+
+    def _construct_condition_metadata(self):
+        """Construct condition step metadata response."""
+        return {"Condition": {"Outcome": self.properties.Outcome}}
+
+    def _construct_fail_metadata(self):
+        """Construct fail step metadata response."""
+        return {"Fail": {"ErrorMessage": self.properties.ErrorMessage}}
 
 
 class _LocalExecutionStatus(enum.Enum):
-    """Placeholder docstring"""
+    """Pipeline execution status."""
 
-    STARTING = "Starting"
     EXECUTING = "Executing"
     SUCCEEDED = "Succeeded"
     FAILED = "Failed"

src/sagemaker/local/local_session.py

@@ -595,7 +595,9 @@ class LocalSession(Session):
     :class:`~sagemaker.session.Session`.
     """
 
-    def __init__(self, boto_session=None, s3_endpoint_url=None, disable_local_code=False):
+    def __init__(
+        self, boto_session=None, default_bucket=None, s3_endpoint_url=None, disable_local_code=False
+    ):
         """Create a Local SageMaker Session.
 
         Args:
@@ -614,7 +616,7 @@ def __init__(self, boto_session=None, s3_endpoint_url=None, disable_local_code=F
         # discourage external use:
         self._disable_local_code = disable_local_code
 
-        super(LocalSession, self).__init__(boto_session)
+        super(LocalSession, self).__init__(boto_session=boto_session, default_bucket=default_bucket)
 
         if platform.system() == "Windows":
             logger.warning("Windows Support for Local Mode is Experimental")
@@ -718,9 +720,12 @@ def __init__(self, fileUri, content_type=None):
 class LocalPipelineSession(LocalSession):
     """Class representing a local session for SageMaker Pipelines executions."""
 
-    def __init__(self, boto_session=None, s3_endpoint_url=None, disable_local_code=False):
+    def __init__(
+        self, boto_session=None, default_bucket=None, s3_endpoint_url=None, disable_local_code=False
+    ):
         super().__init__(
             boto_session=boto_session,
+            default_bucket=default_bucket,
             s3_endpoint_url=s3_endpoint_url,
             disable_local_code=disable_local_code,
         )

src/sagemaker/local/pipeline.py

@@ -15,15 +15,17 @@
 from abc import ABC, abstractmethod
 
 import logging
+import json
 from copy import deepcopy
 from datetime import datetime
 from typing import Dict, List
-from sagemaker.workflow.conditions import ConditionTypeEnum
+from botocore.exceptions import ClientError
 
+from sagemaker.workflow.conditions import ConditionTypeEnum
 from sagemaker.workflow.steps import StepTypeEnum, Step
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.parameters import Parameter
-from sagemaker.workflow.functions import Join, JsonGet
+from sagemaker.workflow.functions import Join, JsonGet, PropertyFile
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.execution_variables import ExecutionVariable, ExecutionVariables
 from sagemaker.workflow.pipeline import PipelineGraph
@@ -116,8 +118,7 @@ def evaluate_pipeline_variable(self, pipeline_variable, step_name):
         elif isinstance(pipeline_variable, ExecutionVariable):
             value = self._evaluate_execution_variable(pipeline_variable)
         elif isinstance(pipeline_variable, JsonGet):
-            # TODO
-            raise NotImplementedError
+            value = self._evaluate_json_get_function(pipeline_variable, step_name)
         else:
             self.execution.update_step_failure(
                 step_name, f"Unrecognized pipeline variable {pipeline_variable.expr}."
@@ -133,7 +134,7 @@ def _evaluate_property_reference(self, pipeline_variable, step_name):
             referenced_step_name = pipeline_variable.step_name
             step_properties = self.execution.step_execution.get(referenced_step_name).properties
             return get_using_dot_notation(step_properties, pipeline_variable.path)
-        except (KeyError, IndexError):
+        except (KeyError, IndexError, TypeError):
             self.execution.update_step_failure(step_name, f"{pipeline_variable.expr} is undefined.")
 
     def _evaluate_execution_variable(self, pipeline_variable):
@@ -154,6 +155,56 @@ def _evaluate_execution_variable(self, pipeline_variable):
             return datetime.now()
         return None
 
+    def _evaluate_json_get_function(self, pipeline_variable, step_name):
+        """Evaluate join function runtime value."""
+        property_file_reference = pipeline_variable.property_file
+        property_file = None
+        if isinstance(property_file_reference, str):
+            processing_step = self.pipeline_dag.step_map[pipeline_variable.step_name]
+            for file in processing_step.property_files:
+                if file.name == property_file_reference:
+                    property_file = file
+                    break
+        elif isinstance(property_file_reference, PropertyFile):
+            property_file = property_file_reference
+        processing_step_response = self.execution.step_execution.get(
+            pipeline_variable.step_name
+        ).properties
+        if (
+            "ProcessingOutputConfig" not in processing_step_response
+            or "Outputs" not in processing_step_response["ProcessingOutputConfig"]
+        ):
+            self.execution.update_step_failure(
+                step_name,
+                f"Step '{pipeline_variable.step_name}' does not yet contain processing outputs.",
+            )
+        processing_output_s3_bucket = None
+        for output in processing_step_response["ProcessingOutputConfig"]["Outputs"]:
+            if output["OutputName"] == property_file.output_name:
+                processing_output_s3_bucket = output["S3Output"]["S3Uri"]
+                break
+        try:
+            file_content = self.sagemaker_session.read_s3_file(
+                processing_output_s3_bucket, property_file.path
+            )
+            file_json = json.loads(file_content)
+            return get_using_dot_notation(file_json, pipeline_variable.json_path)
+        except ClientError as e:
+            self.execution.update_step_failure(
+                step_name,
+                f"Received an error while file reading file '{property_file.path}' from S3: "
+                f"{e.response.get('Code')}: {e.response.get('Message')}",
+            )
+        except json.JSONDecodeError:
+            self.execution.update_step_failure(
+                step_name,
+                f"Contents of property file '{property_file.name}' are not in valid JSON format.",
+            )
+        except (KeyError, IndexError, TypeError):
+            self.execution.update_step_failure(
+                step_name, f"Invalid json path '{pipeline_variable.json_path}'"
+            )
+
 
 class _StepExecutor(ABC):
     """An abstract base class for step executors running steps locally"""

src/sagemaker/local/utils.py

@@ -166,7 +166,7 @@ def get_using_dot_notation(dictionary, keys):
         Nested object within dictionary as defined by "keys"
 
     Raises:
-     KeyError or IndexError if the provided key does not exist in input dictionary
+     KeyError/IndexError/TypeError if the provided key does not exist in input dictionary
     """
     if keys is None:
         return dictionary
@@ -175,14 +175,23 @@ def get_using_dot_notation(dictionary, keys):
     rest = None
     if len(split_keys) > 1:
         rest = split_keys[1]
-    list_accessor = re.search(r"(\w+)\[(\d+)]", key)
-    if list_accessor:
-        key = list_accessor.group(1)
-        list_index = int(list_accessor.group(2))
-        return get_using_dot_notation(dictionary[key][list_index], rest)
-    dict_accessor = re.search(r"(\w+)\[['\"](\S+)['\"]]", key)
-    if dict_accessor:
-        key = dict_accessor.group(1)
-        inner_key = dict_accessor.group(2)
-        return get_using_dot_notation(dictionary[key][inner_key], rest)
-    return get_using_dot_notation(dictionary[key], rest)
+    bracket_accessors = re.findall(r"\[(.+?)]", key)
+    if bracket_accessors:
+        pre_bracket_key = key.split("[", 1)[0]
+        inner_dict = dictionary[pre_bracket_key]
+    else:
+        inner_dict = dictionary[key]
+    for bracket_accessor in bracket_accessors:
+        if (
+            bracket_accessor.startswith("'")
+            and bracket_accessor.endswith("'")
+            or bracket_accessor.startswith('"')
+            and bracket_accessor.endswith('"')
+        ):
+            # key accessor
+            inner_key = bracket_accessor[1:-1]
+        else:
+            # list accessor
+            inner_key = int(bracket_accessor)
+        inner_dict = inner_dict[inner_key]
+    return get_using_dot_notation(inner_dict, rest)