wehub-resource-sync
262 lines
9.1 KiB
Python
262 lines
9.1 KiB
Python
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
|