From 7dac2d707b860b2bb262358777c3a9e67f73922d Mon Sep 17 00:00:00 2001
From: root <root@ip-172-31-59-9.us-west-2.compute.internal>
Date: Tue, 7 Oct 2025 01:02:44 +0000
Subject: [PATCH 1/2] adding static unique_ptr to manage lifecycle of
 PjrtClient

---
 torch_xla/csrc/runtime/runtime.cpp | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/torch_xla/csrc/runtime/runtime.cpp b/torch_xla/csrc/runtime/runtime.cpp
index 4dced2531bd5..ff199831c72d 100644
--- a/torch_xla/csrc/runtime/runtime.cpp
+++ b/torch_xla/csrc/runtime/runtime.cpp
@@ -17,7 +17,7 @@ static std::atomic<bool> g_computation_client_initialized(false);
 // Creates a new instance of a `ComputationClient` (e.g.
 // `PjRtComputationClient`), and initializes the computation client.
 // Can only be called when g_computation_client_initialized is false.
-static absl::StatusOr<ComputationClient * absl_nonnull>
+static absl::StatusOr<std::unique_ptr<ComputationClient>>
 InitializeComputationClient() {
   if (sys_util::GetEnvBool("XLA_DUMP_FATAL_STACK", false)) {
     tsl::testing::InstallStacktraceHandler();
@@ -46,7 +46,7 @@ InitializeComputationClient() {
   // Set only if we actually successfully initialized a client.
   g_computation_client_initialized = true;
 
-  return client.release();
+  return client;
 }
 
 const absl::StatusOr<ComputationClient * absl_nonnull>& GetComputationClient() {
@@ -55,9 +55,20 @@ const absl::StatusOr<ComputationClient * absl_nonnull>& GetComputationClient() {
   // Since we only allow a single initialization, as soon as this function is
   // called, we store the initialization result in this trivially destructible
   // reference.
-  static const auto& maybe_client =
-      *new absl::StatusOr<ComputationClient*>(InitializeComputationClient());
-  return maybe_client;
+  static absl::StatusOr<std::unique_ptr<ComputationClient>> init_result =
+      InitializeComputationClient();
+
+  static absl::StatusOr<ComputationClient* absl_nonnull> maybeClient = []() {
+    if (init_result.ok()) {
+      return absl::StatusOr<ComputationClient * absl_nonnull>(
+          init_result.value().get());
+    } else {
+      // Return the same error from initialization
+      return absl::StatusOr<ComputationClient * absl_nonnull>(
+          init_result.status());
+    }
+  }();
+  return maybeClient;
 }
 
 ComputationClient* GetComputationClientIfInitialized() {

From d77ff9407e6c95d4323eadfe75014149889c1c1a Mon Sep 17 00:00:00 2001
From: root <root@ip-172-31-59-9.us-west-2.compute.internal>
Date: Tue, 7 Oct 2025 01:34:23 +0000
Subject: [PATCH 2/2] lint

---
 torch_xla/csrc/runtime/runtime.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/torch_xla/csrc/runtime/runtime.cpp b/torch_xla/csrc/runtime/runtime.cpp
index ff199831c72d..270155a538d5 100644
--- a/torch_xla/csrc/runtime/runtime.cpp
+++ b/torch_xla/csrc/runtime/runtime.cpp
@@ -58,17 +58,16 @@ const absl::StatusOr<ComputationClient * absl_nonnull>& GetComputationClient() {
   static absl::StatusOr<std::unique_ptr<ComputationClient>> init_result =
       InitializeComputationClient();
 
-  static absl::StatusOr<ComputationClient* absl_nonnull> maybeClient = []() {
+  static absl::StatusOr<ComputationClient* absl_nonnull> maybe_client = []() {
     if (init_result.ok()) {
       return absl::StatusOr<ComputationClient * absl_nonnull>(
           init_result.value().get());
     } else {
-      // Return the same error from initialization
       return absl::StatusOr<ComputationClient * absl_nonnull>(
           init_result.status());
     }
   }();
-  return maybeClient;
+  return maybe_client;
 }
 
 ComputationClient* GetComputationClientIfInitialized() {