From 9f6d289e645cd0e29de5a2db96b7c90d96a6a63b Mon Sep 17 00:00:00 2001 From: David Cavazos Date: Thu, 25 Feb 2021 12:59:19 -0800 Subject: [PATCH] dataflow: use us-central1-f for gpu test --- dataflow/gpu-workers/e2e_test.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dataflow/gpu-workers/e2e_test.py b/dataflow/gpu-workers/e2e_test.py index 064b03ac283..e322491ca2c 100644 --- a/dataflow/gpu-workers/e2e_test.py +++ b/dataflow/gpu-workers/e2e_test.py @@ -25,6 +25,8 @@ PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"] BUCKET_NAME = f"dataflow-gpu-test-{SUFFIX}" IMAGE_NAME = f"gcr.io/{PROJECT}/dataflow/gpu-workers/test-{SUFFIX}:latest" +REGION = "us-central1" +ZONE = "us-central1-f" @pytest.fixture(scope="session") @@ -109,8 +111,6 @@ def test_python_version(image_name: str, configure_docker: None) -> None: def test_end_to_end(bucket_name: str, image_name: str) -> None: # Run the Beam pipeline in Dataflow making sure GPUs are used. gpu_type = "nvidia-tesla-t4" - region = "us-central1" - worker_zone = "us-central1-a" subprocess.run( [ "python", @@ -118,11 +118,11 @@ def test_end_to_end(bucket_name: str, image_name: str) -> None: f"--output-path-prefix=gs://{bucket_name}/outputs/", "--runner=DataflowRunner", f"--project={PROJECT}", - f"--region={region}", + f"--region={REGION}", f"--temp_location=gs://{bucket_name}/temp", "--worker_machine_type=custom-1-13312-ext", f"--worker_harness_container_image={image_name}", - f"--worker_zone={worker_zone}", + f"--worker_zone={ZONE}", f"--experiments=worker_accelerator=type={gpu_type},count=1,install-nvidia-driver", "--experiments=use_runner_v2", ],