Add the pub/sub source (#15)

* added the pub/sub source

* updated the readme

* updated the shards

---------

Co-authored-by: xqhu <xqhu@google.com>

Author: liferoad

Makefile

@@ -119,6 +119,7 @@ ifeq ($(MODEL_ENV), "TORCH")
 	--dataflow_service_option $(SERVICE_OPTIONS) \
 	--number_of_worker_harness_threads 1 \
 	--experiments=disable_worker_container_image_prepull \
+	--experiments=use_pubsub_streaming \
 	--sdk_container_image $(CUSTOM_CONTAINER_IMAGE) \
 	--sdk_location container \
 	--input $(INPUT_DATA) \
@@ -140,6 +141,7 @@ else
 	--dataflow_service_option $(SERVICE_OPTIONS) \
 	--number_of_worker_harness_threads 1 \
 	--experiments=disable_worker_container_image_prepull \
+	--experiments=use_pubsub_streaming \
 	--sdk_container_image $(CUSTOM_CONTAINER_IMAGE) \
 	--sdk_location container \
 	--input $(INPUT_DATA) \

README.md

@@ -268,6 +268,12 @@ When using resnet101 to score 50k images, the job took ~1h and cost ~0.5$ with r
 For `mobilenet_v2`, it cost 0.05$ with ~1h.
 Note the cost and time depends on your job settings and the regions.
 
+### Run the Beam pipeline with the Pub/Sub source
+When `INPUT_DATA` from the `.env` file defines a valid Pub/Sub topic (e.g., `projects/apache-beam-testing/topics/Imagenet_openimage_50k_benchmark`),
+the Beam pipeline is created using the Pub/Sub source with `FixedWindows` and switches to `beam.io.fileio.WriteToFiles` that supports the streaming pipeline.
+We use `shards=0` here since 0 shards is the recommended approach and Dataflow would decide how many files it should write.
+Note that the streaming job will run forever until it is canceled or drained.
+
 ## FAQ
 
 ### Permission error when using any GCP command

src/config.py

@@ -13,10 +13,11 @@
 #  limitations under the License.
 
 # standard libraries
+import re
 from enum import Enum
 
 # third party libraries
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, Field, root_validator, validator
 
 
 class ModelName(str, Enum):
@@ -51,13 +52,23 @@ def validate_fields(cls, values):
         return values
 
 
+def _validate_topic_path(topic_path):
+    pattern = r"projects/.+/topics/.+"
+    return bool(re.match(pattern, topic_path))
+
+
 class SourceConfig(BaseModel):
-    input: str = Field(..., description="the input path to a text file")
+    input: str = Field(..., description="the input path to a text file or a Pub/Sub topic")
     images_dir: str = Field(
         None,
         description="Path to the directory where images are stored."
         "Not required if image names in the input file have absolute path.",
     )
+    streaming: bool = False
+
+    @validator("streaming", pre=True, always=True)
+    def set_streaming(cls, v, values):
+        return _validate_topic_path(values["input"])
 
 
 class SinkConfig(BaseModel):

src/pipeline.py

@@ -17,7 +17,7 @@
 # standard libraries
 import io
 import os
-from typing import Iterable, Iterator, Optional, Tuple
+from typing import Iterable, Iterator, Optional, Tuple, Union
 
 # third party libraries
 import apache_beam as beam
@@ -45,7 +45,9 @@ def get_model_class(model_name: ModelName) -> nn.Module:
     return model_class
 
 
-def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> Tuple[str, Image.Image]:
+def read_image(image_file_name: Union[str, bytes], path_to_dir: Optional[str] = None) -> Tuple[str, Image.Image]:
+    if isinstance(image_file_name, bytes):
+        image_file_name = image_file_name.decode()
     if path_to_dir is not None:
         image_file_name = os.path.join(path_to_dir, image_file_name)
     with FileSystems().open(image_file_name, "r") as file:
@@ -126,14 +128,24 @@ def build_pipeline(pipeline, source_config: SourceConfig, sink_config: SinkConfi
     else:
         raise ValueError("Only support PytorchModelHandler and TFModelHandlerTensor!")
 
-    # read the text file and create the pair of input data with the file name and its image
-    filename_value_pair = (
-        pipeline
-        | "ReadImageNames" >> beam.io.ReadFromText(source_config.input)
-        | "FilterEmptyLines" >> beam.ParDo(filter_empty_lines)
-        | "ReadImageData"
-        >> beam.Map(lambda image_name: read_image(image_file_name=image_name, path_to_dir=source_config.images_dir))
-    )
+    if source_config.streaming:
+        # read the text file path from Pub/Sub and use FixedWindows to group these images
+        # and then run the model inference and store the results into GCS
+        filename_value_pair = (
+            pipeline
+            | "ReadImageNamesFromPubSub" >> beam.io.ReadFromPubSub(topic=source_config.input)
+            | "Window into fixed intervals" >> beam.WindowInto(beam.window.FixedWindows(60 * 5))
+            | "ReadImageData" >> beam.Map(lambda image_name: read_image(image_file_name=image_name))
+        )
+    else:
+        # read the text file and create the pair of input data with the file name and its image
+        filename_value_pair = (
+            pipeline
+            | "ReadImageNames" >> beam.io.ReadFromText(source_config.input)
+            | "FilterEmptyLines" >> beam.ParDo(filter_empty_lines)
+            | "ReadImageData"
+            >> beam.Map(lambda image_name: read_image(image_file_name=image_name, path_to_dir=source_config.images_dir))
+        )
 
     if model_config.model_state_dict_path:
         filename_value_pair = filename_value_pair | "PreprocessImages" >> beam.MapTuple(
@@ -151,7 +163,15 @@ def build_pipeline(pipeline, source_config: SourceConfig, sink_config: SinkConfi
         | "ProcessOutput" >> beam.ParDo(PostProcessor())
     )
 
-    # save the predictions to a text file
-    predictions | "WriteOutputToGCS" >> beam.io.WriteToText(  # pylint: disable=expression-not-assigned
-        sink_config.output, shard_name_template="", append_trailing_newlines=True
-    )
+    # combine all the window results into one text for GCS
+    if source_config.streaming:
+        (
+            predictions
+            | "WriteOutputToGCS"
+            >> beam.io.fileio.WriteToFiles(sink_config.output, shards=0)  # pylint: disable=expression-not-assigned
+        )
+    else:
+        # save the predictions to a text file
+        predictions | "WriteOutputToGCS" >> beam.io.WriteToText(  # pylint: disable=expression-not-assigned
+            sink_config.output, shard_name_template="", append_trailing_newlines=True
+        )

src/run.py

@@ -64,13 +64,6 @@ def run(argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult
     """
     known_args, pipeline_args = parse_known_args(argv)
 
-    pipeline_options = PipelineOptions(pipeline_args)
-    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
-
-    pipeline = test_pipeline
-    if not test_pipeline:
-        pipeline = beam.Pipeline(options=pipeline_options)
-
     # setup configs
     model_config = ModelConfig(
         model_state_dict_path=known_args.model_state_dict_path,
@@ -83,6 +76,14 @@ def run(argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult
     source_config = SourceConfig(input=known_args.input)
     sink_config = SinkConfig(output=known_args.output)
 
+    # setup pipeline
+    pipeline_options = PipelineOptions(pipeline_args, streaming=source_config.streaming)
+    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
+
+    pipeline = test_pipeline
+    if not test_pipeline:
+        pipeline = beam.Pipeline(options=pipeline_options)
+
     # build the pipeline using configs
     build_pipeline(pipeline, source_config=source_config, sink_config=sink_config, model_config=model_config)
 

tests/test_pipeline.py

@@ -48,3 +48,10 @@ def test_build_pipeline_with_tf():
 
     p = beam.Pipeline()
     build_pipeline(p, source_config=source_config, sink_config=sink_config, model_config=model_config)
+
+
+def test_source_config_streaming():
+    source_config = SourceConfig(input=str(DATA_FILE_PATH / "openimage_10.txt"))
+    assert source_config.streaming is False
+    source_config = SourceConfig(input="projects/apache-beam-testing/topics/Imagenet_openimage_50k_benchmark")
+    assert source_config.streaming is True