diff --git a/test/build_image.yaml b/test/build_image.yaml
index 2f7b518af079..c0d2b7539a60 100644
--- a/test/build_image.yaml
+++ b/test/build_image.yaml
@@ -20,8 +20,7 @@ spec:
   entrypoint: build-images
   arguments:
     parameters:
-    - name: commit-sha
-      value: master
+    - name: context-gcs-uri
     - name: bootstrapper-image
     - name: api-image
     - name: frontend-image
@@ -31,7 +30,6 @@ spec:
   - name: build-images
     inputs:
       parameters:
-      - name: commit-sha
       - name: bootstrapper-image
       - name: api-image
       - name: frontend-image
@@ -39,21 +37,21 @@ spec:
       - name: persistenceagent-image
     steps:
     - - name: build-bootstrapper-image
-        template: build-image
+        template: build-image-from-context-archive
         arguments:
           parameters:
-          - name: commit-sha
-            value: "{{inputs.parameters.commit-sha}}"
+          - name: context-gcs-uri
+            value: "{{workflow.parameters.context-gcs-uri}}"
           - name: docker-path
             value: ml-pipeline
           - name: image-name
             value: "{{inputs.parameters.bootstrapper-image}}"
       - name: build-api-server-image
-        template: build-image
+        template: build-image-from-context-archive
         arguments:
           parameters:
-          - name: commit-sha
-            value: "{{inputs.parameters.commit-sha}}"
+          - name: context-gcs-uri
+            value: "{{workflow.parameters.context-gcs-uri}}"
           - name: docker-path
             value: .
           - name: image-name
@@ -61,11 +59,11 @@ spec:
           - name: docker-file
             value: backend/Dockerfile
       - name: build-frontend-image
-        template: build-image
+        template: build-image-from-context-archive
         arguments:
           parameters:
-          - name: commit-sha
-            value: "{{inputs.parameters.commit-sha}}"
+          - name: context-gcs-uri
+            value: "{{workflow.parameters.context-gcs-uri}}"
           - name: docker-path
             value: .
           - name: docker-file
@@ -73,11 +71,11 @@ spec:
           - name: image-name
             value: "{{inputs.parameters.frontend-image}}"
       - name: build-scheduledworkflow-image
-        template: build-image
+        template: build-image-from-context-archive
         arguments:
           parameters:
-          - name: commit-sha
-            value: "{{inputs.parameters.commit-sha}}"
+          - name: context-gcs-uri
+            value: "{{workflow.parameters.context-gcs-uri}}"
           - name: docker-path
             value: .
           - name: image-name
@@ -85,11 +83,11 @@ spec:
           - name: docker-file
             value: backend/Dockerfile.scheduledworkflow
       - name: build-persistenceagent-image
-        template: build-image
+        template: build-image-from-context-archive
         arguments:
           parameters:
-          - name: commit-sha
-            value: "{{inputs.parameters.commit-sha}}"
+          - name: context-gcs-uri
+            value: "{{workflow.parameters.context-gcs-uri}}"
           - name: docker-path
             value: .
           - name: image-name
@@ -127,3 +125,71 @@ spec:
       securityContext:
         privileged: true
       mirrorVolumeMounts: true
+
+  # Build image and upload to GCR
+  - name: build-image-from-context-archive
+    inputs:
+      parameters:
+      # GCS URI prefix pointing to a .tar.gz archive of Docker build context
+      - name: context-gcs-uri
+      # The relative code path to the Dockerfile
+      - name: docker-path
+      # Name of the Docker file to use. "Dockerfile" by default
+      - name: docker-file
+        value: Dockerfile
+      - name: image-name
+    outputs:
+      parameters:
+        - name: image-name
+          valueFrom:
+            path: /outputs/image-name/file
+    container:
+      image: google/cloud-sdk
+      imagePullPolicy: 'Always'
+      env:
+      - name: DOCKER_HOST
+        value: 127.0.0.1
+      command:
+        - bash
+        - -c
+        - |
+          CONTEXT_GCS_URI="{{inputs.parameters.context-gcs-uri}}" #gs://$TEST_RESULT_BUCKET/$PULL_PULL_SHA/source_code_<base_sha>_<date>.tar.gz
+          DOCKER_PATH="{{inputs.parameters.docker-path}}"
+          DOCKER_FILE="{{inputs.parameters.docker-file}}"
+          IMAGE_NAME="{{inputs.parameters.image-name}}"
+
+          BASE_DIR=/image_builder
+          mkdir $BASE_DIR
+          cd $BASE_DIR
+
+          echo "Downloading Docker build context from $CONTEXT_GCS_URI..."
+          downloaded_code_archive_file=$(mktemp)
+          gsutil cp "$CONTEXT_GCS_URI" "$downloaded_code_archive_file"
+          tar -xzf "$downloaded_code_archive_file" --directory .
+
+          echo "Waiting for Docker-in-Docker daemon to start..."
+          until docker ps; do sleep 3; done;
+
+          gcloud auth configure-docker
+
+          if [ "$BUILD_SCRIPT" == "" ]; then
+            echo "Build image ${IMAGE_NAME} using ${DOCKER_PATH}/${DOCKER_FILE}..."
+            docker build -t ${IMAGE_NAME} -f ${DOCKER_PATH}/${DOCKER_FILE} ${DOCKER_PATH}
+          else
+            echo "Build image ${IMAGE_NAME} using ${BUILD_SCRIPT}..."
+            cd $(dirname ${BUILD_SCRIPT})
+            bash $(basename ${BUILD_SCRIPT}) -i ${IMAGE_NAME}
+          fi
+
+          echo "Pushing image ${IMAGE_NAME}..."
+          docker push ${IMAGE_NAME}
+
+          image_name_output_file=/outputs/image-name/file
+          mkdir -p "$(dirname "$image_name_output_file")"
+          echo $IMAGE_NAME > "$image_name_output_file"
+    sidecars:
+    - name: dind
+      image: docker:17.10-dind
+      securityContext:
+        privileged: true
+      mirrorVolumeMounts: true
\ No newline at end of file
diff --git a/test/presubmit-tests.sh b/test/presubmit-tests.sh
index 65e00b90bb4e..7849d67c123d 100755
--- a/test/presubmit-tests.sh
+++ b/test/presubmit-tests.sh
@@ -70,6 +70,16 @@ echo "presubmit test starts"
 # activating the service account
 gcloud auth activate-service-account --key-file="${GOOGLE_APPLICATION_CREDENTIALS}"
 
+#Uploading the source code to GCS:
+local_code_archive_file=$(mktemp)
+date_string=$(TZ=PST8PDT date +%Y-%m-%d_%H-%M-%S_%Z)
+code_archive_prefix="gs://$TEST_RESULT_BUCKET/$PULL_PULL_SHA/source_code"
+remote_code_archive_uri="${code_archive_prefix}_${PULL_BASE_SHA}_${date_string}.tar.gz"
+
+tar -czf "$local_code_archive_file" .
+
+gsutil cp "$local_code_archive_file" "$remote_code_archive_uri"
+
 #Creating a new GKE cluster if needed
 if [ "$CLUSTER_TYPE" == "create-gke" ]; then
   echo "create test cluster"
@@ -116,7 +126,7 @@ kubectl apply -n argo -f https://raw.githubusercontent.com/argoproj/argo/$ARGO_V
 
 echo "submitting argo workflow for commit ${PULL_PULL_SHA}..."
 ARGO_WORKFLOW=`argo submit $(dirname $0)/${WORKFLOW_FILE} \
--p commit-sha="${PULL_PULL_SHA}" \
+-p context-gcs-uri="$remote_code_archive_uri" \
 -p test-results-gcs-dir="${TEST_RESULTS_GCS_DIR}" \
 -p cluster-type="${CLUSTER_TYPE}" \
 -p bootstrapper-image="${GCR_IMAGE_BASE_DIR}/bootstrapper" \