Merge branch 'master' into revert_changes

awslabs · Mar 23, 2021 · 071b05d · 071b05d
2 parents 1dcfbbb + 7c5c8c4
commit 071b05d
Show file tree

Hide file tree

Showing 182 changed files with 14,014 additions and 3,363 deletions.
diff --git a/config/buildspec.yml b/config/buildspec.yml
@@ -10,6 +10,7 @@ env:
     run_pytest_tensorflow: "disable"
     run_pytest_tensorflow2: "disable"
     run_pytest_xgboost: "disable"
+    run_pytest_profiler: "enable"
     run_integration_pytest_pytorch: "disable"
     run_integration_pytest_mxnet: "disable"
     run_integration_pytest_tensorflow: "disable"
@@ -28,26 +29,30 @@ phases:
         - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
         - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
-        - pip install --upgrade pip==19.3.1
+        - pip install --upgrade pip==20.3.3
+        - pip install -q matplotlib==3.3.1 seaborn==0.10.1 nbconvert==5.6.1 papermill==2.1.2 jupyter==1.0.0 scipy==1.5.2 scikit-learn==0.23.2 bokeh==2.2.3 simplejson==3.17.2
         - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --upgrade pyYaml==5.1; else pip install -q pyYaml; fi
         - pip install -q pytest wheel pytest-html pre-commit awscli pytest-cov
 
   pre_build:
     commands:
-      - cd $CODEBUILD_SRC_DIR && pre-commit install && pre-commit run --all-files
+      - cd $CODEBUILD_SRC_DIR
 
   build:
     commands:
       - cd $CODEBUILD_SRC_DIR  && python setup.py bdist_wheel --universal
+      # We do not need to force install smdebug-rules. The container used for PR builds do not have smdebug rules binary.
+      # Force installing rules binary attempts to re-install ipython-genutils which fails on PyTorch Ubuntu 16.04 containers.
       - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal
-      - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --force-reinstall $RULES_CODEBUILD_SRC_DIR/dist/*.whl; else pip install $RULES_CODEBUILD_SRC_DIR/dist/*.whl; fi
-      - cd $CODEBUILD_SRC_DIR  && pip install --force-reinstall dist/*.whl && cd ..
+      - if [ "$run_pytest_xgboost" = "enable" ]; then pip install -U $RULES_CODEBUILD_SRC_DIR/dist/*.whl; else pip install $RULES_CODEBUILD_SRC_DIR/dist/*.whl; fi
+      - cd $CODEBUILD_SRC_DIR  && pip install -U dist/*.whl && cd ..
       - cd $CODEBUILD_SRC_DIR  && chmod +x config/tests.sh && PYTHONPATH=. && ./config/tests.sh  && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
       - pip show smdebug
       - pip show smdebug_rules
       - echo 'Uploading Coverage to CodeCov'
       - bash $CODEBUILD_SRC_DIR/config/codecov.sh
-      - cd $RULES_CODEBUILD_SRC_DIR && chmod +x config/tests.sh && PYTHONPATH=. && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      # - cd $RULES_CODEBUILD_SRC_DIR && chmod +x config/tests.sh && PYTHONPATH=. && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      # TODO: Uncomment the line above when it is safe to do so.
 
   post_build:
     commands:

diff --git a/config/buildspec_build_wheel.yml b/config/buildspec_build_wheel.yml
@@ -15,8 +15,8 @@ phases:
         - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
         - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
         - cd $CODEBUILD_SRC_DIR && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
-        - pip install --upgrade pip==19.3.1
-        - pip install -q pytest wheel pyYaml pytest-html pre-commit pytest-cov
+        - pip install --upgrade pip==20.3.3
+        - pip install -q pytest==5.3.3 wheel pyYaml pytest-html pre-commit pytest-cov
         - pip uninstall -y boto3 && pip uninstall -y aiobotocore && pip uninstall -y botocore
 
   build:

diff --git a/config/buildspec_pre_commit.yml b/config/buildspec_pre_commit.yml
@@ -0,0 +1,26 @@
+version: 0.2
+phases:
+  install:
+    commands:
+        -  . config/change_branch.sh
+        - su && apt-get update
+        - apt-get install sudo -qq -o=Dpkg::Use-Pty=0 # silence output: https://askubuntu.com/a/668859/724247
+        - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
+        - pip install --upgrade pip==19.3.1
+        - pip install -q pre-commit awscli
+
+  pre_build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR && pre-commit install && pre-commit run --all-files
+
+  build:
+    commands:
+      - echo "PRE COMMIT SUCCESSFUL"
+
+  post_build:
+    commands:
+      - . $CODEBUILD_SRC_DIR/config/upload_on_end.sh
+      - rm -rf $CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - rm -rf $RULES_CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 0 ]; then echo "ERROR BUILD FAILED " && exit 1 ; fi
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 1 ]; then echo "INFO BUILD SUCCEEDED !!! " ; fi
diff --git a/config/buildspec_tensorflow_2_3.yml b/config/buildspec_tensorflow_2_3.yml
@@ -0,0 +1,65 @@
+# Build Spec for AWS CodeBuild CI TF 2.3.1 CPU and GPU Containers
+# Target Containers:
+    # CPU Container: 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.3.1-cpu-py37-ubuntu18.04
+    # GPU Container 1: 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04
+    # GPU Container 2: 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu102-ubuntu18.04
+
+# The PR CI Uses:
+  # 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04
+
+version: 0.2
+env:
+  variables:
+    run_pytest_pytorch: "disable"
+    run_pytest_mxnet: "disable"
+    run_pytest_tensorflow: "disable"
+    run_pytest_tensorflow2: "enable"
+    run_pytest_xgboost: "disable"
+    run_pytest_profiler: "enable"
+    run_integration_pytest_pytorch: "disable"
+    run_integration_pytest_mxnet: "disable"
+    run_integration_pytest_tensorflow: "disable"
+    run_integration_pytest_tensorflow2: "enable"
+    run_integration_pytest_xgboost: "disable"
+    # below needs to be enabled
+    zero_code_change_test: "enable"
+    # set code coverage flag
+    code_coverage_smdebug: "true"
+phases:
+  install:
+    commands:
+        -  . config/change_branch.sh
+        - su && apt-get update
+        - apt-get install sudo -qq -o=Dpkg::Use-Pty=0 # silence output: https://askubuntu.com/a/668859/724247
+        - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
+        - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
+        - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
+        - pip install --upgrade pip==19.3.1
+        - pip install -q matplotlib==3.3.1 seaborn==0.10.1 nbconvert==5.6.1 papermill==2.1.2 jupyter==1.0.0 scipy==1.5.2 scikit-learn==0.23.2 bokeh==2.2.3 simplejson==3.17.2 transformers==4.2.1
+        - pip install -q pytest wheel pytest-html pre-commit awscli pytest-cov
+
+  pre_build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR
+
+  build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR  && python setup.py bdist_wheel --universal
+      # We do not need to force install smdebug-rules. The container used for PR builds do not have smdebug rules binary.
+      # Force installing rules binary attempts to re-install ipython-genutils which fails on PyTorch Ubuntu 16.04 containers.
+      - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal
+      - cd $CODEBUILD_SRC_DIR  && pip install --force-reinstall dist/*.whl && cd ..
+      - pip install --force-reinstall numpy==1.18.5
+      - cd $CODEBUILD_SRC_DIR  && chmod +x config/tests.sh && PYTHONPATH=. && ./config/tests.sh  && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      - pip show smdebug
+      - echo 'Uploading Coverage to CodeCov'
+      - bash $CODEBUILD_SRC_DIR/config/codecov.sh
+      - cd $RULES_CODEBUILD_SRC_DIR && chmod +x config/tests.sh && PYTHONPATH=. && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+
+  post_build:
+    commands:
+      - . $CODEBUILD_SRC_DIR/config/upload_on_end.sh
+      - rm -rf $CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - rm -rf $RULES_CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 0 ]; then echo "ERROR BUILD FAILED " && exit 1 ; fi
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 1 ]; then echo "INFO BUILD SUCCEEDED !!! " ; fi
diff --git a/config/buildspec_tensorflow_2_4.yml b/config/buildspec_tensorflow_2_4.yml
@@ -0,0 +1,62 @@
+# Build Spec for AWS CodeBuild CI TF 2.4.x CPU and GPU Containers
+# Containers Used:
+    # Note: The public DLC is not yet available so this buildspec is currently consuming a custom built container
+
+version: 0.2
+env:
+  variables:
+    run_pytest_pytorch: "disable"
+    run_pytest_mxnet: "disable"
+    run_pytest_tensorflow: "disable"
+    run_pytest_tensorflow2: "enable"
+    run_pytest_xgboost: "disable"
+    run_pytest_profiler: "enable"
+    run_integration_pytest_pytorch: "disable"
+    run_integration_pytest_mxnet: "disable"
+    run_integration_pytest_tensorflow: "disable"
+    run_integration_pytest_tensorflow2: "enable"
+    run_integration_pytest_xgboost: "disable"
+    # below needs to be enabled
+    zero_code_change_test: "enable"
+    # set code coverage flag
+    code_coverage_smdebug: "true"
+phases:
+  install:
+    commands:
+        -  . config/change_branch.sh
+        - su && apt-get update
+        - apt-get install sudo -qq -o=Dpkg::Use-Pty=0 # silence output: https://askubuntu.com/a/668859/724247
+        - sudo apt-get update -qq -o=Dpkg::Use-Pty=0
+        - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
+        - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
+        - pip install --upgrade pip==19.3.1
+        - pip install -q matplotlib==3.3.1 seaborn==0.10.1 nbconvert==5.6.1 papermill==2.1.2 jupyter==1.0.0 scipy==1.5.2 scikit-learn==0.23.2 bokeh==2.2.3 simplejson==3.17.2 transformers==4.2.1
+        - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --upgrade pyYaml==5.1; else pip install -q pyYaml; fi
+        - pip install -q pytest wheel pytest-html pre-commit awscli pytest-cov
+
+  pre_build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR
+
+  build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR  && python setup.py bdist_wheel --universal
+      # We do not need to force install smdebug-rules. The container used for PR builds do not have smdebug rules binary.
+      # Force installing rules binary attempts to re-install ipython-genutils which fails on PyTorch Ubuntu 16.04 containers.
+      - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal
+      - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --force-reinstall $RULES_CODEBUILD_SRC_DIR/dist/*.whl; else pip install $RULES_CODEBUILD_SRC_DIR/dist/*.whl; fi
+      - cd $CODEBUILD_SRC_DIR  && pip install --force-reinstall dist/*.whl && cd ..
+      - cd $CODEBUILD_SRC_DIR  && chmod +x config/tests.sh && PYTHONPATH=. && ./config/tests.sh  && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      - pip show smdebug
+      - pip show smdebug_rules
+      - echo 'Uploading Coverage to CodeCov'
+      - bash $CODEBUILD_SRC_DIR/config/codecov.sh
+      - cd $RULES_CODEBUILD_SRC_DIR && chmod +x config/tests.sh && PYTHONPATH=. && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+
+  post_build:
+    commands:
+      - . $CODEBUILD_SRC_DIR/config/upload_on_end.sh
+      - rm -rf $CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - rm -rf $RULES_CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 0 ]; then echo "ERROR BUILD FAILED " && exit 1 ; fi
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 1 ]; then echo "INFO BUILD SUCCEEDED !!! " ; fi
diff --git a/config/buildspec_vanilla_framework_tests.yml b/config/buildspec_vanilla_framework_tests.yml
@@ -7,6 +7,8 @@ env:
     run_pytest_xgboost: "enable"
     run_integration_pytest_pytorch: "enable"
     run_integration_pytest_mxnet: "enable"
+    run_pytest_profiler: "disable"
+    run_integration_pytest_tensorflow2: "disable"
     run_integration_pytest_tensorflow: "enable"
     run_integration_pytest_xgboost: "enable"
     zero_code_change_test: "disable"
@@ -17,9 +19,10 @@ phases:
         - apt-get install sudo -qq -o=Dpkg::Use-Pty=0
         - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
-        - pip install --upgrade pip==19.3.1
-        - pip install -q pytest pytest-cov wheel pyYaml pytest-html keras==2.3.1 mxnet torch xgboost pre-commit tensorflow_datasets torchvision
+        - pip install --upgrade pip==20.3.3
+        - pip install -q -U pytest pytest-cov wheel pyYaml pytest-html keras==2.3.1 mxnet==1.6.0 torch==1.8.0 xgboost==1.3.3 pre-commit tensorflow_datasets==4.0.1 torchvision
         - cd $CODEBUILD_SRC_DIR && chmod +x config/install_smdebug.sh && chmod +x config/check_smdebug_install.sh && ./config/install_smdebug.sh;
+        - pip install --force-reinstall numpy==1.18.5
 
   build:
     commands:

diff --git a/config/buildspec_xgboost_1_2_1.yml b/config/buildspec_xgboost_1_2_1.yml
@@ -0,0 +1,61 @@
+# Build Spec for AWS CodeBuild CI for XGBoost 1.2-1
+# CPU Container Used: 683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-xgboost:1.2-1-cpu-py3
+# GPU Container Used: 683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-xgboost:1.2-1
+
+# Note: The xgboost team maintains a single container for both CPU and GPU
+# There is no special image tag for GPU containers;
+
+version: 0.2
+
+env:
+  variables:
+    run_pytest_pytorch: "disable"
+    run_pytest_mxnet: "disable"
+    run_pytest_tensorflow: "disable"
+    run_pytest_tensorflow2: "disable"
+    run_pytest_profiler: "disable"
+    run_pytest_xgboost: "enable"
+    run_integration_pytest_pytorch: "disable"
+    run_integration_pytest_mxnet: "disable"
+    run_integration_pytest_tensorflow: "disable"
+    run_integration_pytest_tensorflow2: "disable"
+    run_integration_pytest_xgboost: "enable"
+    zero_code_change_test: "enable"
+    code_coverage_smdebug_rules: "true"
+
+phases:
+  install:
+    commands:
+        # The XGBoost container requires an update PUB_KEY to run the apt-get update
+        - if [ "$run_pytest_xgboost" = "enable" ]; then su && apt-key adv --keyserver hkp://keyserver.ubuntu.com:80  --recv 684BA42D; fi
+        - su && apt-key adv --keyserver hkp://keyserver.ubuntu.com:80  --recv 684BA42D && apt-get update
+        - apt-get install sudo -qq -o=Dpkg::Use-Pty=0 # silence output: https://askubuntu.com/a/668859/724247
+        - pip install awscli==1.19.5
+        - . config/change_branch.sh #EXPORTS BRANCHES FOR OTHER REPOS AND CURRENT REPO.
+        - cd $CODEBUILD_SRC_DIR && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
+        - pip install --upgrade pip==19.3.1
+        - pip install -q pytest==6.1.2 pytest-cov==2.10.1 wheel==0.35.1 pyYaml==5.3.1 pytest-html==3.0.0 sagemaker==2.16.3 pre-commit==2.6.0
+        - pip install -q matplotlib==3.3.1 && pip3 install seaborn==0.10.1 nbconvert==5.6.1 papermill==2.1.2 jupyter==1.0.0 scipy==1.5.2 scikit-learn==0.23.2 bokeh==2.2.3
+        - pip uninstall -y boto3 && pip uninstall -y aiobotocore && pip uninstall -y botocore
+        - pip uninstall -y typing # see: https://t.corp.amazon.com/P43975146/overview
+
+  pre_build:
+    commands:
+      - cd $CODEBUILD_SRC_DIR
+      - cd $RULES_CODEBUILD_SRC_DIR
+
+  build:
+    commands:
+      - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal && pip install --upgrade --force-reinstall dist/*.whl && cd ..
+      - cd $CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal && pip install --force-reinstall dist/*.whl && cd ..
+      - cd $RULES_CODEBUILD_SRC_DIR  && chmod +x config/tests.sh && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      - cd $CODEBUILD_SRC_DIR && chmod +x config/tests.sh && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && ./config/tests.sh  && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
+      - echo 'Connect to CodeCov'
+      - bash $CODEBUILD_SRC_DIR/config/codecov.sh
+
+  post_build:
+    commands:
+      - rm -rf $CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - rm -rf $RULES_CODEBUILD_SRC_DIR/upload/$CURRENT_COMMIT_PATH
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 0 ]; then echo "ERROR BUILD FAILED " && exit 1 ; fi
+      - if [ "$CODEBUILD_BUILD_SUCCEEDING" -eq 1 ]; then echo "INFO BUILD SUCCEEDED !!! " ; fi
diff --git a/config/buildspec_zero_code_change.yml b/config/buildspec_zero_code_change.yml
@@ -11,6 +11,7 @@ env:
     run_pytest_tensorflow2: "disable"
     run_pytest_xgboost: "disable"
     run_integration_pytest_pytorch: "disable"
+    run_pytest_profiler: "disable"
     run_integration_pytest_mxnet: "enable"
     run_integration_pytest_tensorflow: "disable"
     run_integration_pytest_tensorflow2: "disable"
@@ -20,11 +21,13 @@ env:
 phases:
   install:
     commands:
+      # The recent update to XGBoost container is requiring to update the PUB_KEY for successfully running the apt-get update
+      - if [ "$run_pytest_xgboost" = "enable" ]; then su && apt-key adv --keyserver hkp://keyserver.ubuntu.com:80  --recv 684BA42D; fi
       - apt-get update
       - apt-get install sudo -qq -o=Dpkg::Use-Pty=0
       - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
       - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
-      - pip install --upgrade pip==19.3.1
+      - pip install --upgrade pip==20.3.3
       - pip install -q pytest wheel pyYaml pytest-html pre-commit awscli pytest-cov
       - cd $CODEBUILD_SRC_DIR && chmod +x config/install_smdebug.sh && chmod +x config/check_smdebug_install.sh && ./config/install_smdebug.sh;
 

diff --git a/config/install_smdebug.sh b/config/install_smdebug.sh
@@ -51,7 +51,7 @@ if [ "$SMDEBUG_S3_BINARY" ]; then
   fi
   echo "Commit hash on sagemaker-debugger repository being used: $CORE_COMMIT"
   cd $CODEBUILD_SRC_DIR && git checkout "$CORE_COMMIT"
-  python setup.py bdist_wheel --universal && pip install --force-reinstall dist/*.whl
+  python setup.py bdist_wheel --universal && pip install -U dist/*.whl
 else
   # if the env var stable_release is not set, then this else block is executed.
   if [ -z "$CORE_COMMIT" ]; then export CORE_COMMIT=$(git log -1 --pretty=%h); fi
@@ -65,15 +65,15 @@ else
   else
     pip install dist/*.whl
   fi
-  cd $CODEBUILD_SRC_DIR && git checkout "$CORE_COMMIT" && python setup.py bdist_wheel --universal && pip install --force-reinstall dist/*.whl
+  cd $CODEBUILD_SRC_DIR && git checkout "$CORE_COMMIT" && python setup.py bdist_wheel --universal && pip install -U dist/*.whl
 fi
 
 if [ "$run_pytest_mxnet" == 'enable' ]; then
   ./config/check_smdebug_install.sh mxnet
 fi
 if [ "$run_pytest_tensorflow" == 'enable' ]; then
   ./config/check_smdebug_install.sh tensorflow
-  pip install tensorflow_datasets
+  pip install tensorflow_datasets==4.0.1
 fi
 if [ "$run_pytest_pytorch" == 'enable' ]; then
   ./config/check_smdebug_install.sh torch

diff --git a/config/profiler/buildspec_profiler_sagemaker_pytorch_1_5_0_integration_tests.yml b/config/profiler/buildspec_profiler_sagemaker_pytorch_1_5_0_integration_tests.yml
@@ -0,0 +1,13 @@
+# Build Spec for AWS CodeBuild CI
+
+version: 0.2
+env:
+  variables:
+    use_current_branch: "true"
+    enable_smdataparallel_tests: "false"
+    force_run_tests: "false"
+    framework: "pytorch"
+phases:
+  build:
+    commands:
+      - chmod +x config/profiler/run_profiler_integration_tests.sh && ./config/profiler/run_profiler_integration_tests.sh
diff --git a/config/profiler/buildspec_profiler_sagemaker_pytorch_1_6_0_integration_tests.yml b/config/profiler/buildspec_profiler_sagemaker_pytorch_1_6_0_integration_tests.yml
@@ -0,0 +1,13 @@
+# Build Spec for AWS CodeBuild CI
+
+version: 0.2
+env:
+  variables:
+    use_current_branch: "true"
+    enable_smdataparallel_tests: "true"
+    force_run_tests: "false"
+    framework: "pytorch"
+phases:
+  build:
+    commands:
+      - chmod +x config/profiler/run_profiler_integration_tests.sh && ./config/profiler/run_profiler_integration_tests.sh
diff --git a/config/profiler/buildspec_profiler_sagemaker_pytorch_1_7_0_integration_tests.yml b/config/profiler/buildspec_profiler_sagemaker_pytorch_1_7_0_integration_tests.yml
@@ -0,0 +1,13 @@
+# Build Spec for AWS CodeBuild CI
+
+version: 0.2
+env:
+  variables:
+    use_current_branch: "true"
+    enable_smdataparallel_tests: "false"
+    force_run_tests: "false"
+    framework: "pytorch"
+phases:
+  build:
+    commands:
+      - chmod +x config/profiler/run_profiler_integration_tests.sh && ./config/profiler/run_profiler_integration_tests.sh