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

Integration test of requirements.txt + script bundle functionality

Author: verdimrc

tests/data/dummy_code_bundle_with_reqs/local_module.py

@@ -0,0 +1,2 @@
+"""A dummy Python module to check importing local files works OK"""
+DUMMY_CONSTANT = 1

tests/data/dummy_code_bundle_with_reqs/main_script.py

@@ -0,0 +1,11 @@
+"""A dummy SageMaker job script testing local imports and requirements.txt installs"""
+
+print("This is the print output from dummy_code_bundle_with_reqs/main_script.py")
+
+print("Trying to import local module...")
+import local_module
+
+print("Trying to import module from requirements.txt...")
+import stepfunctions
+
+print("Done")

tests/data/dummy_code_bundle_with_reqs/requirements.txt

@@ -0,0 +1,3 @@
+# As a test dependency, we'll use the AWS Step Functions Data Science SDK - since it's more of a
+# notebook environment tool than a training/processing job library, so shouldn't be in base images
+stepfunctions

tests/integ/test_sklearn.py renamed to tests/integ/test_processing_sklearn.py

@@ -18,12 +18,13 @@
 import pytest
 import numpy
 
-from sagemaker.sklearn import SKLearn
-from sagemaker.sklearn import SKLearnModel
+from sagemaker.sklearn import SKLearn, SKLearnModel, SKLearnProcessor
 from sagemaker.utils import sagemaker_timestamp, unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
+ROLE = "SageMakerRole"
+
 
 @pytest.fixture(scope="module")
 @pytest.mark.skip(
@@ -45,6 +46,20 @@ def sklearn_training_job(
     sagemaker_session.boto_region_name
 
 
+def test_framework_processing_job_with_deps(
+    sagemaker_session,
+    sklearn_latest_version,
+    sklearn_latest_py_version,
+    cpu_instance_type,
+):
+    return _run_processing_job(
+        sagemaker_session,
+        cpu_instance_type,
+        sklearn_latest_version,
+        sklearn_latest_py_version,
+    )
+
+
 def test_training_with_additional_hyperparameters(
     sagemaker_session,
     sklearn_latest_version,
@@ -57,7 +72,7 @@ def test_training_with_additional_hyperparameters(
 
         sklearn = SKLearn(
             entry_point=script_path,
-            role="SageMakerRole",
+            role=ROLE,
             instance_type=cpu_instance_type,
             framework_version=sklearn_latest_version,
             py_version=sklearn_latest_py_version,
@@ -88,7 +103,7 @@ def test_training_with_network_isolation(
 
         sklearn = SKLearn(
             entry_point=script_path,
-            role="SageMakerRole",
+            role=ROLE,
             instance_type=cpu_instance_type,
             framework_version=sklearn_latest_version,
             py_version=sklearn_latest_py_version,
@@ -145,7 +160,7 @@ def test_deploy_model(
         script_path = os.path.join(DATA_DIR, "sklearn_mnist", "mnist.py")
         model = SKLearnModel(
             model_data,
-            "SageMakerRole",
+            ROLE,
             entry_point=script_path,
             framework_version=sklearn_latest_version,
             sagemaker_session=sagemaker_session,
@@ -198,7 +213,7 @@ def test_failed_training_job(
 
         sklearn = SKLearn(
             entry_point=script_path,
-            role="SageMakerRole",
+            role=ROLE,
             framework_version=sklearn_latest_version,
             py_version=sklearn_latest_py_version,
             instance_count=1,
@@ -215,6 +230,33 @@ def test_failed_training_job(
             sklearn.fit(train_input, job_name=job_name)
 
 
+def _run_processing_job(
+    sagemaker_session, instance_type, sklearn_version, py_version, wait=True
+):
+    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 = SKLearnProcessor(
+            framework_version=sklearn_version,
+            py_version=py_version,
+            role=ROLE,
+            instance_count=1,
+            instance_type=instance_type,
+            sagemaker_session=sagemaker_session,
+            base_job_name="test-sklearn",
+        )
+
+        processor.run(
+            code=entry_point,
+            source_dir=code_path,
+            inputs=[],
+            wait=wait,
+        )
+        return processor.latest_job.name
+
+
 def _run_mnist_training_job(
     sagemaker_session, instance_type, sklearn_version, py_version, wait=True
 ):
@@ -226,7 +268,7 @@ def _run_mnist_training_job(
 
         sklearn = SKLearn(
             entry_point=script_path,
-            role="SageMakerRole",
+            role=ROLE,
             framework_version=sklearn_version,
             py_version=py_version,
             instance_type=instance_type,