dataflow: downgrade to python 3.10

dataflow/conftest.py

@@ -79,7 +79,8 @@ def storage_bucket(name: str) -> str:
         yield bucket.name
 
         # Print all the objects in the bucket before deleting for debugging.
-        logging.info(f"Deleting bucket {bucket.name} with the following contents:")
+        logging.info(
+            f"Deleting bucket {bucket.name} with the following contents:")
         total_files = 0
         total_size = 0
         for blob in bucket.list_blobs():
@@ -159,7 +160,8 @@ def pubsub_topic(name: str, project: str = PROJECT) -> str:
         from google.cloud import pubsub
 
         publisher_client = pubsub.PublisherClient()
-        topic_path = publisher_client.topic_path(project, Utils.hyphen_name(name))
+        topic_path = publisher_client.topic_path(
+            project, Utils.hyphen_name(name))
         topic = publisher_client.create_topic(request={"name": topic_path})
 
         logging.info(f"Created pubsub_topic: {topic.name}")
@@ -169,7 +171,8 @@ def pubsub_topic(name: str, project: str = PROJECT) -> str:
         # library throws an error upon deletion.
         # We use gcloud for a workaround. See also:
         # https://github.com/GoogleCloudPlatform/python-docs-samples/issues/4492
-        cmd = ["gcloud", "pubsub", "--project", project, "topics", "delete", topic.name]
+        cmd = ["gcloud", "pubsub", "--project",
+               project, "topics", "delete", topic.name]
         logging.info(f"{cmd}")
         subprocess.check_call(cmd)
         logging.info(f"Deleted pubsub_topic: {topic.name}")
@@ -224,7 +227,8 @@ def _infinite_publish_job() -> None:
             publisher_client = pubsub.PublisherClient()
             for i in itertools.count():
                 msg = new_msg(i)
-                publisher_client.publish(topic_path, msg.encode("utf-8")).result()
+                publisher_client.publish(
+                    topic_path, msg.encode("utf-8")).result()
                 time.sleep(sleep_sec)
 
         # Start a subprocess in the background to do the publishing.
@@ -281,7 +285,8 @@ def cloud_build_submit(
                     ]
                     logging.info(f"{cmd}")
                     subprocess.check_call(cmd)
-                    logging.info(f"Cloud build finished successfully: {config}")
+                    logging.info(
+                        f"Cloud build finished successfully: {config}")
                     yield f.read()
             except Exception as e:
                 logging.exception(e)
@@ -299,7 +304,8 @@ def cloud_build_submit(
             ]
             logging.info(f"{cmd}")
             subprocess.check_call(cmd)
-            logging.info(f"Created image: gcr.io/{project}/{image_name}:{UUID}")
+            logging.info(
+                f"Created image: gcr.io/{project}/{image_name}:{UUID}")
             yield f"{image_name}:{UUID}"
         else:
             raise ValueError("must specify either `config` or `image_name`")
@@ -317,7 +323,8 @@ def cloud_build_submit(
             ]
             logging.info(f"{cmd}")
             subprocess.check_call(cmd)
-            logging.info(f"Deleted image: gcr.io/{project}/{image_name}:{UUID}")
+            logging.info(
+                f"Deleted image: gcr.io/{project}/{image_name}:{UUID}")
 
     @staticmethod
     def dataflow_job_url(
@@ -424,6 +431,8 @@ def job_is_done() -> bool:
             return False
 
         Utils.wait_until(job_is_done, timeout_sec, poll_interval_sec)
+        assert job_is_done(), (f"Dataflow job is not done after {timeout_sec} seconds\n"
+                               + Utils.dataflow_job_url(job_id, project, region))
 
     @staticmethod
     def dataflow_jobs_cancel(

dataflow/gpu-examples/tensorflow-landsat-prime/Dockerfile

@@ -25,15 +25,15 @@ COPY *.py ./
 RUN apt-get update \
     # Install Python and other system dependencies.
     && apt-get install -y --no-install-recommends \
-        curl g++ python3.11-dev python3-pip python3-distutils \
+        curl g++ python3.10-dev python3.10-venv python3-distutils \
     && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 10 \
+    && update-alternatives --install /usr/bin/python python /usr/bin/python3.10 10 \
+    && curl https://bootstrap.pypa.io/get-pip.py | python \
     # Install the pipeline requirements.
-    && pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir -r requirements.txt \
     && pip check
 
 # Set the entrypoint to Apache Beam SDK worker launcher.
 # Check this matches the apache-beam version in the requirements.txt
-COPY --from=apache/beam_python3.11_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.10_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
 ENTRYPOINT [ "/opt/apache/beam/boot" ]

dataflow/gpu-examples/tensorflow-landsat-prime/noxfile_config.py

@@ -22,10 +22,10 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ Test only on Python 3.11.
+    # > ℹ️ Test only on Python 3.10.
     # > The Python version used is defined by the Dockerfile, so it's redundant
     # > to run multiple tests since they would all be running the same Dockerfile.
-    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.10"],
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.11"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

dataflow/gpu-examples/tensorflow-landsat/Dockerfile

@@ -25,15 +25,15 @@ COPY *.py ./
 RUN apt-get update \
     # Install Python and other system dependencies.
     && apt-get install -y --no-install-recommends \
-        curl g++ python3.11-dev python3-pip python3-distutils \
+    curl g++ python3.10-dev python3.10-venv python3-distutils \
     && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 10 \
+    && update-alternatives --install /usr/bin/python python /usr/bin/python3.10 10 \
+    && curl https://bootstrap.pypa.io/get-pip.py | python \
     # Install the pipeline requirements.
-    && pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir -r requirements.txt \
     && pip check
 
 # Set the entrypoint to Apache Beam SDK worker launcher.
 # Check this matches the apache-beam version in the requirements.txt
-COPY --from=apache/beam_python3.11_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.10_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
 ENTRYPOINT [ "/opt/apache/beam/boot" ]

dataflow/gpu-examples/tensorflow-landsat/noxfile_config.py

@@ -22,10 +22,10 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ Test only on Python 3.11.
+    # > ℹ️ Test only on Python 3.10.
     # > The Python version used is defined by the Dockerfile, so it's redundant
     # > to run multiple tests since they would all be running the same Dockerfile.
-    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.10"],
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.11"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

dataflow/gpu-examples/tensorflow-minimal/Dockerfile

@@ -25,15 +25,15 @@ COPY *.py ./
 RUN apt-get update \
     # Install Python and other system dependencies.
     && apt-get install -y --no-install-recommends \
-        curl g++ python3.11-dev python3-pip python3-distutils \
+        curl g++ python3.10-dev python3.10-venv python3-distutils \
     && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 10 \
+    && update-alternatives --install /usr/bin/python python /usr/bin/python3.10 10 \
+    && curl https://bootstrap.pypa.io/get-pip.py | python \
     # Install the pipeline requirements.
-    && pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir -r requirements.txt \
     && pip check
 
 # Set the entrypoint to Apache Beam SDK worker launcher.
 # Check this matches the apache-beam version in the requirements.txt
-COPY --from=apache/beam_python3.11_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.10_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
 ENTRYPOINT [ "/opt/apache/beam/boot" ]

dataflow/gpu-examples/tensorflow-minimal/noxfile_config.py

@@ -22,10 +22,10 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ Test only on Python 3.11.
+    # > ℹ️ Test only on Python 3.10.
     # > The Python version used is defined by the Dockerfile, so it's redundant
     # > to run multiple tests since they would all be running the same Dockerfile.
-    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.10"],
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.11"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

people-and-planet-ai/geospatial-classification/noxfile_config.py

@@ -22,8 +22,8 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ Test only on Python 3.11.
-    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.10"],
+    # > ℹ️ Test only on Python 3.10.
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.11"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

people-and-planet-ai/timeseries-classification/Dockerfile

@@ -15,25 +15,18 @@
 # NVIDIA CUDA container: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/cuda
 # Supported NVIDIA images: https://gitlab.com/nvidia/container-images/cuda/blob/master/doc/supported-tags.md
 # TensorFlow/CUDA compatibility: https://www.tensorflow.org/install/source#gpu
-FROM nvcr.io/nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
+FROM python:3.10-slim
 
 WORKDIR /pipeline
 
 COPY requirements.txt .
 COPY *.py ./
 
-RUN apt-get update \
-    # Install Python and other system dependencies.
-    && apt-get install -y --no-install-recommends \
-        curl g++ python3.11-dev python3-pip python3-distutils \
-    && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 10 \
-    # Install the pipeline requirements.
-    && pip install --no-cache-dir --upgrade pip \
+RUN pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir -r requirements.txt \
     && pip check
 
 # Set the entrypoint to Apache Beam SDK worker launcher.
 # Check this matches the apache-beam version in the requirements.txt
-COPY --from=apache/beam_python3.11_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.10_sdk:2.46.0 /opt/apache/beam /opt/apache/beam
 ENTRYPOINT [ "/opt/apache/beam/boot" ]

people-and-planet-ai/timeseries-classification/e2e_test.py

@@ -19,7 +19,6 @@
 import time
 import uuid
 
-from google.cloud import aiplatform
 from google.cloud import storage
 from googleapiclient.discovery import build
 import numpy as np
@@ -154,6 +153,7 @@ def service_url(bucket_name: str, container_image: str) -> str:
             f"--project={PROJECT}",
             f"--region={REGION}",
             "--memory=1G",
+            "--timeout=30m",
             f"--set-env-vars=PROJECT={PROJECT}",
             f"--set-env-vars=STORAGE_PATH=gs://{bucket_name}",
             f"--set-env-vars=REGION={REGION}",
@@ -260,42 +260,17 @@ def create_datasets(
 
 
 @pytest.fixture(scope="session")
-def train_model(service_url: str, access_token: str, create_datasets: str) -> str:
+def train_model(service_url: str, access_token: str, create_datasets: str) -> None:
+    logging.info("Training model")
     response = requests.post(
         url=f"{service_url}/train-model",
         headers={"Authorization": f"Bearer {access_token}"},
-        json={"train_epochs": 10, "batch_size": 8},
-    ).json()
-    logging.info(f"train_model response: {response}")
-
-    job_id = response["job_id"]
-    job_url = response["job_url"]
-    logging.info(f"train_model job_id: {job_id}")
-    logging.info(f"train_model job_url: {job_url}")
-
-    # Wait until the model training job finishes.
-    ai_client = aiplatform.gapic.JobServiceClient(
-        client_options={"api_endpoint": "us-central1-aiplatform.googleapis.com"}
+        json={"train_epochs": 10, "batch_size": 8, "sync": True},
     )
-
-    status = None
-    logging.info("Waiting for model to train.")
-    for _ in range(0, TIMEOUT_SEC, POLL_INTERVAL_SEC):
-        # https://googleapis.dev/python/aiplatform/latest/aiplatform_v1/job_service.html
-        job = ai_client.get_custom_job(
-            name=f"projects/{PROJECT}/locations/{REGION}/customJobs/{job_id}"
-        )
-        status = job.state.name
-        if status in VERTEX_AI_FINISHED_STATE:
-            break
-        time.sleep(POLL_INTERVAL_SEC)
-
-    logging.info(f"Model job finished with status {status}")
-    assert status == VERTEX_AI_SUCCESS_STATE, f"job_url: {job_url}"
-    yield job_id
+    logging.info(f"train_model response: {response}")
 
 
-def test_predict(service_url: str, access_token: str, train_model: str) -> None:
+def test_predict(service_url: str, access_token: str, train_model: None) -> None:
     with open("test_data/56980685061237.npz", "rb") as f:
         input_data = pd.DataFrame(np.load(f)["x"])
 

people-and-planet-ai/timeseries-classification/main.py

@@ -100,7 +100,6 @@ def run_train_model() -> dict:
         params = {
             "project": args.get("project", PROJECT),
             "region": args.get("region", REGION),
-            "container_image": args.get("container_image", CONTAINER_IMAGE),
             "train_data_dir": args.get("train_data_dir", TRAIN_DATA_DIR),
             "eval_data_dir": args.get("eval_data_dir", EVAL_DATA_DIR),
             "training_dir": args.get("training_dir", TRAINING_DIR),
@@ -109,13 +108,12 @@ def run_train_model() -> dict:
             "machine_type": args.get("machine_type", DEFAULT_MACHINE_TYPE),
             "gpu_type": args.get("gpu_type", DEFAULT_GPU_TYPE),
             "gpu_count": args.get("gpu_count", DEFAULT_GPU_COUNT),
+            "sync": args.get("sync", False),
         }
-        job_id = train_model.run(**params)
+        train_model.run(**params)
 
         return {
             "method": "train-model",
-            "job_id": job_id,
-            "job_url": f"https://console.cloud.google.com/vertex-ai/locations/{REGION}/training/{job_id}/cpu?project={PROJECT}",
             "params": params,
         }
     except Exception as e:
@@ -148,6 +146,4 @@ def run_predict() -> dict:
 
 
 if __name__ == "__main__":
-    import os
-
     app.run(debug=True, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))

people-and-planet-ai/timeseries-classification/noxfile_config.py

@@ -22,10 +22,10 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ Test only on Python 3.11.
+    # > ℹ️ Test only on Python 3.9.
     # > The Python version used is defined by the Dockerfile, so it's redundant
     # > to run multiple tests since they would all be running the same Dockerfile.
-    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.10"],
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.10", "3.11"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

people-and-planet-ai/timeseries-classification/predict.py

@@ -36,7 +36,7 @@ def predict(model: keras.Model, inputs: Dict[str, np.ndarray]) -> pd.DataFrame:
     }
 
     predictions = model.predict(inputs_batch)
-    return data[trainer.PADDING :].assign(is_fishing=predictions["is_fishing"][0])
+    return data[trainer.PADDING:].assign(is_fishing=predictions["is_fishing"][0])
 
 
 def run(model_dir: str, inputs: Dict[str, List[float]]) -> Dict[str, np.ndarray]: