dataflow: new minimal gpu samples

dataflow/conftest.py

@@ -158,40 +158,64 @@ def _infinite_publish_job() -> None:
         p.terminate()
 
     @staticmethod
-    def container_image(
-        image_path: str,
+    def cloud_build_submit(
+        image_name: Optional[str] = None,
+        config: Optional[str] = None,
+        substitutions: Optional[Dict[str, str]] = None,
         project: str = PROJECT,
-        tag: str = "latest",
-    ) -> str:
-        image_name = f"gcr.io/{project}/{image_path}-{UUID}:{tag}"
+    ) -> None:
+        """Sends a Cloud Build job, if an image_name is provided it will be deleted at teardown."""
         cmd = ["gcloud", "auth", "configure-docker"]
         print(cmd)
-        subprocess.run(cmd, check=True)
-        cmd = [
-            "gcloud",
-            "builds",
-            "submit",
-            f"--project={project}",
-            f"--tag={image_name}",
-            ".",
-        ]
-        print(cmd)
-        subprocess.run(cmd, check=True)
 
-        print(f"container_image: {image_name}")
-        yield image_name
+        if substitutions:
+            cmd_substitutions = [
+                f"--substitutions={','.join([k + '=' + v for k, v in substitutions.items()])}"
+            ]
+        else:
+            cmd_substitutions = []
 
-        cmd = [
-            "gcloud",
-            "container",
-            "images",
-            "delete",
-            image_name,
-            f"--project={project}",
-            "--quiet",
-        ]
-        print(cmd)
         subprocess.run(cmd, check=True)
+        if config:
+            cmd = [
+                "gcloud",
+                "builds",
+                "submit",
+                f"--project={project}",
+                f"--config={config}",
+                *cmd_substitutions,
+            ]
+            print(cmd)
+            subprocess.run(cmd, check=True)
+            yield config
+        elif image_name:
+            cmd = [
+                "gcloud",
+                "builds",
+                "submit",
+                f"--project={project}",
+                f"--tag=gcr.io/{project}/{image_name}-{UUID}:latest",
+                *cmd_substitutions,
+                ".",
+            ]
+            print(cmd)
+            subprocess.run(cmd, check=True)
+            yield f"gcr.io/{project}/{image_name}-{UUID}:latest"
+        else:
+            raise ValueError("must specify either `config` or `image_name`")
+
+        if image_name:
+            cmd = [
+                "gcloud",
+                "container",
+                "images",
+                "delete",
+                f"gcr.io/{project}/{image_name}-{UUID}:latest",
+                f"--project={project}",
+                "--quiet",
+            ]
+            print(cmd)
+            subprocess.run(cmd, check=True)
 
     @staticmethod
     def dataflow_job_id_from_job_name(

dataflow/flex-templates/streaming_beam/e2e_test.py

@@ -13,14 +13,12 @@
 import json
 import time
 
-# `conftest` cannot be imported when running in `nox`, but we still
-# try to import it for the autocomplete when writing the tests.
 try:
+    # `conftest` cannot be imported when running in `nox`, but we still
+    # try to import it for the autocomplete when writing the tests.
     from conftest import Utils
 except ModuleNotFoundError:
-    from typing import Any
-
-    Utils = Any
+    Utils = None
 import pytest
 
 NAME = "dataflow-flex-templates-streaming-beam"
@@ -61,7 +59,7 @@ def pubsub_publisher(utils: Utils, pubsub_topic: str) -> bool:
 
 @pytest.fixture(scope="session")
 def flex_template_image(utils: Utils) -> str:
-    yield from utils.container_image(NAME)
+    yield from utils.cloud_build_submit(NAME)
 
 
 @pytest.fixture(scope="session")

dataflow/gpu-workers/pytorch-minimal/.dockerignore

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

dataflow/gpu-workers/pytorch-minimal/.gcloudignore

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

dataflow/gpu-workers/pytorch-minimal/Dockerfile

@@ -0,0 +1,30 @@
+# Copyright 2021 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.
+
+FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
+
+WORKDIR /pipeline
+
+# Copy the Apache Beam worker files and the pipeline source files.
+COPY --from=apache/beam_python3.8_sdk:2.29.0 /opt/apache/beam /opt/apache/beam
+COPY requirements.txt .
+COPY *.py ./
+
+# Install the pipeline requirements and check that there are no conflicts.
+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.
+ENTRYPOINT [ "/opt/apache/beam/boot" ]