From 9f9bed6c58c9618079c4d2a9f5f5fd3fa3a934a1 Mon Sep 17 00:00:00 2001 From: Tanmay Singh Date: Wed, 28 Oct 2026 10:00:18 +0530 Subject: [PATCH 1/2] [RunInference] Import resource independently of ModelManager `import resource` shared a try block with the ModelManager import, so when ModelManager's optional dependencies (torch) were missing, resource was set to None too and RunInference reported 0 for model memory. Fixes #40854 --- CHANGES.md | 1 + sdks/python/apache_beam/ml/inference/base.py | 5 +++- .../apache_beam/ml/inference/base_test.py | 23 +++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index 59253c35bf87..982d8eb40ab2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -147,6 +147,7 @@ * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). * (Go) Fixed generic registration (for example `register.Function1x1` or `register.DoFn1x1`) panicking on functions and methods that take or return an anonymous struct. Unnamed struct types are now only registered with the schema registry, since they have no type key for the runtime type registry ([#21789](https://github.com/apache/beam/issues/21789)). +* (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 5ee65afa0615..4d5eb4f74c0a 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -68,10 +68,13 @@ try: # pylint: disable=wrong-import-order, wrong-import-position import resource +except ImportError: + resource = None # type: ignore[assignment] +try: + # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.ml.inference.model_manager import ModelManager except ImportError: - resource = None # type: ignore[assignment] ModelManager = None # type: ignore[assignment] _NANOSECOND_TO_MILLISECOND = 1_000_000 diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index dec1c2906cd2..55f90c36dd65 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -23,6 +23,7 @@ import os import pickle import random +import subprocess import sys import tempfile import threading @@ -3153,5 +3154,27 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag( dofn._model.force_reset() +class ResourceImportTest(unittest.TestCase): + @unittest.skipIf( + sys.platform == 'win32', + 'The resource module is not available on Windows.') + def test_resource_available_when_model_manager_import_fails(self): + # Run in a fresh interpreter in which importing model_manager fails, as it + # does when its optional dependencies (e.g. torch) are not installed. + code = '\n'.join([ + 'import sys', + "sys.modules['apache_beam.ml.inference.model_manager'] = None", + 'from apache_beam.ml.inference import base', + 'assert base.ModelManager is None', + 'assert base.resource is not None', + 'assert base._get_current_process_memory_in_bytes() > 0', + ]) + result = subprocess.run([sys.executable, '-c', code], + capture_output=True, + text=True, + check=False) + self.assertEqual(result.returncode, 0, msg=result.stderr) + + if __name__ == '__main__': unittest.main() From d03d8a22d0c5ca531962b693ec7d898ef8057312 Mon Sep 17 00:00:00 2001 From: Tanmay Singh Date: Thu, 29 Oct 2026 17:30:58 +0530 Subject: [PATCH 2/2] Replace subprocess test with a RunInference warning test --- .../apache_beam/ml/inference/base_test.py | 41 ++++++++----------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 55f90c36dd65..d4cfce8ec4f4 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -23,7 +23,6 @@ import os import pickle import random -import subprocess import sys import tempfile import threading @@ -1341,6 +1340,24 @@ def test_timing_metrics(self): self.assertEqual(load_model_latency.result.count, 1) self.assertEqual(load_model_latency.result.mean, 500) + @unittest.skipIf( + sys.platform == 'win32', + 'The resource module is not available on Windows.') + def test_no_resource_warning_when_model_manager_unavailable(self): + # ModelManager needs optional dependencies (e.g. torch). RunInference should + # still measure model memory with `resource` when ModelManager is missing. + with unittest.mock.patch.object(base, 'ModelManager', None), \ + unittest.mock.patch.object(base.logging, 'warning') as mock_warning: + with TestPipeline() as pipeline: + _ = ( + pipeline + | beam.Create([1, 5, 3, 10]) + | base.RunInference(FakeModelHandler())) + messages = [str(call) for call in mock_warning.call_args_list] + self.assertFalse( + any('Resource module is not available' in m for m in messages), + messages) + def test_forwards_batch_args(self): examples = list(range(100)) with TestPipeline('FnApiRunner') as pipeline: @@ -3154,27 +3171,5 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag( dofn._model.force_reset() -class ResourceImportTest(unittest.TestCase): - @unittest.skipIf( - sys.platform == 'win32', - 'The resource module is not available on Windows.') - def test_resource_available_when_model_manager_import_fails(self): - # Run in a fresh interpreter in which importing model_manager fails, as it - # does when its optional dependencies (e.g. torch) are not installed. - code = '\n'.join([ - 'import sys', - "sys.modules['apache_beam.ml.inference.model_manager'] = None", - 'from apache_beam.ml.inference import base', - 'assert base.ModelManager is None', - 'assert base.resource is not None', - 'assert base._get_current_process_memory_in_bytes() > 0', - ]) - result = subprocess.run([sys.executable, '-c', code], - capture_output=True, - text=True, - check=False) - self.assertEqual(result.returncode, 0, msg=result.stderr) - - if __name__ == '__main__': unittest.main()