Merge remote-tracking branch 'verdi/pr-framework-processor-round-02' into feat/fw-processor

Author: athewsey

src/sagemaker/mxnet/processing.py

@@ -0,0 +1,79 @@
+# Copyright 2019-2021 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 code related to MXNet Processors which are used for Processing jobs.
+
+These jobs let customers perform data pre-processing, post-processing, feature engineering,
+data validation, and model evaluation and interpretation on SageMaker.
+"""
+from __future__ import absolute_import
+
+from sagemaker.mxnet.estimator import MXNet
+from sagemaker.processing import FrameworkProcessor
+
+
+class MXNetProcessor(FrameworkProcessor):
+    """Handles Amazon SageMaker processing tasks for jobs using MXNet containers."""
+
+    estimator_cls = MXNet
+
+    def __init__(
+        self,
+        framework_version,  # New arg
+        s3_prefix,  # New arg
+        role,
+        instance_count,
+        instance_type,
+        py_version="py3",  # New kwarg
+        image_uri=None,
+        volume_size_in_gb=30,
+        volume_kms_key=None,
+        output_kms_key=None,
+        max_runtime_in_seconds=None,
+        base_job_name=None,
+        sagemaker_session=None,
+        env=None,
+        tags=None,
+        network_config=None,
+    ):
+        """This processor executes a Python script in a managed MXNet execution environment.
+
+        Unless ``image_uri`` is specified, the MXNet environment is an
+        Amazon-built Docker container that executes functions defined in the supplied
+        ``entry_point`` Python script.
+
+        The arguments have the exact same meaning as in ``FrameworkProcessor``.
+
+        .. tip::
+
+            You can find additional parameters for initializing this class at
+            :class:`~sagemaker.processing.FrameworkProcessor`.
+        """
+        super().__init__(
+            self.estimator_cls,
+            framework_version,
+            s3_prefix,
+            role,
+            instance_count,
+            instance_type,
+            py_version,
+            image_uri,
+            volume_size_in_gb,
+            volume_kms_key,
+            output_kms_key,
+            max_runtime_in_seconds,
+            base_job_name,
+            sagemaker_session,
+            env,
+            tags,
+            network_config,
+        )