[RunInference] Fix WatchFilePattern selecting the newest model

CHANGES.md

@@ -108,6 +108,7 @@
 
 ## Bugfixes
 
+* (Python) Fixed `WatchFilePattern` comparing model files only against their own path, so an older model could be chosen and a poll that saw several models could crash `RunInference` (`AsSingleton` requires one value) ([#40638](https://github.com/apache/beam/issues/40638)).
 * (Python) Fixed BigQuery `FILE_LOADS` silently omitting files when dynamic destinations use different spellings of the same table ([#40001](https://github.com/apache/beam/pull/40001)).
 * (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)).
 * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)).

sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py

@@ -48,9 +48,11 @@
   --file_pattern=<glob_pattern>
 
 file_pattern is path(can contain glob characters), which will be passed to
-WatchContinuously transform for model updates. WatchContinuously watches the
-file_pattern and emits a latest file path, sorted by timestamp. Files that
-are read before and updated with same name will be ignored as an update.
+WatchFilePattern for model updates. WatchFilePattern emits the matching file
+with the newest modification time. If one check finds several models, only
+the newest is emitted. Reusing a filename does not publish an update; upload
+a new file. A file whose modification time is older than the model already
+selected is ignored, so a rollback must be uploaded with a newer timestamp.
 
 The pipeline expects there is at least one file present to match the
 file_pattern before pipeline startup. Presumably, this would be the

sdks/python/apache_beam/ml/inference/base_test.py

@@ -2200,7 +2200,7 @@ def test_runinference_loads_same_model_with_identifier_multi_process_shared(
   def test_run_inference_watch_file_pattern_side_input_label(self):
     pipeline = TestPipeline()
     # label of the WatchPattern transform.
-    side_input_str = 'WatchFilePattern/ApplyGlobalWindow'
+    side_input_str = 'WatchFilePattern/SelectLatestModel/EmitLatestModel'
     from apache_beam.ml.inference.utils import WatchFilePattern
     file_pattern_side_input = (
         pipeline
@@ -2212,7 +2212,7 @@ def test_run_inference_watch_file_pattern_side_input_label(self):
 
   def test_run_inference_watch_file_pattern_keyword_arg_side_input_label(self):
     # label of the WatchPattern transform.
-    side_input_str = 'WatchFilePattern/ApplyGlobalWindow'
+    side_input_str = 'WatchFilePattern/SelectLatestModel/EmitLatestModel'
     pipeline = TestPipeline()
     pcoll = pipeline | 'start' >> beam.Create([1, 2, 3])
     result_pcoll = pcoll | base.RunInference(

sdks/python/apache_beam/ml/inference/utils.py

@@ -21,24 +21,30 @@
 """
 import os
 from collections.abc import Iterable
-from functools import partial
 from typing import Any
 from typing import Optional
 from typing import Union
 
 import apache_beam as beam
+from apache_beam.coders import FloatCoder
 from apache_beam.io.fileio import EmptyMatchTreatment
 from apache_beam.io.fileio import MatchContinuously
+from apache_beam.io.filesystem import FileMetadata
 from apache_beam.ml.inference.base import ModelMetadata
 from apache_beam.ml.inference.base import PredictionResult
 from apache_beam.transforms import trigger
 from apache_beam.transforms import window
-from apache_beam.transforms.userstate import CombiningValueStateSpec
+from apache_beam.transforms.userstate import ReadModifyWriteStateSpec
 from apache_beam.utils.timestamp import MAX_TIMESTAMP
 from apache_beam.utils.timestamp import Timestamp
 
 _START_TIME_STAMP = Timestamp.now()
 
+# Beam keeps DoFn state per key. Keying each file by its own path gives every
+# path a fresh timestamp and never compares models with each other, so every
+# watched file is keyed the same way.
+_GLOBAL_MODEL_KEY = 'WatchFilePattern'
+
 
 def _convert_to_result(
     batch: Iterable,
@@ -60,54 +66,108 @@ def _convert_to_result(
   return [PredictionResult(x, y, model_id) for x, y in zip(batch, predictions)]
 
 
-class _ConvertIterToSingleton(beam.DoFn):
-  """
-  Internal only; No backwards compatibility.
-
-  The MatchContinuously transform examines all files present in a given
-  directory and returns those that have timestamps older than the
-  pipeline's start time. This can produce an Iterable rather than a
-  Singleton. This class only returns the file path when it is first
-  encountered, and it is cached as part of the side input caching mechanism.
-  If the path is seen again, it will not return anything.
-  By doing this, we can ensure that the output of this transform can be wrapped
-  with beam.pvalue.AsSingleton().
+def _as_seconds(timestamp: Union[Timestamp, float]) -> float:
+  return float(timestamp)
+
+
+def _is_newer(
+    candidate: Optional[FileMetadata],
+    incumbent: Optional[FileMetadata]) -> bool:
+  """Return true when candidate should replace incumbent as the latest file.
+
+  The newest modification time wins. Equal timestamps are broken by path so
+  the choice is the same no matter which order the files are merged in.
   """
-  COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum)
+  if candidate is None:
+    return False
+  if incumbent is None:
+    return True
+  candidate_ts = _as_seconds(candidate.last_updated_in_seconds)
+  incumbent_ts = _as_seconds(incumbent.last_updated_in_seconds)
+  if candidate_ts != incumbent_ts:
+    return candidate_ts > incumbent_ts
+  return candidate.path > incumbent.path
+
+
+def _with_global_key(metadata: FileMetadata):
+  return _GLOBAL_MODEL_KEY, metadata
 
-  def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)):
-    counter = count_state.read()
-    if counter == 0:
-      count_state.add(1)
-      yield element[1]
 
+def _model_metadata(path: str) -> ModelMetadata:
+  model_name = os.path.splitext(os.path.basename(path))[0]
+  return ModelMetadata(model_id=path, model_name=model_name)
 
-class _GetLatestFileByTimeStamp(beam.DoFn):
+
+class _LatestFileCombineFn(beam.CombineFn):
+  """Selects the FileMetadata with the newest modification time.
+
+  Internal only; no backwards compatibility.
   """
-  Internal only; No backwards compatibility.
+  def create_accumulator(self):
+    return None
+
+  def add_input(self, accumulator, element):
+    if _is_newer(element, accumulator):
+      return element
+    return accumulator
+
+  def merge_accumulators(self, accumulators):
+    merged = None
+    for accumulator in accumulators:
+      if _is_newer(accumulator, merged):
+        merged = accumulator
+    return merged
+
+  def extract_output(self, accumulator):
+    return accumulator
+
+
+class _EmitLatestModel(beam.DoFn):
+  """Emits a model only when it is newer than every model seen so far.
 
-  This DoFn checks the timestamps of files against the time that the pipeline
-  began running. It returns the files that were modified after the pipeline
-  started. If no such files are found, it returns a default file as fallback.
-   """
-  TIME_STATE = CombiningValueStateSpec(
-      'max', combine_fn=partial(max, default=_START_TIME_STAMP))
+  Internal only; no backwards compatibility.
 
-  def process(self, element, time_state=beam.DoFn.StateParam(TIME_STATE)):
+  Input is ``(key, FileMetadata)`` and every file shares one key, so this
+  state is the newest timestamp across the whole watch, not per path. Files
+  at or before pipeline construction are the initial model and are not
+  emitted. Emitting one value per pane keeps the output usable as a
+  singleton side input.
+  """
+  LATEST_TS = ReadModifyWriteStateSpec('latest_ts', FloatCoder())
+
+  def process(self, element, latest_ts=beam.DoFn.StateParam(LATEST_TS)):
     _, file_metadata = element
-    new_ts = file_metadata.last_updated_in_seconds
-    old_ts = time_state.read()
-    if new_ts > old_ts:
-      time_state.clear()
-      time_state.add(new_ts)
-      model_path = file_metadata.path
-    else:
-      model_path = ''
-
-    model_name = os.path.splitext(os.path.basename(model_path))[0]
-    return [
-        (model_path, ModelMetadata(model_id=model_path, model_name=model_name))
-    ]
+    if file_metadata is None:
+      return
+    new_ts = _as_seconds(file_metadata.last_updated_in_seconds)
+    old_ts = latest_ts.read()
+    if old_ts is None:
+      old_ts = _as_seconds(_START_TIME_STAMP)
+    if new_ts <= old_ts:
+      return
+    latest_ts.write(new_ts)
+    yield _model_metadata(file_metadata.path)
+
+
+class _SelectLatestModel(beam.PTransform):
+  """Reduce watched files to singleton updates of the newest model.
+
+  Internal only; no backwards compatibility.
+  """
+  def expand(self, pcoll) -> beam.PCollection[ModelMetadata]:
+    # Window first so one poll's files are reduced together. RunInference
+    # reads this output with AsSingleton, which rejects a pane that contains
+    # more than one model. The timestamp state above then drops a later file
+    # whose modification time is older than a model already emitted.
+    return (
+        pcoll
+        | 'AttachKey' >> beam.Map(_with_global_key)
+        | 'ApplyGlobalWindow' >> beam.WindowInto(
+            window.GlobalWindows(),
+            trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)),
+            accumulation_mode=trigger.AccumulationMode.DISCARDING)
+        | 'LatestFilePerPane' >> beam.CombinePerKey(_LatestFileCombineFn())
+        | 'EmitLatestModel' >> beam.ParDo(_EmitLatestModel()))
 
 
 class WatchFilePattern(beam.PTransform):
@@ -117,33 +177,37 @@ def __init__(
       interval=360,
       stop_timestamp=MAX_TIMESTAMP,
   ):
-    """
-    Watches a directory for updates to files matching a given file pattern.
+    """Watches a directory for updates to files matching a given file pattern.
+
+    Emits the matching file with the newest modification time, one model per
+    pane, so the result can be read with
+    :class:`~apache_beam.pvalue.AsSingleton`.
 
     Args:
-      file_pattern: The file path to read from as a local file path or a
-        GCS ``gs://`` path. The path can contain glob characters
-        (``*``, ``?``, and ``[...]`` sets).
-        interval: Interval at which to check for files matching file_pattern
-        in seconds.
+      file_pattern: The file path to read from, as a local path or a GCS
+        ``gs://`` path. The path can contain glob characters (``*``, ``?``,
+        and ``[...]`` sets).
+      interval: Interval, in seconds, at which to check for files matching
+        ``file_pattern``.
       stop_timestamp: Timestamp after which no more files will be checked.
 
     **Note**:
 
-    1. Any previously used filenames cannot be reused. If a file is added
-        or updated to a previously used filename, this transform will ignore
-        that update. To trigger a model update, always upload a file with
-        unique name. Do not preserve old modification timestamps when copying
-        or uploading a file (for example, when rolling back to an older model),
-        as files with timestamps older than the pipeline start time are ignored.
-    2. Do not upload more than one model update per polling ``interval``, as
-        matching multiple new files in the same poll can cause undefined
-        ordering or singleton side-input errors.
-    3. Initially, before the pipeline startup time, WatchFilePattern expects
-        at least one file present that matches the file_pattern.
-    4. This transform is supported in streaming mode since
-        MatchContinuously produces an unbounded source. Running in batch
-        mode can lead to undesired results or result in pipeline being stuck.
+    1. Each model update must use a new filename. A path is emitted only once,
+       and ``RunInference`` reloads a model only when that path changes. To
+       roll back, upload the desired model under a new name whose modification
+       time is newer than the model currently in use. A copy that keeps an
+       older timestamp (for example ``cp -p``) is not selected. On GCS the
+       timestamp is the object upload time.
+    2. A poll may observe several matching files. Only the file with the
+       newest modification time is emitted.
+    3. At least one file matching ``file_pattern`` must exist before the
+       pipeline starts. Files whose modification time is at or before pipeline
+       construction are not emitted. The model configured on the
+       ``ModelHandler`` is used until a newer file appears.
+    4. This transform is supported in streaming mode.
+       ``MatchContinuously`` is an unbounded source, and running it in batch
+       mode can stall the pipeline or drop updates.
     """
     self.file_pattern = file_pattern
     self.interval = interval
@@ -157,10 +221,4 @@ def expand(self, pcoll) -> beam.PCollection[ModelMetadata]:
             interval=self.interval,
             stop_timestamp=self.stop_timestamp,
             empty_match_treatment=EmptyMatchTreatment.DISALLOW)
-        | "AttachKey" >> beam.Map(lambda x: (x.path, x))
-        | "GetLatestFileMetaData" >> beam.ParDo(_GetLatestFileByTimeStamp())
-        | "AcceptNewSideInputOnly" >> beam.ParDo(_ConvertIterToSingleton())
-        | 'ApplyGlobalWindow' >> beam.transforms.WindowInto(
-            window.GlobalWindows(),
-            trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)),
-            accumulation_mode=trigger.AccumulationMode.DISCARDING))
+        | 'SelectLatestModel' >> _SelectLatestModel())