Merge pull request aws#36 from athewsey/feat/fw-processor

Fix XGBoostProcessor and add TF integration test

Author: verdimrc

src/sagemaker/mxnet/__init__.py

@@ -10,8 +10,9 @@
 # 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.
-"""Placeholder docstring"""
+"""Classes for using MXNet with Amazon SageMaker."""
 from __future__ import absolute_import  # noqa: F401
 
 from sagemaker.mxnet.estimator import MXNet  # noqa: F401
 from sagemaker.mxnet.model import MXNetModel, MXNetPredictor  # noqa: F401
+from sagemaker.mxnet.processing import MXNetProcessor  # noqa: F401

src/sagemaker/pytorch/__init__.py

@@ -10,8 +10,9 @@
 # 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.
-"""Placeholder docstring"""
+"""Classes for using PyTorch with Amazon SageMaker."""
 from __future__ import absolute_import
 
 from sagemaker.pytorch.estimator import PyTorch  # noqa: F401
 from sagemaker.pytorch.model import PyTorchModel, PyTorchPredictor  # noqa: F401
+from sagemaker.pytorch.processing import PyTorchProcessor  # noqa: F401

src/sagemaker/sklearn/__init__.py

@@ -10,7 +10,7 @@
 # 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.
-"""Placeholder docstring"""
+"""Classes for using Scikit-Learn with Amazon SageMaker."""
 from __future__ import absolute_import
 
 from sagemaker.sklearn.estimator import SKLearn  # noqa: F401

src/sagemaker/tensorflow/__init__.py

@@ -15,3 +15,4 @@
 
 from sagemaker.tensorflow.estimator import TensorFlow  # noqa: F401 (imported but unused)
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor  # noqa: F401
+from sagemaker.tensorflow.processing import TensorFlowProcessor  # noqa: F401

src/sagemaker/xgboost/__init__.py

@@ -10,7 +10,10 @@
 # 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.
-"""Placeholder docstring"""
+"""Classes for using XGBoost with Amazon SageMaker."""
+from __future__ import absolute_import
+
 from sagemaker.xgboost.defaults import XGBOOST_NAME  # noqa: F401
 from sagemaker.xgboost.estimator import XGBoost  # noqa: F401
 from sagemaker.xgboost.model import XGBoostModel, XGBoostPredictor  # noqa: F401
+from sagemaker.xgboost.processing import XGBoostProcessor  # noqa: F401

src/sagemaker/xgboost/processing.py

@@ -21,7 +21,7 @@
 from sagemaker.xgboost.estimator import XGBoost
 
 
-class XGBoostEstimator(FrameworkProcessor):
+class XGBoostProcessor(FrameworkProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using XGBoost containers."""
 
     estimator_cls = XGBoost

tests/conftest.py

@@ -411,4 +411,8 @@ def _parametrize_framework_version_fixtures(metafunc, fixture_prefix, config):
     if fixture_name in metafunc.fixturenames:
         config = config["versions"]
         py_versions = config[latest_version].get("py_versions", config[latest_version].keys())
-        metafunc.parametrize(fixture_name, (sorted(py_versions)[-1],), scope="session")
+        if "repository" in py_versions or "registries" in py_versions:
+            # Config did not specify `py_versions` and is not arranged by py_version. Assume py3
+            metafunc.parametrize(fixture_name, ("py3",), scope="session")
+        else:
+            metafunc.parametrize(fixture_name, (sorted(py_versions)[-1],), scope="session")

tests/integ/test_tf.py

@@ -18,11 +18,11 @@
 
 import pytest
 
-from sagemaker.tensorflow import TensorFlow
+from sagemaker.tensorflow import TensorFlow, TensorFlowProcessor
 from sagemaker.utils import unique_name_from_base, sagemaker_timestamp
 
 import tests.integ
-from tests.integ import kms_utils, timeout
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, kms_utils, timeout
 from tests.integ.retry import retries
 from tests.integ.s3_utils import assert_s3_files_exist
 
@@ -39,6 +39,35 @@
 ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
 
 
+@pytest.mark.release
+def test_framework_processing_job_with_deps(
+    sagemaker_session,
+    instance_type,
+    tensorflow_training_latest_version,
+    tensorflow_training_latest_py_version,
+):
+    with timeout.timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        code_path = os.path.join(DATA_DIR, "dummy_code_bundle_with_reqs")
+        entry_point = "main_script.py"
+
+        processor = TensorFlowProcessor(
+            framework_version=tensorflow_training_latest_version,
+            py_version=tensorflow_training_latest_py_version,
+            role=ROLE,
+            instance_count=1,
+            instance_type=instance_type,
+            sagemaker_session=sagemaker_session,
+            base_job_name="test-tensorflow",
+        )
+
+        processor.run(
+            code=entry_point,
+            source_dir=code_path,
+            inputs=[],
+            wait=True,
+        )
+
+
 def test_mnist_with_checkpoint_config(
     sagemaker_session,
     instance_type,
@@ -51,7 +80,7 @@ def test_mnist_with_checkpoint_config(
     checkpoint_local_path = "/test/checkpoint/path"
     estimator = TensorFlow(
         entry_point=SCRIPT,
-        role="SageMakerRole",
+        role=ROLE,
         instance_count=1,
         instance_type=instance_type,
         sagemaker_session=sagemaker_session,

tests/integ/test_xgboost.py

@@ -0,0 +1,50 @@
+# Copyright 2018-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.
+from __future__ import absolute_import
+
+import os
+import pytest
+from sagemaker.xgboost.processing import XGBoostProcessor
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ.timeout import timeout
+
+ROLE = "SageMakerRole"
+
+
+@pytest.mark.release
+def test_framework_processing_job_with_deps(
+    sagemaker_session,
+    xgboost_latest_version,
+    xgboost_latest_py_version,
+    cpu_instance_type,
+):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        code_path = os.path.join(DATA_DIR, "dummy_code_bundle_with_reqs")
+        entry_point = "main_script.py"
+
+        processor = XGBoostProcessor(
+            framework_version=xgboost_latest_version,
+            py_version=xgboost_latest_py_version,
+            role=ROLE,
+            instance_count=1,
+            instance_type=cpu_instance_type,
+            sagemaker_session=sagemaker_session,
+            base_job_name="test-xgboost",
+        )
+
+        processor.run(
+            code=entry_point,
+            source_dir=code_path,
+            inputs=[],
+            wait=True,
+        )