Error Handling: return status value when loading PjRt dynamic plugin. (#9495)

Author: ysiraichi

test/run_tests.sh

@@ -227,6 +227,7 @@ function run_xla_op_tests2 {
   run_test "$_TEST_DIR/test_assume_pure_spmd.py"
   run_test "$_TEST_DIR/test_assume_pure_torch.py"
   run_test "$_TEST_DIR/test_dynamic_shapes_detector.py"
+  run_test "$_TEST_DIR/test_runtime_client_initialization_error.py"
 }
 
 function run_xla_op_tests3 {

test/test_runtime_client_initialization_error.py

@@ -0,0 +1,35 @@
+import os
+import torch_xla
+import torch_xla.core.xla_env_vars as xenv
+import unittest
+
+
+class TestClientInitializationError(unittest.TestCase):
+
+  def test(self):
+
+    def initialize_client(device):
+      os.environ[xenv.PJRT_DEVICE] = device
+
+      # The message does not change!
+      # After the first call with DUMMY_DEVICE, all other calls will have
+      # "DUMMY_DEVICE" in their message.
+      message = (
+          f"No PjRtPlugin registered for: DUMMY_DEVICE. "
+          f"Make sure the environment variable {xenv.PJRT_DEVICE} is set "
+          "to a correct device name.")
+
+      with self.assertRaisesRegex(RuntimeError, expected_regex=message):
+        torch_xla._XLAC._init_computation_client()
+
+    # Run the initialization function the first time, ending up in an
+    # exception thrown.
+    initialize_client("DUMMY_DEVICE")
+
+    # Even if the device exists, this call should fail, since the result
+    # of the first call is cached.
+    initialize_client("CPU")
+
+
+if __name__ == '__main__':
+  unittest.main()

torch_xla/csrc/runtime/pjrt_registry.cpp

@@ -61,10 +61,21 @@ xla::GpuAllocatorConfig GetGpuAllocatorConfig() {
   return allocator_config;
 }
 
-std::shared_ptr<const PjRtPlugin> GetPjRtPlugin(
+absl::StatusOr<std::shared_ptr<const PjRtPlugin>> GetPjRtPlugin(
     const std::string& device_type) {
-  auto plugin_path = pjrt_plugins_.find(device_type);
-  return plugin_path != pjrt_plugins_.end() ? plugin_path->second : nullptr;
+  auto entry = pjrt_plugins_.find(device_type);
+  if (entry == pjrt_plugins_.end()) {
+    std::string message = absl::StrCat(
+        "No PjRtPlugin registered for: ", device_type,
+        ". Make sure the environment variable ", env::kEnvPjRtDevice,
+        " is set to a correct device name. See "
+        "https://github.com/pytorch/xla/blob/master/docs/source/"
+        "contribute/plugins.md for more information on "
+        "implementing and registering a new "
+        "plugin.");
+    return XLA_ERROR_WITH_LOCATION(absl::FailedPreconditionError(message));
+  }
+  return entry->second;
 }
 
 }  // namespace
@@ -83,10 +94,10 @@ InitializePjRt(const std::string& device_type) {
 
   if (sys_util::GetEnvBool(env::kEnvPjrtDynamicPlugins, false) &&
       device_type != "CPU") {
-    std::shared_ptr<const PjRtPlugin> plugin = GetPjRtPlugin(device_type);
+    TF_VLOG(1) << "Initializing client for PjRt plugin " << device_type;
+    XLA_ASSIGN_OR_RETURN(std::shared_ptr<const PjRtPlugin> plugin,
+                         GetPjRtPlugin(device_type));
     if (plugin) {
-      TF_VLOG(1) << "Initializing client for PjRt plugin " << device_type;
-
       // Init the absl logging to avoid the log spam.
       absl::InitializeLog();
 

torch_xla/csrc/runtime/runtime.cpp

@@ -19,10 +19,6 @@ static std::atomic<bool> g_computation_client_initialized(false);
 // Can only be called when g_computation_client_initialized is false.
 static absl::StatusOr<ComputationClient * absl_nonnull>
 InitializeComputationClient() {
-  ABSL_CHECK(!g_computation_client_initialized)
-      << "InitializeComputationClient() can only be called once.";
-  g_computation_client_initialized = true;
-
   if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
     tsl::testing::InstallStacktraceHandler();
   }
@@ -37,13 +33,19 @@ InitializeComputationClient() {
         absl::FailedPreconditionError("$PJRT_DEVICE is not set."));
   }
 
+  ABSL_CHECK(!g_computation_client_initialized)
+      << "ComputationClient can only be initialized once.";
+
   std::unique_ptr<ComputationClient> client;
   if (use_ifrt) {
     XLA_ASSIGN_OR_RETURN(client, IfrtComputationClient::Create());
   } else {
     XLA_ASSIGN_OR_RETURN(client, PjRtComputationClient::Create());
   }
 
+  // Set only if we actually successfully initialized a client.
+  g_computation_client_initialized = true;
+
   return client.release();
 }
 
@@ -59,7 +61,7 @@ const absl::StatusOr<ComputationClient * absl_nonnull>& GetComputationClient() {
 }
 
 ComputationClient* absl_nonnull GetComputationClientOrDie() {
-  return GetComputationClient().value();
+  return GetValueOrThrow(GetComputationClient());
 }
 
 ComputationClient* GetComputationClientIfInitialized() {