Switched to the Pytorch container to test Beam with Python 3.10 (#24)

Author: liferoad

.github/workflows/docker_test.yml

@@ -49,10 +49,10 @@ jobs:
           registry: "us-docker.pkg.dev"
           username: "oauth2accesstoken"
           password: "${{ steps.auth.outputs.access_token }}"
-      - name: Set up Python 3.8
+      - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: "3.8"
+          python-version: "3.10"
       - name: Init env with the test Beam and docker URI
         run: |
           cp tests/sample.env.tf .env
@@ -61,7 +61,7 @@ jobs:
           make test-latest-env
           sed -i '/CUSTOM_CONTAINER_IMAGE=/d' .env
           echo -e "\n" >> .env
-          echo "CUSTOM_CONTAINER_IMAGE=us-docker.pkg.dev/apache-beam-testing/dataflow-ml-starter/tf_gpu:test-beam" >> .env
+          echo "CUSTOM_CONTAINER_IMAGE=us-docker.pkg.dev/apache-beam-testing/dataflow-ml-starter/pytorch_gpu:test-beam" >> .env
           echo '${{ steps.auth.outputs.access_token }}' | docker login -u oauth2accesstoken --password-stdin https://us-docker.pkg.dev
           make init
       - name: Build and push Docker image

tests/sample.env.pytorch

@@ -0,0 +1,33 @@
+################################################################################
+### PYTHON SDK SETTINGS
+################################################################################
+PYTHON_VERSION=3.10
+BEAM_VERSION=2.60.0
+DOCKERFILE_TEMPLATE=pytorch_gpu.Dockerfile
+DOCKER_CREDENTIAL_REGISTRIES="us-docker.pkg.dev"
+################################################################################
+### GCP SETTINGS
+################################################################################
+PROJECT_ID=apache-beam-testing
+REGION=us-central1
+ZONE=us-central1-f
+DISK_SIZE_GB=50
+MACHINE_TYPE=n1-standard-2
+VM_NAME=beam-ml-starter-gpu
+################################################################################
+### DATAFLOW JOB SETTINGS
+################################################################################
+STAGING_LOCATION=gs://temp-storage-for-perf-tests/loadtests
+TEMP_LOCATION=gs://temp-storage-for-perf-tests/loadtests
+CUSTOM_CONTAINER_IMAGE=us-docker.pkg.dev/apache-beam-testing/dataflow-ml-starter/pytorch_gpu:test
+SERVICE_OPTIONS="worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver"
+################################################################################
+### DATAFLOW JOB MODEL SETTINGS
+################################################################################
+MODEL_STATE_DICT_PATH="gs://apache-beam-ml/models/torchvision.models.mobilenet_v2.pth"
+MODEL_NAME=mobilenet_v2
+################################################################################
+### DATAFLOW JOB INPUT&OUTPUT SETTINGS
+################################################################################
+INPUT_DATA="gs://apache-beam-ml/testing/inputs/openimage_50k_benchmark.txt"
+OUTPUT_DATA="gs://temp-storage-for-end-to-end-tests/temp-storage-for-end-to-end-tests/dataflow-ml-starter/result_gpu.txt"