chore: import upstream snapshot with attribution

python/ray/_private/runtime_env/setup_hook.py

@@ -0,0 +1,261 @@
+import base64
+import logging
+import os
+import traceback
+from typing import Any, Callable, Dict, Optional, Union
+
+import ray
+import ray._private.ray_constants as ray_constants
+import ray.cloudpickle as pickle
+from ray._common.utils import load_class
+from ray._private.function_manager import build_setup_hook_export_entry
+from ray.runtime_env import RuntimeEnv
+
+logger = logging.getLogger(__name__)
+
+RUNTIME_ENV_FUNC_IDENTIFIER = "ray_runtime_env_func::"
+
+
+def get_import_export_timeout():
+    return int(
+        os.environ.get(
+            ray_constants.RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR, "60"
+        )
+    )
+
+
+def decode_function_key(key: bytes) -> str:
+    # b64encode only includes A-Z, a-z, 0-9, + and / characters
+    return RUNTIME_ENV_FUNC_IDENTIFIER + base64.b64encode(key).decode()
+
+
+def _encode_function_key(key: str) -> bytes:
+    assert key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER)
+    return base64.b64decode(key[len(RUNTIME_ENV_FUNC_IDENTIFIER) :])
+
+
+def _raise_setup_hook_conflict(existing_hook_value: str, setup_hook_desc: str) -> None:
+    raise RuntimeError(
+        "Conflicting worker_process_setup_hook: the setup hook env "
+        f"var is already set to '{existing_hook_value}', but "
+        f"runtime_env specifies {setup_hook_desc}."
+    )
+
+
+def export_setup_func_callable(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func, Callable)
+    try:
+        key = worker.function_actor_manager.export_setup_func(
+            setup_func, timeout=get_import_export_timeout()
+        )
+    except Exception as e:
+        raise ray.exceptions.RuntimeEnvSetupError(
+            "Failed to export the setup function."
+        ) from e
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = decode_function_key(key)
+    runtime_env["env_vars"] = env_vars
+    # Note: This field is no-op. We don't have a plugin for the setup hook
+    # because we can implement it simply using an env var.
+    # This field is just for the observability purpose, so we store
+    # the name of the method.
+    runtime_env["worker_process_setup_hook"] = setup_func.__name__
+    return runtime_env
+
+
+def export_setup_func_module(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func_module: str,
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func_module, str)
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = setup_func_module
+    runtime_env["env_vars"] = env_vars
+    return runtime_env
+
+
+def _check_setup_hook_consistency(
+    existing_hook_value: str,
+    setup_func: Union[Callable, str],
+    worker: "ray.Worker",
+) -> None:
+    """Validate that an already-set hook env var is consistent with setup_func.
+
+    When the env var is already populated (e.g. inherited from a job supervisor),
+    we compare it against the `worker_process_setup_hook` field in the runtime_env
+    to detect silent mismatches.
+
+    Args:
+        existing_hook_value: The value of the existing hook env var.
+        setup_func: The setup function or module path.
+        worker: The worker instance.
+
+    Raises:
+        RuntimeError: If a conflict between the existing env var and setup_func is detected.
+    """
+    if isinstance(setup_func, Callable):
+        try:
+            _encode_function_key(existing_hook_value)
+        except Exception:
+            _raise_setup_hook_conflict(
+                existing_hook_value, f"callable '{setup_func.__name__}'"
+            )
+        _check_callable_hooks_match(existing_hook_value, setup_func, worker)
+    elif isinstance(setup_func, str):
+        try:
+            _encode_function_key(existing_hook_value)
+            existing_is_callable_ref = True
+        except Exception:
+            existing_is_callable_ref = False
+        if existing_is_callable_ref or existing_hook_value != setup_func:
+            _raise_setup_hook_conflict(existing_hook_value, f"'{setup_func}'")
+
+
+def _check_callable_hooks_match(
+    existing_hook_value: str,
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> None:
+    """Verify a callable produces the same GCS key as the existing env var."""
+    _, _, expected_key = build_setup_hook_export_entry(
+        setup_func, worker.current_job_id.binary()
+    )
+    expected_env_value = decode_function_key(expected_key)
+
+    if existing_hook_value != expected_env_value:
+        _raise_setup_hook_conflict(
+            existing_hook_value, f"callable '{setup_func.__name__}'"
+        )
+
+
+def upload_worker_process_setup_hook_if_needed(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    """Uploads the worker_process_setup_hook to GCS with a key.
+
+    runtime_env["worker_process_setup_hook"] is converted to a decoded key
+    that can load the worker setup hook function from GCS.
+    i.e., you can use internalKV.Get(runtime_env["worker_process_setup_hook])
+    to access the worker setup hook from GCS.
+
+    Args:
+        runtime_env: The runtime_env. The value will be modified
+            when returned.
+        worker: ray.worker instance.
+
+    Returns:
+        The modified runtime_env with the setup hook processed into an env var.
+    """
+    setup_func = runtime_env.get("worker_process_setup_hook")
+
+    if setup_func is None:
+        return runtime_env
+
+    env_vars = runtime_env.get("env_vars", {})
+    existing_hook = env_vars.get(ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR)
+
+    if existing_hook is not None:
+        # A setup hook is already populated (e.g. inherited from job supervisor).
+        # Validate that it is consistent with the current worker_process_setup_hook.
+        _check_setup_hook_consistency(existing_hook, setup_func, worker)
+        return runtime_env
+
+    if isinstance(setup_func, Callable):
+        return export_setup_func_callable(runtime_env, setup_func, worker)
+    elif isinstance(setup_func, str):
+        return export_setup_func_module(runtime_env, setup_func)
+    else:
+        raise TypeError(
+            "worker_process_setup_hook must be a function, " f"got {type(setup_func)}."
+        )
+
+
+def load_and_execute_setup_hook(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    """Load the setup hook from a given key and execute.
+
+    Args:
+        worker_process_setup_hook_key: The key to import the setup hook
+            from GCS.
+    Returns:
+        An error message if it fails. None if it succeeds.
+    """
+    assert worker_process_setup_hook_key is not None
+    if not worker_process_setup_hook_key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER):
+        return load_and_execute_setup_hook_module(worker_process_setup_hook_key)
+    else:
+        return load_and_execute_setup_hook_func(worker_process_setup_hook_key)
+
+
+def load_and_execute_setup_hook_module(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    try:
+        setup_func = load_class(worker_process_setup_hook_key)
+        setup_func()
+        return None
+    except Exception:
+        error_message = (
+            "Failed to execute the setup hook method, "
+            f"{worker_process_setup_hook_key} "
+            "from ``ray.init(runtime_env="
+            f"{{'worker_process_setup_hook': {worker_process_setup_hook_key}}})``. "
+            "Please make sure the given module exists and is available "
+            "from ray workers. For more details, see the error trace below.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+
+def load_and_execute_setup_hook_func(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    worker = ray._private.worker.global_worker
+    assert worker.connected
+    func_manager = worker.function_actor_manager
+    try:
+        worker_setup_func_info = func_manager.fetch_registered_method(
+            _encode_function_key(worker_process_setup_hook_key),
+            timeout=get_import_export_timeout(),
+        )
+    except Exception:
+        error_message = (
+            "Failed to import setup hook within "
+            f"{get_import_export_timeout()} seconds.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func = pickle.loads(worker_setup_func_info.function)
+    except Exception:
+        error_message = (
+            "Failed to deserialize the setup hook method.\n" f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func()
+    except Exception:
+        error_message = (
+            f"Failed to execute the setup hook method. Function name:"
+            f"{worker_setup_func_info.function_name}\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    return None