[pjrt] ensure client destruction on process exit (#1999)

`torch_xla` doesn't call `PJRT_ClientDestroy` properly. This means that
we are not closing the devices properly.

Recently, this started causing hangs on `n300` boards on subsequent
execution of tests.

This PR introduces a global singleton object which will ensure that we
are properly destroying the client instance on process shutdown. The
singleton serves as a fallback mechanism if the framework doesn't call
`PJRT_ClientDestroy` - like in the case of `torch_xla`.

Additionally, optimizer sub-meshes are now closed after each
compilation; this previously was needed to avoid hangs, but now it
causes them. Leaving the mechanism of persisting optimizer submesh in
the code base, so that we can play with it if needed. Obviously, we need
to dig deep into these issues to fix them properly.

NOTE: this does not solve the case when the process terminates abruptly,
e.g. in case of `SIGSEGV` (segmentation fault). For this, ideally we
would want a fix on `tt-metal` side.

Closes #1824

Author: pilkicTT

pjrt_implementation/inc/api/client_instance.h

@@ -34,6 +34,32 @@ namespace module_builder {
 class ModuleBuilder;
 }
 
+// Singleton class that wraps the PJRT Client Instance.
+// Ensures that we properly destroy the client instance on process termination.
+//
+// NOTE: This is needed since `torch_xla` implementation doesn't call
+// `PJRT_Client_Destroy` API properly and `tt-metal` currently cannot recover
+// (on n300 boards) if we do not properly close all previously opened devices -
+// which is done on client destruction.
+//
+// NOTE: This serves only as a fallback option if `PJRT_Client_Destroy` is not
+// called by the framework.
+class GlobalClientInstanceSingleton {
+public:
+  static ClientInstance *getClientInstance();
+  static PJRT_Error *initClient();
+  static void destroyClient();
+
+private:
+  GlobalClientInstanceSingleton(std::unique_ptr<ClientInstance> client_instance)
+      : m_client_instance(std::move(client_instance)) {}
+
+  bool isInitialized() const { return m_client_instance != nullptr; }
+
+  static GlobalClientInstanceSingleton &getInstance();
+  std::unique_ptr<ClientInstance> m_client_instance;
+};
+
 // Represents PJRT_Client structure and the functionality around it.
 class ClientInstance {
 
@@ -96,6 +122,9 @@ class ClientInstance {
   tt::runtime::Device
   getOrCreateOptimizerSubmesh(const std::vector<uint32_t> &target_mesh_shape);
 
+  // Closes the currently opened optimizer submesh device, if any.
+  void closeOptimizerSubmesh();
+
   // Compiles given mlir program.
   tt_pjrt_status compileMlirProgram(
       const PJRT_Program *mlir_program,

pjrt_implementation/src/api/client_instance.cc

@@ -37,6 +37,37 @@
 
 namespace tt::pjrt {
 
+PJRT_Error *GlobalClientInstanceSingleton::initClient() {
+  std::unique_ptr<ClientInstance> client = std::make_unique<ClientInstance>();
+  PJRT_Error *error = client->initialize();
+  if (error) {
+    return error;
+  }
+
+  GlobalClientInstanceSingleton &singleton_instance = getInstance();
+  singleton_instance.m_client_instance = std::move(client);
+
+  return nullptr;
+}
+
+void GlobalClientInstanceSingleton::destroyClient() {
+  GlobalClientInstanceSingleton &singleton_instance = getInstance();
+  if (singleton_instance.isInitialized()) {
+    singleton_instance.m_client_instance.reset();
+  }
+}
+
+GlobalClientInstanceSingleton &GlobalClientInstanceSingleton::getInstance() {
+  static GlobalClientInstanceSingleton singleton =
+      GlobalClientInstanceSingleton(nullptr);
+  return singleton;
+}
+
+ClientInstance *GlobalClientInstanceSingleton::getClientInstance() {
+  auto &singleton = GlobalClientInstanceSingleton::getInstance();
+  return singleton.m_client_instance.get();
+}
+
 ClientInstance::ClientInstance()
     : m_system_descriptor(nullptr),
       m_module_builder(std::make_unique<module_builder::ModuleBuilder>()),
@@ -58,9 +89,7 @@ ClientInstance::ClientInstance()
 ClientInstance::~ClientInstance() {
   DLOG_F(LOG_DEBUG, "ClientInstance::~ClientInstance");
 
-  if (m_optimizer_submesh.has_value()) {
-    tt::runtime::releaseSubMeshDevice(*m_optimizer_submesh);
-  }
+  closeOptimizerSubmesh();
 
   if (m_parent_mesh.has_value()) {
     tt::runtime::closeMeshDevice(*m_parent_mesh);
@@ -446,6 +475,13 @@ ClientInstance::openMeshDevice(const std::vector<uint32_t> &mesh_shape) {
   return tt::runtime::openMeshDevice(options);
 }
 
+void ClientInstance::closeOptimizerSubmesh() {
+  if (m_optimizer_submesh.has_value()) {
+    tt::runtime::releaseSubMeshDevice(*m_optimizer_submesh);
+    m_optimizer_submesh.reset();
+  }
+}
+
 tt::runtime::Device ClientInstance::getOrCreateOptimizerSubmesh(
     const std::vector<uint32_t> &target_mesh_shape) {
 
@@ -487,23 +523,29 @@ PJRT_Error *onClientCreate(PJRT_Client_Create_Args *args) {
            args->create_options[i].name);
   }
 
-  std::unique_ptr<ClientInstance> client = std::make_unique<ClientInstance>();
-  PJRT_Error *error = client->initialize();
+  PJRT_Error *error = GlobalClientInstanceSingleton::initClient();
+
   if (error) {
+    DLOG_F(ERROR, "Failed to initialize PJRT client instance");
     return error;
   }
 
-  // Successful return.
-  args->client = reinterpret_cast<PJRT_Client *>(client.release());
+  ClientInstance *client_instance =
+      GlobalClientInstanceSingleton::getClientInstance();
+  assert(client_instance != nullptr);
+  args->client = reinterpret_cast<PJRT_Client *>(client_instance);
 
   return nullptr;
 }
 
 PJRT_Error *onClientDestroy(PJRT_Client_Destroy_Args *args) {
   DLOG_F(LOG_DEBUG, "ClientInstance::PJRT_Client_Destroy");
 
-  delete ClientInstance::unwrap(args->client);
-
+  ClientInstance *client_instance = ClientInstance::unwrap(args->client);
+  ClientInstance *global_client_instance =
+      GlobalClientInstanceSingleton::getClientInstance();
+  assert(client_instance == global_client_instance);
+  GlobalClientInstanceSingleton::destroyClient();
   return nullptr;
 }
 

pjrt_implementation/src/api/module_builder/module_builder.cc

@@ -861,7 +861,12 @@ tt_pjrt_status ModuleBuilder::convertFromTTIRToTTNN(
   enableVerboseIRPrinting(ttir_to_ttnn_pm);
 
   // Run the pass manager.
-  if (mlir::failed(ttir_to_ttnn_pm.run(mlir_module.get()))) {
+  mlir::LogicalResult mlir_result = ttir_to_ttnn_pm.run(mlir_module.get());
+
+  // Close the optimizer submesh now that the compilation is complete.
+  client_instance->closeOptimizerSubmesh();
+
+  if (mlir::failed(mlir_result)) {
     DLOG_F(ERROR, "Failed to convert from TTIR to TTNN module");
     return tt_pjrt_status::kInternal;
   }