wehub-resource-sync
751 lines
31 KiB
Python
751 lines
31 KiB
Python
import dis
|
|
import hashlib
|
|
import importlib
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from collections import defaultdict, namedtuple
|
|
from typing import Callable, Optional
|
|
|
|
import ray
|
|
import ray._private.profiling as profiling
|
|
from ray import cloudpickle as pickle
|
|
from ray._common.serialization import pickle_dumps
|
|
from ray._private import ray_constants
|
|
from ray._private.inspect_util import (
|
|
is_class_method,
|
|
is_function_or_method,
|
|
is_static_method,
|
|
)
|
|
from ray._private.ray_constants import KV_NAMESPACE_FUNCTION_TABLE
|
|
from ray._private.utils import (
|
|
check_oversized_function,
|
|
ensure_str,
|
|
format_error_message,
|
|
)
|
|
from ray._raylet import (
|
|
WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS,
|
|
JobID,
|
|
PythonFunctionDescriptor,
|
|
)
|
|
from ray.remote_function import RemoteFunction
|
|
from ray.util.tracing.tracing_helper import _inject_tracing_into_class
|
|
|
|
FunctionExecutionInfo = namedtuple(
|
|
"FunctionExecutionInfo", ["function", "function_name", "max_calls"]
|
|
)
|
|
ImportedFunctionInfo = namedtuple(
|
|
"ImportedFunctionInfo",
|
|
["job_id", "function_id", "function_name", "function", "module", "max_calls"],
|
|
)
|
|
|
|
"""FunctionExecutionInfo: A named tuple storing remote function information."""
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def make_function_table_key(key_type: bytes, job_id: JobID, key: Optional[bytes]):
|
|
if key is None:
|
|
return b":".join([key_type, job_id.hex().encode()])
|
|
else:
|
|
return b":".join([key_type, job_id.hex().encode(), key])
|
|
|
|
|
|
def build_setup_hook_export_entry(
|
|
setup_func: Callable, job_id: JobID
|
|
) -> tuple[bytes, bytes, bytes]:
|
|
"""Compute the exported payload and GCS key for a setup hook callable.
|
|
|
|
Args:
|
|
setup_func: The setup hook function to export.
|
|
job_id: The job ID to export the setup hook for.
|
|
|
|
Returns:
|
|
A tuple of (pickled_function, function_id, key).
|
|
"""
|
|
pickled_function = pickle_dumps(
|
|
setup_func,
|
|
"Cannot serialize the worker_process_setup_hook " f"{setup_func.__name__}",
|
|
)
|
|
function_to_run_id = hashlib.shake_128(pickled_function).digest(
|
|
ray_constants.ID_SIZE
|
|
)
|
|
key = make_function_table_key(
|
|
# This value should match with gcs_function_manager.h.
|
|
# Otherwise, it won't be GC'ed.
|
|
WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS.encode(),
|
|
# b"FunctionsToRun",
|
|
job_id,
|
|
function_to_run_id,
|
|
)
|
|
return pickled_function, function_to_run_id, key
|
|
|
|
|
|
class FunctionActorManager:
|
|
"""A class used to export/load remote functions and actors.
|
|
Attributes:
|
|
_worker: The associated worker that this manager related.
|
|
_functions_to_export: The remote functions to export when
|
|
the worker gets connected.
|
|
_actors_to_export: The actors to export when the worker gets
|
|
connected.
|
|
_function_execution_info: The function_id
|
|
and execution_info.
|
|
_num_task_executions: The function
|
|
execution times.
|
|
imported_actor_classes: The set of actor classes keys (format:
|
|
ActorClass:function_id) that are already in GCS.
|
|
"""
|
|
|
|
def __init__(self, worker: "ray._private.worker.Worker"):
|
|
"""Initialize FunctionActorManager.
|
|
|
|
Args:
|
|
worker: The worker this manager belongs to.
|
|
"""
|
|
self._worker = worker
|
|
self._functions_to_export = []
|
|
self._actors_to_export = []
|
|
# This field is a dictionary that maps function IDs
|
|
# to a FunctionExecutionInfo object. This should only be used on
|
|
# workers that execute remote functions.
|
|
self._function_execution_info = defaultdict(lambda: {})
|
|
self._num_task_executions = defaultdict(lambda: {})
|
|
# A set of all of the actor class keys that have been imported by the
|
|
# import thread. It is safe to convert this worker into an actor of
|
|
# these types.
|
|
self.imported_actor_classes = set()
|
|
self._loaded_actor_classes = {}
|
|
# Deserialize an ActorHandle will call load_actor_class(). If a
|
|
# function closure captured an ActorHandle, the deserialization of the
|
|
# function will be:
|
|
# -> fetch_and_register_remote_function (acquire lock)
|
|
# -> _load_actor_class_from_gcs (acquire lock, too)
|
|
# So, the lock should be a reentrant lock.
|
|
self.lock = threading.RLock()
|
|
|
|
self.execution_infos = {}
|
|
# This is the counter to keep track of how many keys have already
|
|
# been exported so that we can find next key quicker.
|
|
self._num_exported = 0
|
|
# This is to protect self._num_exported when doing exporting
|
|
self._export_lock = threading.Lock()
|
|
|
|
def increase_task_counter(self, function_descriptor):
|
|
function_id = function_descriptor.function_id
|
|
self._num_task_executions[function_id] += 1
|
|
|
|
def get_task_counter(self, function_descriptor):
|
|
function_id = function_descriptor.function_id
|
|
return self._num_task_executions[function_id]
|
|
|
|
def compute_collision_identifier(self, function_or_class: Callable) -> bytes:
|
|
"""The identifier is used to detect excessive duplicate exports.
|
|
The identifier is used to determine when the same function or class is
|
|
exported many times. This can yield false positives.
|
|
Args:
|
|
function_or_class: The function or class to compute an identifier
|
|
for.
|
|
Returns:
|
|
The identifier. Note that different functions or classes can give
|
|
rise to same identifier. However, the same function should
|
|
hopefully always give rise to the same identifier. TODO(rkn):
|
|
verify if this is actually the case. Note that if the
|
|
identifier is incorrect in any way, then we may give warnings
|
|
unnecessarily or fail to give warnings, but the application's
|
|
behavior won't change.
|
|
"""
|
|
import io
|
|
|
|
string_file = io.StringIO()
|
|
dis.dis(function_or_class, file=string_file, depth=2)
|
|
collision_identifier = function_or_class.__name__ + ":" + string_file.getvalue()
|
|
|
|
# Return a hash of the identifier in case it is too large.
|
|
return hashlib.sha256(collision_identifier.encode("utf-8")).digest()
|
|
|
|
def load_function_or_class_from_local(self, module_name, function_or_class_name):
|
|
"""Try to load a function or class in the module from local."""
|
|
module = importlib.import_module(module_name)
|
|
parts = [part for part in function_or_class_name.split(".") if part]
|
|
object = module
|
|
try:
|
|
for part in parts:
|
|
object = getattr(object, part)
|
|
return object
|
|
except Exception:
|
|
return None
|
|
|
|
def export_setup_func(
|
|
self, setup_func: Callable, timeout: Optional[int] = None
|
|
) -> bytes:
|
|
"""Export the setup hook function and return the key."""
|
|
pickled_function, function_to_run_id, key = build_setup_hook_export_entry(
|
|
setup_func, self._worker.current_job_id.binary()
|
|
)
|
|
|
|
check_oversized_function(
|
|
pickled_function, setup_func.__name__, "function", self._worker
|
|
)
|
|
|
|
try:
|
|
self._worker.gcs_client.internal_kv_put(
|
|
key,
|
|
pickle.dumps(
|
|
{
|
|
"job_id": self._worker.current_job_id.binary(),
|
|
"function_id": function_to_run_id,
|
|
"function": pickled_function,
|
|
}
|
|
),
|
|
# overwrite
|
|
True,
|
|
ray_constants.KV_NAMESPACE_FUNCTION_TABLE,
|
|
timeout=timeout,
|
|
)
|
|
except Exception as e:
|
|
logger.exception(
|
|
"Failed to export the setup hook " f"{setup_func.__name__}."
|
|
)
|
|
raise e
|
|
|
|
return key
|
|
|
|
def export(self, remote_function: RemoteFunction) -> None:
|
|
"""Pickle a remote function and export it to redis.
|
|
|
|
Args:
|
|
remote_function: the RemoteFunction object.
|
|
"""
|
|
if self._worker.load_code_from_local:
|
|
function_descriptor = remote_function._function_descriptor
|
|
module_name, function_name = (
|
|
function_descriptor.module_name,
|
|
function_descriptor.function_name,
|
|
)
|
|
# If the function is dynamic, we still export it to GCS
|
|
# even if load_code_from_local is set True.
|
|
if (
|
|
self.load_function_or_class_from_local(module_name, function_name)
|
|
is not None
|
|
):
|
|
return
|
|
function = remote_function._function
|
|
pickled_function = remote_function._pickled_function
|
|
|
|
check_oversized_function(
|
|
pickled_function,
|
|
remote_function._function_name,
|
|
"remote function",
|
|
self._worker,
|
|
)
|
|
key = make_function_table_key(
|
|
b"RemoteFunction",
|
|
self._worker.current_job_id,
|
|
remote_function._function_descriptor.function_id.binary(),
|
|
)
|
|
if self._worker.gcs_client.internal_kv_exists(key, KV_NAMESPACE_FUNCTION_TABLE):
|
|
return
|
|
val = pickle.dumps(
|
|
{
|
|
"job_id": self._worker.current_job_id.binary(),
|
|
"function_id": remote_function._function_descriptor.function_id.binary(), # noqa: E501
|
|
"function_name": remote_function._function_name,
|
|
"module": function.__module__,
|
|
"function": pickled_function,
|
|
"collision_identifier": self.compute_collision_identifier(function),
|
|
"max_calls": remote_function._max_calls,
|
|
}
|
|
)
|
|
self._worker.gcs_client.internal_kv_put(
|
|
key, val, True, KV_NAMESPACE_FUNCTION_TABLE
|
|
)
|
|
|
|
def fetch_registered_method(
|
|
self, key: str, timeout: Optional[int] = None
|
|
) -> Optional[ImportedFunctionInfo]:
|
|
vals = self._worker.gcs_client.internal_kv_get(
|
|
key, KV_NAMESPACE_FUNCTION_TABLE, timeout=timeout
|
|
)
|
|
if vals is None:
|
|
return None
|
|
else:
|
|
vals = pickle.loads(vals)
|
|
fields = [
|
|
"job_id",
|
|
"function_id",
|
|
"function_name",
|
|
"function",
|
|
"module",
|
|
"max_calls",
|
|
]
|
|
return ImportedFunctionInfo._make(vals.get(field) for field in fields)
|
|
|
|
def fetch_and_register_remote_function(self, key):
|
|
"""Import a remote function."""
|
|
remote_function_info = self.fetch_registered_method(key)
|
|
if not remote_function_info:
|
|
return False
|
|
(
|
|
job_id_str,
|
|
function_id_str,
|
|
function_name,
|
|
serialized_function,
|
|
module,
|
|
max_calls,
|
|
) = remote_function_info
|
|
|
|
function_id = ray.FunctionID(function_id_str)
|
|
job_id = ray.JobID(job_id_str)
|
|
max_calls = int(max_calls)
|
|
|
|
# This function is called by ImportThread. This operation needs to be
|
|
# atomic. Otherwise, there is race condition. Another thread may use
|
|
# the temporary function above before the real function is ready.
|
|
with self.lock:
|
|
self._num_task_executions[function_id] = 0
|
|
|
|
try:
|
|
function = pickle.loads(serialized_function)
|
|
except Exception:
|
|
# If an exception was thrown when the remote function was
|
|
# imported, we record the traceback and notify the scheduler
|
|
# of the failure.
|
|
traceback_str = format_error_message(traceback.format_exc())
|
|
|
|
def f(*args, **kwargs):
|
|
raise RuntimeError(
|
|
"The remote function failed to import on the "
|
|
"worker. This may be because needed library "
|
|
"dependencies are not installed in the worker "
|
|
"environment or cannot be found from sys.path "
|
|
f"{sys.path}:\n\n{traceback_str}"
|
|
)
|
|
|
|
# Use a placeholder method when function pickled failed
|
|
self._function_execution_info[function_id] = FunctionExecutionInfo(
|
|
function=f, function_name=function_name, max_calls=max_calls
|
|
)
|
|
|
|
# Log the error message. Log at DEBUG level to avoid overly
|
|
# spamming the log on import failure. The user gets the error
|
|
# via the RuntimeError message above.
|
|
logger.debug(
|
|
"Failed to unpickle the remote function "
|
|
f"'{function_name}' with "
|
|
f"function ID {function_id.hex()}. "
|
|
f"Job ID:{job_id}."
|
|
f"Traceback:\n{traceback_str}. "
|
|
)
|
|
else:
|
|
# The below line is necessary. Because in the driver process,
|
|
# if the function is defined in the file where the python
|
|
# script was started from, its module is `__main__`.
|
|
# However in the worker process, the `__main__` module is a
|
|
# different module, which is `default_worker.py`
|
|
function.__module__ = module
|
|
self._function_execution_info[function_id] = FunctionExecutionInfo(
|
|
function=function, function_name=function_name, max_calls=max_calls
|
|
)
|
|
return True
|
|
|
|
def get_execution_info(
|
|
self, job_id: JobID, function_descriptor: PythonFunctionDescriptor
|
|
) -> FunctionExecutionInfo:
|
|
"""Get the FunctionExecutionInfo of a remote function.
|
|
Args:
|
|
job_id: ID of the job that the function belongs to.
|
|
function_descriptor: The FunctionDescriptor of the function to get.
|
|
Returns:
|
|
A FunctionExecutionInfo object.
|
|
"""
|
|
function_id = function_descriptor.function_id
|
|
# If the function has already been loaded,
|
|
# There's no need to load again
|
|
if function_id in self._function_execution_info:
|
|
return self._function_execution_info[function_id]
|
|
if self._worker.load_code_from_local:
|
|
# Load function from local code.
|
|
if not function_descriptor.is_actor_method():
|
|
# If the function is not able to be loaded,
|
|
# try to load it from GCS,
|
|
# even if load_code_from_local is set True
|
|
if self._load_function_from_local(function_descriptor) is True:
|
|
return self._function_execution_info[function_id]
|
|
# Load function from GCS.
|
|
# Wait until the function to be executed has actually been
|
|
# registered on this worker. We will push warnings to the user if
|
|
# we spend too long in this loop.
|
|
# The driver function may not be found in sys.path. Try to load
|
|
# the function from GCS.
|
|
with profiling.profile("wait_for_function"):
|
|
self._wait_for_function(function_descriptor, job_id)
|
|
try:
|
|
function_id = function_descriptor.function_id
|
|
info = self._function_execution_info[function_id]
|
|
except KeyError as e:
|
|
message = (
|
|
"Error occurs in get_execution_info: "
|
|
"job_id: %s, function_descriptor: %s. Message: %s"
|
|
% (job_id, function_descriptor, e)
|
|
)
|
|
raise KeyError(message)
|
|
return info
|
|
|
|
def _load_function_from_local(self, function_descriptor):
|
|
assert not function_descriptor.is_actor_method()
|
|
function_id = function_descriptor.function_id
|
|
|
|
module_name, function_name = (
|
|
function_descriptor.module_name,
|
|
function_descriptor.function_name,
|
|
)
|
|
|
|
object = self.load_function_or_class_from_local(module_name, function_name)
|
|
if object is not None:
|
|
# Directly importing from local may break function with dynamic ray.remote,
|
|
# such as the _start_controller function utilized for the Ray service.
|
|
if isinstance(object, RemoteFunction):
|
|
function = object._function
|
|
else:
|
|
function = object
|
|
self._function_execution_info[function_id] = FunctionExecutionInfo(
|
|
function=function,
|
|
function_name=function_name,
|
|
max_calls=0,
|
|
)
|
|
self._num_task_executions[function_id] = 0
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def _wait_for_function(
|
|
self,
|
|
function_descriptor: PythonFunctionDescriptor,
|
|
job_id: str,
|
|
timeout: float = 10,
|
|
):
|
|
"""Wait until the function to be executed is present on this worker.
|
|
This method will simply loop until the import thread has imported the
|
|
relevant function. If we spend too long in this loop, that may indicate
|
|
a problem somewhere and we will push an error message to the user.
|
|
If this worker is an actor, then this will wait until the actor has
|
|
been defined.
|
|
|
|
Args:
|
|
function_descriptor: The FunctionDescriptor of the function that
|
|
we want to execute.
|
|
job_id: The ID of the job to push the error message to
|
|
if this times out.
|
|
timeout: Seconds to wait before pushing a warning to the user.
|
|
"""
|
|
start_time = time.time()
|
|
# Only send the warning once.
|
|
warning_sent = False
|
|
while True:
|
|
with self.lock:
|
|
if self._worker.actor_id.is_nil():
|
|
if function_descriptor.function_id in self._function_execution_info:
|
|
break
|
|
else:
|
|
key = make_function_table_key(
|
|
b"RemoteFunction",
|
|
job_id,
|
|
function_descriptor.function_id.binary(),
|
|
)
|
|
if self.fetch_and_register_remote_function(key) is True:
|
|
break
|
|
else:
|
|
assert not self._worker.actor_id.is_nil()
|
|
# Actor loading will happen when execute_task is called.
|
|
assert self._worker.actor_id in self._worker.actors
|
|
break
|
|
|
|
if time.time() - start_time > timeout:
|
|
warning_message = (
|
|
"This worker was asked to execute a function "
|
|
f"that has not been registered ({function_descriptor}, "
|
|
f"node={self._worker.node_ip_address}, "
|
|
f"worker_id={self._worker.worker_id.hex()}, "
|
|
f"pid={os.getpid()}). You may have to restart Ray."
|
|
)
|
|
if not warning_sent:
|
|
logger.error(warning_message)
|
|
ray._private.utils.push_error_to_driver(
|
|
self._worker,
|
|
ray_constants.WAIT_FOR_FUNCTION_PUSH_ERROR,
|
|
warning_message,
|
|
job_id=job_id,
|
|
)
|
|
warning_sent = True
|
|
time.sleep(0.001)
|
|
|
|
def export_actor_class(
|
|
self, Class, actor_creation_function_descriptor, actor_method_names
|
|
):
|
|
if self._worker.load_code_from_local:
|
|
module_name, class_name = (
|
|
actor_creation_function_descriptor.module_name,
|
|
actor_creation_function_descriptor.class_name,
|
|
)
|
|
# If the class is dynamic, we still export it to GCS
|
|
# even if load_code_from_local is set True.
|
|
if (
|
|
self.load_function_or_class_from_local(module_name, class_name)
|
|
is not None
|
|
):
|
|
return
|
|
|
|
# `current_job_id` shouldn't be NIL, unless:
|
|
# 1) This worker isn't an actor;
|
|
# 2) And a previous task started a background thread, which didn't
|
|
# finish before the task finished, and still uses Ray API
|
|
# after that.
|
|
assert not self._worker.current_job_id.is_nil(), (
|
|
"You might have started a background thread in a non-actor "
|
|
"task, please make sure the thread finishes before the "
|
|
"task finishes."
|
|
)
|
|
job_id = self._worker.current_job_id
|
|
key = make_function_table_key(
|
|
b"ActorClass",
|
|
job_id,
|
|
actor_creation_function_descriptor.function_id.binary(),
|
|
)
|
|
serialized_actor_class = pickle_dumps(
|
|
Class,
|
|
f"Could not serialize the actor class "
|
|
f"{actor_creation_function_descriptor.repr}",
|
|
)
|
|
actor_class_info = {
|
|
"class_name": actor_creation_function_descriptor.class_name.split(".")[-1],
|
|
"module": actor_creation_function_descriptor.module_name,
|
|
"class": serialized_actor_class,
|
|
"job_id": job_id.binary(),
|
|
"collision_identifier": self.compute_collision_identifier(Class),
|
|
"actor_method_names": json.dumps(list(actor_method_names)),
|
|
}
|
|
|
|
check_oversized_function(
|
|
actor_class_info["class"],
|
|
actor_class_info["class_name"],
|
|
"actor",
|
|
self._worker,
|
|
)
|
|
|
|
self._worker.gcs_client.internal_kv_put(
|
|
key, pickle.dumps(actor_class_info), True, KV_NAMESPACE_FUNCTION_TABLE
|
|
)
|
|
# TODO(rkn): Currently we allow actor classes to be defined
|
|
# within tasks. I tried to disable this, but it may be necessary
|
|
# because of https://github.com/ray-project/ray/issues/1146.
|
|
|
|
def load_actor_class(
|
|
self,
|
|
job_id: JobID,
|
|
actor_creation_function_descriptor: PythonFunctionDescriptor,
|
|
) -> type:
|
|
"""Load the actor class.
|
|
Args:
|
|
job_id: job ID of the actor.
|
|
actor_creation_function_descriptor: Function descriptor of
|
|
the actor constructor.
|
|
Returns:
|
|
The actor class.
|
|
"""
|
|
function_id = actor_creation_function_descriptor.function_id
|
|
# Check if the actor class already exists in the cache.
|
|
actor_class = self._loaded_actor_classes.get(function_id, None)
|
|
if actor_class is None:
|
|
# Load actor class.
|
|
if self._worker.load_code_from_local:
|
|
# Load actor class from local code first.
|
|
actor_class = self._load_actor_class_from_local(
|
|
actor_creation_function_descriptor
|
|
)
|
|
# If the actor is unable to be loaded
|
|
# from local, try to load it
|
|
# from GCS even if load_code_from_local is set True
|
|
if actor_class is None:
|
|
actor_class = self._load_actor_class_from_gcs(
|
|
job_id, actor_creation_function_descriptor
|
|
)
|
|
|
|
else:
|
|
# Load actor class from GCS.
|
|
actor_class = self._load_actor_class_from_gcs(
|
|
job_id, actor_creation_function_descriptor
|
|
)
|
|
|
|
# Re-inject tracing into the loaded class. This is necessary because
|
|
# cloudpickle doesn't preserve __signature__ attributes on module-level
|
|
# functions. When a class is pickled and unpickled, user-defined methods
|
|
# are looked up from the module, losing the __signature__ that was set by
|
|
# _inject_tracing_into_class during actor creation. Re-injecting tracing
|
|
# ensures the method signatures include _ray_trace_ctx when tracing is
|
|
# enabled, matching the behavior expected by _tracing_actor_method_invocation.
|
|
_inject_tracing_into_class(actor_class)
|
|
|
|
# Save the loaded actor class in cache.
|
|
self._loaded_actor_classes[function_id] = actor_class
|
|
|
|
# Generate execution info for the methods of this actor class.
|
|
module_name = actor_creation_function_descriptor.module_name
|
|
actor_class_name = actor_creation_function_descriptor.class_name
|
|
actor_methods = inspect.getmembers(
|
|
actor_class, predicate=is_function_or_method
|
|
)
|
|
for actor_method_name, actor_method in actor_methods:
|
|
# Actor creation function descriptor use a unique function
|
|
# hash to solve actor name conflict. When constructing an
|
|
# actor, the actor creation function descriptor will be the
|
|
# key to find __init__ method execution info. So, here we
|
|
# use actor creation function descriptor as method descriptor
|
|
# for generating __init__ method execution info.
|
|
if actor_method_name == "__init__":
|
|
method_descriptor = actor_creation_function_descriptor
|
|
else:
|
|
method_descriptor = PythonFunctionDescriptor(
|
|
module_name, actor_method_name, actor_class_name
|
|
)
|
|
method_id = method_descriptor.function_id
|
|
executor = self._make_actor_method_executor(
|
|
actor_method_name, actor_method
|
|
)
|
|
self._function_execution_info[method_id] = FunctionExecutionInfo(
|
|
function=executor,
|
|
function_name=actor_method_name,
|
|
max_calls=0,
|
|
)
|
|
self._num_task_executions[method_id] = 0
|
|
self._num_task_executions[function_id] = 0
|
|
return actor_class
|
|
|
|
def _load_actor_class_from_local(self, actor_creation_function_descriptor):
|
|
"""Load actor class from local code."""
|
|
module_name, class_name = (
|
|
actor_creation_function_descriptor.module_name,
|
|
actor_creation_function_descriptor.class_name,
|
|
)
|
|
|
|
object = self.load_function_or_class_from_local(module_name, class_name)
|
|
|
|
if object is not None:
|
|
if isinstance(object, ray.actor.ActorClass):
|
|
return object.__ray_metadata__.modified_class
|
|
else:
|
|
return ray.actor._modify_class(object)
|
|
else:
|
|
return None
|
|
|
|
def _create_fake_actor_class(
|
|
self, actor_class_name, actor_method_names, traceback_str
|
|
):
|
|
class TemporaryActor:
|
|
async def __dummy_method(self):
|
|
"""Dummy method for this fake actor class to work for async actors.
|
|
Without this method, this temporary actor class fails to initialize
|
|
if the original actor class was async."""
|
|
pass
|
|
|
|
def temporary_actor_method(*args, **kwargs):
|
|
raise RuntimeError(
|
|
f"The actor with name {actor_class_name} "
|
|
"failed to import on the worker. This may be because "
|
|
"needed library dependencies are not installed in the "
|
|
f"worker environment:\n\n{traceback_str}"
|
|
)
|
|
|
|
for method in actor_method_names:
|
|
setattr(TemporaryActor, method, temporary_actor_method)
|
|
|
|
return TemporaryActor
|
|
|
|
def _load_actor_class_from_gcs(self, job_id, actor_creation_function_descriptor):
|
|
"""Load actor class from GCS."""
|
|
key = make_function_table_key(
|
|
b"ActorClass",
|
|
job_id,
|
|
actor_creation_function_descriptor.function_id.binary(),
|
|
)
|
|
|
|
# Fetch raw data from GCS.
|
|
vals = self._worker.gcs_client.internal_kv_get(key, KV_NAMESPACE_FUNCTION_TABLE)
|
|
fields = ["job_id", "class_name", "module", "class", "actor_method_names"]
|
|
if vals is None:
|
|
vals = {}
|
|
else:
|
|
vals = pickle.loads(vals)
|
|
(job_id_str, class_name, module, pickled_class, actor_method_names) = (
|
|
vals.get(field) for field in fields
|
|
)
|
|
|
|
class_name = ensure_str(class_name)
|
|
module_name = ensure_str(module)
|
|
job_id = ray.JobID(job_id_str)
|
|
actor_method_names = json.loads(ensure_str(actor_method_names))
|
|
|
|
actor_class = None
|
|
try:
|
|
with self.lock:
|
|
actor_class = pickle.loads(pickled_class)
|
|
except Exception:
|
|
logger.debug("Failed to load actor class %s.", class_name)
|
|
# If an exception was thrown when the actor was imported, we record
|
|
# the traceback and notify the scheduler of the failure.
|
|
traceback_str = format_error_message(traceback.format_exc())
|
|
# The actor class failed to be unpickled, create a fake actor
|
|
# class instead (just to produce error messages and to prevent
|
|
# the driver from hanging).
|
|
actor_class = self._create_fake_actor_class(
|
|
class_name, actor_method_names, traceback_str
|
|
)
|
|
|
|
# The below line is necessary. Because in the driver process,
|
|
# if the function is defined in the file where the python script
|
|
# was started from, its module is `__main__`.
|
|
# However in the worker process, the `__main__` module is a
|
|
# different module, which is `default_worker.py`
|
|
actor_class.__module__ = module_name
|
|
return actor_class
|
|
|
|
def _make_actor_method_executor(self, method_name: str, method: Callable):
|
|
"""Make an executor that wraps a user-defined actor method.
|
|
The wrapped method updates the worker's internal state and performs any
|
|
necessary checkpointing operations.
|
|
Args:
|
|
method_name: The name of the actor method.
|
|
method: The actor method to wrap. This should be a
|
|
method defined on the actor class and should therefore take an
|
|
instance of the actor as the first argument.
|
|
Returns:
|
|
A function that executes the given actor method on the worker's
|
|
stored instance of the actor. The function also updates the
|
|
worker's internal state to record the executed method.
|
|
"""
|
|
|
|
def actor_method_executor(__ray_actor, *args, **kwargs):
|
|
# Execute the assigned method.
|
|
is_bound = is_class_method(method) or is_static_method(
|
|
type(__ray_actor), method_name
|
|
)
|
|
if is_bound:
|
|
return method(*args, **kwargs)
|
|
else:
|
|
return method(__ray_actor, *args, **kwargs)
|
|
|
|
# Set method_name and method as attributes to the executor closure
|
|
# so we can make decision based on these attributes in task executor.
|
|
# Precisely, asyncio support requires to know whether:
|
|
# - the method is a ray internal method: starts with __ray
|
|
# - the method is a coroutine function: defined by async def
|
|
actor_method_executor.name = method_name
|
|
actor_method_executor.method = method
|
|
|
|
return actor_method_executor
|