Use GCSClient instead of PythonGCSClient

build/wrapper.cc

@@ -257,62 +257,72 @@ JuliaGcsClient::JuliaGcsClient(const std::string &gcs_address) {
 }
 
 Status JuliaGcsClient::Connect() {
-    gcs_client_ = std::make_unique<gcs::PythonGcsClient>(options_);
-    return gcs_client_->Connect();
+    // https://github.com/beacon-biosignals/ray/blob/448a83caf44108fc1bc44fa7c6c358cffcfcb0d7/src/ray/gcs/gcs_client/test/usage_stats_client_test.cc#L34-L41
+    io_service_ = std::make_unique<instrumented_io_context>();
+    io_service_thread_ = std::make_unique<std::thread>([this] {
+        std::unique_ptr<boost::asio::io_service::work> work(
+            new boost::asio::io_service::work(*io_service_));
+        io_service_->run();
+    });
+    gcs_client_ = std::make_unique<gcs::GcsClient>(options_);
+    return gcs_client_->Connect(*io_service_);
+}
+
+void JuliaGcsClient::Disconnect() {
+    if (!gcs_client_){
+        return;
+    }
+    io_service_->stop();
+    io_service_thread_->join();
+    gcs_client_->Disconnect();
+    gcs_client_.reset();
 }
 
-std::string JuliaGcsClient::Get(const std::string &ns,
-                                const std::string &key,
-                                int64_t timeout_ms) {
+std::string JuliaGcsClient::Get(const std::string &ns, const std::string &key) {
     if (!gcs_client_) {
         throw std::runtime_error("GCS client not initialized; did you forget to Connect?");
     }
     std::string value;
-    Status status = gcs_client_->InternalKVGet(ns, key, timeout_ms, value);
+    Status status = gcs_client_->InternalKV().Get(ns, key, value);
     if (!status.ok()) {
         throw std::runtime_error(status.ToString());
     }
     return value;
 }
 
-int JuliaGcsClient::Put(const std::string &ns,
+bool JuliaGcsClient::Put(const std::string &ns,
                         const std::string &key,
                         const std::string &value,
-                        bool overwrite,
-                        int64_t timeout_ms) {
+                        bool overwrite) {
     if (!gcs_client_) {
         throw std::runtime_error("GCS client not initialized; did you forget to Connect?");
     }
-    int added_num;
-    Status status = gcs_client_->InternalKVPut(ns, key, value, overwrite, timeout_ms, added_num);
+    bool added_num;
+    Status status = gcs_client_->InternalKV().Put(ns, key, value, overwrite, added_num);
     if (!status.ok()) {
         throw std::runtime_error(status.ToString());
     }
     return added_num;
 }
 
-std::vector<std::string> JuliaGcsClient::Keys(const std::string &ns,
-                                              const std::string &prefix,
-                                              int64_t timeout_ms) {
+std::vector<std::string> JuliaGcsClient::Keys(const std::string &ns, const std::string &prefix) {
     if (!gcs_client_) {
         throw std::runtime_error("GCS client not initialized; did you forget to Connect?");
     }
     std::vector<std::string> results;
-    Status status = gcs_client_->InternalKVKeys(ns, prefix, timeout_ms, results);
+    Status status = gcs_client_->InternalKV().Keys(ns, prefix, results);
     if (!status.ok()) {
         throw std::runtime_error(status.ToString());
     }
     return results;
 }
 
-bool JuliaGcsClient::Exists(const std::string &ns,
-                            const std::string &key,
-                            int64_t timeout_ms) {
+bool JuliaGcsClient::Exists(const std::string &ns, const std::string &key) {
     if (!gcs_client_) {
         throw std::runtime_error("GCS client not initialized; did you forget to Connect?");
     }
     bool exists;
-    Status status = gcs_client_->InternalKVExists(ns, key, timeout_ms, exists);
+    Status status = gcs_client_->InternalKV().Exists(ns, key, exists);
     if (!status.ok()) {
         throw std::runtime_error(status.ToString());
     }
@@ -715,6 +725,7 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod)
     mod.add_type<JuliaGcsClient>("JuliaGcsClient")
         .constructor<const std::string&>()
         .method("Connect", &JuliaGcsClient::Connect)
+        .method("Disconnect", &JuliaGcsClient::Disconnect)
         .method("Put", &JuliaGcsClient::Put)
         .method("Get", &JuliaGcsClient::Get)
         .method("Keys", &JuliaGcsClient::Keys)

build/wrapper.h

@@ -23,34 +23,42 @@ using ray::core::TaskOptions;
 using ray::core::WorkerType;
 std::string ToString(ray::FunctionDescriptor function_descriptor);
 
-// a wrapper class to manage the IO service + thread that the GcsClient needs.
-// we may want to use the PythonGcsClient however, which does not do async
-// operations in separate threads as far as I can tell...in which case we would
-// not need this wrapper at all.
+// JuliaGCSClient is a wrapper class to manage the IO service + thread that the GcsClient needs.
+// https://github.com/beacon-biosignals/ray/blob/448a83caf44108fc1bc44fa7c6c358cffcfcb0d7/src/ray/gcs/gcs_client/gcs_client.h#L61-L87
+// Note that connection timeout information, etc. is parsed from the RayConfig parameters:
+// https://github.com/ray-project/ray/blob/e060dc2de2251cf883de9b43dda9c73e448e6cbd/src/ray/common/ray_config_def.h
+// these can be overwritten via environment variables before initializing the RayConfig / GCSClient, e.g.
+// ENV["RAY_gcs_server_request_timeout_seconds"] = 10
+// Timeouts that may be useful to configure:
+// - gcs_rpc_server_reconnect_timeout_s (default 5): timeout for connecting to GCSClient
+// - gcs_rpc_server_reconnect_timeout_s (default 60): timeout for reconnecting to GCSClient
+// - gcs_server_request_timeout_seconds (default 60): timeout for fetching from GCSClient
 class JuliaGcsClient {
 public:
     JuliaGcsClient(const ray::gcs::GcsClientOptions &options);
     JuliaGcsClient(const std::string &gcs_address);
 
     ray::Status Connect();
 
-    std::string Get(const std::string &ns,
-                    const std::string &key,
-                    int64_t timeout_ms);
-    int Put(const std::string &ns,
+    void Disconnect();
+
+    // Get, Put, Exists, Keys use methods belonging to an InternalKV field of the GCSClient
+    // https://github.com/beacon-biosignals/ray/blob/448a83caf44108fc1bc44fa7c6c358cffcfcb0d7/src/ray/gcs/gcs_client/accessor.h#L687
+    std::string Get(const std::string &ns, const std::string &key);
+
+    bool Put(const std::string &ns,
             const std::string &key,
-            const std::string &val,
-            bool overwrite,
-            int64_t timeout_ms);
-    std::vector<std::string> Keys(const std::string &ns,
-                                  const std::string &prefix,
-                                  int64_t timeout_ms);
-    bool Exists(const std::string &ns,
-                const std::string &key,
-                int64_t timeout_ms);
-
-    std::unique_ptr<ray::gcs::PythonGcsClient> gcs_client_;
+            const std::string &value,
+            bool overwrite);
+
+    std::vector<std::string> Keys(const std::string &ns, const std::string &prefix);
+
+    bool Exists(const std::string &ns, const std::string &key);
+
+    std::unique_ptr<ray::gcs::GcsClient> gcs_client_;
     ray::gcs::GcsClientOptions options_;
+    std::unique_ptr<instrumented_io_context> io_service_;
+    std::unique_ptr<std::thread> io_service_thread_;
 };
 
 JLCXX_MODULE define_julia_module(jlcxx::Module& mod);

src/function_manager.jl

@@ -55,19 +55,26 @@ end
 # so "jlfun" seems reasonable
 const FUNCTION_MANAGER_NAMESPACE = "jlfun"
 
-Base.@kwdef struct FunctionManager
+Base.@kwdef mutable struct FunctionManager
     gcs_client::ray_jll.JuliaGcsClient
     functions::Dict{String,Any}
+
+    function FunctionManager(gcs_client, functions)
+        fm = new(gcs_client, functions)
+        f(x) = ray_jll.Disconnect(x.gcs_client)
+        return finalizer(f, fm)
+    end
 end
 
 const FUNCTION_MANAGER = Ref{FunctionManager}()
 
 function _init_global_function_manager(gcs_address)
     @info "Connecting function manager to GCS at $gcs_address..."
     gcs_client = ray_jll.JuliaGcsClient(gcs_address)
-    ray_jll.Connect(gcs_client)
+    status = ray_jll.Connect(gcs_client)
+    ray_jll.check_status(status)
+    finalizer(ray_jll.Disconnect, gcs_client)
     FUNCTION_MANAGER[] = FunctionManager(; gcs_client, functions=Dict{String,Any}())
-
     return nothing
 end
 
@@ -82,30 +89,14 @@ function export_function!(fm::FunctionManager, f, job_id=get_job_id())
     @debug "Exporting function to function store:" fd key function_locations
     # DFK: I _think_ the string memory may be mangled if we don't `deepcopy`. Not sure but
     # it can't hurt
-    if ray_jll.Exists(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, deepcopy(key), -1)
+    if ray_jll.Exists(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, deepcopy(key))
         @debug "Function already present in GCS store:" fd key
     else
         @debug "Exporting function to GCS store:" fd key
         val = base64encode(serialize, f)
         check_oversized_function(val, fd)
-        ray_jll.Put(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key, val, true, -1)
-    end
-end
-
-function timedwait_for_function(fm::FunctionManager, fd::ray_jll.JuliaFunctionDescriptor,
-                                job_id=get_job_id(); timeout_s=10)
-    key = function_key(fd, job_id)
-    status = try
-        exists = ray_jll.Exists(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key, timeout_s)
-        exists ? :ok : :timed_out
-    catch e
-        if e isa ErrorException && contains(e.msg, "Deadline Exceeded")
-            return :timed_out
-        else
-            rethrow()
-        end
+        ray_jll.Put(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key, val, true)
     end
-    return status
 end
 
 # XXX: this will error if the function is not found in the store.
@@ -118,7 +109,7 @@ function import_function!(fm::FunctionManager, fd::ray_jll.JuliaFunctionDescript
     return get!(fm.functions, fd.function_hash) do
         key = function_key(fd, job_id)
         @debug "Function not found locally, retrieving from function store" fd key
-        val = ray_jll.Get(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key, -1)
+        val = ray_jll.Get(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key)
         try
             io = IOBuffer()
             iob64 = Base64DecodePipe(io)

src/ray_julia_jll/common.jl

@@ -79,6 +79,16 @@ function check_status(status::Status)
     return nothing
 end
 
+function Connect(body, client::JuliaGcsClient)
+    status = Connect(client)
+    check_status(status)
+    try
+        return body(client)
+    finally
+        Disconnect(client)
+    end
+end
+
 #####
 ##### Function descriptor wrangling
 #####

test/function_manager.jl

@@ -1,6 +1,7 @@
 @testset "function manager" begin
-    using Ray: FunctionManager, export_function!, import_function!, timedwait_for_function
-    using .ray_julia_jll: JuliaGcsClient, Connect, function_descriptor,
+    using Ray: FUNCTION_MANAGER_NAMESPACE, FunctionManager, function_key, export_function!,
+               import_function!
+    using .ray_julia_jll: JuliaGcsClient, Connect, Disconnect, function_descriptor,
                           JuliaFunctionDescriptor, Exists
 
     client = JuliaGcsClient("127.0.0.1:6379")
@@ -14,14 +15,16 @@
     export_function!(fm, f, jobid)
 
     fd = function_descriptor(f)
-    @test timedwait_for_function(fm, fd, jobid; timeout_s=10) == :ok
+    key = function_key(fd, jobid)
+    @test ray_jll.Exists(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, key)
     f2 = import_function!(fm, fd, jobid)
 
     @test f2.(1:10) == f.(1:10)
 
     mfd = function_descriptor(MyMod.f)
     @test_throws ErrorException import_function!(fm, mfd, jobid)
-    @test timedwait_for_function(fm, mfd, jobid; timeout_s=1) == :timed_out
+    mkey = function_key(mfd, jobid)
+    @test !(ray_jll.Exists(fm.gcs_client, FUNCTION_MANAGER_NAMESPACE, mkey))
     export_function!(fm, MyMod.f, jobid)
 
     # can import the function even when it's aliased in another module:

test/ray_julia_jll/gcs_client.jl

@@ -1,44 +1,36 @@
 @testset "GCS client" begin
     using UUIDs
     using .ray_julia_jll: JuliaGcsClient, Connect, Put, Get, Keys, Exists, Status, ok,
-                          ToString
+                          ToString, Disconnect
 
     client = JuliaGcsClient("127.0.0.1:6379")
+    @test isnothing(Disconnect(client))
 
     ns = string("TESTING-", uuid4())
 
     # throws if not connected
-    @test_throws ErrorException Put(client, ns, "computer", "mistaek", false, -1)
-    @test_throws ErrorException Get(client, ns, "computer", -1)
-    @test_throws ErrorException Keys(client, ns, "", -1)
-    @test_throws ErrorException Exists(client, ns, "computer", -1)
-
-    status = Connect(client)
-    @test ok(status)
-    @test ToString(status) == "OK"
-
-    @test Put(client, ns, "computer", "mistaek", false, -1) == 1
-    @test Get(client, ns, "computer", -1) == "mistaek"
-    @test Keys(client, ns, "", -1) == ["computer"]
-    @test Keys(client, ns, "comp", -1) == ["computer"]
-    @test Keys(client, ns, "comppp", -1) == []
-    @test Exists(client, ns, "computer", -1)
-
-    # no overwrite
-    @test Put(client, ns, "computer", "blah", false, -1) == 0
-    @test Get(client, ns, "computer", -1) == "mistaek"
-
-    # overwrite ("added" only increments on new key I think)
-    @test Put(client, ns, "computer", "blah", true, -1) == 0
-    @test Get(client, ns, "computer", -1) == "blah"
-
-    # throw on missing key
-    @test_throws ErrorException Get(client, ns, "none", -1)
-
-    # ideally we'd throw on connect but it returns OK......
-    badclient = JuliaGcsClient("127.0.0.1:6378")
-    status = Connect(badclient)
-
-    # ...but then throws when we try to do anything so at least there's that
-    @test_throws ErrorException Put(badclient, ns, "computer", "mistaek", false, -1)
+    @test_throws ErrorException Put(client, ns, "computer", "mistaek", false)
+    @test_throws ErrorException Get(client, ns, "computer")
+    @test_throws ErrorException Keys(client, ns, "")
+    @test_throws ErrorException Exists(client, ns, "computer")
+
+    Connect(client) do client
+        @test Put(client, ns, "computer", "mistaek", false) == 1
+        @test Get(client, ns, "computer") == "mistaek"
+        @test Keys(client, ns, "") == ["computer"]
+        @test Keys(client, ns, "comp") == ["computer"]
+        @test Keys(client, ns, "comppp") == []
+        @test Exists(client, ns, "computer")
+
+        # no overwrite
+        @test Put(client, ns, "computer", "blah", false) == 0
+        @test Get(client, ns, "computer") == "mistaek"
+
+        # overwrite ("added" only increments on new key I think)
+        @test Put(client, ns, "computer", "blah", true) == 0
+        @test Get(client, ns, "computer") == "blah"
+
+        # throw on missing key
+        @test_throws ErrorException Get(client, ns, "none")
+    end
 end