IFRT proxy: Add some debug information.

PiperOrigin-RevId: 691111539
openxla · Oct 30, 2024 · 8baa106 · 8baa106
1 parent 7efe551
commit 8baa106
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/xla/python/ifrt_proxy/server/BUILD b/xla/python/ifrt_proxy/server/BUILD
@@ -264,6 +264,7 @@ cc_library(
     srcs = ["version.cc"],
     hdrs = ["version.h"],
     deps = [
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

diff --git a/xla/python/ifrt_proxy/server/grpc_service_impl.cc b/xla/python/ifrt_proxy/server/grpc_service_impl.cc
@@ -75,7 +75,8 @@ ::grpc::Status GrpcServiceImpl::IfrtSession(
   const uint64_t session_id =
       next_session_id_.fetch_add(1, std::memory_order_relaxed);
 
-  VLOG(0) << "Starting a new IFRT session with session_id=" << session_id;
+  VLOG(0) << "Starting a new IFRT session with session_id=" << session_id
+          << ", version=" << metadata.version();
 
   // Create a host buffer store for the session.
   auto host_buffer_store =

diff --git a/xla/python/ifrt_proxy/server/version.cc b/xla/python/ifrt_proxy/server/version.cc
@@ -18,6 +18,7 @@
 
 #include <algorithm>
 
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -32,6 +33,13 @@ absl::StatusOr<int> ChooseVersion(int client_min_version,
                                   int server_max_version) {
   const int version = std::min(server_max_version, client_max_version);
 
+  LOG(INFO) << "IFRT proxy: ChooseVersion(client_min_version="
+            << client_min_version
+            << ", client_max_version=" << client_max_version
+            << ", server_min_version=" << server_min_version
+            << ", server_max_version=" << server_max_version << ") is "
+            << version;
+
   if (version < server_min_version || version < client_min_version) {
     return absl::InvalidArgumentError(absl::StrCat(
         "IFRT Proxy client and server failed to agree on the "