dataflow/custom-containers: add readmes and support cli options (GoogleCloudPlatform#7601)

* add readmes and support cli options

* run from run.yaml file

* update tests to run from run.yaml

* launch jobs asynchronously, do not wait until finish

* add googleapiclient library

* remove subprocess unused import

* remove unused substitution

* specify image tag with uuid for tests

* clean commands

Author: David Cavazos

dataflow/custom-containers/miniconda/.gcloudignore

@@ -0,0 +1,5 @@
+# Ignore everything except the source files.
+**/*
+!Dockerfile
+!requirements.txt
+!*.py

dataflow/custom-containers/miniconda/Dockerfile

@@ -18,7 +18,7 @@ FROM continuumio/miniconda3:4.10.3-alpine AS builder
 
 # Create a virtual environment and make it standalone with conda-pack.
 # https://conda.github.io/conda-pack
-RUN conda create -y -n env python=3.8 \
+RUN conda create -y -n env python=3.9 \
     && conda install -y conda-pack \
     && conda-pack -n env -o /tmp/env.tar \
     && mkdir /opt/python \
@@ -28,8 +28,10 @@ RUN conda create -y -n env python=3.8 \
 # -- Main image
 FROM ubuntu:latest
 
+WORKDIR /pipeline
+
 # Set the entrypoint to Apache Beam SDK worker launcher.
-COPY --from=apache/beam_python3.8_sdk:2.37.0 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.9_sdk:2.37.0 /opt/apache/beam /opt/apache/beam
 ENTRYPOINT [ "/opt/apache/beam/boot" ]
 
 # Copy the python installation from the builder stage.

dataflow/custom-containers/miniconda/README.md

@@ -0,0 +1,57 @@
+# Miniconda custom container
+
+This sample runs a Dataflow pipeline where the workers use a custom container with a
+[Miniconda](https://docs.conda.io/en/latest/miniconda.html)
+Python environment.
+
+## Before you begin
+
+Make sure you have followed the
+[Dataflow setup instructions](../../README.md).
+
+You will also need:
+* A Cloud Storage bucket, click [here to create one](https://console.cloud.google.com/storage/create-bucket) if needed.
+
+```sh
+export PROJECT=$(gcloud config get-value project)
+export BUCKET="my-cloud-storage-bucket"
+```
+
+> ℹ️ Make sure your `BUCKET` name does _not_ include the `gs://` prefix.
+
+## Building the Docker image
+
+We use
+[Cloud Build](https://cloud.google.com/build)
+to build the container image for the workers and save it in
+[Container Registry](https://cloud.google.com/container-registry/).
+
+```sh
+export IMAGE=gcr.io/$PROJECT/samples/dataflow-miniconda:latest
+
+gcloud builds submit . --tag=$IMAGE
+```
+
+## Running the Dataflow job
+
+We use Cloud Build to run the [Dataflow](https://cloud.google.com/dataflow) job.
+
+The [`run.yaml`](run.yaml) file contains the command we use to launch the Dataflow job.
+
+> ℹ️ We launch the job using the worker image to make sure the job launches
+> with the same Python version as the workers and all the dependencies installed.
+
+```sh
+# Choose the location where you want to run your Dataflow job.
+# For a list of all supported locations, see:
+#   https://cloud.google.com/dataflow/docs/resources/locations
+export REGION="us-central1"
+
+export JOB_NAME="dataflow-miniconda-$(date +"%F-%H%M%S")"
+export TEMP_LOCATION="gs://$BUCKET/samples/dataflow-miniconda"
+
+gcloud builds submit \
+    --config run.yaml \
+    --substitutions "_JOB_NAME=$JOB_NAME,_REGION=$REGION,_TEMP_LOCATION=$TEMP_LOCATION" \
+    --no-source
+```

dataflow/custom-containers/miniconda/e2e_test.py

@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import subprocess
-
 try:
     # `conftest` cannot be imported when running in `nox`, but we still
     # try to import it for the autocomplete when writing the tests.
@@ -34,21 +32,24 @@ def bucket_name(utils: Utils) -> str:
 
 @pytest.fixture(scope="session")
 def container_image(utils: Utils) -> str:
-    yield from utils.cloud_build_submit(image_name=NAME)
-
-
-def test_tensorflow_minimal(
-    utils: Utils, bucket_name: str, container_image: str
-) -> None:
-    subprocess.check_call(
-        [
-            "python",
-            "main.py",
-            "--runner=DataflowRunner",
-            f"--project={utils.project}",
-            f"--region={utils.region}",
-            f"--temp_location=gs://{bucket_name}",
-            f"--sdk_container_image={container_image}",
-            "--experiment=use_runner_v2",
-        ]
+    yield from utils.cloud_build_submit(NAME)
+
+
+@pytest.fixture(scope="session")
+def run_dataflow_job(utils: Utils, bucket_name: str, container_image: str) -> str:
+    yield from utils.cloud_build_submit(
+        config="run.yaml",
+        substitutions={
+            "_IMAGE": container_image,
+            "_JOB_NAME": utils.hyphen_name(NAME),
+            "_TEMP_LOCATION": f"gs://{bucket_name}/temp",
+            "_REGION": utils.region,
+        },
+        source="--no-source",
     )
+
+
+def test_custom_container_miniconda(utils: Utils, run_dataflow_job: str) -> None:
+    # Wait until the job finishes.
+    job_id = utils.dataflow_job_id(utils.hyphen_name(NAME))
+    utils.dataflow_jobs_wait(job_id)

dataflow/custom-containers/miniconda/main.py

@@ -16,17 +16,21 @@
 
 import logging
 import platform
+from typing import List, Optional
 
 import apache_beam as beam
+from apache_beam.options.pipeline_options import PipelineOptions
 
 
-def run() -> None:
-    with beam.Pipeline() as pipeline:
-        (
-            pipeline
-            | "Create data" >> beam.Create(["Hello", "World!", platform.platform()])
-            | "Print" >> beam.Map(logging.info)
-        )
+def run(beam_args: Optional[List[str]] = None) -> None:
+    beam_options = PipelineOptions(beam_args, save_main_session=True)
+    pipeline = beam.Pipeline(options=beam_options)
+    (
+        pipeline
+        | "Create data" >> beam.Create(["Hello", "World!", platform.platform()])
+        | "Print" >> beam.Map(logging.info)
+    )
+    pipeline.run()
 
 
 if __name__ == "__main__":

dataflow/custom-containers/miniconda/noxfile_config.py

@@ -22,10 +22,10 @@
 
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
-    # > ℹ️ We're opting out of all Python versions except 3.8.
+    # > ℹ️ We're opting out of all Python versions except 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.9", "3.10"],
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.10"],
     # Old samples are opted out of enforcing Python type hints
     # All new samples should feature them
     "enforce_type_hints": True,

dataflow/custom-containers/miniconda/requirements-test.txt

@@ -1,3 +1,4 @@
+google-api-python-client==2.40.0
 google-cloud-storage==1.43.0
 pytest-xdist==2.5.0
 pytest==6.2.4

dataflow/custom-containers/miniconda/requirements.txt

@@ -1 +1 @@
-apache-beam[gcp]==2.36.0
+apache-beam[gcp]==2.37.0

dataflow/custom-containers/miniconda/run.yaml

@@ -0,0 +1,50 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This Cloud Build config runs a Dataflow job using GPUs.
+# We use the same worker image to launch the job.
+# That way we guarantee the same Python version for the workers.
+# It also already has all the requirements installed.
+
+# -----------------------------------------------------------------------------
+# To learn more about this file:
+#   https://cloud.google.com/build/docs/build-config
+#
+# To learn more about Cloud Build variable substitutions:
+#   https://cloud.google.com/build/docs/configuring-builds/substitute-variable-values#using_user-defined_substitutions
+# -----------------------------------------------------------------------------
+
+substitutions:
+  _JOB_NAME: ''
+  _TEMP_LOCATION: ''
+  _REGION: us-central1
+  _IMAGE: samples/dataflow-miniconda:latest
+
+steps:
+- name: gcr.io/$PROJECT_ID/$_IMAGE
+  entrypoint: python
+  args:
+  - /pipeline/main.py
+  - --runner=DataflowRunner
+  - --project=$PROJECT_ID
+  - --region=$_REGION
+  - --job_name=$_JOB_NAME
+  - --temp_location=$_TEMP_LOCATION
+  - --sdk_container_image=gcr.io/$PROJECT_ID/$_IMAGE
+
+options:
+  logging: CLOUD_LOGGING_ONLY
+
+# Use the Compute Engine default service account to launch the job.
+serviceAccount: projects/$PROJECT_ID/serviceAccounts/[email protected]

dataflow/custom-containers/minimal/.gcloudignore

@@ -0,0 +1,5 @@
+# Ignore everything except the source files.
+**/*
+!Dockerfile
+!requirements.txt
+!*.py