Expose KeyValueStore of the distributed service as a python binding.

PiperOrigin-RevId: 448820874
nxp-upstream · May 15, 2022 · 73c1b1e · 73c1b1e
1 parent cd33ee9
commit 73c1b1e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/tensorflow/compiler/xla/python/xla.cc b/tensorflow/compiler/xla/python/xla.cc
@@ -386,7 +386,17 @@ PYBIND11_MODULE(xla_extension, m) {
              std::shared_ptr<DistributedRuntimeClient>>
       distributed_runtime_client(m, "DistributedRuntimeClient");
   distributed_runtime_client.def("connect", &DistributedRuntimeClient::Connect)
-      .def("shutdown", &DistributedRuntimeClient::Shutdown);
+      .def("shutdown", &DistributedRuntimeClient::Shutdown)
+      .def(
+          "blocking_key_value_get",
+          [](DistributedRuntimeClient& client, std::string key,
+             int64_t timeout_in_ms) {
+            return client.BlockingKeyValueGet(
+                key, absl::Milliseconds(timeout_in_ms));
+          },
+          py::arg("key"), py::arg("timeout_in_ms"))
+      .def("key_value_set", &DistributedRuntimeClient::KeyValueSet,
+           py::arg("key"), py::arg("value"));
 
   m.def(
       "get_distributed_runtime_service",

diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py
@@ -43,7 +43,7 @@
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes.
-_version = 66
+_version = 67
 
 # Version number for MLIR:Python components.
 mlir_api_version = 18

diff --git a/tensorflow/compiler/xla/python/xla_extension/__init__.pyi b/tensorflow/compiler/xla/python/xla_extension/__init__.pyi
@@ -495,6 +495,8 @@ class DistributedRuntimeService:
 class DistributedRuntimeClient:
   def connect(self) -> _Status: ...
   def shutdown(self) -> _Status: ...
+  def blocking_key_value_get(self, key: str, timeout_in_ms: int) -> _Status: ...
+  def key_value_set(self, key: str, value: str) -> _Status: ...
 
 def get_distributed_runtime_service(
     address: str,