[AWS SageMaker] Add integration test for sample pipeline train (#3876)

* add integ test for sample pipeline train

* change docker build context integ test

* add spot test and use train ccomponent test for sample train pipeline

* small changes and ran flake8 and black

* address comments

components/aws/sagemaker/codebuild/integration-test.buildspec.yml

@@ -7,18 +7,17 @@ env:
 phases:
   build:
     commands:
-      - cd components/aws
-      - docker build . -f ./sagemaker/tests/integration_tests/Dockerfile -t amazon/integration-test-image --quiet
+      - docker build . -f ./components/aws/sagemaker/tests/integration_tests/Dockerfile -t amazon/integration-test-image --quiet
 
-      - cd sagemaker/codebuild/scripts && export CONTAINER_VARIABLE_ARGUMENTS="$(./construct_environment_array.sh)"
+      - cd components/aws/sagemaker/codebuild/scripts && export CONTAINER_VARIABLE_ARGUMENTS="$(./construct_environment_array.sh)"
 
       # Run the container and copy the results to /tmp
       # Passes all listed host environment variables through to the container
       - docker run --name integration-test-container $(echo $CONTAINER_VARIABLE_ARGUMENTS) amazon/integration-test-image
 
   post_build:
     commands:
-      - docker cp integration-test-container:/tests/integration_tests/integration_tests.log /tmp/results.xml
+      - docker cp integration-test-container:/pipelines/components/aws/sagemaker/tests/integration_tests/integration_tests.log /tmp/results.xml
       - docker rm -f integration-test-container
 
 reports:

components/aws/sagemaker/tests/integration_tests/Dockerfile

@@ -24,7 +24,7 @@ RUN curl -sSL -o /usr/local/bin/argo https://github.com/argoproj/argo/releases/d
  && chmod +x /usr/local/bin/argo
 
 # Copy conda environment early to avoid cache busting
-COPY ./sagemaker/tests/integration_tests/environment.yml environment.yml
+COPY ./components/aws/sagemaker/tests/integration_tests/environment.yml environment.yml
 
 # Create conda environment for running tests and set as start-up environment
 RUN conda env create -f environment.yml
@@ -38,6 +38,7 @@ ENV S3_DATA_BUCKET="kfp-test-data"
 ENV MINIO_LOCAL_PORT=9000
 ENV KFP_NAMESPACE="kubeflow"
 
-COPY ./sagemaker/ .
+RUN mkdir pipelines
+COPY ./ ./pipelines/
 
-ENTRYPOINT [ "/bin/bash", "./tests/integration_tests/scripts/run_integration_tests" ]
+ENTRYPOINT [ "/bin/bash", "./pipelines/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests" ]

components/aws/sagemaker/tests/integration_tests/README.md

@@ -18,8 +18,8 @@
     1. Update the `S3_DATA_BUCKET` parameter with the name of the bucket created earlier.
     1. (Optional) If you have already created an EKS cluster for testing, replace the `EKS_EXISTING_CLUSTER` field with it's name.
 1. Build the image by doing the following:
-    1. Navigate to the `components/aws` directory.
-    1. Run `docker build . -f sagemaker/tests/integration_tests/Dockerfile -t amazon/integration_test`
+    1. Navigate to the root of this github directory.
+    1. Run `docker build . -f components/aws/sagemaker/tests/integration_tests/Dockerfile -t amazon/integration_test`
 1. Run the image, injecting your environment variable files:
     1. Navigate to the `components/aws` directory.
-    1. Run `docker run --env-file sagemaker/tests/integration_tests/.env amazon/integration_test`
+    1. Run `docker run --env-file components/aws/sagemaker/tests/integration_tests/.env amazon/integration_test`

components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py

@@ -11,9 +11,14 @@
     "test_file_dir",
     [
         pytest.param(
-            "resources/config/simple-mnist-training", marks=pytest.mark.canary_test
+            "resources/config/simple-mnist-training",
+            marks=pytest.mark.canary_test
         ),
-        pytest.param("resources/config/fsx-mnist-training", marks=pytest.mark.fsx_test),
+        pytest.param(
+            "resources/config/fsx-mnist-training",
+            marks=pytest.mark.fsx_test
+        ),
+        "resources/config/spot-sample-pipeline-training"
     ],
 )
 def test_trainingjob(

components/aws/sagemaker/tests/integration_tests/resources/config/spot-sample-pipeline-training/config.yaml

@@ -0,0 +1,25 @@
+PipelineDefinition: ../../../../../samples/contrib/aws-samples/simple_train_pipeline/training-pipeline.py
+TestName: sample-pipeline-training
+Timeout: 3600
+ExpectedTrainingImage: ((KMEANS_REGISTRY)).dkr.ecr.((REGION)).amazonaws.com/kmeans:1
+Arguments:
+  region: ((REGION))
+  image: ((KMEANS_REGISTRY)).dkr.ecr.((REGION)).amazonaws.com/kmeans:1
+  hyperparameters:
+    k: "10"
+    feature_dim: "784"
+  channels:
+    - ChannelName: train
+      DataSource:
+        S3DataSource:
+          S3Uri: s3://((DATA_BUCKET))/mnist_kmeans_example/train_data
+          S3DataType: S3Prefix
+          S3DataDistributionType: FullyReplicated
+      CompressionType: None
+      RecordWrapperType: None
+      InputMode: File
+  spot_instance: "True"
+  checkpoint_config:
+    S3Uri: s3://((DATA_BUCKET))/mnist_kmeans_example/train-checkpoints
+  model_artifact_path: s3://((DATA_BUCKET))/mnist_kmeans_example/output
+  role: ((ROLE_ARN))

components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py

@@ -1,7 +1,6 @@
 import kfp
 from kfp import components
 from kfp import dsl
-from kfp.aws import use_aws_secret
 
 sagemaker_train_op = components.load_component_from_file("../../train/component.yaml")
 

components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests

@@ -203,4 +203,6 @@ fi
 
 [ ! -z "${PYTEST_MARKER}" ] && pytest_args+=( -m "${PYTEST_MARKER}" )
 
-cd tests/integration_tests && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n $(nproc)
+DIR_THIS_FILE="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+cd $DIR_THIS_FILE/../ && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n $(nproc)

samples/contrib/aws-samples/simple_train_pipeline/training-pipeline.py

@@ -5,12 +5,17 @@
 
 import kfp
 import json
+import os
 import copy
 from kfp import components
 from kfp import dsl
 from kfp.aws import use_aws_secret
 
-sagemaker_train_op = components.load_component_from_file('../../../../components/aws/sagemaker/train/component.yaml')
+
+cur_file_dir = os.path.dirname(__file__)
+components_dir = os.path.join(cur_file_dir, '../../../../components/aws/sagemaker/')
+
+sagemaker_train_op = components.load_component_from_file(components_dir + '/train/component.yaml')
 
 channelObjList = []