[core] Add the function to cleanup Redis backend.

BUILD.bazel

@@ -2723,6 +2723,7 @@ pyx_library(
         "//:global_state_accessor_lib",
         "//:ray_util",
         "//:raylet_lib",
+        "//:redis_client",
         "//:src/ray/ray_exported_symbols.lds",
         "//:src/ray/ray_version_script.lds",
         "//:stats_lib",

python/ray/_private/gcs_utils.py

@@ -577,3 +577,40 @@ def use_gcs_for_bootstrap():
     This function is included for the purposes of backwards compatibility.
     """
     return True
+
+
+def cleanup_redis_storage(
+    host: str, port: int, password: str, use_ssl: bool, storage_namespace: str
+):
+    """This function is used to cleanup the storage. Before we having
+    a good design for storage backend, it can be used to delete the old
+    data. It support redis cluster and non cluster mode.
+
+    Args:
+       host: The host address of the Redis.
+       port: The port of the Redis.
+       password: The password of the Redis.
+       use_ssl: Whether to encrypt the connection.
+       storage_namespace: The namespace of the storage to be deleted.
+    """
+
+    from ray._raylet import del_key_from_storage  # type: ignore
+
+    if not isinstance(host, str):
+        raise ValueError("Host must be a string")
+
+    if not isinstance(password, str):
+        raise ValueError("Password must be a string")
+
+    if port < 0:
+        raise ValueError(f"Invalid port: {port}")
+
+    if not isinstance(use_ssl, bool):
+        raise TypeError("use_ssl must be a boolean")
+
+    if not isinstance(storage_namespace, str):
+        raise ValueError("storage namespace must be a string")
+
+    # Right now, GCS store all data into a hash set key by storage_namespace.
+    # So we only need to delete the specific key to cleanup the cluster.
+    return del_key_from_storage(host, port, password, use_ssl, storage_namespace)

python/ray/_raylet.pyi

@@ -3,9 +3,9 @@ from typing import Any, Awaitable, TypeVar
 R = TypeVar("R")
 
 
-class ObjectRef(Awaitable[R]):
+class ObjectRef(Awaitable[R]): # type: ignore
     pass
 
 
-class ObjectID(Awaitable[R]):
+class ObjectID(Awaitable[R]): # type: ignore
     pass

python/ray/_raylet.pyx

@@ -109,7 +109,7 @@ from ray.includes.libcoreworker cimport (
 
 from ray.includes.ray_config cimport RayConfig
 from ray.includes.global_state_accessor cimport CGlobalStateAccessor
-
+from ray.includes.global_state_accessor cimport RedisDelKeySync
 from ray.includes.optional cimport (
     optional
 )
@@ -2779,3 +2779,7 @@ cdef void async_callback(shared_ptr[CRayObject] obj,
     py_callback = <object>user_callback
     py_callback(result)
     cpython.Py_DECREF(py_callback)
+
+
+def del_key_from_storage(host, port, password, use_ssl, key):
+    return RedisDelKeySync(host, port, password, use_ssl, key)

python/ray/includes/global_state_accessor.pxd

@@ -2,6 +2,7 @@ from libcpp.string cimport string as c_string
 from libcpp cimport bool as c_bool
 from libcpp.vector cimport vector as c_vector
 from libcpp.memory cimport unique_ptr
+from libc.stdint cimport int32_t as c_int32_t
 from ray.includes.unique_ids cimport (
     CActorID,
     CJobID,
@@ -43,3 +44,65 @@ cdef extern from "ray/gcs/gcs_client/global_state_accessor.h" nogil:
         CRayStatus GetNodeToConnectForDriver(
             const c_string &node_ip_address,
             c_string *node_to_connect)
+
+cdef extern from * namespace "ray::gcs" nogil:
+    """
+    #include <thread>
+    #include "ray/gcs/redis_client.h"
+    namespace ray {
+    namespace gcs {
+
+    class Cleanup {
+      public:
+        Cleanup(std::function<void()> f): f_(f) {}
+        ~Cleanup() { f_(); }
+      private:
+        std::function<void()> f_;
+    };
+
+    bool RedisDelKeySync(const std::string& host,
+                         int32_t port,
+                         const std::string& password,
+                         bool use_ssl,
+                         const std::string& key) {
+      RedisClientOptions options(host, port, password, false, use_ssl);
+      auto cli = std::make_unique<RedisClient>(options);
+
+      instrumented_io_context io_service;
+
+      auto thread = std::make_unique<std::thread>([&]() {
+        boost::asio::io_service::work work(io_service);
+        io_service.run();
+      });
+
+      Cleanup _([&](){
+        io_service.stop();
+        thread->join();
+      });
+
+      auto status = cli->Connect(io_service);
+      if(!status.ok()) {
+        RAY_LOG(ERROR) << "Failed to connect to redis: " << status.ToString();
+        return false;
+      }
+
+      auto context = cli->GetShardContext(key);
+      auto cmd = std::vector<std::string>{"DEL", key};
+      auto reply = context->RunArgvSync(cmd);
+      if(reply->ReadAsInteger() == 1) {
+        RAY_LOG(INFO) << "Successfully deleted " << key;
+        return true;
+      } else {
+        RAY_LOG(ERROR) << "Failed to delete " << key;
+        return false;
+      }
+    }
+
+    }
+    }
+    """
+    c_bool RedisDelKeySync(const c_string& host,
+                           c_int32_t port,
+                           const c_string& password,
+                           c_bool use_ssl,
+                           const c_string& key)

python/ray/includes/global_state_accessor.pxi

@@ -12,6 +12,7 @@ from ray.includes.unique_ids cimport (
 
 from ray.includes.global_state_accessor cimport (
     CGlobalStateAccessor,
+    RedisDelKeySync,
 )
 
 from libcpp.string cimport string as c_string

python/ray/tests/test_gcs_utils.py

@@ -7,6 +7,7 @@
 import grpc
 import pytest
 import ray
+import redis
 from ray._private.gcs_utils import GcsClient
 import ray._private.gcs_utils as gcs_utils
 from ray._private.test_utils import (
@@ -211,6 +212,44 @@ async def check(expect_liveness):
     )
 
 
+@pytest.fixture(params=[True, False])
+def redis_replicas(request, monkeypatch):
+    if request.param:
+        monkeypatch.setenv("TEST_EXTERNAL_REDIS_REPLICAS", "3")
+    yield
+
+
+@pytest.mark.skipif(
+    not enable_external_redis(), reason="Only valid when start with an external redis"
+)
+def test_redis_cleanup(redis_replicas, shutdown_only):
+    addr = ray.init(
+        namespace="a", _system_config={"external_storage_namespace": "c1"}
+    ).address_info["address"]
+    gcs_client = GcsClient(address=addr)
+    gcs_client.internal_kv_put(b"ABC", b"DEF", True, None)
+
+    ray.shutdown()
+    addr = ray.init(
+        namespace="a", _system_config={"external_storage_namespace": "c2"}
+    ).address_info["address"]
+    gcs_client = GcsClient(address=addr)
+    gcs_client.internal_kv_put(b"ABC", b"XYZ", True, None)
+    ray.shutdown()
+    redis_addr = os.environ["RAY_REDIS_ADDRESS"]
+    host, port = redis_addr.split(":")
+    if os.environ.get("TEST_EXTERNAL_REDIS_REPLICAS", "1") != "1":
+        cli = redis.RedisCluster(host, int(port))
+    else:
+        cli = redis.Redis(host, int(port))
+
+    assert set(cli.keys()) == {b"c1", b"c2"}
+    gcs_utils.cleanup_redis_storage(host, int(port), "", False, "c1")
+    assert set(cli.keys()) == {b"c2"}
+    gcs_utils.cleanup_redis_storage(host, int(port), "", False, "c2")
+    assert len(cli.keys()) == 0
+
+
 if __name__ == "__main__":
     import sys
 

src/ray/gcs/redis_client.h

@@ -20,14 +20,13 @@
 #include "ray/common/asio/instrumented_io_context.h"
 #include "ray/common/status.h"
 #include "ray/gcs/asio.h"
+#include "ray/gcs/redis_context.h"
 #include "ray/util/logging.h"
 
 namespace ray {
 
 namespace gcs {
 
-class RedisContext;
-
 class RedisClientOptions {
  public:
   RedisClientOptions(const std::string &ip,