Revert "Use tfrt cpu client (#3898)"

This reverts commit a6ea0f3.

Author: JackCaoG

test/pjrt/test_experimental_pjrt_multi_cpu.py

@@ -77,8 +77,7 @@ function run_xla_backend_mp {
 
 function run_pjrt {
   echo "Running in PjRt runtime: $@"
-  # TODO(darisoy): run these tests with multiple CPU devices, this fails due to TF issue.
-  PJRT_DEVICE=CPU CPU_NUM_DEVICES=1 run_test "$@"
+  PJRT_DEVICE=CPU run_test "$@"
 }
 
 function run_async_scalar {

test/run_tests.sh

@@ -150,7 +150,6 @@ cc_library(
         "//tensorflow/compiler/xla/pjrt:cpu_device",
         "//tensorflow/compiler/xla/pjrt:tpu_client",
         "//tensorflow/compiler/xla/pjrt:pjrt_client",
-        "//tensorflow/compiler/xla/pjrt:tfrt_cpu_pjrt_client",
         "//tensorflow/compiler/xla/rpc:grpc_stub",
         "//tensorflow/compiler/xla/service:cpu_plugin",
         "//tensorflow/compiler/xla/service:platform_util",

third_party/xla_client/BUILD

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

third_party/xla_client/env_vars.cc

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

third_party/xla_client/env_vars.h

@@ -9,7 +9,6 @@
 #include "tensorflow/compiler/xla/literal.h"
 #include "tensorflow/compiler/xla/pjrt/cpu_device.h"
 #include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
-#include "tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.h"
 #include "tensorflow/compiler/xla/pjrt/tpu_client.h"
 #include "tensorflow/compiler/xla/shape.h"
 #include "tensorflow/compiler/xla/xla_client/computation_client.h"
@@ -47,10 +46,7 @@ PjRtComputationClient::PjRtComputationClient() {
   std::string device_type = sys_util::GetEnvString(env::kEnvPjRtDevice, "");
   if (device_type == "CPU") {
     TF_VLOG(1) << "Initializing PjRt CPU client...";
-    bool async = sys_util::GetEnvBool(env::kEnvPjrtAsyncCpuClient, true);
-    int cpu_device_count = sys_util::GetEnvInt(env::kEnvNumCpu, 1);
-    client_ =
-        std::move(xla::GetTfrtCpuClient(async, cpu_device_count).ValueOrDie());
+    client_ = std::move(xla::GetCpuClient(/*asynchronous=*/false).ValueOrDie());
   } else if (device_type == "TPU") {
     TF_VLOG(1) << "Initializing PjRt TPU client...";
     int64_t max_inflight_computations = sys_util::GetEnvInt(

third_party/xla_client/pjrt_computation_client.cc

@@ -26,4 +26,3 @@
 TPU_PROCESS_ADDRESSES = 'TPU_PROCESS_ADDRESSES'
 TPU_VISIBLE_CHIPS = 'TPU_VISIBLE_CHIPS'
 TPU_PROCESS_PORT = 'TPU_PROCESS_PORT'
-PJRT_CPU_ASYNC_CLIENT = 'PJRT_CPU_ASYNC_CLIENT'