diff --git a/tensorflow_quantum/core/ops/load_module.py b/tensorflow_quantum/core/ops/load_module.py index b5002ad84..a7acfbf9a 100644 --- a/tensorflow_quantum/core/ops/load_module.py +++ b/tensorflow_quantum/core/ops/load_module.py @@ -21,26 +21,57 @@ from tensorflow.python.platform import resource_loader +class _LazyLoader: + """Lazily loads a TensorFlow op library on first attribute access. + + This defers the call to `load_library.load_op_library` until the module + is actually used, preventing TensorFlow device initialization at import + time. This allows users to configure TensorFlow devices (e.g., enabling + memory growth) after importing tensorflow_quantum. + """ + + def __init__(self, name): + """Initialize the lazy loader. + + Args: + name: The name of the module, e.g. "_tfq_simulate_ops.so" + """ + self._name = name + self._module = None + + def _load(self): + """Load the module if not already loaded.""" + if self._module is None: + try: + path = resource_loader.get_path_to_datafile(self._name) + self._module = load_library.load_op_library(path) + except: + path = os.path.join(get_python_lib(), + "tensorflow_quantum/core/ops", self._name) + self._module = load_library.load_op_library(path) + return self._module + + def __getattr__(self, name): + """Load the module on first attribute access and delegate.""" + module = self._load() + return getattr(module, name) + + def load_module(name): - """Loads the module with the given name. + """Returns a lazy loader for the module with the given name. - First attempts to load the module as though it was embedded into the binary - using Bazel. If that fails, then it attempts to load the module as though - it was installed in site-packages via PIP. + The actual library loading is deferred until the module is first used. + This prevents TensorFlow device initialization at import time, allowing + users to configure TensorFlow devices after importing tensorflow_quantum. Args: name: The name of the module, e.g. "_tfq_simulate_ops.so" Returns: - A python module containing the Python wrappers for the Ops. + A lazy loader object that behaves like the loaded module but defers + loading until first attribute access. Raises: - RuntimeError: If the library cannot be found. + RuntimeError: If the library cannot be found when first accessed. """ - try: - path = resource_loader.get_path_to_datafile(name) - return load_library.load_op_library(path) - except: - path = os.path.join(get_python_lib(), "tensorflow_quantum/core/ops", - name) - return load_library.load_op_library(path) + return _LazyLoader(name) diff --git a/tensorflow_quantum/core/ops/load_module_test.py b/tensorflow_quantum/core/ops/load_module_test.py new file mode 100644 index 000000000..6428e848f --- /dev/null +++ b/tensorflow_quantum/core/ops/load_module_test.py @@ -0,0 +1,51 @@ +# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for load_module lazy loading functionality.""" +import tensorflow as tf +from absl.testing import parameterized +from tensorflow_quantum.core.ops.load_module import load_module, _LazyLoader + + +class LoadModuleTest(tf.test.TestCase, parameterized.TestCase): + """Tests for load_module function and _LazyLoader class.""" + + def test_load_module_returns_lazy_loader(self): + """Test that load_module returns a _LazyLoader instance.""" + loader = load_module("_tfq_utility_ops.so") + self.assertIsInstance(loader, _LazyLoader) + + def test_lazy_loader_defers_loading(self): + """Test that _LazyLoader does not load the module on construction.""" + loader = _LazyLoader("_tfq_utility_ops.so") + # _module should be None before any attribute access + self.assertIsNone(loader._module) + + def test_lazy_loader_loads_on_attribute_access(self): + """Test that _LazyLoader loads the module on attribute access.""" + loader = load_module("_tfq_utility_ops.so") + # Access an attribute to trigger loading + _ = loader.tfq_append_circuit + # Now _module should be loaded + self.assertIsNotNone(loader._module) + + def test_lazy_loader_attribute_access_works(self): + """Test that attributes from the loaded module are accessible.""" + loader = load_module("_tfq_utility_ops.so") + # Accessing an op should return a callable + self.assertTrue(callable(loader.tfq_append_circuit)) + + +if __name__ == '__main__': + tf.test.main() diff --git a/test_device_config_after_import.py b/test_device_config_after_import.py new file mode 100644 index 000000000..14ce0985b --- /dev/null +++ b/test_device_config_after_import.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python +# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Test that TensorFlow device configuration works after importing TFQ. + +This script verifies the fix for the issue where importing tensorflow_quantum +before setting TensorFlow device configuration (e.g., enabling memory growth) +resulted in a RuntimeError: + + RuntimeError: Physical devices cannot be modified after being initialized + +The fix uses lazy loading of native op libraries to defer TensorFlow device +initialization until the ops are actually used. + +Usage: + python test_device_config_after_import.py +""" + +import sys +import tensorflow as tf + +# Import tensorflow_quantum BEFORE configuring devices. +# This used to trigger device initialization and cause errors. +import tensorflow_quantum as tfq + +# Now try to configure devices - this should work without RuntimeError +gpus = tf.config.list_physical_devices('GPU') +if gpus: + try: + # Try setting memory growth - this would fail before the fix + tf.config.experimental.set_memory_growth(gpus[0], True) + print("SUCCESS: Device configuration after import works!") + print(f" - Configured memory growth for GPU: {gpus[0]}") + except RuntimeError as e: + print(f"FAILED: {e}") + sys.exit(1) +else: + # No GPU available, but we can still test that importing TFQ + # doesn't prematurely initialize devices by checking CPU config + cpus = tf.config.list_physical_devices('CPU') + print("SUCCESS: TFQ import did not prematurely initialize devices!") + print(f" - Available CPUs: {cpus}") + print(" - No GPU available to test memory growth, but import test passed.") + +# Verify TFQ is actually usable after configuration +print(f" - TFQ version: {tfq.__version__}")