Use TFRT CPU client and make parameters modifiable through environment variables.

Author: darisoy

third_party/xla_client/env_vars.cc

@@ -20,6 +20,7 @@ const char* const kEnvTpuvmMode = "TPUVM_MODE";
 const char* const kEnvPjRtDevice = "PJRT_DEVICE";
 const char* const kEnvPjRtTpuMaxInflightComputations =
     "PJRT_TPU_MAX_INFLIGHT_COMPUTATIONS";
+const char* const kEnvAsyncCpuClient = "CPU_ASYNC_CLIENT";
 
 }  // namespace env
 }  // namespace xla

third_party/xla_client/env_vars.h

@@ -20,6 +20,7 @@ extern const char* const kEnvStartService;
 extern const char* const kEnvTpuvmMode;
 extern const char* const kEnvPjRtDevice;
 extern const char* const kEnvPjRtTpuMaxInflightComputations;
+extern const char* const kEnvAsyncCpuClient;
 
 }  // namespace env
 }  // namespace xla

third_party/xla_client/pjrt_computation_client.cc

@@ -46,7 +46,9 @@ PjRtComputationClient::PjRtComputationClient() {
   std::string device_type = sys_util::GetEnvString(env::kEnvPjRtDevice, "");
   if (device_type == "CPU") {
     TF_VLOG(1) << "Initializing PjRt CPU client...";
-    client_ = std::move(xla::GetCpuClient(/*asynchronous=*/false).ValueOrDie());
+    bool async = sys_util::GetEnvBool(env::kEnvAsyncCpuClient, true);
+    int cpu_device_count = sys_util::GetEnvInt(env::kEnvNumCpu, 1);
+    client_ = std::move(xla::GetTfrtCpuClient(async, cpu_device_count).ValueOrDie());
   } else if (device_type == "TPU") {
     TF_VLOG(1) << "Initializing PjRt TPU client...";
     int64_t max_inflight_computations = sys_util::GetEnvInt(