[Coordination Service] Fix pjrt_c_api_gpu_test after introducing TryGet

KV try_get functions should be linked in pjrt_client creation PiperOrigin-RevId: 714149039
openxla · Jan 10, 2025 · ba29286 · ba29286
1 parent 1e81016
commit ba29286
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/xla/pjrt/c/pjrt_c_api_gpu_test.cc b/xla/pjrt/c/pjrt_c_api_gpu_test.cc
@@ -351,6 +351,8 @@ absl::StatusOr<PJRT_Client_Create_Args> BuildCreateArg(
   args.kv_get_user_arg = &kv_callback_data->kv_get_c_func;
   args.kv_put_callback = kv_callback_data->c_kv_put;
   args.kv_put_user_arg = &kv_callback_data->kv_put_c_func;
+  args.kv_try_get_user_arg = &kv_callback_data->kv_try_get_c_func;
+  args.kv_try_get_callback = kv_callback_data->c_kv_try_get;
   args.client = nullptr;
   return args;
 }