diff --git a/CHANGES.md b/CHANGES.md index 44e7568cb8ee..8f5530831266 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -80,6 +80,7 @@ ## Bugfixes +* (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..d4cfce8ec4f4 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -1340,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: