Repository: GoogleCloudPlatform/python-docs-samples (language: python)

Issue:
Add gemma-on-dataflow streaming example
## Description

Adds a code example for running a Dataflow RunInference pipeline using streaming, custom containers, and pubsub

## Checklist
- [ ] I have followed [Sample Guidelines from AUTHORING_GUIDE.MD](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md)
- [ ] README is updated to include [all relevant information](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#readme-file)
- [ ] **Tests** pass:   `nox -s py-3.9` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup))
- [ ] **Lint** pass:   `nox -s lint` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup))
- [ ] These samples need a new **API enabled** in testing projects to pass (let us know which ones)
- [ ] These samples need a new/updated **env vars** in testing projects set to pass (let us know which ones)
- [ ] This sample adds a new sample directory, and I updated the [CODEOWNERS file](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/.github/CODEOWNERS) with the codeowners for this sample
- [ ] This sample adds a new **Product API**, and I updated the [Blunderbuss issue/PR auto-assigner](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/.github/blunderbuss.yml) with the codeowners for this sample
- [ ] Please **merge** this PR for me once it is approved


Discussion / hints:
Is there a persistent but not-publicly-accessible bucket in python-docs-samples-tests project that we can use for test data? would you be able to help put some content there?
Not that I'm aware of. Typically, all tests create their own unique temporary bucket to run tests on. You can copy any necessary files there as part of the test setup.

Here's an example of it.

https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/dataflow/run-inference/tests/e2e_test.py
ok, it should be possible to create one and it would work as long as the bucket not cleaned up. I think Jack's example needs some additional data that cannot be stored with the example itself.
It kind of depends on how thoroughly we test the example, if we're building the container from scratch repeatedly then we need the model downloaded and available on a bucket somewhere. Otherwise, if we just have the container pre-built, then we just need the container in the project's artifact registry. 
> It kind of depends on how thoroughly we test the example

I think the purpose of the suite would be to test the example, to make sure it keeps working and test future changes in an automated way, in case follow-up changes are needed. Rebuilding docker container should be fine, there are test fixtures that can invoke cloud build, fill in necessary params, and clean up the image after the test. There might be some learning curve to learn to use the fixtures but the tests end up pretty simple, see:

https://github.com/GoogleCloudPlatform/python-docs-samples/blob/2b8c938b636614df65eabae17edf229e661ff3b9/dataflow/flex-templates/pipeline_with_dependencies/e2e_test.py#L41C41-L41C55

and

https://github.com/GoogleCloudPlatform/python-docs-samples/blob/2b8c938b636614df65eabae17edf229e661ff3b9/dataflow/conftest.py#L519
Yes, I agree. The test should include building the image since we find many issues with those regularly. In terms of data, it can be anything as simple as just passing a handful of hardcoded elements just to make sure that it's able to process them, and check that it goes through end to end, not running out of memory, etc.

And yes, the fixtures can be a little tricky at first, but once you get used to them, setting up the tests can be very simple. If you have any questions on how the fixtures work, feel free to ask. Here are the [fixture docs](https://docs.pytest.org/en/6.2.x/fixture.html) for some background, but in general we use them to setup and cleanup any Cloud resources.

Each test should generate a random unique ID and all resources should include that so that there are no clashes between concurrent tests since they all run in the same project. The fixtures in the [dataflow `conftest.py`](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/dataflow/conftest.py) should hopefully contain most of the necessary resources for most Dataflow tests, otherwise feel free to edit or add new fixtures as necessary.

Testing streaming pipelines is definitely trickier since they run forever, so we have to wait for them to be ready, then pass some data, check the outputs, and then stop the pipeline. It should be fairly similar to this one:

https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/dataflow/run-inference/tests/e2e_test.py
@davidcavazos I've gotten the first run at the e2e test written, but we're to the point where we need the Gemma model artifacts staged on GCS as part of the process. 
```

ERROR: Could not find a version that satisfies the requirement keras_nlp==0.8.2 (from versions: 0.0.0, 0.0.1, 0.0.2, 0.1.0, 0.1.1, 0.2.0.dev1, 0.2.0.dev2, 0.2.0, 0.3.0, 0.3.1.dev1, 0.3.1, 0.4.0.dev0, 0.4.0, 0.4.1.dev0, 0.4.1, 0.5.0.dev0, 0.5.0.dev4, 0.5.0.dev5, 0.5.0, 0.5.1, 0.5.2, 0.6.0.dev0, 0.6.0, 0.6.1.dev0, 0.6.1)
ERROR: No matching distribution found for keras_nlp==0.8.2
```

You probably need to configure this test to run using a newer version of Python (not 3.8). See: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/c35d01f37918e31c4adf29e63e8061eb0a5f1779/dataflow/flex-templates/pipeline_with_dependencies/noxfile_config.py#L22 for an example how to use 3.11.


`AttributeError: 'NoneType' object has no attribute 'run_cmd'`

It's one of those annoying ways in which conftest files work in pytest. It has a weird way of discovering where it is since it's on a parent directory. I've had some similar issues recently, try the following:

1. Move the test into the same directory as the `main.py`, you can then remove the `tests/` subdirectory (no need for an `__init__.py` file).
2. Remove the whole `try-except` from importing the `conftest`.
3. Add the following to your `noxfile_config.py`: `"envs": {"PYTHONPATH": ".."}`
4. Use the following instructions for the test.

```py
"""End-to-end tests.

1. Set your project.
    export GOOGLE_CLOUD_PROJECT="my-project-id"

2. To use an existing bucket, set it without the 'gs://' prefix.
    export GOOGLE_CLOUD_BUCKET="my-bucket-name"

3. Change directory to the sample location.
    cd dataflow/gemma

4. Set the PYTHONPATH to where the conftest is located.
    export PYTHONPATH=..

OPTION A: Run tests with pytest, you can use -k to run specific tests
    python -m venv env
    source env/bin/activate
    pip install -r requirements.txt -r requirements-test.txt
    pip check
    python -m pytest --verbose -s

OPTION B: Run tests with nox
    pip install nox
    nox -s py-3.10

NOTE: For the tests to find the conftest in the testing infrastructure,
      add the PYTHONPATH to the "env" in your noxfile_config.py file.
"""
```
Tests are finally green

Fix patch:
diff --git a/dataflow/gemma-flex-template/Dockerfile b/dataflow/gemma-flex-template/Dockerfile
new file mode 100644
index 000000000000..f63a462cba4a
--- /dev/null
+++ b/dataflow/gemma-flex-template/Dockerfile
@@ -0,0 +1,46 @@
+# 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 uses Ubuntu with Python 3.10 and comes with CUDA drivers for
+# GPU use.
+ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime
+
+FROM ${SERVING_BUILD_IMAGE}
+
+WORKDIR /workspace
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends apt-utils curl wget git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade pip \
+    && git clone https://github.com/google/gemma_pytorch.git \
+    && pip install --no-cache-dir -r requirements.txt ./gemma_pytorch
+
+# Copy SDK entrypoint binary from Apache Beam image, which makes it possible to
+# use the image as SDK container image.
+# The Beam version should match the version specified in requirements.txt
+COPY --from=apache/beam_python3.10_sdk:2.56.0 /opt/apache/beam /opt/apache/beam
+
+# Copy Flex Template launcher binary from the launcher image, which makes it
+# possible to use the image as a Flex Template base image.
+COPY --from=gcr.io/dataflow-templates-base/python310-template-launcher-base:latest /opt/google/dataflow/python_template_launcher /opt/google/dataflow/python_template_launcher
+
+# Copy the model directory downloaded from Kaggle and the pipeline code.
+COPY pytorch_model pytorch_model
+COPY custom_model_gemma.py custom_model_gemma.py
+
+ENV FLEX_TEMPLATE_PYTHON_PY_FILE="custom_model_gemma.py"
+
+# Set the entrypoint to Apache Beam SDK launcher.
+ENTRYPOINT ["/opt/apache/beam/boot"]
\ No newline at end of file
diff --git a/dataflow/gemma-flex-template/README.md b/dataflow/gemma-flex-template/README.md
new file mode 100644
index 000000000000..8ade42a8c460
--- /dev/null
+++ b/dataflow/gemma-flex-template/README.md
@@ -0,0 +1,179 @@
+# RunInference on Dataflow streaming with Gemma and Flex Templates
+
+Gemma is a family of lightweight, state-of-the-art open models built from research and technology used to create the Gemini models.
+You can use Gemma models in your Apache Beam inference pipelines with the `RunInference` transform.
+
+This example demonstrates how to use a Gemma model running on Pytorch in a streaming Dataflow pipeline that has Pub/Sub sources and sinks. This pipeline
+is deployed by using Flex Templates.
+
+For more information about using RunInference, see [Get started with AI/ML pipelines](https://beam.apache.org/documentation/ml/overview/) in the Apache Beam documentation.
+
+## Before you begin
+
+Make sure you have followed the [Dataflow setup instructions](/dataflow/README.md).
+
+Follow the steps in this section to create the necessary environment to run this workflow.
+
+### Enable Google Cloud services
+
+This workflow uses multiple Google Cloud products, including Dataflow, Pub/Sub, Google Cloud Storage, and Artifact Registry. Before you start the workflow, create a Google Cloud project that has the following services enabled:
+
+* Dataflow
+* Pub/Sub
+* Compute Engine
+* Cloud Logging
+* Google Cloud Storage
+* Google Cloud Storage JSON
+* Cloud Build
+* Datastore
+* Cloud Resource Manager
+* Artifact Registry
+
+Using these services incurs billing charges.
+
+Your Google Cloud project also needs to have Nvidia L4 GPU quota. For more information, see [GPU quota](https://cloud.google.com/compute/resource-usage#gpu_quota) in the Google Cloud documentation.
+
+### Download and save the model
+
+Save a version of the Gemma 2B model. Downloaded the model from [Kaggle](https://www.kaggle.com/models/google/gemma/pyTorch/1.1-2b-it). This download is a `.tar.gz` archive. Extract the archive into a directory and name it `pytorch_model`. 
+
+### Create a cloud storage bucket
+
+Click [here to create a GCS bucket](https://console.cloud.google.com/storage/create-bucket) or run the following commands:
+
+```sh
+export GCS_BUCKET="your--bucket"
+gsutil mb gs://$GCS_BUCKET
+```
+
+Make sure your GCS bucket name does __not__ include the `gs://` prefix
+
+### Create a custom container
+
+To build a custom container, use Docker. This repository contains a Dockerfile that you can use to build your custom container. To build and push a container to Artifact Registry by using Docker or Cloud Build, follow the instructions in the [Build and push the image](https://cloud.google.com/dataflow/docs/guides/build-container-image#build_and_push_the_image) section of "Build custom container images for Dataflow" in the Google Cloud documentation.
+
+### Create Pub/Sub topics for input and output
+
+To create your Pub/Sub source and sink, follow the instructions in [Create a Pub/Sub topic](https://cloud.google.com/pubsub/docs/create-topic#pubsub_create_topic-Console) in the Google Cloud documentation. For this example, create two topics, one input topic and one output topic. Follow the instructions in [Create pull subscriptions](https://cloud.google.com/pubsub/docs/create-subscription) to create a pull subscription for each of the two topics you jsut created. The input subscription allows you to provide input to the pipeline, while the output subscription will allow you to see the output from the pipeline during and after execution.
+
+## Code overview
+
+This section provides details about the custom model handler and the formatting `DoFn` used in this example.
+
+### Custom model handler
+This example defines a custom model handler that loads the model. The model handler constructs a configuration object and loads the model's checkpoint from the local filesystem.
+Because this approach differs from the PyTorch model loading process followed in the Beam PyTorch model handler, a custom implementation is necessary. 
+
+To customize the behavior of the handler, implement the following methods: `load_model`, `validate_inference_args`, and `share_model_across_processes`.
+
+The PyTorch implementation of the Gemma models has a `generate` method
+that generates text based on a prompt. To route the prompts correctly, use this function in the `run_inference` function.
+
+```py
+class GemmaPytorchModelHandler(ModelHandler[str, PredictionResult,
+                                            GemmaForCausalLM]):
+    def __init__(self,
+                 model_variant: str,
+                 checkpoint_path: str,
+                 tokenizer_path: str,
+                 device: Optional[str] = 'cpu'):
+        """ Implementation of the ModelHandler interface for Gemma-on-Pytorch
+        using text as input.
+
+        Example Usage::
+
+          pcoll | RunInference(GemmaPytorchHandler())
+
+        Args:
+          model_variant: The Gemma model name.
+          checkpoint_path: the path to a local copy of gemma model weights.
+          tokenizer_path: the path to a local copy of the gemma tokenizer
+          device: optional. the device to run inference on. can be either
+            'cpu' or 'gpu', defaults to cpu. 
+        """
+        model_config = get_config_for_2b(
+        ) if "2b" in model_variant else get_config_for_7b()
+        model_config.tokenizer = tokenizer_path
+        model_config.quant = 'quant' in model_variant
+        model_config.tokenizer = tokenizer_path
+
+        self._model_config = model_config
+        self._checkpoint_path = checkpoint_path
+        if device == 'GPU':
+            logging.info("Device is set to CUDA")
+            self._device = torch.device('cuda')
+        else:
+            logging.info("Device is set to CPU")
+            self._device = torch.device('cpu')
+        self._env_vars = {}
+
+    def load_model(self) -> GemmaForCausalLM:
+        """Loads and initializes a model for processing."""
+        torch.set_default_dtype(self._model_config.get_dtype())
+        model = GemmaForCausalLM(self._model_config)
+        model.load_weights(self._checkpoint_path)
+        model = model.to(self._device).eval()
+        return model
+```
+
+### Formatting DoFn
+
+The output from a keyed model handler is a tuple of the form `(key, PredictionResult)`. To format that output into a string before sending it to the answer Pub/Sub topic, use an extra `DoFn`.
+
+```py
+| "Format output" >> beam.Map(
+    lambda response: json.dumps(
+        {"input": response.example, "outputs": response.inference}
+    )
+)
+```
+
+## Build the Flex Template
+Run the following code from the directory to build the Dataflow flex template.
+
+- Replace `$GCS_BUCKET` with a Google Cloud Storage bucket.
+- Set `SDK_CONTAINER_IMAGE` to the name of the Docker image created previously.
+- `$PROJECT` is the Google Cloud project that you created previously. 
+
+```sh
+gcloud dataflow flex-template build gs://$GCS_BUCKET/config.json \
+  --image $SDK_CONTAINER_IMAGE \
+  --sdk-language "PYTHON" \
+  --metadata-file metadata.json \
+  --project $PROJECT
+```
+
+## Start the pipeline
+To start the Dataflow streaming job, run the following code from the directory. Replace `$TEMPLATE_FILE`, `$REGION`, `$GCS_BUCKET`, `$INPUT_SUBSCRIPTION`, `$OUTPUT_TOPIC`, `$SDK_CONTAINER_IMAGE`, and `$PROJECT` with the Google Cloud project resources you created previously. Ensure that `$INPUT_SUBSCRIPTION` and `$OUTPUT_TOPIC` are the fully qualified subscription and topic names, respectively. It might take as much as 30 minutes for the worker to start up and to begin accepting messages from the input Pub/Sub topic. 
+
+```sh
+gcloud dataflow flex-template run "gemma-flex-`date +%Y%m%d-%H%M%S`" \
+  --template-file-gcs-location $TEMPLATE_FILE \
+  --region $REGION \
+  --temp-location gs://$GCS_BUCKET/tmp \
+  --staging-location gs://$GCS_BUCKET \
+  --parameters messages_subscription=$INPUT_SUBSCRIPTION \
+  --parameters responses_topic=$OUTPUT_TOPIC \
+  --parameters device="GPU" \
+  --parameters sdk_container_image=$SDK_CONTAINER_IMAGE \
+  --additional-experiments "worker_accelerator=type:nvidia-l4;count:1;install-nvidia-driver" \
+  --project $PROJECT \
+  --worker-machine-type "g2-standard-4"
+```
+
+## Send a prompt to the model and check the response
+
+In the Google Cloud console, navigate to the Pub/Sub topics page, and then select your input topic. On the **Messages** tab, click **Publish Message**. Add a message for the Dataflow job to pick up and pass through the model. For example, your input message could be "Tell the the sentiment of the following sentence: I like pineapple on pizza."
+
+The Dataflow job outputs the response to the Pub/Sub sink topic. To check the response from the model, you can manually pull messages from the destination topic. For more information, see [Publish messages](https://cloud.google.com/pubsub/docs/publisher#publish-messages) in the Google Cloud documentation.
+
+## Clean up resources
+
+To avoid incurring charges to your Google Cloud account for the resources used in this example, clean up the resources that you created.
+
+
+*   Cancel the streaming Dataflow job.
+*   **Optional**: Archive the streaming Dataflow job.
+*   Delete the Pub/Sub topic and subscriptions.
+*   Delete the custom container from Artifact Registry.
+*   Delete the created GCS bucket.
\ No newline at end of file
diff --git a/dataflow/gemma-flex-template/custom_model_gemma.py b/dataflow/gemma-flex-template/custom_model_gemma.py
new file mode 100644
index 000000000000..0b6230ca56cb
--- /dev/null
+++ b/dataflow/gemma-flex-template/custom_model_gemma.py
@@ -0,0 +1,176 @@
+# Copyright 2024 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 collections.abc import Iterable, Sequence
+import json
+import logging
+
+import apache_beam as beam
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.options.pipeline_options import PipelineOptions
+
+from gemma.config import get_config_for_2b
+from gemma.config import get_config_for_7b
+from gemma.model import GemmaForCausalLM
+
+import torch
+
+
+class GemmaPytorchModelHandler(ModelHandler[str, PredictionResult, GemmaForCausalLM]):
+    def __init__(
+        self,
+        model_variant: str,
+        checkpoint_path: str,
+        tokenizer_path: str,
+        device: str | None = "cpu",
+    ):
+        """Implementation of the ModelHandler interface for Gemma-on-Pytorch
+        using text as input.
+
+        Example Usage::
+
+          pcoll | RunInference(GemmaPytorchHandler())
+
+        Args:
+          model_variant: The Gemma model name.
+          checkpoint_path: the path to a local copy of gemma model weights.
+          tokenizer_path: the path to a local copy of the gemma tokenizer
+          device: optional. the device to run inference on. can be either
+            'cpu' or 'gpu', defaults to cpu.
+        """
+        model_config = (
+            get_config_for_2b() if "2b" in model_variant else get_config_for_7b()
+        )
+        model_config.tokenizer = tokenizer_path
+        model_config.quant = "quant" in model_variant
+        model_config.tokenizer = tokenizer_path
+
+        self._model_config = model_config
+        self._checkpoint_path = checkpoint_path
+        if device == "GPU":
+            logging.info("Device is set to CUDA")
+            self._device = torch.device("cuda")
+        else:
+            logging.info("Device is set to CPU")
+            self._device = torch.device("cpu")
+        self._env_vars = {}
+
+    def share_model_across_processes(self) -> bool:
+        """Allows us to load a model only once per worker VM, decreasing
+        pipeline memory requirements.
+        """
+        return True
+
+    def load_model(self) -> GemmaForCausalLM:
+        """Loads and initializes a model for processing."""
+        torch.set_default_dtype(self._model_config.get_dtype())
+        model = GemmaForCausalLM(self._model_config)
+        model.load_weights(self._checkpoint_path)
+        model = model.to(self._device).eval()
+        return model
+
+    def run_inference(
+        self,
+        batch: Sequence[str],
+        model: GemmaForCausalLM,
+        inference_args: dict | None = None,
+    ) -> Iterable[PredictionResult]:
+        """Runs inferences on a batch of text strings.
+
+        Args:
+          batch: A sequence of examples as text strings.
+          model: The Gemma model being used.
+          inference_args: Any additional arguments for an inference.
+
+        Returns:
+          An Iterable of type PredictionResult.
+        """
+        result = model.generate(prompts=batch, device=self._device)
+        predictions = [result]
+        return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--messages_subscription",
+        required=True,
+        help="Pub/Sub subscription for input text messages",
+    )
+    parser.add_argument(
+        "--responses_topic",
+        required=True,
+        help="Pub/Sub topic for output text responses",
+    )
+    parser.add_argument(
+        "--model_variant",
+        required=False,
+        default="gemma-2b-it",
+        help="name of the gemma variant being used",
+    )
+    parser.add_argument(
+        "--checkpoint_path",
+        required=False,
+        default="pytorch_model/gemma-2b-it.ckpt",
+        help="path to the Gemma model weights in the custom worker container",
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        required=False,
+        default="pytorch_model/tokenizer.model",
+        help="path to the Gemma tokenizer in the custom worker container",
+    )
+    parser.add_argument(
+        "--device",
+        required=False,
+        default="cpu",
+        help="device to run the model on",
+    )
+
+    args, beam_args = parser.parse_known_args()
+
+    config = get_config_for_2b()
+
+    logging.getLogger().setLevel(logging.INFO)
+    beam_options = PipelineOptions(
+        beam_args,
+        save_main_session=True,
+        streaming=True,
+    )
+
+    handler = GemmaPytorchModelHandler(
+        model_variant=args.model_variant,
+        checkpoint_path=args.checkpoint_path,
+        tokenizer_path=args.tokenizer_path,
+        device=args.device,
+    )
+
+    with beam.Pipeline(options=beam_options) as pipeline:
+        _ = (
+            pipeline
+            | "Subscribe to Pub/Sub" >> beam.io.ReadFromPubSub(subscription=args.messages_subscription)
+            | "Decode" >> beam.Map(lambda msg: msg.decode("utf-8"))
+            | "RunInference Gemma" >> RunInference(handler)
+            | "Format output" >> beam.Map(
+                lambda response: json.dumps(
+                    {"input": response.example, "outputs": response.inference}
+                )
+            )
+            | "Encode" >> beam.Map(lambda msg: msg.encode("utf-8"))
+            | "Publish to Pub/Sub" >> beam.io.gcp.pubsub.WriteToPubSub(topic=args.responses_topic)
+        )
diff --git a/dataflow/gemma-flex-template/metadata.json b/dataflow/gemma-flex-template/metadata.json
new file mode 100644
index 000000000000..df15d87bf79e
--- /dev/null
+++ b/dataflow/gemma-flex-template/metadata.json
@@ -0,0 +1,27 @@
+{
+    "name": "Streaming Gemma-on-Dataflow",
+    "description": "A flex template configuring a streaming Gemma pipeline using pytorch and pub/sub",
+    "parameters": [
+      {
+        "name": "messages_subscription",
+        "label": "Input messages subscription",
+        "helpText": "The pub/sub subscription to pull inputs from"
+      },
+      {
+        "name": "responses_topic",
+        "label": "Output messages topic",
+        "helpText": "The pub/sub topic to push outputs to"
+      },
+      {
+        "name": "device",
+        "label": "Device",
+        "helpText": "The device to run inference on. Can be 'CPU' or 'GPU'."
+      },
+      {
+        "name": "disk_size_gb",
+        "label": "disk_size_gb",
+        "helpText": "disk_size_gb for worker",
+        "isOptional": true
+      }
+    ]
+  }
\ No newline at end of file
diff --git a/dataflow/gemma-flex-template/noxfile_config.py b/dataflow/gemma-flex-template/noxfile_config.py
new file mode 100644
index 000000000000..d483d3f479d2
--- /dev/null
+++ b/dataflow/gemma-flex-template/noxfile_config.py
@@ -0,0 +1,26 @@
+# Copyright 2024 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 is a test configuration file. It is not a part of the sample.
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    # Opting out of all Python versions except 3.10.
+    # The Python version used is defined by the Dockerfile and the job
+    # submission enviornment must match.
+    "ignored_versions": ["2.7", "3.6", "3.7", "3.8", "3.9", "3.11", "3.12"],
+    "envs": {
+        "PYTHONPATH": ".."
+    },
+}
diff --git a/dataflow/gemma-flex-template/requirements.txt b/dataflow/gemma-flex-template/requirements.txt
new file mode 100644
index 000000000000..5aea06d762a3
--- /dev/null
+++ b/dataflow/gemma-flex-template/requirements.txt
@@ -0,0 +1,9 @@
+# For reproducible builds, it is better to also include transitive dependencies: 
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/c93accadf3bd29e9c3166676abb2c95564579c5e/dataflow/flex-templates/pipeline_with_dependencies/requirements.txt#L22, 
+# but for simplicity of this example, we are only including the top-level dependencies.
+apache_beam[gcp]==2.58.0
+immutabledict==4.2.0
+
+# Also required, please download and install gemma_pytorch.
+#   git clone https://github.com/google/gemma_pytorch.git
+#   pip install ./gemma_pytorch


Test patch:
diff --git a/dataflow/conftest.py b/dataflow/conftest.py
index edfd386daf1f..fce16f48f471 100644
--- a/dataflow/conftest.py
+++ b/dataflow/conftest.py
@@ -794,6 +794,7 @@ def dataflow_flex_template_run(
         parameters: dict[str, str] = {},
         project: str = PROJECT,
         region: str = REGION,
+        additional_experiments: dict[str,str] = {},
     ) -> str:
         import yaml
 
@@ -815,6 +816,11 @@ def dataflow_flex_template_run(
             for name, value in {
                 **parameters,
             }.items()
+        ] + [
+            f"--additional-experiments={name}={value}"
+            for name, value in {
+                **additional_experiments,
+            }.items()
         ]
         logging.info(f"{cmd}")
 
@@ -825,7 +831,7 @@ def dataflow_flex_template_run(
         logging.info(f">> {Utils.dataflow_job_url(job_id, project, region)}")
         yield job_id
 
-        Utils.dataflow_jobs_cancel(job_id)
+        Utils.dataflow_jobs_cancel(job_id, region=region)
 
     @staticmethod
     def dataflow_extensible_template_run(
@@ -849,7 +855,7 @@ def dataflow_extensible_template_run(
             f"--gcs-location={template_path}",
             f"--project={project}",
             f"--region={region}",
-            f"--staging-location=gs://{bucket_name}/staging",
+            f"--temp-location=gs://{bucket_name}/staging",
         ] + [
             f"--parameters={name}={value}"
             for name, value in {
diff --git a/dataflow/gemma-flex-template/e2e_test.py b/dataflow/gemma-flex-template/e2e_test.py
new file mode 100644
index 000000000000..f95f78ec089d
--- /dev/null
+++ b/dataflow/gemma-flex-template/e2e_test.py
@@ -0,0 +1,162 @@
+# Copyright 2024 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.
+"""End-to-end tests.
+
+1. Set your project.
+    export GOOGLE_CLOUD_PROJECT="my-project-id"
+
+2. To use an existing bucket, set it without the 'gs://' prefix.
+    export GOOGLE_CLOUD_BUCKET="my-bucket-name"
+
+3. Change directory to the sample location.
+    cd dataflow/gemma-flex-template
+
+4. Set the PYTHONPATH to where the conftest is located.
+    export PYTHONPATH=..
+
+OPTION A: Run tests with pytest, you can use -k to run specific tests
+    python -m venv env
+    source env/bin/activate
+    pip install -r requirements.txt -r requirements-test.txt
+    pip check
+    python -m pytest --verbose -s
+
+OPTION B: Run tests with nox
+    pip install nox
+    nox -s py-3.10
+
+NOTE: For the tests to find the conftest in the testing infrastructure,
+      add the PYTHONPATH to the "env" in your noxfile_config.py file.
+"""
+
+from collections.abc import Callable
+
+import conftest  # python-docs-samples/dataflow/conftest.py
+from conftest import Utils
+
+import pytest
+
+DATAFLOW_MACHINE_TYPE = "g2-standard-4"
+GEMMA_GCS = "gs://perm-dataflow-gemma-example-testdata/pytorch_model"
+NAME = "dataflow/gemma-flex-template/streaming"
+
+# There are limited resources in us-central1, so change to a known good region.
+REGION = "us-west1"
+
+
+@pytest.fixture(scope="session")
+def test_name() -> str:
+    # Many fixtures expect a fixture called `test_name`, so be sure to define it!
+    return "dataflow/gemma-flex-template"
+
+
+@pytest.fixture(scope="session")
+def bucket_name(utils: Utils) -> str:
+    yield from utils.storage_bucket(NAME)
+
+
+@pytest.fixture(scope="session")
+def messages_topic(pubsub_topic: Callable[[str], str]) -> str:
+    return pubsub_topic("messages")
+
+
+@pytest.fixture(scope="session")
+def messages_subscription(
+    pubsub_subscription: Callable[[str, str], str], messages_topic: str
+) -> str:
+    return pubsub_subscription("messages", messages_topic)
+
+
+@pytest.fixture(scope="session")
+def responses_topic(pubsub_topic: Callable[[str], str]) -> str:
+    return pubsub_topic("responses")
+
+
+@pytest.fixture(scope="session")
+def responses_subscription(
+    pubsub_subscription: Callable[[str, str], str], responses_topic: str
+) -> str:
+    return pubsub_subscription("responses", responses_topic)
+
+
+@pytest.fixture(scope="session")
+def flex_template_image(utils: Utils) -> str:
+    conftest.run_cmd("gsutil", "cp", "-r", GEMMA_GCS, ".")
+    yield from utils.cloud_build_submit(NAME)
+
+
+@pytest.fixture(scope="session")
+def flex_template_path(utils: Utils, bucket_name: str, flex_template_image: str) -> str:
+    yield from utils.dataflow_flex_template_build(bucket_name, flex_template_image)
+
+
+@pytest.fixture(scope="session")
+def dataflow_job(
+    utils: Utils,
+    project: str,
+    location: str,
+    bucket_name: str,
+    flex_template_image: str,
+    flex_template_path: str,
+    messages_subscription: str,
+    responses_topic: str,
+) -> str:
+    yield from utils.dataflow_flex_template_run(
+        job_name=NAME,
+        template_path=flex_template_path,
+        bucket_name=bucket_name,
+        project=project,
+        region=REGION,
+        parameters={
+            "messages_subscription": messages_subscription,
+            "responses_topic": responses_topic,
+            "device": "GPU",
+            "sdk_container_image": f"gcr.io/{project}/{flex_template_image}",
+            "machine_type": f"{DATAFLOW_MACHINE_TYPE}",
+            "disk_size_gb": "50",
+        },
+        additional_experiments={
+            "worker_accelerator": "type:nvidia-l4;count:1;install-nvidia-driver",
+        },
+    )
+
+
+@pytest.mark.timeout(5400)
+def test_pipeline_dataflow(
+    project: str,
+    location: str,
+    dataflow_job: str,
+    messages_topic: str,
+    responses_subscription: str,
+) -> None:
+    print(f"Waiting for the Dataflow workers to start: {dataflow_job}")
+    conftest.wait_until(
+        lambda: conftest.dataflow_num_workers(project, REGION, dataflow_job) > 0,
+        "workers are running",
+    )
+    num_workers = conftest.dataflow_num_workers(project, REGION, dataflow_job)
+    print(f"Dataflow job num_workers: {num_workers}")
+
+    messages = ["This is a test for a Python sample."]
+    conftest.pubsub_publish(messages_topic, messages)
+
+    print(f"Waiting for messages on {responses_subscription}")
+    responses = conftest.pubsub_wait_for_messages(responses_subscription)
+
+    print(f"Received {len(responses)} responses(s)")
+
+    for msg in responses:
+        print(f"- {type(msg)} - {msg}")
+
+    assert responses, "expected at least one response"
diff --git a/dataflow/gemma-flex-template/requirements-test.txt b/dataflow/gemma-flex-template/requirements-test.txt
new file mode 100644
index 000000000000..50c76d34a122
--- /dev/null
+++ b/dataflow/gemma-flex-template/requirements-test.txt
@@ -0,0 +1,7 @@
+google-cloud-aiplatform==1.62.0
+google-cloud-dataflow-client==0.8.12
+google-cloud-pubsub==2.23.0
+google-cloud-storage==2.18.2
+pytest==8.3.2
+pytest-timeout==2.3.1
+pyyaml==6.0.2