Improved abstraction on Processor objects when handling local session. Other minor changes.

Author: EC2 Default User

src/sagemaker/local/entities.py

@@ -75,41 +75,30 @@ def start(self, processing_inputs, processing_output_config, environment, proces
             if "DatasetDefinition" in item:
                 raise RuntimeError("DatasetDefinition is not currently supported in Local Mode")
 
-            if "S3Input" in item and item["S3Input"]:
-                data_uri = item["S3Input"]["S3Uri"]
-            else:
+            try:
+                s3_input = item["S3Input"]
+            except KeyError:
                 raise ValueError("Processing input must have a valid ['S3Input']")
 
-            if item["S3Input"]["S3InputMode"]:
-                input_mode = item["S3Input"]["S3InputMode"]
-            else:
-                raise ValueError("Processing input must have a valid ['S3InputMode']")
-
-            item["DataUri"] = data_uri
-
-            if (
-                "S3DataDistributionType" in item["S3Input"]
-                and item["S3Input"]["S3DataDistributionType"] != "FullyReplicated"
-            ):
+            item["DataUri"] = s3_input["S3Uri"]
 
+            if "S3InputMode" in s3_input and s3_input["S3InputMode"] != "File":
                 raise RuntimeError(
-                    "DataDistribution: %s is not currently supported in Local Mode"
-                    % item["S3Input"]["S3DataDistributionType"]
+                    "S3InputMode: %s is not currently supported in Local Mode"
+                    % s3_input["S3InputMode"]
                 )
 
-            if input_mode != "File":
+            if ("S3DataDistributionType" in s3_input
+                and s3_input["S3DataDistributionType"] != "FullyReplicated"):
                 raise RuntimeError(
-                    "S3InputMode: %s is not currently supported in Local Mode" % input_mode
+                    "DataDistribution: %s is not currently supported in Local Mode"
+                    % s3_input["S3DataDistributionType"]
                 )
 
-            if (
-                "S3CompressionType" in item["S3Input"]
-                and item["S3Input"]["S3CompressionType"] != "None"
-            ):
-
+            if "S3CompressionType" in s3_input and s3_input["S3CompressionType"] != "None":
                 raise RuntimeError(
                     "CompressionType: %s is not currently supported in Local Mode"
-                    % item["S3Input"]["S3CompressionType"]
+                    % s3_input["S3CompressionType"]
                 )
 
         if processing_output_config and "Outputs" in processing_output_config:
@@ -121,14 +110,15 @@ def start(self, processing_inputs, processing_output_config, environment, proces
                         "FeatureStoreOutput is not currently supported in Local Mode"
                     )
 
-                if "S3Output" in item and item["S3Output"]:
-                    upload_mode = item["S3Output"]["S3UploadMode"]
-                else:
-                    raise ValueError("Please specify a valid ['S3Output'] when using Local Mode.")
+                try:
+                    s3_output = item["S3Output"]
+                except KeyError:
+                    raise ValueError("Processing output must have a valid ['S3Output']")
 
-                if upload_mode != "EndOfJob":
+                if s3_output["S3UploadMode"] != "EndOfJob":
                     raise RuntimeError(
-                        "UploadMode: %s is not currently supported in Local Mode." % upload_mode
+                        "UploadMode: %s is not currently supported in Local Mode."
+                        % s3_output["S3UploadMode"]
                     )
 
         self.start_time = datetime.datetime.now()
@@ -149,7 +139,8 @@ def start(self, processing_inputs, processing_output_config, environment, proces
     def describe(self):
         """Describes a local processing job.
 
-        Returns: An object describing the processing job.
+        Returns:
+            An object describing the processing job.
         """
 
         response = {

src/sagemaker/local/image.py

@@ -117,7 +117,7 @@ def process(
 
         Args:
             processing_inputs (dict): The processing input specification.
-            processing_output_config: The processing output configuration specification.
+            processing_output_config (dict): The processing output configuration specification.
             environment (dict): The environment collection for the processing job.
             processing_job_name (str): Name of the local processing job being run.
         """
@@ -158,15 +158,18 @@ def process(
         except RuntimeError as e:
             # _stream_output() doesn't have the command line. We will handle the exception
             # which contains the exit code and append the command line to it.
-            msg = "Failed to run: %s, %s" % (compose_command, str(e))
-            raise RuntimeError(msg)
+            msg = f"Failed to run: {compose_command}"
+            raise RuntimeError(msg) from e
         finally:
             # Uploading processing outputs back to Amazon S3.
             self._upload_processing_outputs(data_dir, processing_output_config)
 
-            # Deleting temporary directories.
-            dirs_to_delete = [shared_dir, data_dir]
-            self._cleanup(dirs_to_delete)
+            try:
+                # Deleting temporary directories.
+                dirs_to_delete = [shared_dir, data_dir]
+                self._cleanup(dirs_to_delete)
+            except OSError:
+                pass
 
         # Print our Job Complete line to have a similar experience to training on SageMaker where
         # you see this line at the end.
@@ -910,10 +913,11 @@ def _check_output(cmd, *popenargs, **kwargs):
 
 
 def _create_processing_config_file_directories(root, host):
-    """
+    """Creates the directory for the processing config files.
+
     Args:
-        root:
-        host:
+        root: The root path.
+        host: The current host.
     """
     for d in ["config"]:
         os.makedirs(os.path.join(root, host, d))

src/sagemaker/processing.py

@@ -115,14 +115,10 @@ def __init__(
         self.arguments = None
 
         if self.instance_type in ("local", "local_gpu"):
-            self.sagemaker_session = sagemaker_session or LocalSession()
-            if not isinstance(self.sagemaker_session, LocalSession):
-                raise RuntimeError(
-                    "instance_type local or local_gpu is only supported with an"
-                    "instance of LocalSession"
-                )
-        else:
-            self.sagemaker_session = sagemaker_session or Session()
+            if not isinstance(sagemaker_session, LocalSession):
+                sagemaker_session = LocalSession()
+
+        self.sagemaker_session = sagemaker_session or Session()
 
     def run(
         self,

src/sagemaker/sklearn/processing.py

@@ -18,7 +18,6 @@
 from __future__ import absolute_import
 
 from sagemaker import image_uris, Session
-from sagemaker.local import LocalSession
 from sagemaker.processing import ScriptProcessor
 from sagemaker.sklearn import defaults
 
@@ -84,17 +83,9 @@ def __init__(
         if not command:
             command = ["python3"]
 
-        if instance_type in ("local", "local_gpu"):
-            session = sagemaker_session or LocalSession()
-            if not isinstance(session, LocalSession):
-                raise RuntimeError(
-                    "instance_type local or local_gpu is only supported with an"
-                    "instance of LocalSession"
-                )
-        else:
-            session = sagemaker_session or Session()
-
+        session = sagemaker_session or Session()
         region = session.boto_region_name
+
         image_uri = image_uris.retrieve(
             defaults.SKLEARN_NAME, region, version=framework_version, instance_type=instance_type
         )

src/sagemaker/spark/processing.py

@@ -35,7 +35,6 @@
 from sagemaker.local.image import _ecr_login_if_needed, _pull_image
 from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
 from sagemaker.s3 import S3Uploader
-from sagemaker.local import LocalSession
 from sagemaker.session import Session
 from sagemaker.spark import defaults
 
@@ -145,16 +144,7 @@ def __init__(
         self.history_server = None
         self._spark_event_logs_s3_uri = None
 
-        if instance_type in ("local", "local_gpu"):
-            session = sagemaker_session or LocalSession()
-            if not isinstance(session, LocalSession):
-                raise RuntimeError(
-                    "instance_type local or local_gpu is only supported with an"
-                    "instance of LocalSession"
-                )
-        else:
-            session = sagemaker_session or Session()
-
+        session = sagemaker_session or Session()
         region = session.boto_region_name
 
         self.image_uri = self._retrieve_image_uri(

tests/integ/test_local_mode.py

@@ -59,7 +59,7 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
 
 
 @pytest.fixture(scope="module")
-def image_uri(
+def sklearn_image_uri(
     sklearn_latest_version,
     sklearn_latest_py_version,
     cpu_instance_type,
@@ -355,12 +355,12 @@ def test_local_processing_sklearn(sagemaker_local_session, sklearn_latest_versio
 
 
 @pytest.mark.local_mode
-def test_local_processing_script_processor(sagemaker_local_session, image_uri):
+def test_local_processing_script_processor(sagemaker_local_session, sklearn_image_uri):
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 
     script_processor = ScriptProcessor(
         role="SageMakerRole",
-        image_uri=image_uri,
+        image_uri=sklearn_image_uri,
         command=["python3"],
         instance_count=1,
         instance_type="local",
@@ -419,6 +419,6 @@ def test_local_processing_script_processor(sagemaker_local_session, image_uri):
         "python3",
         "/opt/ml/processing/input/code/dummy_script.py",
     ]
-    assert job_description["AppSpecification"]["ImageUri"] == image_uri
+    assert job_description["AppSpecification"]["ImageUri"] == sklearn_image_uri
 
     assert job_description["Environment"] == {"DUMMY_ENVIRONMENT_VARIABLE": "dummy-value"}