[core] Expose Redis getter to Python and use to retrieve session name

BUILD.bazel

@@ -2531,6 +2531,7 @@ pyx_library(
  deps = [
  "//:core_worker_lib",
  "//:global_state_accessor_lib",
+ "//:gcs_server_lib",
  "//:raylet_lib",
  "//:redis_client",
  "//:src/ray/ray_exported_symbols.lds",

python/ray/_private/node.py

@@ -24,8 +24,9 @@
 import ray._private.services
 import ray._private.utils
 from ray._private import storage
-from ray._raylet import GcsClient
+from ray._raylet import GcsClient, get_session_key_from_storage
 from ray._private.resource_spec import ResourceSpec
+from ray._private.services import serialize_config
 from ray._private.utils import open_log, try_to_create_directory, try_to_symlink
 
 # Logger for this module. It should be configured at the entry point
@@ -177,9 +178,15 @@ def __init__(
 
  # Register the temp dir.
  if head:
- # date including microsecond
- date_str = datetime.datetime.today().strftime("%Y-%m-%d_%H-%M-%S_%f")
- self._session_name = f"session_{date_str}_{os.getpid()}"
+ # We expect this the first time we initialize a cluster, but not during
+ # subsequent restarts of the head node.
+ maybe_key = self.check_persisted_session_name()
+ if maybe_key is None:
+ # date including microsecond
+ date_str = datetime.datetime.today().strftime("%Y-%m-%d_%H-%M-%S_%f")
+ self._session_name = f"session_{date_str}_{os.getpid()}"
+ else:
+ self._session_name = ray._private.utils.decode(maybe_key)
  else:
  if ray_params.session_name is None:
  assert not self._default_worker
@@ -317,6 +324,41 @@ def __init__(
  self.validate_ip_port(self.gcs_address)
  self._record_stats()
 
+ def check_persisted_session_name(self):
+ if self._ray_params.external_addresses is None:
+ return None
+ self._redis_address = self._ray_params.external_addresses[0]
+ # Address is ip:port or redis:https://ip:port
+ parts = self._redis_address.split(":https://", 1)
+ enable_redis_ssl = False
+ if len(parts) == 1:
+ redis_ip_address, redis_port = parts[0].rsplit(":", 1)
+ else:
+ # rediss for SSL
+ if len(parts) != 2 or parts[0] not in ("redis", "rediss"):
+ raise ValueError(
+ f"Invalid redis address {self._redis_address}."
+ "Expected format is ip:port or redis:https://ip:port, "
+ "or rediss:https://ip:port for SSL."
+ )
+ redis_ip_address, redis_port = parts[1].rsplit(":", 1)
+ if parts[0] == "rediss":
+ enable_redis_ssl = True
+ if int(redis_port) < 0:
+ raise ValueError(
+ f"Invalid Redis port provided: {redis_port}."
+ "The port must be a non-negative integer."
+ )
+
+ return get_session_key_from_storage(
+ redis_ip_address,
+ int(redis_port),
+ self._ray_params.redis_password,
+ enable_redis_ssl,
+ serialize_config(self._config),
+ b"session_name",
+ )
+
  @staticmethod
  def validate_ip_port(ip_port):
  """Validates the address is in the ip:port format"""
@@ -1175,12 +1217,22 @@ def _write_cluster_info_to_kv(self):
 
  ray_usage_lib.put_cluster_metadata(self.get_gcs_client())
  # Make sure GCS is up.
- self.get_gcs_client().internal_kv_put(
+ added = self.get_gcs_client().internal_kv_put(
  b"session_name",
  self._session_name.encode(),
- True,
+ False,
  ray_constants.KV_NAMESPACE_SESSION,
  )
+ if not added:
+ curr_val = self.get_gcs_client().internal_kv_get(
+ b"session_name", ray_constants.KV_NAMESPACE_SESSION
+ )
+ assert curr_val != self._session_name, (
+ f"Session name {self._session_name} does not match "
+ f"persisted value {curr_val}. Perhaps there was an "
+ f"error connecting to Redis."
+ )
+
  self.get_gcs_client().internal_kv_put(
  b"session_dir",
  self._session_dir.encode(),
@@ -1215,13 +1267,9 @@ def start_head_processes(self):
  logger.debug(
  f"Process STDOUT and STDERR is being " f"redirected to {self._logs_dir}."
  )
- assert self._redis_address is None
  assert self._gcs_address is None
  assert self._gcs_client is None
 
- if self._ray_params.external_addresses is not None:
- self._redis_address = self._ray_params.external_addresses[0]
-
  self.start_gcs_server()
  assert self.get_gcs_client() is not None
  self._write_cluster_info_to_kv()

python/ray/_raylet.pyx

@@ -156,7 +156,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.global_state_accessor cimport RedisDelKeySync, RedisGetKeySync
 from ray.includes.optional cimport (
  optional, nullopt
 )
@@ -4574,3 +4574,25 @@ cdef void async_callback(shared_ptr[CRayObject] obj,
 
 def del_key_from_storage(host, port, password, use_ssl, key):
  return RedisDelKeySync(host, port, password, use_ssl, key)
+
+
+def get_session_key_from_storage(host, port, password, use_ssl, config, key):
+ """
+ Get the session key from the storage.
+ Intended to be used for session_name only.
+ Args:
+ host: The address of the owner (caller) of the
+ generator task.
+ port: The task ID of the generator task.
+ password: The redis password.
+ use_ssl: Whether to use SSL.
+ config: The Ray config. Used to get storage namespace.
+ key: The key to retrieve.
+ """
+ cdef:
+ c_string data
+ result = RedisGetKeySync(host, port, password, use_ssl, config, key, &data)
+ if result:
+ return data
+ else:
+ return None

python/ray/includes/global_state_accessor.pxd

@@ -46,6 +46,73 @@ cdef extern from "ray/gcs/gcs_client/global_state_accessor.h" nogil:
  const c_string &node_ip_address,
  c_string *node_to_connect)
 
+cdef extern from * namespace "ray::gcs" nogil:
+ """
+ #include <thread>
+ #include "ray/gcs/gcs_server/store_client_kv.h"
+ namespace ray {
+ namespace gcs {
+
+ bool RedisGetKeySync(const std::string& host,
+ int32_t port,
+ const std::string& password,
+ bool use_ssl,
+ const std::string& config,
+ const std::string& key,
+ std::string* data) {
+ InitShutdownRAII ray_log_shutdown_raii(ray::RayLog::StartRayLog,
+ ray::RayLog::ShutDownRayLog,
+ "ray_init",
+ ray::RayLogLevel::WARNING,
+ "" /* log_dir */);
+
+ RedisClientOptions options(host, port, password, false, use_ssl);
+
+ std::string config_list;
+ RAY_CHECK(absl::Base64Unescape(config, &config_list));
+ RayConfig::instance().initialize(config_list);
+
+ instrumented_io_context io_service;
+
+ auto redis_client = std::make_shared<RedisClient>(options);
+ auto status = redis_client->Connect(io_service);
+
+ if(!status.ok()) {
+ RAY_LOG(ERROR) << "Failed to connect to redis: " << status.ToString();
+ return false;
+ }
+
+ auto cli = std::make_unique<StoreClientInternalKV>(
+ std::make_unique<RedisStoreClient>(std::move(redis_client)));
+
+ bool ret_val = false;
+ cli->Get("session", key, [&](std::optional<std::string> result) {
+ if (result.has_value()) {
+ *data = result.value();
+ ret_val = true;
+ } else {
+ RAY_LOG(INFO) << "Failed to retrieve the key " << key
+ << " from persistent storage.";
+ ret_val = false;
+ }
+ });
+ io_service.run_for(std::chrono::milliseconds(1000));
+
+ return ret_val;
+ }
+
+ }
+ }
+ """
+ c_bool RedisGetKeySync(const c_string& host,
+ c_int32_t port,
+ const c_string& password,
+ c_bool use_ssl,
+ const c_string& config,
+ const c_string& key,
+ c_string* data)
+
+
 cdef extern from * namespace "ray::gcs" nogil:
  """
  #include <thread>

python/ray/tests/BUILD

@@ -225,6 +225,7 @@ py_test_module_list(
 py_test_module_list(
  files = [
  "test_gcs_ha_e2e.py",
+ "test_gcs_ha_e2e_2.py",
  "test_memory_pressure.py",
  "test_node_labels.py",
  ],

python/ray/tests/test_advanced_9.py

@@ -363,9 +363,8 @@ def test_redis_not_available(monkeypatch, call_ray_stop_only):
  capture_output=True,
  )
  assert "Could not establish connection to Redis" in p.stderr.decode()
- assert "Please check" in p.stderr.decode()
- assert "gcs_server.out for details" in p.stderr.decode()
- assert "RuntimeError: Failed to start GCS" in p.stderr.decode()
+ assert "Please check " in p.stderr.decode()
+ assert "redis storage is alive or not." in p.stderr.decode()
 
 
 @pytest.mark.skipif(not enable_external_redis(), reason="Only valid in redis env")

python/ray/tests/test_gcs_fault_tolerance.py

@@ -868,6 +868,36 @@ def check_raylet_healthy():
  sleep(1)
 
 
+def test_session_name(ray_start_cluster):
+ # Kill GCS and check that raylets kill themselves when not backed by Redis,
+ # and stay alive when backed by Redis.
+ # Raylets should kill themselves due to cluster ID mismatch in the
+ # non-persisted case.
+ cluster = ray_start_cluster
+ cluster.add_node()
+ cluster.wait_for_nodes()
+
+ head_node = cluster.head_node
+ session_dir = head_node.get_session_dir_path()
+
+ gcs_server_process = head_node.all_processes["gcs_server"][0].process
+ gcs_server_pid = gcs_server_process.pid
+ cluster.remove_node(head_node, allow_graceful=False)
+ # Wait to prevent the gcs server process becoming zombie.
+ gcs_server_process.wait()
+ wait_for_pid_to_exit(gcs_server_pid, 1000)
+
+ # Add head node back
+ cluster.add_node()
+ head_node = cluster.head_node
+ new_session_dir = head_node.get_session_dir_path()
+
+ if not enable_external_redis():
+ assert session_dir != new_session_dir
+ else:
+ assert session_dir == new_session_dir
+
+
 @pytest.mark.parametrize(
  "ray_start_regular_with_external_redis",
  [

python/ray/tests/test_gcs_ha_e2e.py

@@ -1,6 +1,8 @@
-import pytest
 import sys
 from time import sleep
+
+import pytest
+
 from ray._private.test_utils import wait_for_condition
 from ray.tests.conftest_docker import * # noqa
 

python/ray/tests/test_gcs_ha_e2e_2.py

@@ -0,0 +1,55 @@
+import pytest
+import sys
+from time import sleep
+from ray._private.test_utils import wait_for_condition
+from ray.tests.conftest_docker import * # noqa
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="Only works on linux.")
+def test_ray_session_name_preserved(docker_cluster):
+ get_nodes_script = """
+import ray
+ray.init("auto")
+print(ray._private.worker._global_node.session_name)
+"""
+ head, worker = docker_cluster
+
+ def get_session_name(to_head=True):
+ if to_head:
+ output = head.exec_run(cmd=f"python -c '{get_nodes_script}'")
+ else:
+ output = worker.exec_run(cmd=f"python -c '{get_nodes_script}'")
+ session_name = output.output.decode().strip().split("\n")[-1]
+ print("Output: ", output.output.decode().strip().split("\n"))
+ assert output.exit_code == 0
+ return session_name
+
+ # Make sure two nodes are alive
+ wait_for_condition(get_session_name, to_head=True)
+ session_name_head = get_session_name(to_head=True)
+ wait_for_condition(get_session_name, to_head=False)
+ session_name_worker = get_session_name(to_head=False)
+ assert session_name_head == session_name_worker
+ print("head killed")
+ head.kill()
+
+ sleep(2)
+
+ head.restart()
+
+ wait_for_condition(get_session_name, to_head=True)
+ session_name_head_after_restart = get_session_name(to_head=True)
+ wait_for_condition(get_session_name, to_head=False)
+ session_name_worker_after_restart = get_session_name(to_head=False)
+ assert session_name_worker_after_restart == session_name_head_after_restart
+ assert session_name_head == session_name_head_after_restart
+ assert session_name_worker_after_restart == session_name_worker
+
+
+if __name__ == "__main__":
+ import os
+
+ if os.environ.get("PARALLEL_CI"):
+ sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+ else:
+ sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/test_output.py

@@ -15,6 +15,16 @@
 )
 
 
+def test_ray_init_no_redis_logs():
+ script = """
+import ray
+
+ray.init()
+ """
+ out_str = run_string_as_driver(script)
+ assert "redis" not in out_str, out_str
+
+
 def test_dedup_logs():
  script = """
 import ray