Repository: google/jax
Pull request #20003 (2024-02-27 23:27:24)

Support mocking number of GPUs in CUDA plugin.
Support mocking number of GPUs in CUDA plugin.

Also move reading jax config value to be right before the client is created. Previously they were read before calling register_plugin, which happens during import and before any call of jax.config.update.

The decorator in mock_gpu_test was used wrongly. jtu.run_on_devices will create the client before jax.config.update is called, which is not desired. Remove the decorator will not fail CPU/TPU tests because the mesh will check the num_shard and the number of devices in the client and skip it if it does not match.

generate_pjrt_gpu_plugin_options is only used in places that do not require compatibility so do not need to update xla_client version.



Patch:
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
index 89344e2ee..e2c4c1d97 100644
--- a/jax/_src/xla_bridge.py
+++ b/jax/_src/xla_bridge.py
@@ -485,6 +485,20 @@ def discover_pjrt_plugins() -> None:
                          "calling %s.initialize()", plugin_module_name)
 
 
+def _options_from_jax_configs(plugin_name):
+  if plugin_name != "cuda":
+    return {}
+
+  options = {}
+  visible_devices = CUDA_VISIBLE_DEVICES.value
+  if visible_devices != 'all':
+    options['visible_devices'] = [int(x) for x in visible_devices.split(',')]
+  options['enable_mock_nccl'] = _USE_MOCK_GPU_CLIENT.value
+  if options['enable_mock_nccl']:
+    options['num_nodes'] = _MOCK_NUM_GPUS.value
+  return options
+
+
 # TODO(b/261345120): decide on a public name and expose a public method which is
 # an alias of this method.
 def register_plugin(
@@ -509,15 +523,19 @@ def register_plugin(
   def factory():
     if not xla_client.pjrt_plugin_initialized(plugin_name):
       xla_client.initialize_pjrt_plugin(plugin_name)
-
+    updated_options = {}
+    if options is not None:
+      updated_options.update(options)
+    updated_options.update(_options_from_jax_configs(plugin_name))
     if distributed.global_state.client is None:
-      return xla_client.make_c_api_client(plugin_name, options, None)
+      return xla_client.make_c_api_client(plugin_name, updated_options, None)
+
     distribute_options = {
         'node_id': distributed.global_state.process_id,
         'num_nodes': distributed.global_state.num_processes,
     }
     if options is not None:
-      distribute_options.update(options)
+      distribute_options.update(updated_options)
     return xla_client.make_c_api_client(
         plugin_name, distribute_options, distributed.global_state.client
     )
diff --git a/jax_plugins/cuda/__init__.py b/jax_plugins/cuda/__init__.py
index bbd65ab55..d8f35fc33 100644
--- a/jax_plugins/cuda/__init__.py
+++ b/jax_plugins/cuda/__init__.py
@@ -75,9 +75,7 @@ def initialize():
   if path is None:
     return
 
-  options = xla_client.generate_pjrt_gpu_plugin_options(
-      xb.CUDA_VISIBLE_DEVICES.value
-  )
+  options = xla_client.generate_pjrt_gpu_plugin_options()
   c_api = xb.register_plugin(
       'cuda', priority=500, library_path=str(path), options=options
   )


Test patch:
diff --git a/tests/mock_gpu_test.py b/tests/mock_gpu_test.py
index ebc32ba25..b955f0398 100644
--- a/tests/mock_gpu_test.py
+++ b/tests/mock_gpu_test.py
@@ -27,7 +27,6 @@ import numpy as np
 config.parse_flags_with_absl()
 
 
-@jtu.run_on_devices('gpu')
 class MockGPUTest(jtu.JaxTestCase):
 
   def setUp(self):
diff --git a/tests/xla_bridge_test.py b/tests/xla_bridge_test.py
index a3d78f1c2..0409187a3 100644
--- a/tests/xla_bridge_test.py
+++ b/tests/xla_bridge_test.py
@@ -194,7 +194,7 @@ class XlaBridgeTest(jtu.JaxTestCase):
     self.assertIn("name2", xb._backend_factories)
     self.assertEqual(registration.priority, 400)
     self.assertTrue(registration.experimental)
-    mock_make.assert_called_once_with("name1", None, None)
+    mock_make.assert_called_once_with("name1", {}, None)
 
   def test_register_plugin_with_config(self):
     test_json_file_path = os.path.join(