add default tag integration test (#28)

* add default tag integration test

Author: guoqiao1992

Makefile

@@ -84,13 +84,28 @@ test-sagemaker: install-sdk build-tests
 	# History server tests can't run in parallel since they use the same container name.
 	pytest -s -vv test/integration/history \
 	--repo=$(DEST_REPO) --tag=$(VERSION) --durations=0 \
+	--spark-version=$(SPARK_VERSION)
+	--framework_version=$(FRAMEWORK_VERSION) \
 	--role $(ROLE) \
 	--image_uri $(IMAGE_URI) \
 	--region ${REGION} \
 	--domain ${AWS_DOMAIN}
 	# OBJC_DISABLE_INITIALIZE_FORK_SAFETY: https://github.com/ansible/ansible/issues/32499#issuecomment-341578864
 	OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES pytest --workers auto -s -vv test/integration/sagemaker \
 	--repo=$(DEST_REPO) --tag=$(VERSION) --durations=0 \
+	--spark-version=$(SPARK_VERSION)
+	--framework_version=$(FRAMEWORK_VERSION) \
+	--role $(ROLE) \
+	--image_uri $(IMAGE_URI) \
+	--region ${REGION} \
+	--domain ${AWS_DOMAIN}
+
+# This is included in a separate target because it will be run only in prod stage
+test-prod:
+	pytest -s -vv test/integration/tag \
+	--repo=$(DEST_REPO) --tag=$(VERSION) --durations=0 \
+	--spark-version=$(SPARK_VERSION)
+	--framework_version=$(FRAMEWORK_VERSION) \
 	--role $(ROLE) \
 	--image_uri $(IMAGE_URI) \
 	--region ${REGION} \

test/integration/conftest.py

@@ -23,6 +23,8 @@ def pytest_addoption(parser) -> str:
     parser.addoption("--region", default="us-west-2")
     parser.addoption("--repo")
     parser.addoption("--tag")
+    parser.addoption("--spark-version")
+    parser.addoption("--framework-version")
     parser.addoption("--domain", default="amazonaws.com")
 
 
@@ -60,6 +62,18 @@ def tag(request) -> str:
     return request.config.getoption("--tag")
 
 
+@pytest.fixture(scope="session")
+def spark_version(request) -> str:
+    """Return Docker image framework_version to use in tests."""
+    return request.config.getoption("--spark-version")
+
+
+@pytest.fixture(scope="session")
+def framework_version(request) -> str:
+    """Return Docker image framework_version to use in tests."""
+    return request.config.getoption("--framework-version")
+
+
 @pytest.fixture(scope="session")
 def domain(request) -> str:
     """Return AWS domain"""

test/integration/local/test_multinode_container.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 import subprocess
+
 import pytest
 
 

test/integration/prod/test_default_tag.py

@@ -0,0 +1,60 @@
+# Copyright 2020 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 datetime import datetime
+
+from sagemaker.s3 import S3Downloader, S3Uploader
+from sagemaker.spark.processing import PySparkProcessor
+
+
+def test_sagemaker_spark_processor_default_tag(spark_version, role, sagemaker_session, sagemaker_client):
+    """Test that spark processor works with default tag"""
+    spark = PySparkProcessor(
+        base_job_name="sm-spark-py",
+        framework_version=spark_version,
+        role=role,
+        instance_count=1,
+        instance_type="ml.c5.xlarge",
+        max_runtime_in_seconds=1200,
+        sagemaker_session=sagemaker_session,
+    )
+    bucket = spark.sagemaker_session.default_bucket()
+    timestamp = datetime.now().isoformat()
+    output_data_uri = "s3://{}/spark/output/sales/{}".format(bucket, timestamp)
+    spark_event_logs_key_prefix = "spark/spark-events/{}".format(timestamp)
+    spark_event_logs_s3_uri = "s3://{}/{}".format(bucket, spark_event_logs_key_prefix)
+
+    with open("test/resources/data/files/data.jsonl") as data:
+        body = data.read()
+        input_data_uri = "s3://{}/spark/input/data.jsonl".format(bucket)
+        S3Uploader.upload_string_as_file_body(
+            body=body, desired_s3_uri=input_data_uri, sagemaker_session=sagemaker_session
+        )
+
+    spark.run(
+        submit_app="test/resources/code/python/hello_py_spark/hello_py_spark_app.py",
+        submit_py_files=["test/resources/code/python/hello_py_spark/hello_py_spark_udfs.py"],
+        arguments=["--input", input_data_uri, "--output", output_data_uri],
+        spark_event_logs_s3_uri=spark_event_logs_s3_uri,
+        wait=True,
+    )
+
+    processing_job = spark.latest_job
+    waiter = sagemaker_client.get_waiter("processing_job_completed_or_stopped")
+    waiter.wait(
+        ProcessingJobName=processing_job.job_name,
+        # poll every 15 seconds. timeout after 15 minutes.
+        WaiterConfig={"Delay": 15, "MaxAttempts": 60},
+    )
+
+    output_contents = S3Downloader.list(output_data_uri, sagemaker_session=sagemaker_session)
+    assert len(output_contents) != 0

test/integration/sagemaker/test_spark.py

@@ -91,11 +91,10 @@ def configuration() -> list:
     return configuration
 
 
-def test_sagemaker_pyspark_multinode(tag, role, image_uri, configuration, sagemaker_session, region, sagemaker_client):
+def test_sagemaker_pyspark_multinode(role, image_uri, configuration, sagemaker_session, region, sagemaker_client):
     """Test that basic multinode case works on 32KB of data"""
     spark = PySparkProcessor(
         base_job_name="sm-spark-py",
-        framework_version=tag,
         image_uri=image_uri,
         role=role,
         instance_count=2,
@@ -168,11 +167,10 @@ def test_sagemaker_pyspark_multinode(tag, role, image_uri, configuration, sagema
 # TODO: similar integ test case for SSE-KMS. This would require test infrastructure bootstrapping a KMS key.
 # Currently, Spark jobs can read data encrypted with SSE-KMS (assuming the execution role has permission),
 # however our Hadoop version (2.8.5) does not support writing data with SSE-KMS (enabled in version 3.0.0).
-def test_sagemaker_pyspark_sse_s3(tag, role, image_uri, sagemaker_session, region, sagemaker_client):
+def test_sagemaker_pyspark_sse_s3(role, image_uri, sagemaker_session, region, sagemaker_client):
     """Test that Spark container can read and write S3 data encrypted with SSE-S3 (default AES256 encryption)"""
     spark = PySparkProcessor(
         base_job_name="sm-spark-py",
-        framework_version=tag,
         image_uri=image_uri,
         role=role,
         instance_count=2,
@@ -212,11 +210,10 @@ def test_sagemaker_pyspark_sse_s3(tag, role, image_uri, sagemaker_session, regio
     assert len(output_contents) != 0
 
 
-def test_sagemaker_scala_jar_multinode(tag, role, image_uri, configuration, sagemaker_session, sagemaker_client):
+def test_sagemaker_scala_jar_multinode(role, image_uri, configuration, sagemaker_session, sagemaker_client):
     """Test SparkJarProcessor using Scala application jar with external runtime dependency jars staged by SDK"""
     spark = SparkJarProcessor(
         base_job_name="sm-spark-scala",
-        framework_version=tag,
         image_uri=image_uri,
         role=role,
         instance_count=2,

test/unit/test_bootstrapper.py

@@ -14,6 +14,7 @@
 from unittest.mock import MagicMock, Mock, PropertyMock, call, mock_open, patch
 
 import pytest
+
 from smspark.bootstrapper import Bootstrapper
 from smspark.config import Configuration
 from smspark.defaults import default_resource_config

test/unit/test_cli.py

@@ -19,6 +19,7 @@
 import click
 import pytest
 from click.testing import CliRunner
+
 from smspark.cli import submit, submit_main
 from smspark.errors import InputError
 from smspark.job import ProcessingJobManager

test/unit/test_history_server_cli.py

@@ -10,15 +10,16 @@
 # 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.
-import pytest
-
 from dataclasses import dataclass
 from typing import Type, Union
 from unittest.mock import patch
+
+import pytest
 from click.testing import CliRunner
+
 from smspark.cli import submit
-from smspark.job import ProcessingJobManager
 from smspark.history_server_cli import run_history_server
+from smspark.job import ProcessingJobManager
 
 ARGS_FORMAT = "--event-logs-s3-uri {} --remote-domain-name {}"
 EVENT_LOGS_S3_URI = "s3://bucket"

test/unit/test_history_server_utils.py

@@ -16,6 +16,7 @@
 from unittest.mock import MagicMock, call, mock_open, patch
 
 import pytest
+
 from smspark.bootstrapper import Bootstrapper
 from smspark.errors import InputError
 from smspark.history_server_utils import (

test/unit/test_spark_event_logs_publisher.py

@@ -15,6 +15,7 @@
 from unittest.mock import call, mock_open, patch
 
 import pytest
+
 from smspark.spark_event_logs_publisher import (
     CONFIG_ENABLE_EVENT_LOG,
     CONFIG_EVENT_LOG_DIR_FORMAT,