Create test + fix.

Author: ysiraichi

BUILD

@@ -77,6 +77,7 @@ test_suite(
         "//test/cpp:test_replication",
         "//test/cpp:test_tensor",
         "//test/cpp:test_xla_sharding",
+        "//torch_xla/csrc/runtime:runtime_test",
         "//torch_xla/csrc/runtime:pjrt_computation_client_test",
         # "//torch_xla/csrc/runtime:ifrt_computation_client_test",
     ],

torch_xla/csrc/runtime/BUILD

@@ -524,3 +524,12 @@ ptxla_cc_test(
 #         "@xla//xla/tools:hlo_module_loader",
 #     ],
 # )
+
+ptxla_cc_test(
+    name = "runtime_test",
+    srcs = ["runtime_test.cpp"],
+    deps = [
+        ":runtime",
+        "@tsl//tsl/platform:test_main",
+    ],
+)

torch_xla/csrc/runtime/runtime.cpp

@@ -12,8 +12,8 @@ namespace torch_xla::runtime {
 
 std::atomic<bool> g_computation_client_initialized(false);
 
-// Creates a new instance of a `ComputationClient` (e.g. `PjRtComputationClient`),
-// and initializes the computation client
+// Creates a new instance of a `ComputationClient` (e.g.
+// `PjRtComputationClient`), and initializes the computation client
 static absl::StatusOr<absl_nonnull std::unique_ptr<ComputationClient>>
 InitializeComputationClient() {
   if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
@@ -22,9 +22,11 @@ InitializeComputationClient() {
 
   std::unique_ptr<ComputationClient> client;
 
-  // Disable IFRT right now as it currently crashes.
+  // TODO: enable IFRT once it's not crashing anymore.
+  // Ref: https://github.com/pytorch/xla/pull/8267
+  //
   // static bool use_ifrt = sys_util::GetEnvBool("XLA_USE_IFRT", false);
-  static bool use_ifrt = false;
+  static const bool use_ifrt = false;
   if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") != "") {
     if (use_ifrt) {
       client = std::make_unique<IfrtComputationClient>();
@@ -40,7 +42,8 @@ InitializeComputationClient() {
 }
 
 absl::StatusOr<ComputationClient*> GetComputationClient() {
-  static absl::StatusOr<absl_nonnull std::unique_ptr<ComputationClient>> maybeClient = InitializeComputationClient();
+  static absl::StatusOr<absl_nonnull std::unique_ptr<ComputationClient>>
+      maybeClient = InitializeComputationClient();
 
   if (!maybeClient.ok()) {
     return maybeClient.status();

torch_xla/csrc/runtime/runtime_test.cpp

@@ -0,0 +1,31 @@
+#include "torch_xla/csrc/runtime/runtime.h"
+
+#include <gtest/gtest.h>
+
+namespace torch_xla::runtime {
+
+TEST(RuntimeTest, NullComputationClient) {
+  auto client = GetComputationClientIfInitialized();
+  EXPECT_EQ(client, nullptr);
+}
+
+TEST(RuntimeTest, GetComputationClientSuccess) {
+  ComputationClient* client;
+
+  client = GetComputationClientIfInitialized();
+  EXPECT_EQ(client, nullptr);
+
+  // Initialize the ComputationClient.
+  // Check all the APIs return the same valid ComputationClient.
+
+  client = GetComputationClientOrDie();
+  EXPECT_NE(client, nullptr);
+
+  auto status = GetComputationClient();
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(client, status.value());
+
+  EXPECT_EQ(client, GetComputationClientIfInitialized());
+}
+
+}  // namespace torch_xla::runtime