From 0c97816162d1dd82363cd8d1738c9cc7d1bb7c56 Mon Sep 17 00:00:00 2001 From: Amogh Atreya Date: Fri, 23 Oct 2026 23:49:03 -0700 Subject: [PATCH] [RunInference] Fix WatchFilePattern selecting the newest model --- CHANGES.md | 1 + ...h_image_classification_with_side_inputs.py | 8 +- .../apache_beam/ml/inference/base_test.py | 4 +- sdks/python/apache_beam/ml/inference/utils.py | 198 +++++++++++------- .../apache_beam/ml/inference/utils_test.py | 159 ++++++++------ 5 files changed, 233 insertions(+), 137 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..85352c710bd6 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py b/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py index 903bf38ab279..a0d424d6847c 100644 --- a/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py +++ b/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py @@ -48,9 +48,11 @@ --file_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 diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index dec1c2906cd2..edd34c8ea13e 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/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( diff --git a/sdks/python/apache_beam/ml/inference/utils.py b/sdks/python/apache_beam/ml/inference/utils.py index 8cbe836e22dd..e48aebb56873 100644 --- a/sdks/python/apache_beam/ml/inference/utils.py +++ b/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()) diff --git a/sdks/python/apache_beam/ml/inference/utils_test.py b/sdks/python/apache_beam/ml/inference/utils_test.py index 66499a5a6f48..af0677e39a48 100644 --- a/sdks/python/apache_beam/ml/inference/utils_test.py +++ b/sdks/python/apache_beam/ml/inference/utils_test.py @@ -21,82 +21,117 @@ import apache_beam as beam from apache_beam.io.filesystem import FileMetadata from apache_beam.ml.inference import utils +from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +def _file(path, seconds_from_start): + return FileMetadata( + path, + 10, + last_updated_in_seconds=utils._START_TIME_STAMP + seconds_from_start) + + +def _model_ids(pcoll, label='ModelIds'): + return pcoll | label >> beam.Map(lambda model: model.model_id) + + class WatchFilePatternTest(unittest.TestCase): - def test_latest_file_by_timestamp_default_value(self): - # match continuously returns the files in sorted timestamp order. - main_input_pcoll = [ - FileMetadata( - 'path1.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP - 20), - FileMetadata( - 'path2.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP - 10) + def test_latest_file_combine_fn_picks_newest_regardless_of_order(self): + older = _file('older.py', 10) + newer = _file('newer.py', 50) + initial = _file('initial.py', -5) + combine = utils._LatestFileCombineFn() + + def accumulate(files): + accumulator = combine.create_accumulator() + for metadata in files: + accumulator = combine.add_input(accumulator, metadata) + return accumulator + + merged = combine.merge_accumulators( + [accumulate([older]), accumulate([newer, initial])]) + self.assertEqual('newer.py', combine.extract_output(merged).path) + swapped = combine.merge_accumulators( + [accumulate([newer, initial]), accumulate([older])]) + self.assertEqual('newer.py', combine.extract_output(swapped).path) + + def test_latest_file_combine_fn_breaks_timestamp_ties_by_path(self): + combine = utils._LatestFileCombineFn() + first = combine.add_input(combine.create_accumulator(), _file('b.py', 10)) + second = combine.add_input(combine.create_accumulator(), _file('a.py', 10)) + merged = combine.merge_accumulators([first, second]) + self.assertEqual('b.py', combine.extract_output(merged).path) + + def test_files_at_or_before_pipeline_start_are_not_emitted(self): + files = [_file('path1.py', -20), _file('path2.py', -10), _file('now.py', 0)] + with TestPipeline() as p: + emitted = _model_ids(p | beam.Create(files) | utils._SelectLatestModel()) + assert_that(emitted, equal_to([])) + + def test_file_newer_than_pipeline_start_is_emitted(self): + with TestPipeline() as p: + emitted = _model_ids( + p | beam.Create([_file('path1.py', 10)]) | utils._SelectLatestModel()) + assert_that(emitted, equal_to(['path1.py'])) + + def test_one_poll_emits_only_the_newest_model(self): + # Newest file is not last, and another post-start file is present, so + # neither input order nor "every new path" can select the right model. + files = [ + _file('path4.py', 20), + _file('path1.py', -20), + _file('path3.py', 10), + _file('path2.py', -10), ] with TestPipeline() as p: - files_pc = ( + updates = p | 'Files' >> beam.Create(files) | utils._SelectLatestModel() + assert_that(_model_ids(updates), equal_to(['path4.py']), label='emitted') + singleton = ( p - | beam.Create(main_input_pcoll) - | beam.Map(lambda x: (x.path, x)) - | beam.ParDo(utils._GetLatestFileByTimeStamp()) - | beam.Map(lambda x: x[0])) - assert_that(files_pc, equal_to(['', ''])) - - def test_latest_file_with_timestamp_after_pipeline_construction_time(self): - main_input_pcoll = [ - FileMetadata( - 'path1.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP + 10) - ] + | 'Probe' >> beam.Create([None]) + | beam.Map( + lambda _, model: model.model_id, + beam.pvalue.AsSingleton(updates))) + assert_that(singleton, equal_to(['path4.py']), label='singleton') + + def test_older_update_does_not_replace_a_newer_model(self): + # Shared key, processed in input order. The second file is still newer + # than pipeline start, but older than the model already accepted. + files = [_file('newer.py', 50), _file('older.py', 10)] with TestPipeline() as p: - files_pc = ( + emitted = _model_ids( p - | beam.Create(main_input_pcoll) - | beam.Map(lambda x: (x.path, x)) - | beam.ParDo(utils._GetLatestFileByTimeStamp()) - | beam.Map(lambda x: x[0])) - assert_that(files_pc, equal_to(['path1.py'])) - - def test_emitting_singleton_output(self): - # match continuously returns the files in sorted timestamp order. - main_input_pcoll = [ - FileMetadata( - 'path1.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP - 20), - # returns default - FileMetadata( - 'path2.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP - 10), - # returns default - FileMetadata( - 'path3.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP + 10), - FileMetadata( - 'path4.py', - 10, - last_updated_in_seconds=utils._START_TIME_STAMP + 20) - ] - # returns path3.py + | beam.Create(files) + | beam.Map(lambda metadata: (utils._GLOBAL_MODEL_KEY, metadata)) + | beam.ParDo(utils._EmitLatestModel())) + assert_that(emitted, equal_to(['newer.py'])) + def test_newer_update_replaces_an_older_model(self): + files = [_file('older.py', 10), _file('newer.py', 50)] with TestPipeline() as p: - files_pc = ( + emitted = _model_ids( p - | beam.Create(main_input_pcoll) - | beam.Map(lambda x: (x.path, x)) - | beam.ParDo(utils._GetLatestFileByTimeStamp()) - | beam.ParDo(utils._ConvertIterToSingleton()) - | beam.Map(lambda x: x[0])) - assert_that(files_pc, equal_to(['', 'path3.py', 'path4.py'])) + | beam.Create(files) + | beam.Map(lambda metadata: (utils._GLOBAL_MODEL_KEY, metadata)) + | beam.ParDo(utils._EmitLatestModel()), + label='UpdatedModelIds') + assert_that(emitted, equal_to(['older.py', 'newer.py'])) + + def test_later_poll_does_not_emit_an_older_model(self): + options = PipelineOptions(flags=['--streaming']) + with TestPipeline(options=options) as p: + stream = ( + TestStream().advance_watermark_to(0).add_elements( + [_file('newer.py', 50)]).advance_processing_time(5).add_elements([ + _file('older.py', 10) + ]).advance_processing_time(5).advance_watermark_to_infinity()) + emitted = _model_ids( + p | stream | utils._SelectLatestModel(), label='PolledModelIds') + assert_that(emitted, equal_to(['newer.py'])) if __name__ == '__main__':