Improved docstrings, handled DatasetDefinition and FeatureStoreOutput, other minor changes.

Author: EC2 Default User

src/sagemaker/local/entities.py

@@ -38,16 +38,16 @@
 HEALTH_CHECK_TIMEOUT_LIMIT = 120
 
 
-class _LocalProcessingJob(object):
+class _LocalProcessingJob:
     """Defines and starts a local processing job."""
 
     _STARTING = "Starting"
     _PROCESSING = "Processing"
     _COMPLETED = "Completed"
-    _states = ["Starting", "Processing", "Completed"]
 
     def __init__(self, container):
-        """
+        """Creates a local processing job.
+
         Args:
             container: the local container object.
         """
@@ -61,16 +61,21 @@ def __init__(self, container):
         self.environment = None
 
     def start(self, processing_inputs, processing_output_config, environment, processing_job_name):
-        """
+        """Starts a local processing job.
+
         Args:
             processing_inputs: The processing input configuration.
             processing_output_config: The processing input configuration.
             environment: The collection of environment variables passed to the job.
             processing_job_name: The processing job name.
         """
+        self.state = self._STARTING
 
         for item in processing_inputs:
-            if item["S3Input"]:
+            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:
                 raise ValueError("Processing input must have a valid ['S3Input']")
@@ -111,10 +116,15 @@ def start(self, processing_inputs, processing_output_config, environment, proces
             processing_outputs = processing_output_config["Outputs"]
 
             for item in processing_outputs:
-                if item["S3Output"]:
+                if "FeatureStoreOutput" in item:
+                    raise RuntimeError(
+                        "FeatureStoreOutput is not currently supported in Local Mode"
+                    )
+
+                if "S3Output" in item and item["S3Output"]:
                     upload_mode = item["S3Output"]["S3UploadMode"]
                 else:
-                    raise ValueError("Processing output must have a valid ['S3Output']")
+                    raise ValueError("Please specify a valid ['S3Output'] when using Local Mode.")
 
                 if upload_mode != "EndOfJob":
                     raise RuntimeError(
@@ -137,7 +147,10 @@ def start(self, processing_inputs, processing_output_config, environment, proces
         self.state = self._COMPLETED
 
     def describe(self):
-        """Describes a local processing job."""
+        """Describes a local processing job.
+
+        Returns: An object describing the processing job.
+        """
 
         response = {
             "ProcessingJobArn": self.processing_job_name,

src/sagemaker/local/image.py

@@ -120,8 +120,6 @@ def process(
             processing_output_config: 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.
-
-        Returns:
         """
 
         self.container_root = self._create_tmp_folder()
@@ -383,8 +381,6 @@ def write_processing_config_files(
         This method writes the hyperparameters, resources and input data
         configuration files.
 
-        Returns: None
-
         Args:
             host (str): Host to write the configuration for
             environment (dict): Environment variable collection.
@@ -509,11 +505,15 @@ def _prepare_training_volumes(
         return volumes
 
     def _prepare_processing_volumes(self, data_dir, processing_inputs, processing_output_config):
-        """
+        """Prepares local container volumes for the processing job.
+
         Args:
-            data_dir:
-            processing_inputs:
-            processing_output_config:
+            data_dir: The local data directory.
+            processing_inputs: The configuration of processing inputs.
+            processing_output_config: The configuration of processing outputs.
+
+        Returns:
+            The volumes configuration.
         """
         shared_dir = os.path.join(self.container_root, "shared")
         volumes = []
@@ -524,9 +524,6 @@ def _prepare_processing_volumes(self, data_dir, processing_inputs, processing_ou
             uri = item["DataUri"]
             input_container_dir = item["S3Input"]["LocalPath"]
 
-            # input_dir = os.path.join(data_dir, "input", input_name)
-            # os.makedirs(input_dir)
-
             data_source = sagemaker.local.data.get_data_source_instance(uri, self.sagemaker_session)
             volumes.append(_Volume(data_source.get_root_dir(), input_container_dir))
 
@@ -545,10 +542,11 @@ def _prepare_processing_volumes(self, data_dir, processing_inputs, processing_ou
         return volumes
 
     def _upload_processing_outputs(self, data_dir, processing_output_config):
-        """
+        """Uploads processing outputs to Amazon S3.
+
         Args:
-            data_dir:
-            processing_output_config:
+            data_dir: The local data directory.
+            processing_output_config: The processing output configuration.
         """
         if processing_output_config and "Outputs" in processing_output_config:
             for item in processing_output_config["Outputs"]:
@@ -711,7 +709,7 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
 
         host_config = {
             "image": self.image,
-            "container_name": container_name_prefix + "-" + host,
+            "container_name": f"{container_name_prefix}-{host}",
             "stdin_open": True,
             "tty": True,
             "volumes": [v.map for v in optml_volumes],

src/sagemaker/local/local_session.py

@@ -75,7 +75,7 @@ def create_processing_job(
         ProcessingOutputConfig=None,
         **kwargs
     ):
-        """Create a processing job in Local Mode
+        """Creates a processing job in Local Mode
 
         Args:
           ProcessingJobName(str): local processing job name.
@@ -128,7 +128,7 @@ def create_processing_job(
         LocalSagemakerClient._processing_jobs[ProcessingJobName] = processing_job
 
     def describe_processing_job(self, ProcessingJobName):
-        """Describe a local processing job.
+        """Describes a local processing job.
 
         Args:
           ProcessingJobName(str): Processing job name to describe.

tests/unit/test_local_session.py

@@ -244,6 +244,107 @@ def test_create_processing_job_invalid_upload_mode(process, LocalSession):
         )
 
 
+@patch("sagemaker.local.image._SageMakerContainer.process")
+@patch("sagemaker.local.local_session.LocalSession")
+def test_create_processing_job_invalid_processing_input(process, LocalSession):
+    local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
+
+    instance_count = 2
+    image = "my-docker-image:1.0"
+
+    app_spec = {"ImageUri": image}
+    resource_config = {"ClusterConfig": {"InstanceCount": instance_count, "InstanceType": "local"}}
+    environment = {"Var1": "Value1"}
+    processing_inputs = [
+        {
+            "InputName": "input1",
+            "DatasetDefinition": {
+                "AthenaDatasetDefinition": {
+                    "Catalog": "cat1",
+                    "Database": "db1",
+                    "OutputS3Uri": "s3://bucket_name/prefix/",
+                    "QueryString": "SELECT * FROM SOMETHING",
+                },
+                "DataDistributionType": "FullyReplicated",
+                "InputMode": "File",
+                "LocalPath": "/opt/ml/processing/input/athena",
+            },
+        }
+    ]
+    processing_output_config = {
+        "Outputs": [
+            {
+                "OutputName": "output1",
+                "S3Output": {
+                    "LocalPath": "/opt/ml/processing/output/output1",
+                    "S3Uri": "s3://some-bucket/some-path/output1",
+                    "S3UploadMode": "Continuous",
+                },
+            }
+        ]
+    }
+    with pytest.raises(RuntimeError):
+        local_sagemaker_client.create_processing_job(
+            "my-processing-job",
+            app_spec,
+            resource_config,
+            environment,
+            processing_inputs,
+            processing_output_config,
+        )
+
+
+@patch("sagemaker.local.image._SageMakerContainer.process")
+@patch("sagemaker.local.local_session.LocalSession")
+def test_create_processing_job_invalid_processing_output(process, LocalSession):
+    local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
+
+    instance_count = 2
+    image = "my-docker-image:1.0"
+
+    app_spec = {"ImageUri": image}
+    resource_config = {"ClusterConfig": {"InstanceCount": instance_count, "InstanceType": "local"}}
+    environment = {"Var1": "Value1"}
+    processing_inputs = [
+        {
+            "InputName": "input1",
+            "S3Input": {
+                "LocalPath": "/opt/ml/processing/input/input1",
+                "S3Uri": "s3://some-bucket/some-path/input1",
+                "S3DataDistributionType": "FullyReplicated",
+                "S3InputMode": "File",
+            },
+        },
+        {
+            "InputName": "input2",
+            "S3Input": {
+                "LocalPath": "/opt/ml/processing/input/input2",
+                "S3Uri": "s3://some-bucket/some-path/input2",
+                "S3DataDistributionType": "FullyReplicated",
+                "S3CompressionType": "None",
+                "S3InputMode": "File",
+            },
+        },
+    ]
+    processing_output_config = {
+        "Outputs": [
+            {
+                "OutputName": "output1",
+                "FeatureStoreOutput": {"FeatureGroupName": "Group1"},
+            }
+        ]
+    }
+    with pytest.raises(RuntimeError):
+        local_sagemaker_client.create_processing_job(
+            "my-processing-job",
+            app_spec,
+            resource_config,
+            environment,
+            processing_inputs,
+            processing_output_config,
+        )
+
+
 @patch("sagemaker.local.local_session.LocalSession")
 def test_describe_invalid_processing_job(*args):
     local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()