diff --git a/lib/python3.12/site-packages/ray/dag/__init__.py b/lib/python3.12/site-packages/ray/dag/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bc9970899cf41e4755e4d6c424b768d2357d78f2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/__init__.py @@ -0,0 +1,46 @@ +from ray.dag.dag_node import DAGNode +from ray.dag.function_node import FunctionNode +from ray.dag.class_node import ( + ClassNode, + ClassMethodNode, +) +from ray.dag.collective_node import CollectiveOutputNode +from ray.dag.input_node import ( + InputNode, + InputAttributeNode, + DAGInputData, +) +from ray.dag.output_node import MultiOutputNode +from ray.dag.dag_operation_future import DAGOperationFuture, GPUFuture +from ray.dag.constants import ( + PARENT_CLASS_NODE_KEY, + PREV_CLASS_METHOD_CALL_KEY, + BIND_INDEX_KEY, + IS_CLASS_METHOD_OUTPUT_KEY, + COLLECTIVE_OPERATION_KEY, + DAGNODE_TYPE_KEY, +) +from ray.dag.vis_utils import plot +from ray.dag.context import DAGContext + +__all__ = [ + "ClassNode", + "ClassMethodNode", + "CollectiveOutputNode", + "DAGNode", + "DAGOperationFuture", + "FunctionNode", + "GPUFuture", + "InputNode", + "InputAttributeNode", + "DAGInputData", + "PARENT_CLASS_NODE_KEY", + "PREV_CLASS_METHOD_CALL_KEY", + "BIND_INDEX_KEY", + "IS_CLASS_METHOD_OUTPUT_KEY", + "COLLECTIVE_OPERATION_KEY", + "DAGNODE_TYPE_KEY", + "plot", + "MultiOutputNode", + "DAGContext", +] diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a42ea65f8068200db7c3e5afc6915550c692be42 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f99f7930994ab183d49c03a407022b03c911b1ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/class_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/class_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d8d5b3fe57cff46154f1c2a0114bdb07c474215 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/class_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/collective_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/collective_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71919c4e09ba87c8775c40e801006b2b2006b7a6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/collective_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/conftest.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/conftest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..faf524535349e1cc459f4e0fe888345c829a014b Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/conftest.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..705a071204ffb9fb610f3ffa1d4c9916c9be1dd5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/context.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..117a0b49cb83702d0e293f86e8a06c37eba2db52 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ef7e6c5e4a07950da5f418e578c51830d95046f Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node_operation.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node_operation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24036c00967048749873879e55b7164a01f603b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_node_operation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/dag_operation_future.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_operation_future.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ddca0e7776537f1ab7fb14613aea01a2c8cbed3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/dag_operation_future.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/format_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/format_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be83e0edbd928e4f46f9b5aab767546c0c80422e Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/format_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/function_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/function_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c229af767d4c1428d08b833ddfef2b6408a79e5f Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/function_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/input_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/input_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bce4dbb9caabdb0177019a48f3af135125392ead Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/input_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/output_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/output_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34cdbcb024f6a330698bc1466d85aea00c8ef2d9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/output_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/py_obj_scanner.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/py_obj_scanner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c35a57530a3101002f734cb72e5b8524f1518ed9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/py_obj_scanner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd382919f3b7168447956e5bb5bb3603acce5870 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/__pycache__/vis_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/__pycache__/vis_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b334bf4f1cc2a3255d0deec3bc843bc4b9ba7ef9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/__pycache__/vis_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/base.py b/lib/python3.12/site-packages/ray/dag/base.py new file mode 100644 index 0000000000000000000000000000000000000000..4153866cdeeba83e32becb86ab42ab2544646028 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/base.py @@ -0,0 +1,8 @@ +"""This module defines the base class for object scanning and gets rid of +reference cycles.""" +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class DAGNodeBase: + """Common base class for a node in a Ray task graph.""" diff --git a/lib/python3.12/site-packages/ray/dag/class_node.py b/lib/python3.12/site-packages/ray/dag/class_node.py new file mode 100644 index 0000000000000000000000000000000000000000..1a5b78e8e7066546759773aed070217b8dac5bdb --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/class_node.py @@ -0,0 +1,320 @@ +from typing import Any, Dict, List, Optional, Tuple, Union +from weakref import ReferenceType + +import ray +from ray.dag.constants import ( + BIND_INDEX_KEY, + IS_CLASS_METHOD_OUTPUT_KEY, + PARENT_CLASS_NODE_KEY, + PREV_CLASS_METHOD_CALL_KEY, +) +from ray.dag.dag_node import DAGNode +from ray.dag.format_utils import get_dag_node_str +from ray.dag.input_node import InputNode +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class ClassNode(DAGNode): + """Represents an actor creation in a Ray task DAG.""" + + def __init__( + self, + cls, + cls_args, + cls_kwargs, + cls_options, + other_args_to_resolve=None, + ): + self._body = cls + self._last_call: Optional["ClassMethodNode"] = None + super().__init__( + cls_args, + cls_kwargs, + cls_options, + other_args_to_resolve=other_args_to_resolve, + ) + + if self._contains_input_node(): + raise ValueError( + "InputNode handles user dynamic input the DAG, and " + "cannot be used as args, kwargs, or other_args_to_resolve " + "in ClassNode constructor because it is not available at " + "class construction or binding time." + ) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return ClassNode( + self._body, + new_args, + new_kwargs, + new_options, + other_args_to_resolve=new_other_args_to_resolve, + ) + + def _execute_impl(self, *args, **kwargs): + """Executor of ClassNode by ray.remote() + + Args and kwargs are to match base class signature, but not in the + implementation. All args and kwargs should be resolved and replaced + with value in bound_args and bound_kwargs via bottom-up recursion when + current node is executed. + """ + return ( + ray.remote(self._body) + .options(**self._bound_options) + .remote(*self._bound_args, **self._bound_kwargs) + ) + + def _contains_input_node(self) -> bool: + """Check if InputNode is used in children DAGNodes with current node + as the root. + """ + children_dag_nodes = self._get_all_child_nodes() + for child in children_dag_nodes: + if isinstance(child, InputNode): + return True + return False + + def __getattr__(self, method_name: str): + # User trying to call .bind() without a bind class method + if method_name == "bind" and "bind" not in dir(self._body): + raise AttributeError(f".bind() cannot be used again on {type(self)} ") + # Raise an error if the method is invalid. + getattr(self._body, method_name) + call_node = _UnboundClassMethodNode(self, method_name, {}) + return call_node + + def __str__(self) -> str: + return get_dag_node_str(self, str(self._body)) + + +class _UnboundClassMethodNode(object): + def __init__(self, actor: ClassNode, method_name: str, options: dict): + # TODO(sang): Theoretically, We should use weakref cuz it is + # a circular dependency but when I used weakref, it fails + # because we cannot serialize the weakref. + self._actor = actor + self._method_name = method_name + self._options = options + + def bind(self, *args, **kwargs): + other_args_to_resolve = { + PARENT_CLASS_NODE_KEY: self._actor, + PREV_CLASS_METHOD_CALL_KEY: self._actor._last_call, + } + + node = ClassMethodNode( + self._method_name, + args, + kwargs, + self._options, + other_args_to_resolve=other_args_to_resolve, + ) + self._actor._last_call = node + return node + + def __getattr__(self, attr: str): + if attr == "remote": + raise AttributeError( + ".remote() cannot be used on ClassMethodNodes. Use .bind() instead " + "to express an symbolic actor call." + ) + else: + return self.__getattribute__(attr) + + def options(self, **options): + self._options = options + return self + + +class _ClassMethodOutput: + """Represents a class method output in a Ray function DAG.""" + + def __init__(self, class_method_call: "ClassMethodNode", output_idx: int): + # The upstream class method call that returns multiple values. + self._class_method_call = class_method_call + # The output index of the return value from the upstream class method call. + self._output_idx = output_idx + + @property + def class_method_call(self) -> "ClassMethodNode": + return self._class_method_call + + @property + def output_idx(self) -> int: + return self._output_idx + + +@DeveloperAPI +class ClassMethodNode(DAGNode): + """Represents an actor method invocation in a Ray function DAG.""" + + def __init__( + self, + method_name: str, + method_args: Tuple[Any], + method_kwargs: Dict[str, Any], + method_options: Dict[str, Any], + other_args_to_resolve: Dict[str, Any], + ): + self._bound_args = method_args or [] + self._bound_kwargs = method_kwargs or {} + self._bound_options = method_options or {} + self._method_name: str = method_name + # Parse other_args_to_resolve and assign to variables + self._parent_class_node: Union[ + ClassNode, ReferenceType["ray._private.actor.ActorHandle"] + ] = other_args_to_resolve.get(PARENT_CLASS_NODE_KEY) + # Used to track lineage of ClassMethodCall to preserve deterministic + # submission and execution order. + self._prev_class_method_call: Optional[ + ClassMethodNode + ] = other_args_to_resolve.get(PREV_CLASS_METHOD_CALL_KEY, None) + # The index/order when bind() is called on this class method + self._bind_index: Optional[int] = other_args_to_resolve.get( + BIND_INDEX_KEY, None + ) + # Represent if the ClassMethodNode is a class method output. If True, + # the node is a placeholder for a return value from the ClassMethodNode + # that returns multiple values. If False, the node is a class method call. + self._is_class_method_output: bool = other_args_to_resolve.get( + IS_CLASS_METHOD_OUTPUT_KEY, False + ) + # Represents the return value from the upstream ClassMethodNode that + # returns multiple values. If the node is a class method call, this is None. + self._class_method_output: Optional[_ClassMethodOutput] = None + if self._is_class_method_output: + # Set the upstream ClassMethodNode and the output index of the return + # value from `method_args`. + self._class_method_output = _ClassMethodOutput( + method_args[0], method_args[1] + ) + + # The actor creation task dependency is encoded as the first argument, + # and the ordering dependency as the second, which ensures they are + # executed prior to this node. + super().__init__( + method_args, + method_kwargs, + method_options, + other_args_to_resolve=other_args_to_resolve, + ) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return ClassMethodNode( + self._method_name, + new_args, + new_kwargs, + new_options, + other_args_to_resolve=new_other_args_to_resolve, + ) + + def _execute_impl(self, *args, **kwargs): + """Executor of ClassMethodNode by ray.remote() + + Args and kwargs are to match base class signature, but not in the + implementation. All args and kwargs should be resolved and replaced + with value in bound_args and bound_kwargs via bottom-up recursion when + current node is executed. + """ + if self.is_class_method_call: + method_body = getattr(self._parent_class_node, self._method_name) + # Execute with bound args. + return method_body.options(**self._bound_options).remote( + *self._bound_args, + **self._bound_kwargs, + ) + else: + assert self._class_method_output is not None + return self._bound_args[0][self._class_method_output.output_idx] + + def __str__(self) -> str: + return get_dag_node_str(self, f"{self._method_name}()") + + def __repr__(self) -> str: + return self.__str__() + + def get_method_name(self) -> str: + return self._method_name + + def _get_bind_index(self) -> int: + return self._bind_index + + def _get_remote_method(self, method_name): + method_body = getattr(self._parent_class_node, method_name) + return method_body + + def _get_actor_handle(self) -> Optional["ray.actor.ActorHandle"]: + if not isinstance(self._parent_class_node, ray.actor.ActorHandle): + return None + return self._parent_class_node + + @property + def num_returns(self) -> int: + """ + Return the number of return values from the class method call. If the + node is a class method output, return the number of return values from + the upstream class method call. + """ + + if self.is_class_method_call: + num_returns = self._bound_options.get("num_returns", None) + if num_returns is None: + method = self._get_remote_method(self._method_name) + num_returns = method.__getstate__()["num_returns"] + return num_returns + else: + assert self._class_method_output is not None + return self._class_method_output.class_method_call.num_returns + + @property + def is_class_method_call(self) -> bool: + """ + Return True if the node is a class method call, False if the node is a + class method output. + """ + return not self._is_class_method_output + + @property + def is_class_method_output(self) -> bool: + """ + Return True if the node is a class method output, False if the node is a + class method call. + """ + return self._is_class_method_output + + @property + def class_method_call(self) -> Optional["ClassMethodNode"]: + """ + Return the upstream class method call that returns multiple values. If + the node is a class method output, return None. + """ + + if self._class_method_output is None: + return None + return self._class_method_output.class_method_call + + @property + def output_idx(self) -> Optional[int]: + """ + Return the output index of the return value from the upstream class + method call that returns multiple values. If the node is a class method + call, return None. + """ + + if self._class_method_output is None: + return None + return self._class_method_output.output_idx diff --git a/lib/python3.12/site-packages/ray/dag/collective_node.py b/lib/python3.12/site-packages/ray/dag/collective_node.py new file mode 100644 index 0000000000000000000000000000000000000000..03609b20cc2eccb51e20ce951a3834f5e49caa0d --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/collective_node.py @@ -0,0 +1,307 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +if TYPE_CHECKING: + import torch + +import ray +from ray.dag import ( + ClassMethodNode, + DAGNode, +) +from ray.dag.constants import COLLECTIVE_OPERATION_KEY, IS_CLASS_METHOD_OUTPUT_KEY +from ray.experimental.channel import ChannelContext +from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType +from ray.experimental.util.types import ( + AllGatherOp, + AllReduceOp, + ReduceScatterOp, + _CollectiveOp, +) +from ray.util.annotations import DeveloperAPI + + +class _CollectiveOperation: + """ + Represent metadata for a collective communicator collective operation. + + Args: + inputs: A list of lists of DAGNode. Each nested list inside + of inputs should contain exactly one object per actor. + If multiple nested lists are provided, then the order of + actors should be the same for each nested list. + op: The collective operation to perform. + transport: The transport to use for the collective operation. + + Requirements: + 1. Input nodes are unique. + 2. Actor handles are unique. + 3. Actor handles match the custom communicator group if specified. + """ + + def __init__( + self, + inputs: List[List[DAGNode]], + op: _CollectiveOp, + transport: Optional[Union[str, Communicator]] = None, + ): + self._actor_handles: List["ray.actor.ActorHandle"] = [] + for i, input_nodes in enumerate(inputs): + # Check non-empty input list + if len(input_nodes) == 0: + nested_list_error_msg = f" at index {i}" if len(inputs) > 1 else "" + raise ValueError( + f"Expected non-empty input list{nested_list_error_msg}." + ) + + # Check input nodes are DAGNode + if not all(isinstance(node, DAGNode) for node in input_nodes): + nested_list_error_msg = ( + f" at list at index {i}" if len(inputs) > 1 else "" + ) + raise ValueError( + f"Expected all input nodes to be DAGNode{nested_list_error_msg}, " + f"but got {input_nodes}." + ) + + # Check unique input nodes + if len(set(input_nodes)) != len(input_nodes): + duplicates = [ + input_node + for input_node in input_nodes + if input_nodes.count(input_node) > 1 + ] + nested_list_error_msg = ( + f" at list at index {i}" if len(inputs) > 1 else "" + ) + raise ValueError( + f"Expected unique input nodes{nested_list_error_msg}, but found duplicates: " + f"{duplicates}" + ) + + current_actor_handles = [] + for input_node in input_nodes: + actor_handle = input_node._get_actor_handle() + if actor_handle is None: + nested_list_error_msg = ( + f" at list at index {i}" if len(inputs) > 1 else "" + ) + raise ValueError( + f"Expected an actor handle from the input node{nested_list_error_msg}" + ) + current_actor_handles.append(actor_handle) + + # Check unique actor handles + if len(set(current_actor_handles)) != len(current_actor_handles): + invalid_input_nodes = [ + input_node + for input_node in input_nodes + if current_actor_handles.count(input_node._get_actor_handle()) > 1 + ] + nested_list_error_msg = ( + f" at list at index {i}" if len(inputs) > 1 else "" + ) + raise ValueError( + f"Expected unique actor handles{nested_list_error_msg}, " + "but found duplicate actor handles from input nodes: " + f"{invalid_input_nodes}" + ) + + if i == 0: + first_actor_handles = current_actor_handles + + # Check all lists of DAGNode have the same number of nodes + if len(inputs[0]) != len(inputs[i]): + raise ValueError( + f"Expected all input lists to have the same number of nodes. " + f"List at index 0 has length {len(inputs[0])}, but list at " + f"index {i} has length {len(inputs[i])}." + ) + + # Check all lists of DAGNode have same set of actor handles + if set(first_actor_handles) != set(current_actor_handles): + raise ValueError( + f"Expected all input lists to have the same set of actor handles. " + f"List at index 0 has actors {set(first_actor_handles)}, but list at " + f"index {i} has actors {set(current_actor_handles)}." + ) + + # Check all lists of DAGNode have same order of actor handles + for j, (first, current) in enumerate( + zip(first_actor_handles, current_actor_handles) + ): + if first != current: + raise ValueError( + f"Expected all input lists to have the same order of actor handles. " + f"List at index 0 has actor {first} at position {j}, but list at " + f"index {i} has actor {current} at position {j}." + ) + self._actor_handles = current_actor_handles + + self._op = op + if transport is None: + transport = TorchTensorType.ACCELERATOR + self._type_hint = TorchTensorType(transport=transport, _direct_return=True) + if isinstance(transport, Communicator): + if set(transport.get_actor_handles()) != set(self._actor_handles): + raise ValueError( + "Expected actor handles to match the custom communicator group" + ) + + def __str__(self) -> str: + return ( + f"CollectiveOperation(" + f"_actor_handles={self._actor_handles}, " + f"_op={self._op}, " + f"_type_hint={self._type_hint})" + ) + + @property + def actor_handles(self) -> List["ray.actor.ActorHandle"]: + return self._actor_handles + + @property + def type_hint(self) -> TorchTensorType: + return self._type_hint + + def get_communicator(self) -> Communicator: + if self._type_hint.communicator_id is not None: + ctx = ChannelContext.get_current() + communicator = ctx.communicators[self._type_hint.communicator_id] + elif self._type_hint.get_custom_communicator() is not None: + communicator = self._type_hint.get_custom_communicator() + else: + raise ValueError("Expected a communicator group") + return communicator + + def execute( + self, *send_buf: "torch.Tensor" + ) -> Union["torch.Tensor", Tuple["torch.Tensor", ...]]: + """ + Call the collective operation on the input tensor(s). Output tensor(s) are + allocated and returned. + + Args: + *send_buf: A variable number of torch tensors to send to the collective + operation. The tensors have the same order as the input nodes. + + Returns: + A torch tensor or a tuple of torch tensors containing the results of the + collective operation. The output tensors have the same length and order + as the input node list of the actor of this operation. + """ + import torch + + if not all(isinstance(t, torch.Tensor) for t in send_buf): + raise ValueError("Expected a torch tensor for each input node") + + communicator = self.get_communicator() + if isinstance(self._op, AllGatherOp): + assert len(send_buf) == 1 + t = send_buf[0] + world_size = len(self._actor_handles) + recv_buf = torch.empty( + (t.shape[0] * world_size, *t.shape[1:]), + dtype=t.dtype, + device=t.device, + ) + communicator.allgather(t, recv_buf) + elif isinstance(self._op, AllReduceOp): + if len(send_buf) == 1: + t = send_buf[0] + recv_buf = torch.empty_like(t) + communicator.allreduce(t, recv_buf, self._op.reduceOp) + else: + if not all(t.dtype == send_buf[0].dtype for t in send_buf): + raise ValueError( + "Expected all input tensors to have the same dtype, " + f"but got {[t.dtype for t in send_buf]}" + ) + + def unflatten_from(flat_buf, bufs): + views = [] + offset = 0 + for t in bufs: + numel = t.numel() + t = flat_buf[offset : offset + numel].view(t.shape) + views.append(t) + offset += numel + return tuple(views) + + flat_buf = torch.nn.utils.parameters_to_vector(send_buf) + communicator.allreduce(flat_buf, flat_buf, self._op.reduceOp) + recv_buf = unflatten_from(flat_buf, send_buf) + elif isinstance(self._op, ReduceScatterOp): + assert len(send_buf) == 1 + t = send_buf[0] + world_size = len(self._actor_handles) + if t.shape[0] % world_size != 0: + raise ValueError( + "Expected the first dimension of the input tensor to be divisible " + f"by the world size {world_size}" + ) + recv_buf = torch.empty( + (t.shape[0] // world_size, *t.shape[1:]), + dtype=t.dtype, + device=t.device, + ) + communicator.reducescatter(t, recv_buf, self._op.reduceOp) + return recv_buf + + +@DeveloperAPI +class CollectiveOutputNode(ClassMethodNode): + """Represent an output node from a communicator collective operation in a Ray DAG.""" + + def __init__( + self, + method_name: str, + method_args: Tuple[ + DAGNode, + ], + method_kwargs: Dict[str, Any], + method_options: Dict[str, Any], + other_args_to_resolve: Dict[str, Any], + ): + # Parse the input node(s). + self._inputs = method_args + # Parse the collective operation. + self._collective_op: _CollectiveOperation = other_args_to_resolve.get( + COLLECTIVE_OPERATION_KEY, None + ) + self._is_class_method_output: bool = other_args_to_resolve.get( + IS_CLASS_METHOD_OUTPUT_KEY, False + ) + if self._collective_op is None and not self._is_class_method_output: + raise ValueError("Expected a collective operation") + + super().__init__( + method_name, + method_args, + method_kwargs, + method_options, + other_args_to_resolve, + ) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return CollectiveOutputNode( + self._method_name, + new_args, + new_kwargs, + new_options, + other_args_to_resolve=new_other_args_to_resolve, + ) + + def _execute_impl(self, *args, **kwargs): + raise NotImplementedError( + "CollectiveOutputNode is only supported with dag.experimental_compile()" + ) + + @property + def collective_op(self) -> _CollectiveOperation: + return self._collective_op diff --git a/lib/python3.12/site-packages/ray/dag/compiled_dag_node.py b/lib/python3.12/site-packages/ray/dag/compiled_dag_node.py new file mode 100644 index 0000000000000000000000000000000000000000..2204cb40ab7bf24d5278f7acb3aac284b573c099 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/compiled_dag_node.py @@ -0,0 +1,3312 @@ +import asyncio +import logging +import threading +import time +import traceback +import uuid +import weakref +from collections import defaultdict +from contextlib import nullcontext +from dataclasses import asdict, dataclass +from typing import ( + Any, + Dict, + List, + Optional, + Set, + Tuple, + Union, +) + +import ray +import ray.exceptions +from ray.dag.constants import ( + RAY_CGRAPH_ENABLE_NVTX_PROFILING, + RAY_CGRAPH_ENABLE_TORCH_PROFILING, + RAY_CGRAPH_VISUALIZE_SCHEDULE, +) +from ray.dag.dag_node_operation import ( + _build_dag_node_operation_graph, + _DAGNodeOperation, + _DAGNodeOperationType, + _DAGOperationGraphNode, + _extract_execution_schedule, + _generate_actor_to_execution_schedule, + _generate_overlapped_execution_schedule, + _visualize_execution_schedule, +) +from ray.dag.dag_operation_future import DAGOperationFuture, GPUFuture, ResolvedFuture +from ray.exceptions import ( + RayCgraphCapacityExceeded, + RayChannelError, + RayChannelTimeoutError, + RayTaskError, +) +from ray.experimental.channel import ( + AwaitableBackgroundReader, + AwaitableBackgroundWriter, + ChannelContext, + ChannelInterface, + ChannelOutputType, + CompiledDAGArgs, + CompositeChannel, + IntraProcessChannel, + ReaderInterface, + SynchronousReader, + SynchronousWriter, + WriterInterface, +) +from ray.experimental.channel.accelerator_context import AcceleratorContext +from ray.experimental.channel.auto_transport_type import ( + AutoTransportType, + TypeHintResolver, +) +from ray.experimental.channel.cached_channel import CachedChannel +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.shared_memory_channel import ( + SharedMemoryType, +) +from ray.experimental.channel.torch_tensor_accelerator_channel import ( + _destroy_communicator, + _init_communicator, +) +from ray.experimental.channel.torch_tensor_type import TorchTensorType +from ray.experimental.compiled_dag_ref import ( + CompiledDAGFuture, + CompiledDAGRef, + _process_return_vals, +) +from ray.util.annotations import DeveloperAPI +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +logger = logging.getLogger(__name__) + +# Keep tracking of every compiled dag created during the lifetime of +# this process. It tracks them as weakref meaning when the compiled dag +# is GC'ed, it is automatically removed from here. It is used to teardown +# compiled dags at interpreter shutdown time. +_compiled_dags = weakref.WeakValueDictionary() + + +# Relying on __del__ doesn't work well upon shutdown because +# the destructor order is not guaranteed. We call this function +# upon `ray.worker.shutdown` which is registered to atexit handler +# so that teardown is properly called before objects are destructed. +def _shutdown_all_compiled_dags(): + global _compiled_dags + for _, compiled_dag in _compiled_dags.items(): + # Kill DAG actors to avoid hanging during shutdown if the actor tasks + # cannot be cancelled. + compiled_dag.teardown(kill_actors=True) + _compiled_dags = weakref.WeakValueDictionary() + + +def _check_unused_dag_input_attributes( + output_node: "ray.dag.MultiOutputNode", input_attributes: Set[str] +) -> Set[str]: + """ + Helper function to check that all input attributes are used in the DAG. + For example, if the user creates an input attribute by calling + InputNode()["x"], we ensure that there is a path from the + InputAttributeNode corresponding to "x" to the DAG's output. If an + input attribute is not used, throw an error. + + Args: + output_node: The starting node for the traversal. + input_attributes: A set of attributes accessed by the InputNode. + """ + from ray.dag import InputAttributeNode + + used_attributes = set() + visited_nodes = set() + stack: List["ray.dag.DAGNode"] = [output_node] + + while stack: + current_node = stack.pop() + if current_node in visited_nodes: + continue + visited_nodes.add(current_node) + + if isinstance(current_node, InputAttributeNode): + used_attributes.add(current_node.key) + + stack.extend(current_node._upstream_nodes) + + unused_attributes = input_attributes - used_attributes + if unused_attributes: + unused_attributes_str = ", ".join(str(key) for key in unused_attributes) + input_attributes_str = ", ".join(str(key) for key in input_attributes) + unused_phrase = "is unused" if len(unused_attributes) == 1 else "are unused" + + raise ValueError( + "Compiled Graph expects input to be accessed " + f"using all of attributes {input_attributes_str}, " + f"but {unused_attributes_str} {unused_phrase}. " + "Ensure all input attributes are used and contribute " + "to the computation of the Compiled Graph output." + ) + + +@DeveloperAPI +def do_allocate_channel( + self, + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + typ: ChannelOutputType, + driver_actor_id: Optional[str] = None, +) -> ChannelInterface: + """Generic actor method to allocate an output channel. + + Args: + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + typ: The output type hint for the channel. + driver_actor_id: If this channel is read by a driver and that driver is an + actual actor, this will be the actor ID of that driver actor. + + Returns: + The allocated channel. + """ + # None means it is called from a driver. + writer: Optional["ray.actor.ActorHandle"] = None + try: + writer = ray.get_runtime_context().current_actor + except RuntimeError: + # This is the driver so there is no current actor handle. + pass + + output_channel = typ.create_channel( + writer, + reader_and_node_list, + driver_actor_id, + ) + return output_channel + + +@DeveloperAPI +def do_exec_tasks( + self, + tasks: List["ExecutableTask"], + schedule: List[_DAGNodeOperation], + overlap_gpu_communication: bool = False, +) -> None: + """A generic actor method to begin executing the operations belonging to an + actor. This runs an infinite loop to execute each _DAGNodeOperation in the + order specified by the schedule. It exits only if the actor dies or an + exception is thrown. + + Args: + tasks: the executable tasks corresponding to the actor methods. + schedule: A list of _DAGNodeOperation that should be executed in order. + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance. + """ + try: + for task in tasks: + task.prepare(overlap_gpu_communication=overlap_gpu_communication) + + if RAY_CGRAPH_ENABLE_NVTX_PROFILING: + assert ( + not RAY_CGRAPH_ENABLE_TORCH_PROFILING + ), "NVTX and torch profiling cannot be enabled at the same time." + try: + import nvtx + except ImportError: + raise ImportError( + "Please install nvtx to enable nsight profiling. " + "You can install it by running `pip install nvtx`." + ) + nvtx_profile = nvtx.Profile() + nvtx_profile.enable() + + if RAY_CGRAPH_ENABLE_TORCH_PROFILING: + assert ( + not RAY_CGRAPH_ENABLE_NVTX_PROFILING + ), "NVTX and torch profiling cannot be enabled at the same time." + + import torch + + torch_profile = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + on_trace_ready=torch.profiler.tensorboard_trace_handler( + "compiled_graph_torch_profiles" + ), + ) + torch_profile.start() + logger.info("Torch profiling started") + + done = False + while True: + if done: + break + for operation in schedule: + done = tasks[operation.exec_task_idx].exec_operation( + self, operation.type, overlap_gpu_communication + ) + if done: + break + + if RAY_CGRAPH_ENABLE_NVTX_PROFILING: + nvtx_profile.disable() + + if RAY_CGRAPH_ENABLE_TORCH_PROFILING: + torch_profile.stop() + logger.info("Torch profiling stopped") + except Exception: + logging.exception("Compiled DAG task exited with exception") + raise + + +@DeveloperAPI +def do_profile_tasks( + self, + tasks: List["ExecutableTask"], + schedule: List[_DAGNodeOperation], + overlap_gpu_communication: bool = False, +) -> None: + """A generic actor method similar to `do_exec_tasks`, but with profiling enabled. + + Args: + tasks: the executable tasks corresponding to the actor methods. + schedule: A list of _DAGNodeOperation that should be executed in order. + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance. + """ + try: + for task in tasks: + task.prepare(overlap_gpu_communication=overlap_gpu_communication) + + if not hasattr(self, "__ray_cgraph_events"): + self.__ray_cgraph_events = [] + + done = False + while True: + if done: + break + for operation in schedule: + start_t = time.perf_counter() + task = tasks[operation.exec_task_idx] + done = task.exec_operation( + self, operation.type, overlap_gpu_communication + ) + end_t = time.perf_counter() + + self.__ray_cgraph_events.append( + _ExecutableTaskRecord( + actor_classname=self.__class__.__name__, + actor_name=ray.get_runtime_context().get_actor_name(), + actor_id=ray.get_runtime_context().get_actor_id(), + method_name=task.method_name, + bind_index=task.bind_index, + operation=operation.type.value, + start_t=start_t, + end_t=end_t, + ) + ) + + if done: + break + except Exception: + logging.exception("Compiled DAG task exited with exception") + raise + + +@DeveloperAPI +def do_cancel_executable_tasks(self, tasks: List["ExecutableTask"]) -> None: + # CUDA events should be destroyed before other CUDA resources. + for task in tasks: + task.destroy_cuda_event() + for task in tasks: + task.cancel() + + +def _wrap_exception(exc): + backtrace = ray._private.utils.format_error_message( + "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)), + task_exception=True, + ) + wrapped = RayTaskError( + function_name="do_exec_tasks", + traceback_str=backtrace, + cause=exc, + ) + return wrapped + + +def _get_comm_group_id(type_hint: ChannelOutputType) -> Optional[str]: + """ + Get the communicator group ID from the type hint. If the type hint does not + require communicator, return None. + + Args: + type_hint: The type hint of the channel. + + Returns: + The communicator group ID if the type hint requires communicator, + otherwise None. + """ + if type_hint.requires_accelerator(): + assert isinstance(type_hint, TorchTensorType) + return type_hint.communicator_id + return None + + +def _device_context_manager(): + """ + Return a context manager for executing communication operations + (i.e., READ and WRITE). For accelerator operations, the context manager + uses the proper cuda device from channel context, otherwise, + nullcontext will be returned. + """ + if not ChannelContext.get_current().torch_available: + return nullcontext() + + import torch + + from ray.experimental.channel.accelerator_context import AcceleratorContext + + device = ChannelContext.get_current().torch_device + + if device.type == "cuda" and not torch.cuda.is_available(): + # In the case of mocked NCCL, we may get a device with type "cuda" + # but CUDA is not available. We return nullcontext() in that case, + # otherwise torch raises a runtime error if the cuda device context + # manager is used. + # TODO(rui): consider better mocking NCCL to support device context. + return nullcontext() + + return AcceleratorContext.get().get_device_context(device) + + +@DeveloperAPI +class CompiledTask: + """Wraps the normal Ray DAGNode with some metadata.""" + + def __init__(self, idx: int, dag_node: "ray.dag.DAGNode"): + """ + Args: + idx: A unique index into the original DAG. + dag_node: The original DAG node created by the user. + """ + self.idx = idx + self.dag_node = dag_node + + # Dict from task index to actor handle for immediate downstream tasks. + self.downstream_task_idxs: Dict[int, "ray.actor.ActorHandle"] = {} + # Case 1: The task represents a ClassMethodNode. + # + # Multiple return values are written to separate `output_channels`. + # `output_idxs` represents the tuple index of the output value for + # multiple returns in a tuple. If an output index is None, it means + # the complete return value is written to the output channel. + # Otherwise, the return value is a tuple and the index is used + # to extract the value to be written to the output channel. + # + # Case 2: The task represents an InputNode. + # + # `output_idxs` can be an integer or a string to retrieve the + # corresponding value from `args` or `kwargs` in the DAG's input. + self.output_channels: List[ChannelInterface] = [] + self.output_idxs: List[Optional[Union[int, str]]] = [] + # The DAGNodes that are arguments to this task. + # This is used for lazy resolution of the arguments' type hints. + self.arg_nodes: List["ray.dag.DAGNode"] = [] + # idxs of possible ClassMethodOutputNodes if they exist, used for visualization + self.output_node_idxs: List[int] = [] + + @property + def args(self) -> Tuple[Any]: + return self.dag_node.get_args() + + @property + def kwargs(self) -> Dict[str, Any]: + return self.dag_node.get_kwargs() + + @property + def num_readers(self) -> int: + return len(self.downstream_task_idxs) + + @property + def arg_type_hints(self) -> List["ChannelOutputType"]: + return [arg_node.type_hint for arg_node in self.arg_nodes] + + def __str__(self) -> str: + return f""" + Node: {self.dag_node} + Arguments: {self.args} + Output: {self.output_channels} + """ + + +class _ExecutableTaskInput: + """Represents an input to an ExecutableTask. + + Args: + input_variant: either an unresolved input (when type is ChannelInterface) + , or a resolved input value (when type is Any) + channel_idx: if input_variant is an unresolved input, this is the index + into the input channels list. + """ + + def __init__( + self, + input_variant: Union[ChannelInterface, Any], + channel_idx: Optional[int], + ): + self.input_variant = input_variant + self.channel_idx = channel_idx + + def resolve(self, channel_results: Any) -> Any: + """ + Resolve the input value from the channel results. + + Args: + channel_results: The results from reading the input channels. + """ + + if isinstance(self.input_variant, ChannelInterface): + value = channel_results[self.channel_idx] + else: + value = self.input_variant + return value + + +@DeveloperAPI +class ExecutableTask: + """A task that can be executed in a compiled DAG, and it + corresponds to an actor method. + """ + + def __init__( + self, + task: "CompiledTask", + resolved_args: List[Any], + resolved_kwargs: Dict[str, Any], + ): + """ + Args: + task: The CompiledTask that this ExecutableTask corresponds to. + resolved_args: The arguments to the method. Arguments that are + not Channels will get passed through to the actor method. + If the argument is a channel, it will be replaced by the + value read from the channel before the method executes. + resolved_kwargs: The keyword arguments to the method. Currently, we + do not support binding kwargs to other DAG nodes, so the values + of the dictionary cannot be Channels. + """ + from ray.dag import CollectiveOutputNode + + self.method_name = task.dag_node.get_method_name() + self.bind_index = task.dag_node._get_bind_index() + self.output_channels = task.output_channels + self.output_idxs = task.output_idxs + self.input_type_hints: List[ChannelOutputType] = task.arg_type_hints + self.output_type_hint: ChannelOutputType = task.dag_node.type_hint + + # The accelerator collective operation. + self.collective_op: Optional["ray.dag.CollectiveOperation"] = None + if isinstance(task.dag_node, CollectiveOutputNode): + self.collective_op = task.dag_node.collective_op + + self.input_channels: List[ChannelInterface] = [] + self.task_inputs: List[_ExecutableTaskInput] = [] + self.resolved_kwargs: Dict[str, Any] = resolved_kwargs + # A unique index which can be used to index into `idx_to_task` to get + # the corresponding task. + self.task_idx = task.idx + + # Reverse map for input_channels: maps an input channel to + # its index in input_channels. + input_channel_to_idx: dict[ChannelInterface, int] = {} + + for arg in resolved_args: + if isinstance(arg, ChannelInterface): + channel = arg + if channel in input_channel_to_idx: + # The same channel was added before, so reuse the index. + channel_idx = input_channel_to_idx[channel] + else: + # Add a new channel to the list of input channels. + self.input_channels.append(channel) + channel_idx = len(self.input_channels) - 1 + input_channel_to_idx[channel] = channel_idx + + task_input = _ExecutableTaskInput(arg, channel_idx) + else: + task_input = _ExecutableTaskInput(arg, None) + self.task_inputs.append(task_input) + + # Currently DAGs do not support binding kwargs to other DAG nodes. + for val in self.resolved_kwargs.values(): + assert not isinstance(val, ChannelInterface) + + # Input reader to read input data from upstream DAG nodes. + self.input_reader: ReaderInterface = SynchronousReader(self.input_channels) + # Output writer to write output data to downstream DAG nodes. + self.output_writer: WriterInterface = SynchronousWriter( + self.output_channels, self.output_idxs + ) + # The intermediate future for a READ or COMPUTE operation, + # and `wait()` must be called to get the actual result of the operation. + # The result of a READ operation will be used by a COMPUTE operation, + # and the result of a COMPUTE operation will be used by a WRITE operation. + self._intermediate_future: Optional[DAGOperationFuture] = None + + def cancel(self): + """ + Close all the input channels and the output channel. The exact behavior + depends on the type of channel. Typically, it will release the resources + used by the channels. + """ + self.input_reader.close() + self.output_writer.close() + + def destroy_cuda_event(self): + """ + If this executable task has created a GPU future that is not yet waited on, + that future is in the channel context cache. Remove the future from the cache + and destroy its CUDA event. + """ + GPUFuture.remove_gpu_future(self.task_idx) + + def prepare(self, overlap_gpu_communication: bool = False): + """ + Prepare the task for execution. The `exec_operation` function can only + be called after `prepare` has been called. + + Args: + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance + """ + for typ_hint in self.input_type_hints: + typ_hint.register_custom_serializer() + self.output_type_hint.register_custom_serializer() + self.input_reader.start() + self.output_writer.start() + + # Stream context type are different between different accelerators. + # Type hint is not applicable here. + self._send_stream = nullcontext() + self._recv_stream = nullcontext() + if not overlap_gpu_communication: + return + + # Set up send_stream and recv_stream when overlap_gpu_communication + # is configured + if self.output_type_hint.requires_accelerator(): + comm_group_id = _get_comm_group_id(self.output_type_hint) + comm_group = ChannelContext.get_current().communicators.get(comm_group_id) + assert comm_group is not None + self._send_stream = comm_group.send_stream + if self.input_type_hints: + for type_hint in self.input_type_hints: + if type_hint.requires_accelerator(): + comm_group_id = _get_comm_group_id(type_hint) + comm_group = ChannelContext.get_current().communicators.get( + comm_group_id + ) + assert comm_group is not None + if not isinstance(self._recv_stream, nullcontext): + assert self._recv_stream == comm_group.recv_stream, ( + "Currently all torch tensor input channels of a " + "Compiled Graph task should use the same recv cuda stream." + ) + self._recv_stream = comm_group.recv_stream + + def wrap_and_set_intermediate_future( + self, val: Any, wrap_in_gpu_future: bool + ) -> None: + """ + Wrap the value in a `DAGOperationFuture` and store to the intermediate future. + The value corresponds to result of a READ or COMPUTE operation. + + If wrap_in_gpu_future is True, the value will be wrapped in a GPUFuture, + Otherwise, the future will be a ResolvedFuture. + + Args: + val: The value to wrap in a future. + wrap_in_gpu_future: Whether to wrap the value in a GPUFuture. + """ + assert self._intermediate_future is None + + if wrap_in_gpu_future: + future = GPUFuture(val, self.task_idx) + else: + future = ResolvedFuture(val) + self._intermediate_future = future + + def reset_and_wait_intermediate_future(self) -> Any: + """ + Reset the intermediate future and wait for the result. + + The wait does not block the CPU because: + - If the future is a ResolvedFuture, the result is immediately returned. + - If the future is a GPUFuture, the result is only waited by the current + CUDA stream, and the CPU is not blocked. + + Returns: + The result of a READ or COMPUTE operation from the intermediate future. + """ + future = self._intermediate_future + self._intermediate_future = None + return future.wait() + + def _read(self, overlap_gpu_communication: bool) -> bool: + """ + Read input data from upstream DAG nodes and cache the intermediate result. + + Args: + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance. + + Returns: + True if system error occurs and exit the loop; otherwise, False. + """ + assert self._intermediate_future is None + exit = False + try: + input_data = self.input_reader.read() + # When overlap_gpu_communication is enabled, wrap the result in + # a GPUFuture so that this read operation (communication) can + # be overlapped with computation. + self.wrap_and_set_intermediate_future( + input_data, + wrap_in_gpu_future=overlap_gpu_communication, + ) + except RayChannelError: + # Channel closed. Exit the loop. + exit = True + return exit + + def _compute( + self, + overlap_gpu_communication: bool, + class_handle, + ) -> bool: + """ + Retrieve the intermediate result from the READ operation and perform the + computation. Then, cache the new intermediate result. The caller must ensure + that the last operation executed is READ so that the function retrieves the + correct intermediate result. + + Args: + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance. + class_handle: An instance of the class to which the actor belongs. For + example, the type of `class_handle` is if the + actor belongs to the `class Worker` class. + Returns: + True if system error occurs and exit the loop; otherwise, False. + """ + input_data = self.reset_and_wait_intermediate_future() + try: + _process_return_vals(input_data, return_single_output=False) + except Exception as exc: + # Previous task raised an application-level exception. + # Propagate it and skip the actual task. We don't need to wrap the + # exception in a RayTaskError here because it has already been wrapped + # by the previous task. + self.wrap_and_set_intermediate_future( + exc, wrap_in_gpu_future=overlap_gpu_communication + ) + return False + + resolved_inputs = [] + for task_input in self.task_inputs: + resolved_inputs.append(task_input.resolve(input_data)) + + if self.collective_op is not None: + # Run an accelerator collective operation. + method = self.collective_op.execute + else: + # Run an actor method. + method = getattr(class_handle, self.method_name) + try: + output_val = method(*resolved_inputs, **self.resolved_kwargs) + except Exception as exc: + output_val = _wrap_exception(exc) + + # When overlap_gpu_communication is enabled, wrap the result in a GPUFuture + # so that this compute operation can be overlapped with communication. + self.wrap_and_set_intermediate_future( + output_val, wrap_in_gpu_future=overlap_gpu_communication + ) + return False + + def _write(self) -> bool: + """ + Retrieve the intermediate result from the COMPUTE operation and write to its + downstream DAG nodes. The caller must ensure that the last operation executed + is COMPUTE so that the function retrieves the correct intermediate result. + + Returns: + True if system error occurs and exit the loop; otherwise, False. + """ + output_val = self.reset_and_wait_intermediate_future() + exit = False + try: + self.output_writer.write(output_val) + except RayChannelError: + # Channel closed. Exit the loop. + exit = True + return exit + + def exec_operation( + self, + class_handle, + op_type: _DAGNodeOperationType, + overlap_gpu_communication: bool = False, + ) -> bool: + """ + An ExecutableTask corresponds to a DAGNode. It consists of three + operations: READ, COMPUTE, and WRITE, which should be executed in + order to ensure that each operation can read the correct intermediate + result. + Args: + class_handle: The handle of the class to which the actor belongs. + op_type: The type of the operation. Possible types are READ, + COMPUTE, and WRITE. + overlap_gpu_communication: Whether to overlap GPU communication with + computation during DAG execution to improve performance. + Returns: + True if the next operation should not be executed; otherwise, False. + """ + if op_type == _DAGNodeOperationType.READ: + with _device_context_manager(): + with self._recv_stream: + return self._read(overlap_gpu_communication) + elif op_type == _DAGNodeOperationType.COMPUTE: + return self._compute(overlap_gpu_communication, class_handle) + elif op_type == _DAGNodeOperationType.WRITE: + with _device_context_manager(): + with self._send_stream: + return self._write() + + +@dataclass +class _ExecutableTaskRecord: + actor_classname: str + actor_name: str + actor_id: str + method_name: str + bind_index: int + operation: str + start_t: float + end_t: float + + def to_dict(self): + return asdict(self) + + +@DeveloperAPI +class CompiledDAG: + """Experimental class for accelerated execution. + + This class should not be called directly. Instead, create + a ray.dag and call experimental_compile(). + + See REP https://github.com/ray-project/enhancements/pull/48 for more + information. + """ + + @ray.remote(num_cpus=0) + class DAGDriverProxyActor: + """ + To support the driver as a reader, the output writer needs to be able to invoke + remote functions on the driver. This is necessary so that the output writer can + create a reader ref on the driver node, and later potentially create a larger + reader ref on the driver node if the channel backing store needs to be resized. + However, remote functions cannot be invoked on the driver. + + A Compiled Graph creates an actor from this class when the DAG is initialized. + The actor is on the same node as the driver. This class has an empty + implementation, though it serves as a way for the output writer to invoke remote + functions on the driver node. + """ + + pass + + def __init__( + self, + submit_timeout: Optional[float] = None, + buffer_size_bytes: Optional[int] = None, + enable_asyncio: bool = False, + max_inflight_executions: Optional[int] = None, + max_buffered_results: Optional[int] = None, + overlap_gpu_communication: Optional[bool] = None, + default_communicator: Optional[Union[Communicator, str]] = "create", + ): + """ + Args: + submit_timeout: The maximum time in seconds to wait for execute() calls. + None means using default timeout (DAGContext.submit_timeout), + 0 means immediate timeout (immediate success or timeout without + blocking), -1 means infinite timeout (block indefinitely). + buffer_size_bytes: The initial buffer size in bytes for messages + that can be passed between tasks in the DAG. The buffers will + be automatically resized if larger messages are written to the + channel. + enable_asyncio: Whether to enable asyncio. If enabled, caller must + be running in an event loop and must use `execute_async` to + invoke the DAG. Otherwise, the caller should use `execute` to + invoke the DAG. + max_inflight_executions: The maximum number of in-flight executions that + can be submitted via `execute` or `execute_async` before consuming + the output using `ray.get()`. If the caller submits more executions, + `RayCgraphCapacityExceeded` is raised. + max_buffered_results: The maximum number of results that can be + buffered at the driver. If more results are buffered, + `RayCgraphCapacityExceeded` is raised. Note that + when result corresponding to an execution is retrieved + (by calling `ray.get()` on a `CompiledDAGRef` or + `CompiledDAGRef` or await on a `CompiledDAGFuture), results + corresponding to earlier executions that have not been retrieved + yet are buffered. + overlap_gpu_communication: (experimental) Whether to overlap GPU + communication with computation during DAG execution. If True, the + communication and computation can be overlapped, which can improve + the performance of the DAG execution. If None, the default value + will be used. + _default_communicator: The default communicator to use to transfer + tensors. Three types of values are valid. (1) Communicator: + For p2p operations, this is the default communicator + to use for nodes annotated with `with_tensor_transport()` and when + shared memory is not the desired option (e.g., when transport="accelerator", + or when transport="auto" for communication between two different GPUs). + For collective operations, this is the default communicator to use + when a custom communicator is not specified. + (2) "create": for each collective operation without a custom communicator + specified, a communicator is created and initialized on its involved actors, + or an already created communicator is reused if the set of actors is the same. + For all p2p operations without a custom communicator specified, it reuses + an already created collective communicator if the p2p actors are a subset. + Otherwise, a new communicator is created. + (3) None: a ValueError will be thrown if a custom communicator is not specified. + + Returns: + Channel: A wrapper around ray.ObjectRef. + """ + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + + self._enable_asyncio: bool = enable_asyncio + self._fut_queue = asyncio.Queue() + self._max_inflight_executions = max_inflight_executions + if self._max_inflight_executions is None: + self._max_inflight_executions = ctx.max_inflight_executions + self._max_buffered_results = max_buffered_results + if self._max_buffered_results is None: + self._max_buffered_results = ctx.max_buffered_results + self._dag_id = uuid.uuid4().hex + self._submit_timeout: Optional[float] = submit_timeout + if self._submit_timeout is None: + self._submit_timeout = ctx.submit_timeout + self._get_timeout: Optional[float] = ctx.get_timeout + self._buffer_size_bytes: Optional[int] = buffer_size_bytes + if self._buffer_size_bytes is None: + self._buffer_size_bytes = ctx.buffer_size_bytes + self._overlap_gpu_communication: Optional[bool] = overlap_gpu_communication + if self._overlap_gpu_communication is None: + self._overlap_gpu_communication = ctx.overlap_gpu_communication + self._create_default_communicator = False + if isinstance(default_communicator, str): + if default_communicator == "create": + self._create_default_communicator = True + default_communicator = None + else: + raise ValueError( + "The only allowed string for default_communicator is 'create', " + f"got {default_communicator}" + ) + elif default_communicator is not None and not isinstance( + default_communicator, Communicator + ): + raise ValueError( + "The default_communicator must be None, a string, or a Communicator, " + f"got {type(default_communicator)}" + ) + self._default_communicator: Optional[Communicator] = default_communicator + + # Dict from passed-in communicator to set of type hints that refer to it. + self._communicator_to_type_hints: Dict[ + Communicator, + Set["ray.experimental.channel.torch_tensor_type.TorchTensorType"], + ] = defaultdict(set) + # Dict from set of actors to created communicator ID. + # These communicators are created by Compiled Graph, rather than passed in. + # Communicators are only created when self._create_default_communicator is True. + self._actors_to_created_communicator_id: Dict[ + Tuple["ray.actor.ActorHandle"], str + ] = {} + + # Set of actors involved in P2P communication using an unresolved communicator. + self._p2p_actors_with_unresolved_communicators: Set[ + "ray.actor.ActorHandle" + ] = set() + # Set of DAG nodes involved in P2P communication using an unresolved communicator. + self._p2p_dag_nodes_with_unresolved_communicators: Set[ + "ray.dag.DAGNode" + ] = set() + # Set of collective operations using an unresolved communicator. + self._collective_ops_with_unresolved_communicators: Set[ + "ray.dag.collective_node._CollectiveOperation" + ] = set() + + self._default_type_hint: ChannelOutputType = SharedMemoryType( + buffer_size_bytes=self._buffer_size_bytes, + # We conservatively set num_shm_buffers to _max_inflight_executions. + # It means that the DAG can be underutilized, but it guarantees there's + # no false positive timeouts. + num_shm_buffers=self._max_inflight_executions, + ) + if not isinstance(self._buffer_size_bytes, int) or self._buffer_size_bytes <= 0: + raise ValueError( + "`buffer_size_bytes` must be a positive integer, found " + f"{self._buffer_size_bytes}" + ) + + # Used to ensure that the future returned to the + # caller corresponds to the correct DAG output. I.e. + # order of futures added to fut_queue should match the + # order of inputs written to the DAG. + self._dag_submission_lock = asyncio.Lock() + + # idx -> CompiledTask. + self.idx_to_task: Dict[int, "CompiledTask"] = {} + # DAGNode -> idx. + self.dag_node_to_idx: Dict["ray.dag.DAGNode", int] = {} + # idx counter. + self.counter: int = 0 + + # Attributes that are set during preprocessing. + # Preprocessing identifies the input node and output node. + self.input_task_idx: Optional[int] = None + self.output_task_idx: Optional[int] = None + # List of task indices that are input attribute nodes. + self.input_attr_task_idxs: List[int] = [] + # Denotes whether execute/execute_async returns a list of refs/futures. + self._returns_list: bool = False + # Number of expected positional args and kwargs that may be passed to + # dag.execute. + self._input_num_positional_args: Optional[int] = None + self._input_kwargs: Tuple[str, ...] = None + + # Cached attributes that are set during compilation. + self.dag_input_channels: Optional[List[ChannelInterface]] = None + self.dag_output_channels: Optional[List[ChannelInterface]] = None + self._dag_submitter: Optional[WriterInterface] = None + self._dag_output_fetcher: Optional[ReaderInterface] = None + + # ObjectRef for each worker's task. The task is an infinite loop that + # repeatedly executes the method specified in the DAG. + self.worker_task_refs: Dict["ray.actor.ActorHandle", "ray.ObjectRef"] = {} + self.actor_to_tasks: Dict[ + "ray.actor.ActorHandle", List["CompiledTask"] + ] = defaultdict(list) + # Mapping from actor handle to its GPU IDs. + # This is used for type hint resolution for with_tensor_transport("auto"). + self.actor_to_gpu_ids: Dict["ray.actor.ActorHandle", List[str]] = {} + self.actor_to_executable_tasks: Dict[ + "ray.actor.ActorHandle", List["ExecutableTask"] + ] = {} + # Mapping from the actor handle to the execution schedule which is a list + # of operations to be executed. + self.actor_to_execution_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGNodeOperation] + ] = defaultdict(list) + # Mapping from the actor handle to the node ID that the actor is on. + # A None actor handle means the actor is the driver. + self.actor_to_node_id: Dict[Optional["ray.actor.ActorHandle"], str] = {} + # The index of the current execution. It is incremented each time + # the DAG is executed. + self._execution_index: int = -1 + # The maximum index of finished executions. + # All results with higher indexes have not been generated yet. + self._max_finished_execution_index: int = -1 + # execution_index -> {channel_index -> result} + self._result_buffer: Dict[int, Dict[int, Any]] = defaultdict(dict) + # channel to possible inner channel + self._channel_dict: Dict[ChannelInterface, ChannelInterface] = {} + + def _create_proxy_actor() -> "ray.actor.ActorHandle": + # Creates the driver actor on the same node as the driver. + # + # To support the driver as a reader, the output writer needs to be able to + # invoke remote functions on the driver (e.g., to create the reader ref, to + # create a reader ref for a larger object when the channel backing store is + # resized, etc.). The driver actor serves as a way for the output writer + # to invoke remote functions on the driver node. + return CompiledDAG.DAGDriverProxyActor.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + ray.get_runtime_context().get_node_id(), soft=False + ) + ).remote() + + self._proxy_actor = _create_proxy_actor() + # Set to True when `teardown` API is called. + self._is_teardown = False + # Execution index to set of channel indices for CompiledDAGRefs + # or CompiledDAGFuture whose destructor has been called. A "None" + # channel index means there is only one channel, and its destructor + # has been called. + self._destructed_ref_idxs: Dict[int, Set[Optional[int]]] = dict() + # Execution index to set of channel indices for CompiledDAGRefs + # or CompiledDAGFuture whose get() has been called. A "None" + # channel index means there is only one channel, and its get() + # has been called. + self._got_ref_idxs: Dict[int, Set[Optional[int]]] = dict() + + @property + def is_teardown(self) -> bool: + return self._is_teardown + + def get_id(self) -> str: + """ + Get the unique ID of the compiled DAG. + """ + return self._dag_id + + def __str__(self) -> str: + return f"CompiledDAG({self._dag_id})" + + def _add_node(self, node: "ray.dag.DAGNode") -> None: + idx = self.counter + self.idx_to_task[idx] = CompiledTask(idx, node) + self.dag_node_to_idx[node] = idx + self.counter += 1 + + def _preprocess(self) -> None: + """Before compiling, preprocess the DAG to build an index from task to + upstream and downstream tasks, and to set the input and output node(s) + of the DAG. + + This function is idempotent. + """ + from ray.dag import ( + ClassMethodNode, + CollectiveOutputNode, + DAGNode, + FunctionNode, + InputAttributeNode, + InputNode, + MultiOutputNode, + ) + + self.input_task_idx, self.output_task_idx = None, None + + input_attributes: Set[str] = set() + # Find the input node and input attribute nodes in the DAG. + for idx, task in self.idx_to_task.items(): + if isinstance(task.dag_node, InputNode): + assert self.input_task_idx is None, "More than one InputNode found" + self.input_task_idx = idx + # handle_unused_attributes: + # Save input attributes in a set. + input_node = task.dag_node + input_attributes.update(input_node.input_attribute_nodes.keys()) + elif isinstance(task.dag_node, InputAttributeNode): + self.input_attr_task_idxs.append(idx) + + # Find the (multi-)output node to the DAG. + for idx, task in self.idx_to_task.items(): + if idx == self.input_task_idx or isinstance( + task.dag_node, InputAttributeNode + ): + continue + if ( + len(task.downstream_task_idxs) == 0 + and task.dag_node.is_cgraph_output_node + ): + assert self.output_task_idx is None, "More than one output node found" + self.output_task_idx = idx + + assert self.output_task_idx is not None + output_node = self.idx_to_task[self.output_task_idx].dag_node + # Add an MultiOutputNode to the end of the DAG if it's not already there. + if not isinstance(output_node, MultiOutputNode): + output_node = MultiOutputNode([output_node]) + self._add_node(output_node) + self.output_task_idx = self.dag_node_to_idx[output_node] + else: + self._returns_list = True + + # TODO: Support no-input DAGs (use an empty object to signal). + if self.input_task_idx is None: + raise NotImplementedError( + "Compiled DAGs currently require exactly one InputNode" + ) + + # Whether the DAG binds directly to the InputNode(), versus binding to + # a positional arg or kwarg of the input. For example, a.foo.bind(inp) + # instead of a.foo.bind(inp[0]) or a.foo.bind(inp.key). + direct_input: Optional[bool] = None + # Collect the set of InputNode keys bound to DAG node args. + input_positional_args: Set[int] = set() + input_kwargs: Set[str] = set() + # Set of tasks with annotation of with_tensor_transport("auto"). + # These only correspond to ClassMethodNodes, but not InputNodes + # or InputAttributeNodes. + auto_transport_tasks: Set["CompiledTask"] = set() + + # For each task node, set its upstream and downstream task nodes. + # Also collect the set of tasks that produce torch.tensors. + for task_idx, task in self.idx_to_task.items(): + dag_node = task.dag_node + if not ( + isinstance(dag_node, InputNode) + or isinstance(dag_node, InputAttributeNode) + or isinstance(dag_node, MultiOutputNode) + or isinstance(dag_node, ClassMethodNode) + ): + if isinstance(dag_node, FunctionNode): + # TODO(swang): Support non-actor tasks. + raise NotImplementedError( + "Compiled DAGs currently only support actor method nodes" + ) + else: + raise ValueError(f"Found unsupported node of type {type(dag_node)}") + + if isinstance(dag_node, ClassMethodNode) and dag_node.is_class_method_call: + actor_handle = dag_node._get_actor_handle() + if actor_handle is None: + raise ValueError( + "Compiled DAGs can only bind methods to an actor " + "that is already created with Actor.remote()" + ) + + if actor_handle not in self.actor_to_gpu_ids: + self.actor_to_gpu_ids[actor_handle] = CompiledDAG._get_gpu_ids( + actor_handle + ) + + if isinstance(dag_node.type_hint, AutoTransportType): + auto_transport_tasks.add(task) + + # Collect actors for accelerator P2P methods. + if dag_node.type_hint.requires_accelerator(): + self._track_communicator_usage(dag_node, {actor_handle}) + # Collect accelerator collective operations. + if isinstance(dag_node, CollectiveOutputNode): + self._track_communicator_usage( + dag_node, + set(dag_node._collective_op.actor_handles), + collective_op=True, + ) + assert not self._overlap_gpu_communication, ( + "Currently, the overlap_gpu_communication option is not " + "supported for accelerator collective operations. Please set " + "overlap_gpu_communication=False." + ) + elif isinstance(dag_node, InputNode) or isinstance( + dag_node, InputAttributeNode + ): + if dag_node.type_hint.requires_accelerator(): + raise ValueError( + "DAG inputs cannot be transferred via accelerator because " + "the driver cannot participate in the communicator group" + ) + if isinstance(dag_node.type_hint, AutoTransportType): + # Currently driver on GPU is not supported, so we always + # use shared memory to transfer tensors. + dag_node.type_hint = TorchTensorType( + device=dag_node.type_hint.device + ) + + if type(dag_node.type_hint) is ChannelOutputType: + # No type hint specified by the user. Replace + # with the default type hint for this DAG. + dag_node.type_hint = self._default_type_hint + + for _, val in task.kwargs.items(): + if isinstance(val, DAGNode): + raise ValueError( + "Compiled DAG currently does not support binding to " + "other DAG nodes as kwargs" + ) + + for _, arg in enumerate(task.args): + if not isinstance(arg, DAGNode): + continue + upstream_node_idx = self.dag_node_to_idx[arg] + upstream_task = self.idx_to_task[upstream_node_idx] + downstream_actor_handle = None + if ( + isinstance(dag_node, ClassMethodNode) + and dag_node.is_class_method_call + ): + downstream_actor_handle = dag_node._get_actor_handle() + + # Add upstream node as the argument nodes of this task, whose + # type hints may be updated when resolved lazily. + task.arg_nodes.append(upstream_task.dag_node) + + if isinstance(upstream_task.dag_node, InputAttributeNode): + # Record all of the keys used to index the InputNode. + # During execution, we will check that the user provides + # the same args and kwargs. + if isinstance(upstream_task.dag_node.key, int): + input_positional_args.add(upstream_task.dag_node.key) + elif isinstance(upstream_task.dag_node.key, str): + input_kwargs.add(upstream_task.dag_node.key) + else: + raise ValueError( + "InputNode() can only be indexed using int " + "for positional args or str for kwargs." + ) + + if direct_input is not None and direct_input: + raise ValueError( + "All tasks must either use InputNode() " + "directly, or they must index to specific args or " + "kwargs." + ) + direct_input = False + + # If the upstream node is an InputAttributeNode, treat the + # DAG's input node as the actual upstream node + upstream_task = self.idx_to_task[self.input_task_idx] + + elif isinstance(upstream_task.dag_node, InputNode): + if direct_input is not None and not direct_input: + raise ValueError( + "All tasks must either use InputNode() directly, " + "or they must index to specific args or kwargs." + ) + direct_input = True + + upstream_task.downstream_task_idxs[task_idx] = downstream_actor_handle + + if upstream_task.dag_node.type_hint.requires_accelerator(): + # Here we are processing the args of the DAGNode, so track + # downstream actors only, upstream actor is already tracked + # when processing the DAGNode itself. + self._track_communicator_usage( + upstream_task.dag_node, + {downstream_actor_handle}, + ) + # Check that all specified input attributes, e.g., InputNode()["x"], + # are used in the DAG. + _check_unused_dag_input_attributes(output_node, input_attributes) + + self._check_leaf_nodes() + + self._resolve_auto_transport(auto_transport_tasks) + + self._init_communicators() + + if direct_input: + self._input_num_positional_args = 1 + elif not input_positional_args: + self._input_num_positional_args = 0 + else: + self._input_num_positional_args = max(input_positional_args) + 1 + self._input_kwargs = tuple(input_kwargs) + + def _init_communicators(self) -> None: + """ + Initialize communicators for the DAG. + """ + + # First, initialize communicators that are passed in by the user. + for communicator, type_hints in self._communicator_to_type_hints.items(): + communicator_id = _init_communicator( + communicator.get_actor_handles(), + communicator, + self._overlap_gpu_communication, + ) + for type_hint in type_hints: + type_hint.set_communicator_id(communicator_id) + + # Second, get registered accelerator context if any. + accelerator_module_name = AcceleratorContext.get().module_name + accelerator_communicator_cls = AcceleratorContext.get().communicator_cls + + # Then, create communicators for collective operations. + # Reuse an already created communicator for the same set of actors. + for collective_op in self._collective_ops_with_unresolved_communicators: + if not self._create_default_communicator: + raise ValueError( + "Communicator creation is not allowed for collective operations." + ) + # using tuple to preserve the order of actors for collective operations + actors = tuple(collective_op.actor_handles) + if actors in self._actors_to_created_communicator_id: + communicator_id = self._actors_to_created_communicator_id[actors] + else: + communicator_id = _init_communicator( + list(actors), + None, + self._overlap_gpu_communication, + accelerator_module_name, + accelerator_communicator_cls, + ) + self._actors_to_created_communicator_id[actors] = communicator_id + collective_op.type_hint.set_communicator_id(communicator_id) + + # Finally, create a communicator for P2P operations. + # Reuse an already created collective op communicator when p2p actors + # are a subset of the actors in the collective op communicator. + p2p_communicator_id = None + if self._p2p_actors_with_unresolved_communicators: + for ( + actors, + communicator_id, + ) in self._actors_to_created_communicator_id.items(): + if self._p2p_actors_with_unresolved_communicators.issubset(actors): + p2p_communicator_id = communicator_id + break + if p2p_communicator_id is None: + p2p_communicator_id = _init_communicator( + list(self._p2p_actors_with_unresolved_communicators), + None, + self._overlap_gpu_communication, + accelerator_module_name, + accelerator_communicator_cls, + ) + for dag_node in self._p2p_dag_nodes_with_unresolved_communicators: + dag_node.type_hint.set_communicator_id(p2p_communicator_id) + + def _track_communicator_usage( + self, + dag_node: "ray.dag.DAGNode", + actors: Set["ray.actor.ActorHandle"], + collective_op: bool = False, + ) -> None: + """ + Track the usage of a communicator. + + This method first determines the communicator to use: if a custom + communicator is specified, use it; if not and a default communicator + is available, use it; otherwise, it records necessary information to + create a new communicator later. + + This method also performs validation checks on the passed-in communicator. + + Args: + dag_node: The DAG node that uses the communicator, this is the node + that has the `with_tensor_transport()` type hint for p2p communication, + or a `CollectiveOutputNode` for collective operations. + actors: The full or partial set of actors that use the communicator. + This method should be called one or multiple times so that all actors + of the communicator are tracked. + collective_op: Whether the communicator is used for a collective operation. + """ + if None in actors: + raise ValueError("Driver cannot participate in the communicator group.") + if collective_op: + type_hint = dag_node._collective_op.type_hint + else: + type_hint = dag_node.type_hint + communicator = type_hint.get_custom_communicator() + + if communicator is None: + if ( + self._default_communicator is None + and not self._create_default_communicator + ): + if dag_node._original_type_hint is not None: + assert isinstance(dag_node._original_type_hint, AutoTransportType) + raise ValueError( + f"with_tensor_transport(transport='auto') is used for DAGNode {dag_node}, " + "This requires specifying a default communicator or 'create' for " + "_default_communicator when calling experimental_compile()." + ) + raise ValueError( + f"DAGNode {dag_node} has no custom communicator specified. " + "Please specify a custom communicator for the DAGNode using " + "`with_tensor_transport()`, or specify a communicator or 'create' for " + "_default_communicator when calling experimental_compile()." + ) + communicator = self._default_communicator + + if communicator is None: + if collective_op: + self._collective_ops_with_unresolved_communicators.add( + dag_node._collective_op + ) + else: + self._p2p_dag_nodes_with_unresolved_communicators.add(dag_node) + self._p2p_actors_with_unresolved_communicators.update(actors) + else: + if collective_op: + if set(communicator.get_actor_handles()) != actors: + raise ValueError( + "The passed-in communicator must have the same set " + "of actors as the collective operation. " + f"The passed-in communicator has actors {communicator.get_actor_handles()} " + f"while the collective operation has actors {actors}." + ) + else: + if not actors.issubset(set(communicator.get_actor_handles())): + raise ValueError( + "The passed-in communicator must include all of the actors " + "used in the P2P operation. " + f"The passed-in communicator has actors {communicator.get_actor_handles()} " + f"while the P2P operation has actors {actors}." + ) + self._communicator_to_type_hints[communicator].add(type_hint) + + def _resolve_auto_transport( + self, + auto_transport_tasks: Set["CompiledTask"], + ) -> None: + """ + Resolve the auto transport type hint for the DAG. + """ + type_hint_resolver = TypeHintResolver(self.actor_to_gpu_ids) + # Resolve AutoChannelType type hints and track the actors that use accelerator. + # This is needed so that the communicator group can be initialized for + # these actors that use accelerator. + for task in auto_transport_tasks: + writer = task.dag_node._get_actor_handle() + readers = task.downstream_task_idxs.values() + writer_and_node = (writer, self._get_node_id(writer)) + reader_and_node_list = [ + (reader, self._get_node_id(reader)) for reader in readers + ] + # Update the type hint to the resolved one. This is needed because + # the resolved type hint's `register_custom_serializer` will be called + # in preparation for channel I/O. + task.dag_node.type_hint = type_hint_resolver.resolve( + task.dag_node.type_hint, + writer_and_node, + reader_and_node_list, + ) + if task.dag_node.type_hint.requires_accelerator(): + self._track_communicator_usage( + task.dag_node, + set(readers).union({writer}), + ) + + def _check_leaf_nodes(self) -> None: + """ + Check if there are leaf nodes in the DAG and raise an error if there are. + """ + from ray.dag import ( + ClassMethodNode, + DAGNode, + ) + + leaf_nodes: List[DAGNode] = [] + for _, task in self.idx_to_task.items(): + if not isinstance(task.dag_node, ClassMethodNode): + continue + if ( + len(task.downstream_task_idxs) == 0 + and not task.dag_node.is_cgraph_output_node + ): + leaf_nodes.append(task.dag_node) + # Leaf nodes are not allowed because the exception thrown by the leaf + # node will not be propagated to the driver. + if len(leaf_nodes) != 0: + raise ValueError( + "Compiled DAG doesn't support leaf nodes, i.e., nodes that don't have " + "downstream nodes and are not output nodes. There are " + f"{len(leaf_nodes)} leaf nodes in the DAG. Please add the outputs of " + f"{[leaf_node.get_method_name() for leaf_node in leaf_nodes]} to the " + f"the MultiOutputNode." + ) + + @staticmethod + def _get_gpu_ids(actor_handle: "ray.actor.ActorHandle") -> List[str]: + """ + Get the GPU IDs of an actor handle. + """ + accelerator_ids = ray.get( + actor_handle.__ray_call__.remote( + lambda self: ray.get_runtime_context().get_accelerator_ids() + ) + ) + return accelerator_ids.get("GPU", []) + + def _get_node_id(self, actor_handle: Optional["ray.actor.ActorHandle"]) -> str: + """ + Get the node ID of an actor handle and cache it. + + Args: + actor_handle: The actor handle, or None if the actor handle is the + driver. + Returns: + The node ID of the actor handle or driver. + """ + if actor_handle in self.actor_to_node_id: + return self.actor_to_node_id[actor_handle] + node_id = None + if actor_handle == self._proxy_actor or actor_handle is None: + node_id = ray.get_runtime_context().get_node_id() + else: + node_id = ray.get( + actor_handle.__ray_call__.remote( + lambda self: ray.get_runtime_context().get_node_id() + ) + ) + self.actor_to_node_id[actor_handle] = node_id + return node_id + + def _get_or_compile( + self, + ) -> None: + """Compile an execution path. This allocates channels for adjacent + tasks to send/receive values. An infinite task is submitted to each + actor in the DAG that repeatedly receives from input channel(s) and + sends to output channel(s). + + This function is idempotent and will cache the previously allocated + channels. After calling this function, _dag_submitter and + _dag_output_fetcher will be set and can be used to invoke and fetch + outputs for the DAG. + """ + from ray.dag import ( + ClassMethodNode, + DAGNode, + InputAttributeNode, + InputNode, + MultiOutputNode, + ) + + if self.input_task_idx is None: + self._preprocess() + assert self.input_task_idx is not None + + if self._dag_submitter is not None: + assert self._dag_output_fetcher is not None + return + + frontier = [self.input_task_idx] + visited = set() + # Create output buffers. This loop does a breadth-first search through the DAG. + while frontier: + cur_idx = frontier.pop(0) + if cur_idx in visited: + continue + visited.add(cur_idx) + + task = self.idx_to_task[cur_idx] + if ( + isinstance(task.dag_node, ClassMethodNode) + and task.dag_node.is_class_method_call + ): + # Create output buffers for the actor method. + assert len(task.output_channels) == 0 + # `output_to_readers` stores the reader tasks for each output of + # the current node. If the current node returns one output, the + # readers are the downstream nodes of the current node. If the + # current node returns multiple outputs, the readers of each + # output are the downstream nodes of the ClassMethodNode that + # is a class method output. + output_to_readers: Dict[CompiledTask, List[CompiledTask]] = defaultdict( + list + ) + for idx in task.downstream_task_idxs: + downstream_task = self.idx_to_task[idx] + downstream_node = downstream_task.dag_node + if ( + isinstance(downstream_node, ClassMethodNode) + and downstream_node.is_class_method_output + ): + output_to_readers[downstream_task] = [ + self.idx_to_task[idx] + for idx in downstream_task.downstream_task_idxs + ] + else: + if task not in output_to_readers: + output_to_readers[task] = [] + output_to_readers[task].append(downstream_task) + fn = task.dag_node._get_remote_method("__ray_call__") + for output, readers in output_to_readers.items(): + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]] = [] + # Use reader_handles_set to deduplicate readers on the + # same actor, because with CachedChannel each actor will + # only read from the upstream channel once. + reader_handles_set = set() + read_by_multi_output_node = False + for reader in readers: + if isinstance(reader.dag_node, MultiOutputNode): + read_by_multi_output_node = True + # inserting at 0 to make sure driver is first reader as + # expected by CompositeChannel read + reader_and_node_list.insert( + 0, + ( + self._proxy_actor, + self._get_node_id(self._proxy_actor), + ), + ) + else: + reader_handle = reader.dag_node._get_actor_handle() + if reader_handle not in reader_handles_set: + reader_handle = reader.dag_node._get_actor_handle() + reader_and_node_list.append( + (reader_handle, self._get_node_id(reader_handle)) + ) + reader_handles_set.add(reader_handle) + + # if driver is an actual actor, gets driver actor id + driver_actor_id = ( + ray.get_runtime_context().get_actor_id() + if read_by_multi_output_node + else None + ) + # Create an output channel for each output of the current node. + output_channel = ray.get( + fn.remote( + do_allocate_channel, + reader_and_node_list, + task.dag_node.type_hint, + driver_actor_id, + ) + ) + output_idx = None + downstream_node = output.dag_node + if ( + isinstance(downstream_node, ClassMethodNode) + and downstream_node.is_class_method_output + ): + output_idx = downstream_node.output_idx + task.output_channels.append(output_channel) + task.output_idxs.append(output_idx) + task.output_node_idxs.append(self.dag_node_to_idx[downstream_node]) + actor_handle = task.dag_node._get_actor_handle() + assert actor_handle is not None + self.actor_to_tasks[actor_handle].append(task) + elif ( + isinstance(task.dag_node, ClassMethodNode) + and task.dag_node.is_class_method_output + ): + task_node = task.dag_node + upstream_node = task_node.class_method_call + assert upstream_node + upstream_task = self.idx_to_task[self.dag_node_to_idx[upstream_node]] + for i in range(len(upstream_task.output_channels)): + if upstream_task.output_idxs[i] == task_node.output_idx: + task.output_channels.append(upstream_task.output_channels[i]) + task.output_idxs.append(upstream_task.output_idxs[i]) + assert len(task.output_channels) == 1 + elif isinstance(task.dag_node, InputNode): + # A dictionary that maps an InputNode or InputAttributeNode to its + # readers and the node on which the reader is running. Use `set` to + # deduplicate readers on the same actor because with CachedChannel + # each actor will only read from the shared memory once. + input_node_to_reader_and_node_set: Dict[ + Union[InputNode, InputAttributeNode], + Set[Tuple["ray.actor.ActorHandle", str]], + ] = defaultdict(set) + + for idx in task.downstream_task_idxs: + reader_task = self.idx_to_task[idx] + assert isinstance(reader_task.dag_node, ClassMethodNode) + reader_handle = reader_task.dag_node._get_actor_handle() + reader_node_id = self._get_node_id(reader_handle) + for arg in reader_task.args: + if isinstance(arg, InputAttributeNode) or isinstance( + arg, InputNode + ): + input_node_to_reader_and_node_set[arg].add( + (reader_handle, reader_node_id) + ) + + # A single channel is responsible for sending the same data to + # corresponding consumers. Therefore, we create a channel for + # each InputAttributeNode, or a single channel for the entire + # input data if there are no InputAttributeNodes. + task.output_channels = [] + for input_dag_node in input_node_to_reader_and_node_set: + reader_and_node_list = list( + input_node_to_reader_and_node_set[input_dag_node] + ) + + output_channel = do_allocate_channel( + self, + reader_and_node_list, + input_dag_node.type_hint, + None, + ) + task.output_channels.append(output_channel) + task.output_idxs.append( + None + if isinstance(input_dag_node, InputNode) + else input_dag_node.key + ) + + # Update the InputAttributeNode's `output_channels`, which is + # used to determine whether to create a CachedChannel. + if isinstance(input_dag_node, InputAttributeNode): + input_attr_idx = self.dag_node_to_idx[input_dag_node] + input_attr_task = self.idx_to_task[input_attr_idx] + input_attr_task.output_channels.append(output_channel) + assert len(input_attr_task.output_channels) == 1 + else: + assert isinstance(task.dag_node, InputAttributeNode) or isinstance( + task.dag_node, MultiOutputNode + ) + + for idx in task.downstream_task_idxs: + frontier.append(idx) + + # Validate input channels for tasks that have not been visited + for node_idx, task in self.idx_to_task.items(): + if ( + node_idx == self.input_task_idx + or node_idx == self.output_task_idx + or isinstance(task.dag_node, InputAttributeNode) + ): + continue + if node_idx not in visited: + has_at_least_one_channel_input = False + for arg in task.args: + if isinstance(arg, DAGNode): + has_at_least_one_channel_input = True + if not has_at_least_one_channel_input: + raise ValueError( + "Compiled DAGs require each task to take a ray.dag.InputNode " + "or at least one other DAGNode as an input. " + "Invalid task node:\n" + f"{task.dag_node}\n" + "Please bind the task to proper DAG nodes." + ) + + from ray.dag.constants import RAY_CGRAPH_ENABLE_DETECT_DEADLOCK + + if RAY_CGRAPH_ENABLE_DETECT_DEADLOCK and self._detect_deadlock(): + raise ValueError( + "This DAG cannot be compiled because it will deadlock on accelerator " + "calls. If you believe this is a false positive, please disable " + "the graph verification by setting the environment variable " + "RAY_CGRAPH_ENABLE_DETECT_DEADLOCK to 0 and file an issue at " + "https://github.com/ray-project/ray/issues/new/." + ) + + input_task = self.idx_to_task[self.input_task_idx] + self.dag_input_channels = input_task.output_channels + assert self.dag_input_channels is not None + + # Create executable tasks for each actor + for actor_handle, tasks in self.actor_to_tasks.items(): + # Dict from arg to the set of tasks that consume it. + arg_to_consumers: Dict[DAGNode, Set[CompiledTask]] = defaultdict(set) + + # Step 1: populate `arg_to_consumers` and perform some validation. + for task in tasks: + has_at_least_one_channel_input = False + for arg in task.args: + if isinstance(arg, DAGNode): + has_at_least_one_channel_input = True + arg_to_consumers[arg].add(task) + arg_idx = self.dag_node_to_idx[arg] + upstream_task = self.idx_to_task[arg_idx] + assert len(upstream_task.output_channels) == 1 + arg_channel = upstream_task.output_channels[0] + assert arg_channel is not None + # TODO: Support no-input DAGs (use an empty object to signal). + if not has_at_least_one_channel_input: + raise ValueError( + "Compiled DAGs require each task to take a " + "ray.dag.InputNode or at least one other DAGNode as an " + "input" + ) + + # Step 2: create cached channels if needed + + # Dict from original channel to the channel to be used in execution. + # The value of this dict is either the original channel or a newly + # created CachedChannel (if the original channel is read more than once). + for arg, consumers in arg_to_consumers.items(): + arg_idx = self.dag_node_to_idx[arg] + upstream_task = self.idx_to_task[arg_idx] + assert len(upstream_task.output_channels) == 1 + arg_channel = upstream_task.output_channels[0] + assert arg_channel is not None + if len(consumers) > 1: + self._channel_dict[arg_channel] = CachedChannel( + len(consumers), + arg_channel, + ) + else: + self._channel_dict[arg_channel] = arg_channel + + # Step 3: create executable tasks for the actor + executable_tasks = [] + for task in tasks: + resolved_args: List[Any] = [] + for arg in task.args: + if isinstance(arg, DAGNode): + arg_idx = self.dag_node_to_idx[arg] + upstream_task = self.idx_to_task[arg_idx] + assert len(upstream_task.output_channels) == 1 + arg_channel = upstream_task.output_channels[0] + assert arg_channel is not None + arg_channel = self._channel_dict[arg_channel] + resolved_args.append(arg_channel) + else: + # Constant arg + resolved_args.append(arg) + executable_task = ExecutableTask( + task, + resolved_args, + task.kwargs, + ) + executable_tasks.append(executable_task) + # Sort executable tasks based on their bind index, i.e., submission order + # so that they will be executed in that order. + executable_tasks.sort(key=lambda task: task.bind_index) + self.actor_to_executable_tasks[actor_handle] = executable_tasks + + from ray.dag.constants import RAY_CGRAPH_ENABLE_PROFILING + + if RAY_CGRAPH_ENABLE_PROFILING: + exec_task_func = do_profile_tasks + else: + exec_task_func = do_exec_tasks + + # Build an execution schedule for each actor + self.actor_to_execution_schedule = self._build_execution_schedule() + for actor_handle, executable_tasks in self.actor_to_executable_tasks.items(): + self.worker_task_refs[actor_handle] = actor_handle.__ray_call__.options( + concurrency_group="_ray_system" + ).remote( + exec_task_func, + executable_tasks, + self.actor_to_execution_schedule[actor_handle], + self._overlap_gpu_communication, + ) + + assert self.output_task_idx is not None + self.dag_output_channels = [] + for output in self.idx_to_task[self.output_task_idx].args: + assert isinstance(output, DAGNode) + output_idx = self.dag_node_to_idx[output] + task = self.idx_to_task[output_idx] + assert len(task.output_channels) == 1 + self.dag_output_channels.append(task.output_channels[0]) + + # Register custom serializers for input, input attribute, and output nodes. + self._register_input_output_custom_serializer() + + assert self.dag_input_channels + assert self.dag_output_channels + assert [ + output_channel is not None for output_channel in self.dag_output_channels + ] + # If no MultiOutputNode was specified during the DAG creation, there is only + # one output. Return a single output channel instead of a list of + # channels. + if not self._returns_list: + assert len(self.dag_output_channels) == 1 + + # Driver should ray.put on input, ray.get/release on output + self._monitor = self._monitor_failures() + input_task = self.idx_to_task[self.input_task_idx] + if self._enable_asyncio: + self._dag_submitter = AwaitableBackgroundWriter( + self.dag_input_channels, + input_task.output_idxs, + is_input=True, + ) + self._dag_output_fetcher = AwaitableBackgroundReader( + self.dag_output_channels, + self._fut_queue, + ) + else: + self._dag_submitter = SynchronousWriter( + self.dag_input_channels, input_task.output_idxs, is_input=True + ) + self._dag_output_fetcher = SynchronousReader(self.dag_output_channels) + + self._dag_submitter.start() + self._dag_output_fetcher.start() + + def _generate_dag_operation_graph_node( + self, + ) -> Dict["ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]]]: + """ + Generate READ, COMPUTE, and WRITE operations for each DAG node. + + Returns: + A dictionary that maps an actor handle to a list of lists of + _DAGOperationGraphNode. For the same actor, the index of the + outer list corresponds to the index of the ExecutableTask in + the list of `executable_tasks` in `actor_to_executable_tasks`, + i.e. `exec_task_idx`. In the inner list, the order of operations + is READ, COMPUTE, and WRITE. + + Example: + { + actor1: [ + [READ COMPUTE WRITE] # exec_task_idx 0 + [READ COMPUTE WRITE] # exec_task_idx 1 + ] + } + """ + from ray.dag.collective_node import CollectiveOutputNode + + assert self.idx_to_task + assert self.actor_to_executable_tasks + + actor_to_operation_nodes: Dict[ + "ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]] + ] = defaultdict(list) + + for actor_handle, executable_tasks in self.actor_to_executable_tasks.items(): + for exec_task_idx, exec_task in enumerate(executable_tasks): + # Divide a DAG node into three _DAGOperationGraphNodes: READ, COMPUTE, + # and WRITE. Each _DAGOperationGraphNode has a _DAGNodeOperation. + task_idx = exec_task.task_idx + dag_node = self.idx_to_task[task_idx].dag_node + method_name = exec_task.method_name + actor_handle = dag_node._get_actor_handle() + requires_accelerator_read = False + for upstream_node in dag_node._upstream_nodes: + if upstream_node.type_hint.requires_accelerator(): + requires_accelerator_read = True + break + requires_accelerator_compute = isinstance( + dag_node, CollectiveOutputNode + ) + requires_accelerator_write = dag_node.type_hint.requires_accelerator() + + read_node = _DAGOperationGraphNode( + _DAGNodeOperation( + exec_task_idx, _DAGNodeOperationType.READ, method_name + ), + task_idx, + actor_handle, + requires_accelerator_read, + ) + compute_node = _DAGOperationGraphNode( + _DAGNodeOperation( + exec_task_idx, _DAGNodeOperationType.COMPUTE, method_name + ), + task_idx, + actor_handle, + requires_accelerator_compute, + ) + write_node = _DAGOperationGraphNode( + _DAGNodeOperation( + exec_task_idx, _DAGNodeOperationType.WRITE, method_name + ), + task_idx, + actor_handle, + requires_accelerator_write, + ) + + actor_to_operation_nodes[actor_handle].append( + [read_node, compute_node, write_node] + ) + + return actor_to_operation_nodes + + def _build_execution_schedule( + self, + ) -> Dict["ray.actor.ActorHandle", List[_DAGNodeOperation]]: + """ + Generate an execution schedule for each actor. The schedule is a list of + _DAGNodeOperation. + + Step 1: Generate a DAG node operation graph. Refer to the functions + `_generate_dag_operation_graph_node` and `_build_dag_node_operation_graph` + for more details. + + Step 2: Topological sort + + It is possible to have multiple _DAGOperationGraphNodes with zero in-degree. + Refer to the function `_select_next_nodes` for the logic of selecting nodes. + + Then, put the selected nodes into the corresponding actors' schedules. + + The schedule should be intuitive to users, meaning that the execution should + perform operations in ascending order of `bind_index` as much as possible. + + [Example]: + + See `test_execution_schedule` for more examples. + + Returns: + actor_to_execution_schedule: A dictionary that maps an actor handle to + the execution schedule which is a list of operations to be executed. + """ + # Step 1: Build a graph of _DAGOperationGraphNode + actor_to_operation_nodes = self._generate_dag_operation_graph_node() + graph = _build_dag_node_operation_graph( + self.idx_to_task, actor_to_operation_nodes + ) + # Step 2: Generate an execution schedule for each actor using topological sort + actor_to_execution_schedule = _generate_actor_to_execution_schedule(graph) + + # Step 3: Overlap GPU communication for the execution schedule if configured + actor_to_overlapped_schedule = None + if self._overlap_gpu_communication: + actor_to_overlapped_schedule = _generate_overlapped_execution_schedule( + actor_to_execution_schedule + ) + + if RAY_CGRAPH_VISUALIZE_SCHEDULE: + _visualize_execution_schedule( + actor_to_execution_schedule, actor_to_overlapped_schedule, graph + ) + + if actor_to_overlapped_schedule is not None: + return _extract_execution_schedule(actor_to_overlapped_schedule) + else: + return _extract_execution_schedule(actor_to_execution_schedule) + + def _detect_deadlock(self) -> bool: + """ + TODO (kevin85421): Avoid false negatives. + + Currently, a compiled graph may deadlock if there are accelerator channels, + and the readers have control dependencies on the same actor. For example: + + actor1.a ---> actor2.f1 + | + ---> actor2.f2 + + The control dependency between `actor2.f1` and `actor2.f2` is that `f1` should + run before `f2`. If `actor1.a` writes to `actor2.f2` before `actor2.f1`, a + deadlock will occur. + + Currently, the execution schedule is not granular enough to detect this + deadlock. + + Returns: + True if a deadlock is detected; otherwise, False. + """ + logger.debug("Deadlock detection has not been implemented yet.") + return False + + def _monitor_failures(self): + get_outer = weakref.ref(self) + + class Monitor(threading.Thread): + def __init__(self): + super().__init__(daemon=True) + self.name = "CompiledGraphMonitorThread" + # Lock to make sure that we only perform teardown for this DAG + # once. + self._in_teardown_lock = threading.Lock() + self._teardown_done = False + + def _outer_ref_alive(self) -> bool: + if get_outer() is None: + logger.error( + "CompiledDAG has been destructed before teardown. " + "This should not occur please report an issue at " + "https://github.com/ray-project/ray/issues/new/.", + stack_info=True, + ) + return False + return True + + def wait_teardown(self, kill_actors: bool = False): + outer = get_outer() + if not self._outer_ref_alive(): + return + + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + teardown_timeout = ctx.teardown_timeout + for actor, ref in outer.worker_task_refs.items(): + timeout = False + try: + ray.get(ref, timeout=teardown_timeout) + except ray.exceptions.GetTimeoutError: + msg = ( + f"Compiled DAG actor {actor} is still running " + f"{teardown_timeout}s after teardown()." + ) + if kill_actors: + msg += ( + " Force-killing actor. " + "Increase RAY_CGRAPH_teardown_timeout if you want " + "teardown to wait longer." + ) + ray.kill(actor) + else: + msg += ( + " Teardown may hang. " + "Call teardown with kill_actors=True if force kill " + "is desired." + ) + + logger.warning(msg) + timeout = True + except Exception: + # We just want to check that the task has finished so + # we don't care if the actor task ended in an + # exception. + pass + + if not timeout: + continue + + try: + ray.get(ref) + except Exception: + pass + + if kill_actors: + # In the previous loop, we allow the actor tasks to exit first. + # Now, we force kill the actors if not yet. + for actor in outer.worker_task_refs: + logger.info(f"Killing actor: {actor}") + ray.kill(actor) + + def teardown(self, kill_actors: bool = False): + with self._in_teardown_lock: + if self._teardown_done: + return + + outer = get_outer() + if not self._outer_ref_alive(): + return + + logger.info("Tearing down compiled DAG") + outer._dag_submitter.close() + outer._dag_output_fetcher.close() + + for actor in outer.actor_to_executable_tasks.keys(): + logger.info(f"Cancelling compiled worker on actor: {actor}") + # Cancel all actor loops in parallel. + cancel_refs = [ + actor.__ray_call__.remote(do_cancel_executable_tasks, tasks) + for actor, tasks in outer.actor_to_executable_tasks.items() + ] + for cancel_ref in cancel_refs: + try: + ray.get(cancel_ref, timeout=30) + except RayChannelError: + # Channel error happens when a channel is closed + # or timed out. In this case, do not log. + pass + except Exception: + logger.exception("Error cancelling worker task") + pass + + for ( + communicator_id + ) in outer._actors_to_created_communicator_id.values(): + _destroy_communicator(communicator_id) + + logger.info("Waiting for worker tasks to exit") + self.wait_teardown(kill_actors=kill_actors) + + logger.info("Teardown complete") + self._teardown_done = True + + def run(self): + try: + outer = get_outer() + if not self._outer_ref_alive(): + return + ray.get(list(outer.worker_task_refs.values())) + except KeyboardInterrupt: + logger.info( + "Received KeyboardInterrupt, tearing down with kill_actors=True" + ) + self.teardown(kill_actors=True) + except Exception as e: + logger.debug(f"Handling exception from worker tasks: {e}") + self.teardown() + + monitor = Monitor() + monitor.start() + return monitor + + def _raise_if_too_many_inflight_executions(self): + num_inflight_executions = ( + self._execution_index - self._max_finished_execution_index + ) + if num_inflight_executions >= self._max_inflight_executions: + raise ray.exceptions.RayCgraphCapacityExceeded( + "The compiled graph can't have more than " + f"{self._max_inflight_executions} in-flight executions, and you " + f"currently have {num_inflight_executions} in-flight executions. " + "Retrieve an output using ray.get before submitting more requests or " + "increase `_max_inflight_executions`. " + "`dag.experimental_compile(_max_inflight_executions=...)`" + ) + + def _has_execution_results( + self, + execution_index: int, + ) -> bool: + """Check whether there are results corresponding to the given execution + index stored in self._result_buffer. This helps avoid fetching and + caching results again. + + Args: + execution_index: The execution index corresponding to the result. + + Returns: + Whether the result for the given index has been fetched and cached. + """ + return execution_index in self._result_buffer + + def _cache_execution_results( + self, + execution_index: int, + result: Any, + ): + """Cache execution results in self._result_buffer. Results are converted + to dictionary format to allow efficient element removal and calculation of + the buffer size. This can only be called once per execution index. + + Args: + execution_index: The execution index corresponding to the result. + result: The results from all channels to be cached. + """ + if not self._has_execution_results(execution_index): + for chan_idx, res in enumerate(result): + # avoid caching for any CompiledDAGRef that has already been destructed. + if not ( + execution_index in self._destructed_ref_idxs + and chan_idx in self._destructed_ref_idxs[execution_index] + ): + self._result_buffer[execution_index][chan_idx] = res + + def _get_execution_results( + self, execution_index: int, channel_index: Optional[int] + ) -> List[Any]: + """Retrieve execution results from self._result_buffer and return the result. + Results are converted back to original list format ordered by output channel + index. + + Args: + execution_index: The execution index to retrieve results from. + channel_index: The index of the output channel corresponding to the result. + Channel indexing is consistent with the order of + self.dag_output_channels. None means that the result wraps outputs from + all output channels. + + Returns: + The execution result corresponding to the given execution index and channel + index. + """ + # Although CompiledDAGRef and CompiledDAGFuture guarantee that the same + # execution index and channel index combination will not be requested multiple + # times and therefore self._result_buffer will always have execution_index as + # a key, we still do a sanity check to avoid misuses. + assert execution_index in self._result_buffer + + if channel_index is None: + # Convert results stored in self._result_buffer back to original + # list representation + result = [ + kv[1] + for kv in sorted( + self._result_buffer.pop(execution_index).items(), + key=lambda kv: kv[0], + ) + ] + else: + result = [self._result_buffer[execution_index].pop(channel_index)] + + if execution_index not in self._got_ref_idxs: + self._got_ref_idxs[execution_index] = set() + self._got_ref_idxs[execution_index].add(channel_index) + self._clean_up_buffers(execution_index) + return result + + def _delete_execution_results(self, execution_index: int, channel_index: int): + """ + Delete the execution results for the given execution index and channel index. + This method should be called when a CompiledDAGRef or CompiledDAGFuture is + destructed. + + Note that this method maintains metadata for the deleted execution results, + and only actually deletes the buffers lazily when the buffer is not needed + anymore. + + Args: + execution_index: The execution index to destruct results from. + channel_index: The index of the output channel corresponding to the result. + """ + if execution_index not in self._destructed_ref_idxs: + self._destructed_ref_idxs[execution_index] = set() + self._destructed_ref_idxs[execution_index].add(channel_index) + self._clean_up_buffers(execution_index) + + def _try_release_result_buffer(self, execution_index: int): + """ + Try to release the result buffer for the given execution index. + """ + + should_release = False + got_channel_idxs = self._got_ref_idxs.get(execution_index, set()) + if None in got_channel_idxs: + assert len(got_channel_idxs) == 1, ( + "when None exists in got_channel_idxs, it means all channels, and " + "it should be the only value in the set", + ) + should_release = True + else: + destructed_channel_idxs = self._destructed_ref_idxs.get( + execution_index, set() + ) + processed_channel_idxs = got_channel_idxs.union(destructed_channel_idxs) + # No more processing is needed for this execution index. + should_release = processed_channel_idxs == set( + range(len(self.dag_output_channels)) + ) + + if not should_release: + return False + + self._result_buffer.pop(execution_index, None) + self._destructed_ref_idxs.pop(execution_index, None) + self._got_ref_idxs.pop(execution_index, None) + return True + + def _try_release_native_buffer( + self, idx_to_release: int, timeout: Optional[float] = None + ) -> bool: + """ + Try to release the native buffer for the given execution index. + + Args: + idx_to_release: The execution index to release buffers from. + timeout: The maximum time in seconds to wait for the release. + + Returns: + Whether the buffers have been released. + """ + if idx_to_release != self._max_finished_execution_index + 1: + # Native buffer can only be released for the next execution index. + return False + + destructed_channel_idxs = self._destructed_ref_idxs.get(idx_to_release, set()) + should_release = False + if None in destructed_channel_idxs: + assert len(destructed_channel_idxs) == 1, ( + "when None exists in destructed_channel_idxs, it means all channels, " + "and it should be the only value in the set", + ) + should_release = True + elif len(destructed_channel_idxs) == len(self.dag_output_channels): + should_release = True + + if not should_release: + return False + + # refs corresponding to idx_to_release are all destructed, + # and they are never fetched or cached. + assert idx_to_release not in self._result_buffer + assert idx_to_release not in self._got_ref_idxs + + try: + self._dag_output_fetcher.release_channel_buffers(timeout) + except RayChannelTimeoutError as e: + raise RayChannelTimeoutError( + "Releasing native buffers corresponding to a stale CompiledDAGRef " + "is taking a long time. If this is expected, increase " + f"RAY_CGRAPH_get_timeout which is currently {self._get_timeout} " + "seconds. Otherwise, this may indicate that the execution " + "is hanging." + ) from e + self._destructed_ref_idxs.pop(idx_to_release) + + return True + + def _try_release_buffer( + self, idx_to_release: int, timeout: Optional[float] = None + ) -> bool: + """ + Try to release the buffer for the given execution index. + First try to release the native buffer, then try to release the result buffer. + + Args: + idx_to_release: The execution index to release buffers from. + timeout: The maximum time in seconds to wait for the release. + + Returns: + Whether the native buffer or result buffer has been released. + """ + if self._try_release_native_buffer(idx_to_release, timeout): + # Releasing native buffer means the corresponding execution result + # is consumed (and discarded). + self._max_finished_execution_index += 1 + return True + return self._try_release_result_buffer(idx_to_release) + + def _try_release_buffers(self): + """ + Repeatedly release buffer if possible. + + This method starts from _max_finished_execution_index + 1 and tries to release + as many buffers as possible. If a native buffer is released, + _max_finished_execution_index will be incremented. + """ + timeout = self._get_timeout + while True: + start_time = time.monotonic() + if not self._try_release_buffer( + self._max_finished_execution_index + 1, timeout + ): + break + + if timeout != -1: + timeout -= time.monotonic() - start_time + timeout = max(timeout, 0) + + def _clean_up_buffers(self, idx_to_release: int): + """ + Clean up native and result buffers. + + This method: + 1. Tries to release the buffer for the given execution index. + This index is the specific one that requires a clean up, + e.g., right after get() is called or a CompiledDAGRef/CompiledDAGFuture + is destructed. + 2. Tries to release all buffers starting from _max_finished_execution_index + 1. + This step is to clean up buffers that are no longer needed. + + Args: + idx_to_release: The execution index that requires a clean up, + e.g., right after get() is called or a CompiledDAGRef/CompiledDAGFuture + is destructed. + """ + self._try_release_buffer(idx_to_release) + self._try_release_buffers() + + def _execute_until( + self, + execution_index: int, + channel_index: Optional[int] = None, + timeout: Optional[float] = None, + ): + """Repeatedly execute this DAG until the given execution index and + buffer results for all CompiledDagRef's. + If the DAG has already been executed up to the given index, it will do nothing. + + Note: If this comes across execution indices for which the corresponding + CompiledDAGRef's have been destructed, it will release the buffer and not + cache the result. + + Args: + execution_index: The execution index to execute until. + channel_index: The index of the output channel to get the result from. + Channel indexing is consistent with the order of + self.dag_output_channels. None means wrapping results from all output + channels into a single list. + timeout: The maximum time in seconds to wait for the execution. + None means using default timeout (DAGContext.get_timeout), + 0 means immediate timeout (immediate success or timeout without + blocking), -1 means infinite timeout (block indefinitely). + + TODO(rui): catch the case that user holds onto the CompiledDAGRefs + """ + if timeout is None: + timeout = self._get_timeout + while self._max_finished_execution_index < execution_index: + if len(self._result_buffer) >= self._max_buffered_results: + raise RayCgraphCapacityExceeded( + "The compiled graph can't have more than " + f"{self._max_buffered_results} buffered results, and you " + f"currently have {len(self._result_buffer)} buffered results. " + "Call `ray.get()` on CompiledDAGRef's (or await on " + "CompiledDAGFuture's) to retrieve results, or increase " + f"`_max_buffered_results` if buffering is desired, note that " + "this will increase driver memory usage." + ) + start_time = time.monotonic() + + # Fetch results from each output channel up to execution_index and cache + # them separately to enable individual retrieval + # If a CompiledDagRef for a specific execution index has been destructed, + # release the channel buffers for that execution index instead of caching + try: + if not self._try_release_native_buffer( + self._max_finished_execution_index + 1, timeout + ): + result = self._dag_output_fetcher.read(timeout) + self._cache_execution_results( + self._max_finished_execution_index + 1, + result, + ) + # We have either released the native buffer or fetched and + # cached the result buffer, therefore we always increment + # _max_finished_execution_index. + self._max_finished_execution_index += 1 + except RayChannelTimeoutError as e: + raise RayChannelTimeoutError( + "If the execution is expected to take a long time, increase " + f"RAY_CGRAPH_get_timeout which is currently {self._get_timeout} " + "seconds. Otherwise, this may indicate that the execution is " + "hanging." + ) from e + + if timeout != -1: + timeout -= time.monotonic() - start_time + timeout = max(timeout, 0) + + def execute( + self, + *args, + **kwargs, + ) -> Union[CompiledDAGRef, List[CompiledDAGRef]]: + """Execute this DAG using the compiled execution path. + + Args: + args: Args to the InputNode. + kwargs: Kwargs to the InputNode + + Returns: + A list of Channels that can be used to read the DAG result. + + Raises: + RayChannelTimeoutError: If the execution does not complete within + self._submit_timeout seconds. + + NOTE: Not thread-safe due to _execution_index etc. + """ + if self._enable_asyncio: + raise ValueError("Use execute_async if enable_asyncio=True") + + self._get_or_compile() + + self._check_inputs(args, kwargs) + if len(args) == 1 and len(kwargs) == 0: + # When serializing a tuple, the Ray serializer invokes pickle5, which adds + # several microseconds of overhead. One common case for Compiled Graphs is + # passing a single argument (oftentimes of of type `bytes`, which requires + # no serialization). To avoid imposing this overhead on this common case, we + # create a fast path for this case that avoids pickle5. + inp = args[0] + else: + inp = CompiledDAGArgs(args=args, kwargs=kwargs) + + # We want to release any buffers we can at this point based on the + # max_finished_execution_index so that the number of inflight executions + # is up to date. + self._try_release_buffers() + self._raise_if_too_many_inflight_executions() + try: + self._dag_submitter.write(inp, self._submit_timeout) + except RayChannelTimeoutError as e: + raise RayChannelTimeoutError( + "If the execution is expected to take a long time, increase " + f"RAY_CGRAPH_submit_timeout which is currently {self._submit_timeout} " + "seconds. Otherwise, this may indicate that execution is hanging." + ) from e + + self._execution_index += 1 + + if self._returns_list: + ref = [ + CompiledDAGRef(self, self._execution_index, channel_index) + for channel_index in range(len(self.dag_output_channels)) + ] + else: + ref = CompiledDAGRef(self, self._execution_index) + + return ref + + def _check_inputs(self, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None: + """ + Helper method to check that the DAG args provided by the user during + execution are valid according to the defined DAG. + """ + if len(args) != self._input_num_positional_args: + raise ValueError( + "dag.execute() or dag.execute_async() must be " + f"called with {self._input_num_positional_args} positional args, got " + f"{len(args)}" + ) + + for kwarg in self._input_kwargs: + if kwarg not in kwargs: + raise ValueError( + "dag.execute() or dag.execute_async() " + f"must be called with kwarg `{kwarg}`" + ) + + async def execute_async( + self, + *args, + **kwargs, + ) -> Union[CompiledDAGFuture, List[CompiledDAGFuture]]: + """Execute this DAG using the compiled execution path. + + NOTE: Not thread-safe. + + Args: + args: Args to the InputNode. + kwargs: Kwargs to the InputNode. + + Returns: + A list of Channels that can be used to read the DAG result. + """ + if not self._enable_asyncio: + raise ValueError("Use execute if enable_asyncio=False") + + self._get_or_compile() + self._check_inputs(args, kwargs) + async with self._dag_submission_lock: + if len(args) == 1 and len(kwargs) == 0: + # When serializing a tuple, the Ray serializer invokes pickle5, which + # adds several microseconds of overhead. One common case for accelerated + # DAGs is passing a single argument (oftentimes of of type `bytes`, + # which requires no serialization). To avoid imposing this overhead on + # this common case, we create a fast path for this case that avoids + # pickle5. + inp = args[0] + else: + inp = CompiledDAGArgs(args=args, kwargs=kwargs) + + self._raise_if_too_many_inflight_executions() + await self._dag_submitter.write(inp) + # Allocate a future that the caller can use to get the result. + fut = asyncio.Future() + await self._fut_queue.put(fut) + + self._execution_index += 1 + + if self._returns_list: + fut = [ + CompiledDAGFuture(self, self._execution_index, fut, channel_index) + for channel_index in range(len(self.dag_output_channels)) + ] + else: + fut = CompiledDAGFuture(self, self._execution_index, fut) + + return fut + + def _visualize_ascii(self) -> str: + """ + Visualize the compiled graph in + ASCII format with directional markers. + + This function generates an ASCII visualization of a Compiled Graph, + where each task node is labeled, + and edges use `<` and `>` markers to show data flow direction. + + This method is called by: + - `compiled_dag.visualize(format="ascii")` + + + + High-Level Algorithm: + - Topological Sorting: Sort nodes topologically to organize + them into layers based on dependencies. + - Grid Initialization: Set up a 2D grid canvas with dimensions based + on the number of layers and the maximum number of nodes per layer. + - Node Placement: Position each node on the grid according to its + layer and relative position within that layer. + Spacing is added for readability, and directional markers (`<` and `>`) + are added to edges to show input/output flow clearly. + + This method should be called + **after** compiling the graph with `experimental_compile()`. + + Returns: + ASCII representation of the CG with Nodes Information, + Edges Information and Graph Built. + + Limitations: + - Note: This is only used for quick visualization for small graphs. + For complex graph (i.e. more than 20 tasks), please use graphviz. + - Scale: Works best for smaller CGs (typically fewer than 20 tasks). + Larger CGs may result in dense, less readable ASCII + outputs due to limited space for node and edge rendering. + - Shape: Ideal for relatively shallow CGs with clear dependency paths. + For deep, highly branched or densely connected CGs, + readability may suffer. + - Edge Overlap: In cases with high fan-out (i.e., nodes with many children) + or fan-in (nodes with many parents), edge lines may intersect or overlap + in the ASCII visualization, potentially obscuring some connections. + - Multi-output Tasks: Multi-output tasks can be visualized, but positioning + may cause line breaks or overlap when a task has multiple outputs that + feed into nodes at varying depths. + + Example: + Basic Visualization: + ```python + # Print the CG structure in ASCII format + print(compiled_dag.visualize(format="ascii")) + ``` + + Example of Ordered Visualization (task is build in order + to reduce line intersection): + ```python + with InputNode() as i: + o1, o2, o3 = a.return_three.bind(i) + o4 = b.echo.bind(o1) + o5 = b.echo.bind(o2) + o6, o7 = b.return_two.bind(o3) + dag = MultiOutputNode([o4, o5, o6, o7]) + + compiled_dag = dag.experimental_compile() + compiled_dag.visualize(format="ascii",view=True) + + + # Output: + # 0:InputNode + # | + # 1:Actor_54777d:return_three + # |---------------------------->|---------------------------->| # noqa + # 2:Output[0] 3:Output[1] 4:Output[2] # noqa + # | | | # noqa + # 5:Actor_c927c9:echo 6:Actor_c927c9:echo 7:Actor_c927c9:return_two # noqa + # | | |---------------------------->| # noqa + # | | 9:Output[0] 10:Output[1] # noqa + # |<----------------------------|-----------------------------|-----------------------------| # noqa + # 8:MultiOutputNode + ``` + + Example of Anti-pattern Visualization (There are intersections): + # We can swtich the nodes ordering to reduce intersections, i.e. swap o2 and o3 + ```python + with InputNode() as i: + o1, o2, o3 = a.return_three.bind(i) + o4 = b.echo.bind(o1) + o5 = b.echo.bind(o3) + o6, o7 = b.return_two.bind(o2) + dag = MultiOutputNode([o4, o5, o6, o7]) + compiled_dag = dag.experimental_compile() + compiled_dag.visualize(format="ascii",view=True) + + # Output (Nodes 5, 7, 9, 10 should connect to Node 8): + # 0:InputNode + # | + # 1:Actor_84835a:return_three + # |---------------------------->|---------------------------->| # noqa + # 2:Output[0] 3:Output[1] 4:Output[2] # noqa + # | | | # noqa + # 5:Actor_02a6a1:echo 6:Actor_02a6a1:return_two 7:Actor_02a6a1:echo # noqa + # | |---------------------------->| # noqa + # | 9:Output[0] 10:Output[1] # noqa + # |<----------------------------------------------------------| # noqa + # 8:MultiOutputNode + ``` + """ + + from ray.dag import ( + ClassMethodNode, + DAGNode, + InputAttributeNode, + InputNode, + MultiOutputNode, + ) + + # Check that the DAG has been compiled + if not hasattr(self, "idx_to_task") or not self.idx_to_task: + raise ValueError( + "The DAG must be compiled before calling 'visualize()'. " + "Please call 'experimental_compile()' first." + ) + + # Check that each CompiledTask has a valid dag_node + for idx, task in self.idx_to_task.items(): + if not hasattr(task, "dag_node") or not isinstance(task.dag_node, DAGNode): + raise ValueError( + f"Task at index {idx} does not have a valid 'dag_node'. " + "Ensure that 'experimental_compile()' completed successfully." + ) + + from collections import defaultdict, deque + + # Create adjacency list representation of the DAG + # Adjacency list for DAG; maps a node index to its downstream nodes. + adj_list: Dict[int, List[int]] = defaultdict(list) + # Indegree count for topological sorting; maps a node index to its indegree. + indegree: Dict[int, int] = defaultdict(int) + + # Tracks whether a node is a multi-output node. + is_multi_output: Dict[int, bool] = defaultdict(bool) + # Maps child node indices to their parent node indices. + child2parent: Dict[int, int] = defaultdict(int) + ascii_visualization = "" + # Node information; maps a node index to its descriptive label. + node_info: Dict[int, str] = {} + # Edge information; tuples of (upstream_index, downstream_index, edge_label). + edge_info: List[Tuple[int, int, str]] = [] + + for idx, task in self.idx_to_task.items(): + dag_node = task.dag_node + label = f"Task {idx} " + + # Determine the type and label of the node + if isinstance(dag_node, InputNode): + label += "InputNode" + elif isinstance(dag_node, InputAttributeNode): + label += f"InputAttributeNode[{dag_node.key}]" + elif isinstance(dag_node, MultiOutputNode): + label += "MultiOutputNode" + elif isinstance(dag_node, ClassMethodNode): + if dag_node.is_class_method_call: + method_name = dag_node.get_method_name() + actor_handle = dag_node._get_actor_handle() + actor_id = ( + actor_handle._actor_id.hex()[:6] if actor_handle else "unknown" + ) + label += f"Actor: {actor_id}... Method: {method_name}" + elif dag_node.is_class_method_output: + label += f"ClassMethodOutputNode[{dag_node.output_idx}]" + else: + label += "ClassMethodNode" + else: + label += type(dag_node).__name__ + + node_info[idx] = label + + for arg_index, arg in enumerate(dag_node.get_args()): + if isinstance(arg, DAGNode): + upstream_task_idx = self.dag_node_to_idx[arg] + + # Get the type hint for this argument + if arg_index < len(task.arg_type_hints): + if task.arg_type_hints[arg_index].requires_accelerator(): + type_hint = "Accelerator" + else: + type_hint = type(task.arg_type_hints[arg_index]).__name__ + else: + type_hint = "UnknownType" + + adj_list[upstream_task_idx].append(idx) + indegree[idx] += 1 + edge_info.append((upstream_task_idx, idx, type_hint)) + + width_adjust = 0 + for upstream_task_idx, child_idx_list in adj_list.items(): + # Mark as multi-output if the node has more than one output path + if len(child_idx_list) > 1: + for child in child_idx_list: + is_multi_output[child] = True + child2parent[child] = upstream_task_idx + width_adjust = max(width_adjust, len(child_idx_list)) + + # Topological sort to determine layers + layers = defaultdict(list) + zero_indegree = deque([idx for idx in self.idx_to_task if indegree[idx] == 0]) + layer_index = 0 + + while zero_indegree: + next_layer = deque() + while zero_indegree: + task_idx = zero_indegree.popleft() + layers[layer_index].append(task_idx) + for downstream in adj_list[task_idx]: + indegree[downstream] -= 1 + if indegree[downstream] == 0: + next_layer.append(downstream) + zero_indegree = next_layer + layer_index += 1 + + # Print detailed node information + ascii_visualization += "Nodes Information:\n" + for idx, info in node_info.items(): + ascii_visualization += f'{idx} [label="{info}"] \n' + + # Print edges + ascii_visualization += "\nEdges Information:\n" + for upstream_task, downstream_task, type_hint in edge_info: + if type_hint == "Accelerator": + edgs_channel = "+++" + else: + edgs_channel = "---" + ascii_visualization += ( + f"{upstream_task} {edgs_channel}>" f" {downstream_task}\n" + ) + + # Add the legend to the output + ascii_visualization += "\nLegend:\n" + ascii_visualization += "+++> : Represents Accelerator-type data channels\n" + ascii_visualization += "---> : Represents Shared Memory data channels\n" + + # Find the maximum width (number of nodes in any layer) + max_width = max(len(layer) for layer in layers.values()) + width_adjust + height = len(layers) + + # Build grid for ASCII visualization + grid = [[" " for _ in range(max_width * 20)] for _ in range(height * 2 - 1)] + + # Place nodes in the grid with more details + task_to_pos = {} + for layer_num, layer_tasks in layers.items(): + layer_y = layer_num * 2 # Every second row is for nodes + for col_num, task_idx in enumerate(layer_tasks): + task = self.idx_to_task[task_idx] + task_info = f"{task_idx}:" + + # Determine if it's an actor method or a regular task + if isinstance(task.dag_node, ClassMethodNode): + if task.dag_node.is_class_method_call: + method_name = task.dag_node.get_method_name() + actor_handle = task.dag_node._get_actor_handle() + actor_id = ( + actor_handle._actor_id.hex()[:6] + if actor_handle + else "unknown" + ) + task_info += f"Actor_{actor_id}:{method_name}" + elif task.dag_node.is_class_method_output: + task_info += f"Output[{task.dag_node.output_idx}]" + else: + task_info += "UnknownMethod" + else: + task_info += type(task.dag_node).__name__ + + adjust_col_num = 0 + if task_idx in is_multi_output: + adjust_col_num = layers[layer_num - 1].index(child2parent[task_idx]) + col_x = (col_num + adjust_col_num) * 30 # Every 30th column for spacing + # Place the task information into the grid + for i, char in enumerate(task_info): + if col_x + i < len(grid[0]): # Ensure we don't overflow the grid + grid[layer_y][col_x + i] = char + + task_to_pos[task_idx] = (layer_y, col_x) + + # Connect the nodes with lines + for upstream_task, downstream_tasks in adj_list.items(): + upstream_y, upstream_x = task_to_pos[upstream_task] + for downstream_task in downstream_tasks: + downstream_y, downstream_x = task_to_pos[downstream_task] + + # Draw vertical line + for y in range(upstream_y + 1, downstream_y): + if grid[y][upstream_x] == " ": + grid[y][upstream_x] = "|" + + # Draw horizontal line with directional arrows + if upstream_x != downstream_x: + for x in range( + min(upstream_x, downstream_x) + 1, + max(upstream_x, downstream_x), + ): + grid[downstream_y - 1][x] = ( + "-" + if grid[downstream_y - 1][x] == " " + else grid[downstream_y - 1][x] + ) + + # Add arrows to indicate flow direction + if downstream_x > upstream_x: + grid[downstream_y - 1][downstream_x - 1] = ">" + else: + grid[downstream_y - 1][downstream_x + 1] = "<" + + # Draw connection to the next task + grid[downstream_y - 1][downstream_x] = "|" + + # Ensure proper multi-output task connection + for idx, task in self.idx_to_task.items(): + if isinstance(task.dag_node, MultiOutputNode): + output_tasks = task.dag_node.get_args() + for i, output_task in enumerate(output_tasks): + if isinstance(output_task, DAGNode): + output_task_idx = self.dag_node_to_idx[output_task] + if output_task_idx in task_to_pos: + output_y, output_x = task_to_pos[output_task_idx] + grid[output_y - 1][output_x] = "|" + + # Convert grid to string for printing + ascii_visualization += "\nGraph Built:\n" + ascii_visualization += "\n".join("".join(row) for row in grid) + + return ascii_visualization + + def get_channel_details( + self, channel: ChannelInterface, downstream_actor_id: str + ) -> str: + """ + Get details about outer and inner channel types and channel ids + based on the channel and the downstream actor ID. + Used for graph visualization. + Args: + channel: The channel to get details for. + downstream_actor_id: The downstream actor ID. + Returns: + A string with details about the channel based on its connection + to the actor provided. + """ + channel_details = type(channel).__name__ + # get outer channel + if channel in self._channel_dict and self._channel_dict[channel] != channel: + channel = self._channel_dict[channel] + channel_details += f"\n{type(channel).__name__}" + if type(channel) is CachedChannel: + channel_details += f", {channel._channel_id[:6]}..." + # get inner channel + if ( + type(channel) is CompositeChannel + and downstream_actor_id in channel._channel_dict + ): + inner_channel = channel._channel_dict[downstream_actor_id] + channel_details += f"\n{type(inner_channel).__name__}" + if type(inner_channel) is IntraProcessChannel: + channel_details += f", {inner_channel._channel_id[:6]}..." + return channel_details + + def visualize( + self, + filename="compiled_graph", + format="png", + view=False, + channel_details=False, + ) -> str: + """ + Visualize the compiled graph by showing tasks and their dependencies. + This method should be called **after** the graph has been compiled using + `experimental_compile()`. + + Args: + filename: For non-ASCII formats, the output file name (without extension). + For ASCII format, the visualization will be printed to the console, + and this argument is ignored. + format: The format of the output file (e.g., 'png', 'pdf', 'ascii'). + view: For non-ASCII formats: Whether to open the file with the default + viewer. For ASCII format: Whether to print the visualization and return + None or return the ascii visualization string directly. + channel_details: If True, adds channel details to edges. + + Returns: + The string representation of the compiled graph. For Graphviz-based formats + (e.g., 'png', 'pdf', 'jpeg'), returns the Graphviz DOT string representation + of the compiled graph. For ASCII format, returns the ASCII string + representation of the compiled graph. + + Raises: + ValueError: If the graph is empty or not properly compiled. + ImportError: If the `graphviz` package is not installed. + + """ + if format == "ascii": + if channel_details: + raise ValueError( + "Parameters 'channel_details' are" + " not compatible with 'ascii' format." + ) + ascii_visualiztion_str = self._visualize_ascii() + if view: + print(ascii_visualiztion_str) + return ascii_visualiztion_str + try: + import graphviz + except ImportError: + raise ImportError( + "Please install graphviz to visualize the compiled graph. " + "You can install it by running `pip install graphviz`." + ) + from ray.dag import ( + ClassMethodNode, + DAGNode, + InputAttributeNode, + InputNode, + MultiOutputNode, + ) + + # Check that the DAG has been compiled + if not hasattr(self, "idx_to_task") or not self.idx_to_task: + raise ValueError( + "The DAG must be compiled before calling 'visualize()'. " + "Please call 'experimental_compile()' first." + ) + + # Check that each CompiledTask has a valid dag_node + for idx, task in self.idx_to_task.items(): + if not hasattr(task, "dag_node") or not isinstance(task.dag_node, DAGNode): + raise ValueError( + f"Task at index {idx} does not have a valid 'dag_node'. " + "Ensure that 'experimental_compile()' completed successfully." + ) + + # Dot file for debugging + dot = graphviz.Digraph(name="compiled_graph", format=format) + # Give every actor a unique color, colors between 24k -> 40k tested as readable + # other colors may be too dark, especially when wrapping back around to 0 + actor_id_to_color = defaultdict( + lambda: f"#{((len(actor_id_to_color) * 2000 + 24000) % 0xFFFFFF):06X}" + ) + # Add nodes with task information + for idx, task in self.idx_to_task.items(): + dag_node = task.dag_node + # Initialize the label and attributes + label = f"Task {idx}\n" + shape = "oval" # Default shape + style = "filled" + fillcolor = "" + + # Handle different types of dag_node + if isinstance(dag_node, InputNode): + label += "InputNode" + shape = "rectangle" + fillcolor = "lightblue" + elif isinstance(dag_node, InputAttributeNode): + label += f"InputAttributeNode[{dag_node.key}]" + shape = "rectangle" + fillcolor = "lightblue" + elif isinstance(dag_node, MultiOutputNode): + label += "MultiOutputNode" + shape = "rectangle" + fillcolor = "yellow" + elif isinstance(dag_node, ClassMethodNode): + if dag_node.is_class_method_call: + # Class Method Call Node + method_name = dag_node.get_method_name() + actor = dag_node._get_actor_handle() + if actor: + class_name = ( + actor._ray_actor_creation_function_descriptor.class_name + ) + actor_id = actor._actor_id.hex() + label += f"Actor: {class_name}\n" + label += f"ID: {actor_id[:6]}...\n" + label += f"Method: {method_name}" + fillcolor = actor_id_to_color[actor_id] + else: + label += f"Method: {method_name}" + fillcolor = "lightgreen" + shape = "oval" + elif dag_node.is_class_method_output: + # Class Method Output Node + label += f"ClassMethodOutputNode[{dag_node.output_idx}]" + shape = "rectangle" + fillcolor = "orange" + else: + # Unexpected ClassMethodNode + label += "ClassMethodNode" + shape = "diamond" + fillcolor = "red" + else: + # Unexpected node type + label += type(dag_node).__name__ + shape = "diamond" + fillcolor = "red" + + # Add the node to the graph with attributes + dot.node(str(idx), label, shape=shape, style=style, fillcolor=fillcolor) + channel_type_str = ( + ( + type(dag_node.type_hint).__name__ + if dag_node.type_hint + else "UnknownType" + ) + + "\n" + if channel_details + else None + ) + + # This logic is built on the assumption that there will only be multiple + # output channels if the task has multiple returns + # case: task with one output + if len(task.output_channels) == 1: + for downstream_node in task.dag_node._downstream_nodes: + downstream_idx = self.dag_node_to_idx[downstream_node] + edge_label = None + if channel_details: + edge_label = channel_type_str + edge_label += self.get_channel_details( + task.output_channels[0], + ( + downstream_node._get_actor_handle()._actor_id.hex() + if type(downstream_node) is ClassMethodNode + else self._proxy_actor._actor_id.hex() + ), + ) + dot.edge(str(idx), str(downstream_idx), label=edge_label) + # case: multi return, output channels connect to class method output nodes + elif len(task.output_channels) > 1: + assert len(task.output_idxs) == len(task.output_channels) + for output_channel, downstream_idx in zip( + task.output_channels, task.output_node_idxs + ): + edge_label = None + if channel_details: + edge_label = channel_type_str + edge_label += self.get_channel_details( + output_channel, + task.dag_node._get_actor_handle()._actor_id.hex(), + ) + dot.edge(str(idx), str(downstream_idx), label=edge_label) + if type(task.dag_node) is InputAttributeNode: + # Add an edge from the InputAttributeNode to the InputNode + dot.edge(str(self.input_task_idx), str(idx)) + dot.render(filename, view=view) + return dot.source + + def _register_input_output_custom_serializer(self): + """ + Register custom serializers for input, input attribute, and output nodes. + """ + assert self.input_task_idx is not None + assert self.output_task_idx is not None + + # Register custom serializers for input node. + input_task = self.idx_to_task[self.input_task_idx] + input_task.dag_node.type_hint.register_custom_serializer() + + # Register custom serializers for input attribute nodes. + for input_attr_task_idx in self.input_attr_task_idxs: + input_attr_task = self.idx_to_task[input_attr_task_idx] + input_attr_task.dag_node.type_hint.register_custom_serializer() + + # Register custom serializers for output nodes. + for output in self.idx_to_task[self.output_task_idx].args: + output.type_hint.register_custom_serializer() + + def teardown(self, kill_actors: bool = False): + """ + Teardown and cancel all actor tasks for this DAG. After this + function returns, the actors should be available to execute new tasks + or compile a new DAG. + + Note: This method is automatically called when the CompiledDAG is destructed + or the script exits. However, this should be explicitly called before compiling + another graph on the same actors. Python may not garbage collect the + CompiledDAG object immediately when you may expect. + """ + if self._is_teardown: + return + + monitor = getattr(self, "_monitor", None) + if monitor is not None: + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + monitor.teardown(kill_actors=kill_actors) + monitor.join(timeout=ctx.teardown_timeout) + # We do not log a warning here if the thread is still alive because + # wait_teardown already logs upon teardown_timeout. + + self._is_teardown = True + + def __del__(self): + self.teardown() + + +@DeveloperAPI +def build_compiled_dag_from_ray_dag( + dag: "ray.dag.DAGNode", + submit_timeout: Optional[float] = None, + buffer_size_bytes: Optional[int] = None, + enable_asyncio: bool = False, + max_inflight_executions: Optional[int] = None, + max_buffered_results: Optional[int] = None, + overlap_gpu_communication: Optional[bool] = None, + default_communicator: Optional[Union[Communicator, str]] = "create", +) -> "CompiledDAG": + compiled_dag = CompiledDAG( + submit_timeout, + buffer_size_bytes, + enable_asyncio, + max_inflight_executions, + max_buffered_results, + overlap_gpu_communication, + default_communicator, + ) + + def _build_compiled_dag(node): + compiled_dag._add_node(node) + return node + + root = dag._find_root() + root.traverse_and_apply(_build_compiled_dag) + compiled_dag._get_or_compile() + global _compiled_dags + _compiled_dags[compiled_dag.get_id()] = compiled_dag + return compiled_dag diff --git a/lib/python3.12/site-packages/ray/dag/conftest.py b/lib/python3.12/site-packages/ray/dag/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..a6a1a22b89a8d92895357ec851d5181854b3e22f --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/conftest.py @@ -0,0 +1,17 @@ +import os + +import pytest + +import ray + +TEST_NAMESPACE = "ray_dag_test_namespace" + + +@pytest.fixture(scope="session") +def shared_ray_instance(): + # Remove ray address for test ray cluster in case we have + # lingering RAY_ADDRESS="http://127.0.0.1:8265" from previous local job + # submissions. + if "RAY_ADDRESS" in os.environ: + del os.environ["RAY_ADDRESS"] + yield ray.init(num_cpus=16, namespace=TEST_NAMESPACE, log_to_driver=True) diff --git a/lib/python3.12/site-packages/ray/dag/constants.py b/lib/python3.12/site-packages/ray/dag/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..8e35e1140f1eab2b08d5aca52b81f11f00ef2ad9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/constants.py @@ -0,0 +1,40 @@ +import os + +# Reserved keys used to handle ClassMethodNode in Ray DAG building. +PARENT_CLASS_NODE_KEY = "parent_class_node" +PREV_CLASS_METHOD_CALL_KEY = "prev_class_method_call" +BIND_INDEX_KEY = "bind_index" +IS_CLASS_METHOD_OUTPUT_KEY = "is_class_method_output" + +# Reserved keys used to handle CollectiveOutputNode in Ray DAG building. +COLLECTIVE_OPERATION_KEY = "collective_operation" + +# Reserved key to distinguish DAGNode type and avoid collision with user dict. +DAGNODE_TYPE_KEY = "__dag_node_type__" + +# Feature flag to turn off the deadlock detection. +RAY_CGRAPH_ENABLE_DETECT_DEADLOCK = ( + os.environ.get("RAY_CGRAPH_ENABLE_DETECT_DEADLOCK", "1") == "1" +) + +# Feature flag to turn on profiling. +RAY_CGRAPH_ENABLE_PROFILING = os.environ.get("RAY_CGRAPH_ENABLE_PROFILING", "0") == "1" + +# Feature flag to turn on NVTX (NVIDIA Tools Extension Library) profiling. +# With this flag, Compiled Graph uses nvtx to automatically annotate and profile +# function calls during each actor's execution loop. +# This cannot be used together with RAY_CGRAPH_ENABLE_TORCH_PROFILING. +RAY_CGRAPH_ENABLE_NVTX_PROFILING = ( + os.environ.get("RAY_CGRAPH_ENABLE_NVTX_PROFILING", "0") == "1" +) + +# Feature flag to turn on torch profiling. +# This cannot be used together with RAY_CGRAPH_ENABLE_NVTX_PROFILING. +RAY_CGRAPH_ENABLE_TORCH_PROFILING = ( + os.environ.get("RAY_CGRAPH_ENABLE_TORCH_PROFILING", "0") == "1" +) + +# Feature flag to turn on visualization of the execution schedule. +RAY_CGRAPH_VISUALIZE_SCHEDULE = ( + os.environ.get("RAY_CGRAPH_VISUALIZE_SCHEDULE", "0") == "1" +) diff --git a/lib/python3.12/site-packages/ray/dag/context.py b/lib/python3.12/site-packages/ray/dag/context.py new file mode 100644 index 0000000000000000000000000000000000000000..89fb981eb01941e2e9b563fd984fe5c920aa6368 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/context.py @@ -0,0 +1,108 @@ +import os +import threading +from dataclasses import dataclass +from typing import Optional + +from ray.util.annotations import DeveloperAPI + +# The context singleton on this process. +_default_context: "Optional[DAGContext]" = None +_context_lock = threading.Lock() + +DEFAULT_SUBMIT_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_submit_timeout", 10)) +DEFAULT_GET_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_get_timeout", 10)) +DEFAULT_TEARDOWN_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_teardown_timeout", 30)) +DEFAULT_READ_ITERATION_TIMEOUT_S = float( + os.environ.get("RAY_CGRAPH_read_iteration_timeout_s", 0.1) +) +# Default buffer size is 1MB. +DEFAULT_BUFFER_SIZE_BYTES = int(os.environ.get("RAY_CGRAPH_buffer_size_bytes", 1e6)) +# The default number of in-flight executions that can be submitted before consuming the +# output. +DEFAULT_MAX_INFLIGHT_EXECUTIONS = int( + os.environ.get("RAY_CGRAPH_max_inflight_executions", 10) +) + +# The default number of results that can be buffered at the driver. +DEFAULT_MAX_BUFFERED_RESULTS = int( + os.environ.get("RAY_CGRAPH_max_buffered_results", 1000) +) + +DEFAULT_OVERLAP_GPU_COMMUNICATION = bool( + os.environ.get("RAY_CGRAPH_overlap_gpu_communication", 0) +) + + +@DeveloperAPI +@dataclass +class DAGContext: + """Global settings for Ray DAG. + + You can configure parameters in the DAGContext by setting the environment + variables, `RAY_CGRAPH_` (e.g., `RAY_CGRAPH_buffer_size_bytes`) or Python. + + Examples: + >>> from ray.dag import DAGContext + >>> DAGContext.get_current().buffer_size_bytes + 1000000 + >>> DAGContext.get_current().buffer_size_bytes = 500 + >>> DAGContext.get_current().buffer_size_bytes + 500 + + Args: + submit_timeout: The maximum time in seconds to wait for execute() + calls. + get_timeout: The maximum time in seconds to wait when retrieving + a result from the DAG during `ray.get`. This should be set to a + value higher than the expected time to execute the entire DAG. + teardown_timeout: The maximum time in seconds to wait for the DAG to + cleanly shut down. + read_iteration_timeout: The timeout in seconds for each read iteration + that reads one of the input channels. If the timeout is reached, the + read operation will be interrupted and will try to read the next + input channel. It must be less than or equal to `get_timeout`. + buffer_size_bytes: The initial buffer size in bytes for messages + that can be passed between tasks in the DAG. The buffers will + be automatically resized if larger messages are written to the + channel. + max_inflight_executions: The maximum number of in-flight executions that + can be submitted via `execute` or `execute_async` before consuming + the output using `ray.get()`. If the caller submits more executions, + `RayCgraphCapacityExceeded` is raised. + overlap_gpu_communication: (experimental) Whether to overlap GPU + communication with computation during DAG execution. If True, the + communication and computation can be overlapped, which can improve + the performance of the DAG execution. + """ + + submit_timeout: int = DEFAULT_SUBMIT_TIMEOUT_S + get_timeout: int = DEFAULT_GET_TIMEOUT_S + teardown_timeout: int = DEFAULT_TEARDOWN_TIMEOUT_S + read_iteration_timeout: float = DEFAULT_READ_ITERATION_TIMEOUT_S + buffer_size_bytes: int = DEFAULT_BUFFER_SIZE_BYTES + max_inflight_executions: int = DEFAULT_MAX_INFLIGHT_EXECUTIONS + max_buffered_results: int = DEFAULT_MAX_BUFFERED_RESULTS + overlap_gpu_communication: bool = DEFAULT_OVERLAP_GPU_COMMUNICATION + + def __post_init__(self): + if self.read_iteration_timeout > self.get_timeout: + raise ValueError( + "RAY_CGRAPH_read_iteration_timeout_s " + f"({self.read_iteration_timeout}) must be less than or equal to " + f"RAY_CGRAPH_get_timeout ({self.get_timeout})" + ) + + @staticmethod + def get_current() -> "DAGContext": + """Get or create a singleton context. + + If the context has not yet been created in this process, it will be + initialized with default settings. + """ + global _default_context + + with _context_lock: + if _default_context is None: + _default_context = DAGContext() + + return _default_context diff --git a/lib/python3.12/site-packages/ray/dag/dag_node.py b/lib/python3.12/site-packages/ray/dag/dag_node.py new file mode 100644 index 0000000000000000000000000000000000000000..72eb29e93c7a130ca97594e96b661d3df3456e14 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/dag_node.py @@ -0,0 +1,705 @@ +import asyncio +import copy +import uuid +from itertools import chain +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + TypeVar, + Union, +) + +import ray +from ray.dag.base import DAGNodeBase +from ray.dag.compiled_dag_node import build_compiled_dag_from_ray_dag +from ray.dag.py_obj_scanner import _PyObjScanner +from ray.experimental.channel import ChannelOutputType +from ray.experimental.channel.auto_transport_type import AutoTransportType +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.torch_tensor_type import TorchTensorType +from ray.experimental.util.types import Device +from ray.util.annotations import DeveloperAPI + +T = TypeVar("T") + + +@DeveloperAPI +class DAGNode(DAGNodeBase): + """Abstract class for a node in a Ray task graph. + + A node has a type (e.g., FunctionNode), data (e.g., function options and + body), arguments (Python values, DAGNodes, and DAGNodes nested within Python + argument values) and options (Ray API .options() used for function, class + or class method) + """ + + def __init__( + self, + args: Tuple[Any], + kwargs: Dict[str, Any], + options: Dict[str, Any], + other_args_to_resolve: Dict[str, Any], + ): + """ + args: + args (Tuple[Any]): Bound node arguments. + ex: func_or_class.bind(1) + kwargs (Dict[str, Any]): Bound node keyword arguments. + ex: func_or_class.bind(a=1) + options (Dict[str, Any]): Bound node options arguments. + ex: func_or_class.options(num_cpus=2) + other_args_to_resolve (Dict[str, Any]): Bound kwargs to resolve + that's specific to subclass implementation without exposing + as args in base class, example: ClassMethodNode + """ + self._bound_args: Tuple[Any] = args or [] + self._bound_kwargs: Dict[str, Any] = kwargs or {} + self._bound_options: Dict[str, Any] = options or {} + self._bound_other_args_to_resolve: Optional[Dict[str, Any]] = ( + other_args_to_resolve or {} + ) + + # The list of nodes that use this DAG node as an argument. + self._downstream_nodes: List["DAGNode"] = [] + + # UUID that is not changed over copies of this node. + self._stable_uuid = uuid.uuid4().hex + + # Indicates whether this DAG node contains nested DAG nodes. + # Nested DAG nodes are allowed in traditional DAGs but not + # in Ray Compiled Graphs, except for MultiOutputNode. + self._args_contain_nested_dag_node = False + + # The list of nodes that this DAG node uses as an argument. + self._upstream_nodes: List["DAGNode"] = self._collect_upstream_nodes() + + # Cached values from last call to execute() + self.cache_from_last_execute = {} + + self._type_hint: ChannelOutputType = ChannelOutputType() + + # If the original type hint is an AutoTransportType, we make a copy + # here when it is resolved to the actual type, as additional debugging + # information. Otherwise, it is None. + self._original_type_hint: Optional[ChannelOutputType] = None + + # Whether this node calls `experimental_compile`. + self.is_cgraph_output_node = False + + def _collect_upstream_nodes(self) -> List["DAGNode"]: + """ + Retrieve upstream nodes and update their downstream dependencies. + + Currently, the DAG assumes that all DAGNodes in `args`, `kwargs`, and + `other_args_to_resolve` are upstream nodes. However, Ray Compiled Graphs + builds the upstream/downstream relationship based only on args. Be cautious + when persisting DAGNodes in `other_args_to_resolve` and kwargs in the future. + + TODO (kevin85421): Currently, the upstream nodes and downstream nodes have + circular references. Therefore, it relies on the garbage collector to clean + them up instead of reference counting. We should consider using weak references + to avoid circular references. + """ + upstream_nodes: List["DAGNode"] = [] + + # Ray Compiled Graphs do not allow nested DAG nodes in arguments. + # Specifically, a DAGNode should not be placed inside any type of + # container. However, we only know if this is a compiled graph + # when calling `experimental_compile`. Therefore, we need to check + # in advance if the arguments contain nested DAG nodes and raise + # an error after compilation. + assert hasattr(self._bound_args, "__iter__") + for arg in self._bound_args: + if isinstance(arg, DAGNode): + upstream_nodes.append(arg) + else: + scanner = _PyObjScanner() + dag_nodes = scanner.find_nodes(arg) + upstream_nodes.extend(dag_nodes) + scanner.clear() + self._args_contain_nested_dag_node = len(dag_nodes) > 0 + + scanner = _PyObjScanner() + other_upstream_nodes: List["DAGNode"] = scanner.find_nodes( + [ + self._bound_kwargs, + self._bound_other_args_to_resolve, + ] + ) + upstream_nodes.extend(other_upstream_nodes) + scanner.clear() + # Update dependencies. + for upstream_node in upstream_nodes: + upstream_node._downstream_nodes.append(self) + return upstream_nodes + + def with_tensor_transport( + self, + transport: Optional[Union[str, Communicator]] = "auto", + device: Literal["default", "cpu", "gpu", "cuda"] = "default", + _static_shape: bool = False, + _direct_return: bool = False, + ): + """ + Configure the torch tensor transport for this node. + + Args: + transport: Specifies the tensor transport mechanism. + - "accelerator": Tensors are communicated using accelerator-specific backends + (e.g., NCCL, XLA, or vendor-provided transport). This is the recommended option + for most use cases, as it supports extensibility and future hardware backends. + - "nccl": Tensors are passed explicitly via NCCL. This option is kept for + backwards compatibility and may be removed in the future. Use "accelerator" + instead unless you have legacy requirements. + - "shm": Tensors are passed via host shared memory and gRPC. Typically used + when accelerator-based transport is unavailable or not suitable. + - "auto" (default): The system automatically selects the appropriate transport + mechanism based on the sender and receiver, usually preferring accelerator-based + transport when available. + device: The target device to use for the tensor transport. + "default": The tensor will maintain its original device placement from the sender + "cpu": The tensor will be explicitly moved to CPU device in the receiver + "gpu" or "cuda": The tensor will be explicitly moved to GPU device in the receiver + _static_shape: A hint indicating whether the shape(s) and dtype(s) + of tensor(s) contained in this value always remain the same + across different executions of the DAG. If this is True, the + transport will be more efficient. + _direct_return: Whether the tensor is sent directly or inside of + other data. If a "nccl" transport is used, this allows the + sender and receiver to eliminate performance overhead from + an additional data transfer. + """ + try: + device = Device(device) + except ValueError: + valid_devices = ", ".join(f"'{d.value}'" for d in Device) + raise ValueError( + f"Invalid device '{device}'. Valid options are: {valid_devices}." + ) + if transport == "auto": + self._type_hint = AutoTransportType( + device=device, + _static_shape=_static_shape, + _direct_return=_direct_return, + ) + elif transport == "nccl": + self._type_hint = TorchTensorType( + transport="accelerator", + device=device, + _static_shape=_static_shape, + _direct_return=_direct_return, + ) + elif transport == "accelerator": + self._type_hint = TorchTensorType( + transport="accelerator", + device=device, + _static_shape=_static_shape, + _direct_return=_direct_return, + ) + elif transport == "shm": + self._type_hint = TorchTensorType( + device=device, + _static_shape=_static_shape, + _direct_return=_direct_return, + ) + else: + if not isinstance(transport, Communicator): + raise ValueError( + f"Invalid transport type: {transport}. " + "Transport must be one of 'auto', 'nccl', 'shm', 'accelerator' or " + "an instance of Communicator type." + ) + self._type_hint = TorchTensorType( + transport=transport, + device=device, + _static_shape=_static_shape, + _direct_return=_direct_return, + ) + return self + + @property + def type_hint(self) -> ChannelOutputType: + return self._type_hint + + @type_hint.setter + def type_hint(self, type_hint: ChannelOutputType) -> None: + if isinstance(self._type_hint, AutoTransportType): + self._original_type_hint = self._type_hint + self._type_hint = type_hint + + def get_args(self) -> Tuple[Any]: + """Return the tuple of arguments for this node.""" + + return self._bound_args + + def get_kwargs(self) -> Dict[str, Any]: + """Return the dict of keyword arguments for this node.""" + + return self._bound_kwargs.copy() + + def get_options(self) -> Dict[str, Any]: + """Return the dict of options arguments for this node.""" + + return self._bound_options.copy() + + def get_other_args_to_resolve(self) -> Dict[str, Any]: + """Return the dict of other args to resolve arguments for this node.""" + return self._bound_other_args_to_resolve.copy() + + def get_stable_uuid(self) -> str: + """Return stable uuid for this node. + 1) Generated only once at first instance creation + 2) Stable across pickling, replacement and JSON serialization. + """ + return self._stable_uuid + + async def get_object_refs_from_last_execute(self) -> Dict[str, Any]: + """Gets cached object refs from the last call to execute(). + + After this DAG is executed through execute(), retrieves a map between node + UUID to a reference to the return value of the default executor on that node. + """ + cache = {} + for node_uuid, value in self.cache_from_last_execute.items(): + if isinstance(value, asyncio.Task): + cache[node_uuid] = await value + else: + cache[node_uuid] = value + + return cache + + def clear_cache(self): + self.cache_from_last_execute = {} + + def experimental_compile( + self, + _submit_timeout: Optional[float] = None, + _buffer_size_bytes: Optional[int] = None, + enable_asyncio: bool = False, + _max_inflight_executions: Optional[int] = None, + _max_buffered_results: Optional[int] = None, + _overlap_gpu_communication: Optional[bool] = None, + _default_communicator: Optional[Union[Communicator, str]] = "create", + ) -> "ray.dag.CompiledDAG": + """Compile an accelerated execution path for this DAG. + + Args: + _submit_timeout: The maximum time in seconds to wait for execute() calls. + None means using default timeout, 0 means immediate timeout + (immediate success or timeout without blocking), -1 means + infinite timeout (block indefinitely). + _buffer_size_bytes: The initial buffer size in bytes for messages + that can be passed between tasks in the DAG. The buffers will + be automatically resized if larger messages are written to the + channel. + enable_asyncio: Whether to enable asyncio for this DAG. + _max_inflight_executions: The maximum number of in-flight executions that + can be submitted via `execute` or `execute_async` before consuming + the output using `ray.get()`. If the caller submits more executions, + `RayCgraphCapacityExceeded` is raised. + _max_buffered_results: The maximum number of results that can be + buffered at the driver. If more than this number of results + are buffered, `RayCgraphCapacityExceeded` is raised. Note that + when result corresponding to an execution is retrieved + (by calling `ray.get()` on a `CompiledDAGRef` or + `CompiledDAGRef` or await on a `CompiledDAGFuture`), results + corresponding to earlier executions that have not been retrieved + yet are buffered. + _overlap_gpu_communication: (experimental) Whether to overlap GPU + communication with computation during DAG execution. If True, the + communication and computation can be overlapped, which can improve + the performance of the DAG execution. If None, the default value + will be used. + _default_communicator: The default communicator to use to transfer + tensors. Three types of values are valid. (1) Communicator: + For p2p operations, this is the default communicator + to use for nodes annotated with `with_tensor_transport()` and when + shared memory is not the desired option (e.g., when transport="nccl", + or when transport="auto" for communication between two different GPUs). + For collective operations, this is the default communicator to use + when a custom communicator is not specified. + (2) "create": for each collective operation without a custom communicator + specified, a communicator is created and initialized on its involved actors, + or an already created communicator is reused if the set of actors is the same. + For all p2p operations without a custom communicator specified, it reuses + an already created collective communicator if the p2p actors are a subset. + Otherwise, a new communicator is created. + (3) None: a ValueError will be thrown if a custom communicator is not specified. + + Returns: + A compiled DAG. + """ + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + if _buffer_size_bytes is None: + _buffer_size_bytes = ctx.buffer_size_bytes + + # Validate whether this DAG node has already been compiled. + if self.is_cgraph_output_node: + raise ValueError( + "It is not allowed to call `experimental_compile` on the same DAG " + "object multiple times no matter whether `teardown` is called or not. " + "Please reuse the existing compiled DAG or create a new one." + ) + # Whether this node is an output node in the DAG. We cannot determine + # this in the constructor because the output node is determined when + # `experimental_compile` is called. + self.is_cgraph_output_node = True + return build_compiled_dag_from_ray_dag( + self, + _submit_timeout, + _buffer_size_bytes, + enable_asyncio, + _max_inflight_executions, + _max_buffered_results, + _overlap_gpu_communication, + _default_communicator, + ) + + def execute( + self, *args, _ray_cache_refs: bool = False, **kwargs + ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]: + """Execute this DAG using the Ray default executor _execute_impl(). + + Args: + _ray_cache_refs: If true, stores the default executor's return values + on each node in this DAG in a cache. These should be a mix of: + - ray.ObjectRefs pointing to the outputs of method and function nodes + - Serve handles for class nodes + - resolved values representing user input at runtime + """ + + def executor(node): + return node._execute_impl(*args, **kwargs) + + result = self.apply_recursive(executor) + if _ray_cache_refs: + self.cache_from_last_execute = executor.cache + return result + + def _get_toplevel_child_nodes(self) -> List["DAGNode"]: + """Return the list of nodes specified as top-level args. + + For example, in `f.remote(a, [b])`, only `a` is a top-level arg. + + This list of nodes are those that are typically resolved prior to + task execution in Ray. This does not include nodes nested within args. + For that, use ``_get_all_child_nodes()``. + """ + + # we use List instead of Set here because the hash key of the node + # object changes each time we create it. So if using Set here, the + # order of returned children can be different if we create the same + # nodes and dag one more time. + children = [] + for a in self.get_args(): + if isinstance(a, DAGNode): + if a not in children: + children.append(a) + for a in self.get_kwargs().values(): + if isinstance(a, DAGNode): + if a not in children: + children.append(a) + for a in self.get_other_args_to_resolve().values(): + if isinstance(a, DAGNode): + if a not in children: + children.append(a) + return children + + def _get_all_child_nodes(self) -> List["DAGNode"]: + """Return the list of nodes referenced by the args, kwargs, and + args_to_resolve in current node, even they're deeply nested. + + Examples: + f.remote(a, [b]) -> [a, b] + f.remote(a, [b], key={"nested": [c]}) -> [a, b, c] + """ + + scanner = _PyObjScanner() + # we use List instead of Set here, reason explained + # in `_get_toplevel_child_nodes`. + children = [] + for n in scanner.find_nodes( + [ + self._bound_args, + self._bound_kwargs, + self._bound_other_args_to_resolve, + ] + ): + if n not in children: + children.append(n) + scanner.clear() + return children + + def _apply_and_replace_all_child_nodes( + self, fn: "Callable[[DAGNode], T]" + ) -> "DAGNode": + """Apply and replace all immediate child nodes using a given function. + + This is a shallow replacement only. To recursively transform nodes in + the DAG, use ``apply_recursive()``. + + Args: + fn: Callable that will be applied once to each child of this node. + + Returns: + New DAGNode after replacing all child nodes. + """ + + replace_table = {} + # CloudPickler scanner object for current layer of DAGNode. Same + # scanner should be use for a full find & replace cycle. + scanner = _PyObjScanner() + # Find all first-level nested DAGNode children in args. + # Update replacement table and execute the replace. + for node in scanner.find_nodes( + [ + self._bound_args, + self._bound_kwargs, + self._bound_other_args_to_resolve, + ] + ): + if node not in replace_table: + replace_table[node] = fn(node) + new_args, new_kwargs, new_other_args_to_resolve = scanner.replace_nodes( + replace_table + ) + scanner.clear() + + # Return updated copy of self. + return self._copy( + new_args, new_kwargs, self.get_options(), new_other_args_to_resolve + ) + + def apply_recursive(self, fn: "Callable[[DAGNode], T]") -> T: + """Apply callable on each node in this DAG in a bottom-up tree walk. + + Args: + fn: Callable that will be applied once to each node in the + DAG. It will be applied recursively bottom-up, so nodes can + assume the fn has been applied to their args already. + + Returns: + Return type of the fn after application to the tree. + """ + + if not type(fn).__name__ == "_CachingFn": + + class _CachingFn: + def __init__(self, fn): + self.cache = {} + self.fn = fn + self.fn.cache = self.cache + self.input_node_uuid = None + + def __call__(self, node: "DAGNode"): + from ray.dag.input_node import InputNode + + if node._stable_uuid not in self.cache: + self.cache[node._stable_uuid] = self.fn(node) + if isinstance(node, InputNode): + if not self.input_node_uuid: + self.input_node_uuid = node._stable_uuid + elif self.input_node_uuid != node._stable_uuid: + raise AssertionError( + "Each DAG should only have one unique InputNode." + ) + return self.cache[node._stable_uuid] + + fn = _CachingFn(fn) + else: + if self._stable_uuid in fn.cache: + return fn.cache[self._stable_uuid] + + return fn( + self._apply_and_replace_all_child_nodes( + lambda node: node.apply_recursive(fn) + ) + ) + + def traverse_and_apply(self, fn: "Callable[[DAGNode], T]"): + """ + Traverse all nodes in the connected component of the DAG that contains + the `self` node, and apply the given function to each node. + """ + visited = set() + queue = [self] + cgraph_output_node: Optional[DAGNode] = None + + while queue: + node = queue.pop(0) + if node._args_contain_nested_dag_node: + self._raise_nested_dag_node_error(node._bound_args) + + if node not in visited: + if node.is_cgraph_output_node: + # Validate whether there are multiple nodes that call + # `experimental_compile`. + if cgraph_output_node is not None: + raise ValueError( + "The DAG was compiled more than once. The following two " + "nodes call `experimental_compile`: " + f"(1) {cgraph_output_node}, (2) {node}" + ) + cgraph_output_node = node + fn(node) + visited.add(node) + """ + Add all unseen downstream and upstream nodes to the queue. + This function should be called by the root of the DAG. However, + in some invalid cases, some nodes may not be descendants of the + root. Therefore, we also add upstream nodes to the queue so that + a meaningful error message can be raised when the DAG is compiled. + + ``` + with InputNode() as inp: + dag = MultiOutputNode([a1.inc.bind(inp), a2.inc.bind(1)]) + ``` + + In the above example, `a2.inc` is not a descendant of inp. If we only + add downstream nodes to the queue, the `a2.inc` node will not be visited + , and the error message will be hard to understand, such as a key error + in the compiled DAG. + """ + for neighbor in chain.from_iterable( + [node._downstream_nodes, node._upstream_nodes] + ): + if neighbor not in visited: + queue.append(neighbor) + + def _raise_nested_dag_node_error(self, args): + """ + Raise an error for nested DAGNodes in Ray Compiled Graphs. + + Args: + args: The arguments of the DAGNode. + """ + for arg in args: + if isinstance(arg, DAGNode): + continue + else: + scanner = _PyObjScanner() + dag_nodes = scanner.find_nodes([arg]) + scanner.clear() + if len(dag_nodes) > 0: + raise ValueError( + f"Found {len(dag_nodes)} DAGNodes from the arg {arg} " + f"in {self}. Please ensure that the argument is a " + "single DAGNode and that a DAGNode is not allowed to " + "be placed inside any type of container." + ) + raise AssertionError( + "A DAGNode's args should contain nested DAGNodes as args, " + "but none were found during the compilation process. This is a " + "Ray internal error. Please report this issue to the Ray team." + ) + + def _find_root(self) -> "DAGNode": + """ + Return the root node of the DAG. The root node must be an InputNode. + """ + from ray.dag.input_node import InputNode + + node = self + while not isinstance(node, InputNode): + if len(node._upstream_nodes) == 0: + raise ValueError( + "No InputNode found in the DAG: when traversing upwards, " + f"no upstream node was found for {node}." + ) + node = node._upstream_nodes[0] + return node + + def apply_functional( + self, + source_input_list: Any, + predicate_fn: Callable, + apply_fn: Callable, + ): + """ + Apply a given function to DAGNodes in source_input_list, and return + the replaced inputs without mutating or coping any DAGNode. + + Args: + source_input_list: Source inputs to extract and apply function on + all children DAGNode instances. + predicate_fn: Applied on each DAGNode instance found and determine + if we should apply function to it. Can be used to filter node + types. + apply_fn: Function to apply on the node on bound attributes. Example:: + + apply_fn = lambda node: node._get_serve_deployment_handle( + node._deployment, node._bound_other_args_to_resolve + ) + + Returns: + replaced_inputs: Outputs of apply_fn on DAGNodes in + source_input_list that passes predicate_fn. + """ + replace_table = {} + scanner = _PyObjScanner() + for node in scanner.find_nodes(source_input_list): + if predicate_fn(node) and node not in replace_table: + replace_table[node] = apply_fn(node) + + replaced_inputs = scanner.replace_nodes(replace_table) + scanner.clear() + + return replaced_inputs + + def _execute_impl( + self, *args, **kwargs + ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]: + """Execute this node, assuming args have been transformed already.""" + raise NotImplementedError + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ) -> "DAGNode": + """Return a copy of this node with the given new args.""" + raise NotImplementedError + + def _copy( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ) -> "DAGNode": + """Return a copy of this node with the given new args.""" + instance = self._copy_impl( + new_args, new_kwargs, new_options, new_other_args_to_resolve + ) + instance._stable_uuid = self._stable_uuid + instance._type_hint = copy.deepcopy(self._type_hint) + instance._original_type_hint = copy.deepcopy(self._original_type_hint) + return instance + + def __getstate__(self): + """Required due to overriding `__getattr__` else pickling fails.""" + return self.__dict__ + + def __setstate__(self, d: Dict[str, Any]): + """Required due to overriding `__getattr__` else pickling fails.""" + self.__dict__.update(d) + + def __getattr__(self, attr: str): + if attr == "bind": + raise AttributeError(f".bind() cannot be used again on {type(self)} ") + elif attr == "remote": + raise AttributeError( + f".remote() cannot be used on {type(self)}. To execute the task " + "graph for this node, use .execute()." + ) + else: + return self.__getattribute__(attr) diff --git a/lib/python3.12/site-packages/ray/dag/dag_node_operation.py b/lib/python3.12/site-packages/ray/dag/dag_node_operation.py new file mode 100644 index 0000000000000000000000000000000000000000..5a192e9f5da23c160ff3ad2fa7833ee54afeb483 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/dag_node_operation.py @@ -0,0 +1,853 @@ +import copy +import heapq +import logging +from collections import defaultdict +from enum import Enum +from functools import total_ordering +from typing import Dict, List, Optional, Set, Tuple + +import ray + +logger = logging.getLogger(__name__) + + +class _DAGNodeOperationType(Enum): + """ + There are three types of operations that a DAG node can perform: + 1. READ: Read from an input channel. + 2. COMPUTE: Execute the method corresponding to the node. + 3. WRITE: Write to an output channel. + """ + + READ = "READ" + COMPUTE = "COMPUTE" + WRITE = "WRITE" + + def viz_str(self): + """ + A string representation of the operation type to be used in visualization. + + The result string is a single character because conciseness is preferred. + """ + if self == _DAGNodeOperationType.READ: + return "R" + elif self == _DAGNodeOperationType.COMPUTE: + return "C" + elif self == _DAGNodeOperationType.WRITE: + return "W" + assert False, f"Unknown operation type: {self}" + + +class _DAGNodeOperation: + def __init__( + self, + exec_task_idx: int, + operation_type: _DAGNodeOperationType, + method_name: Optional[str] = None, + ): + """ + Args: + exec_task_idx: The index of the task that this operation belongs to + in the actor's ExecutableTask list. The index is not the same + as bind_index because there may be more tasks bound to an actor + than tasks that appear in the current compiled DAG. + operation_type: The type of operation to perform. + method_name: The name of the method that this operation originates + from. This is only for visualization and debugging purposes. + """ + self.exec_task_idx = exec_task_idx + self.type = operation_type + self.method_name = method_name + + def __repr__(self): + return ( + f"_DAGNodeOperation(" + f"exec_task_idx: {self.exec_task_idx}, " + f"type: {self.type}, " + f"method_name: {self.method_name})" + ) + + def viz_str(self): + """ + A string representation of the node to be used in visualization. + """ + return f"[{self.exec_task_idx}] {self.method_name} {self.type.viz_str()}" + + def __hash__(self): + return hash((self.exec_task_idx, self.type)) + + def __eq__(self, other): + # An operation is uniquely identified by its `exec_task_idx` and type. + # `method_name` is only for debugging purposes. + return self.exec_task_idx == other.exec_task_idx and self.type == other.type + + +@total_ordering +class _DAGOperationGraphNode: + def __init__( + self, + operation: _DAGNodeOperation, + task_idx: int, + actor_handle: "ray.actor.ActorHandle", + requires_accelerator: bool, + ): + """ + _DAGOperationGraphNode represents a node in the DAG operation graph. + It contains information about the node's in-degree, out-degree, edges, + and the operation it performs. + + Args: + operation: The operation that this node performs. The operation + can be a READ, COMPUTE, or WRITE operation. + task_idx: A unique index which can be used to index into + `CompiledDAG.idx_to_task` to get the corresponding task. + actor_handle: The actor handle to which this operation belongs. + requires_accelerator: Whether this operation requires accelerator. + """ + self.operation = operation + self.task_idx = task_idx + self.actor_handle = actor_handle + self.requires_accelerator = requires_accelerator + # The in_edges and out_edges are dicts of tuples to strings. + # Each tuple (the key) contains an integer `task_idx`, which can be + # used to index into `idx_to_task` to get the corresponding task, + # and a `_DAGNodeOperationType`, which can be READ, COMPUTE, or WRITE. + # The string (the value) is the visualization information of the edge, + # it is a tuple of a label of the edge and a boolean indicating whether + # the edge is a control dependency. + self.in_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {} + self.out_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {} + # The synchronous nodes are all the nodes that belong to the same accelerator + # operation. Each node is represented by a tuple of its task idx and type. + self.sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set() + # The pending synchronous nodes are the nodes that are pending to be executed, + # i.e., their in-degrees are zero. When a synchronous node is pending, it + # will be added to the pending synchronous nodes of all the nodes in the + # accelerator operation. + self.pending_sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set() + + def __repr__(self): + return ( + f"_DAGOperationGraphNode(" + f"operation: {self.operation}, " + f"task_idx: {self.task_idx}, " + f"actor_id: {self.actor_handle._ray_actor_id}, " + f"requires_accelerator: {self.requires_accelerator})" + ) + + def __lt__(self, other: "_DAGOperationGraphNode"): + """ + This function defines the order of the nodes in the priority queue used in + `_select_next_nodes`. The priority queue is a min-heap, so the node with + higher priority is considered "less than" the other node. + """ + if self.is_accelerator_op != other.is_accelerator_op: + # When one node is an accelerator operation and the other is not, + # prioritize the accelerator operation. + return self.is_accelerator_op + else: + # When either both nodes are accelerator operations or both nodes + # are not accelerator operations, prioritize the earlier task within + # the same actor and load balance tasks across actors. The tie is + # broken by the `task_idx`. + return (self.operation.exec_task_idx, self.task_idx) < ( + other.operation.exec_task_idx, + other.task_idx, + ) + + def __eq__(self, other: "_DAGOperationGraphNode"): + """ + Two operations are equal only when they have the same `exec_task_idx` and `type` + and belong to the same actor. + """ + return ( + self.actor_handle == other.actor_handle + and self.operation.exec_task_idx == other.operation.exec_task_idx + and self.operation.type == other.operation.type + ) + + def __hash__(self): + """ + An operation is uniquely identified by its `task_idx` and type. + """ + return hash((self.operation, self.task_idx)) + + @property + def in_degree(self) -> int: + return len(self.in_edges) + + @property + def is_ready(self) -> bool: + """ + If a node is not an accelerator operation, it is ready when it has a zero + in-degree. + If it is an accelerator operation, it is ready when all the nodes in the + operation have zero in-degrees. + """ + return self.in_degree == 0 and ( + len(self.pending_sync_idxs) == len(self.sync_idxs) + ) + + @property + def is_read(self) -> bool: + return self.operation.type == _DAGNodeOperationType.READ + + @property + def is_accelerator_read(self) -> bool: + """ + A node is an accelerator read if it is a read node and requires accelerator. + """ + return ( + self.operation.type == _DAGNodeOperationType.READ + and self.requires_accelerator + ) + + @property + def is_accelerator_compute(self) -> bool: + """ + A node is an accelerator compute if it is a compute node and requires accelerator. + """ + return ( + self.operation.type == _DAGNodeOperationType.COMPUTE + and self.requires_accelerator + ) + + @property + def is_accelerator_write(self) -> bool: + """ + A node is an accelerator write if it is a write node and requires accelerator. + """ + return ( + self.operation.type == _DAGNodeOperationType.WRITE + and self.requires_accelerator + ) + + @property + def is_accelerator_op(self) -> bool: + return ( + self.is_accelerator_read + or self.is_accelerator_compute + or self.is_accelerator_write + ) + + def viz_str(self): + """ + A string representation of the node to be used in visualization. + """ + return self.operation.viz_str() + + @property + def _actor_id(self): + return self.actor_handle._ray_actor_id.hex() + + +def _add_edge( + from_node: _DAGOperationGraphNode, + to_node: _DAGOperationGraphNode, + label: str = "", + control_dependency: bool = False, +): + """ + Add an edge from `from_node` to `to_node`. + + Args: + from_node: The node from which the edge originates. + to_node: The node to which the edge points. + label: The label of the edge. This will be used to annotate the edge + in the visualization of the execution schedule. + """ + from_node.out_edges[(to_node.task_idx, to_node.operation.type)] = ( + label, + control_dependency, + ) + to_node.in_edges[(from_node.task_idx, from_node.operation.type)] = ( + label, + control_dependency, + ) + + +def _update_pending_sync_idxs( + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], + node: _DAGOperationGraphNode, +) -> None: + """ + Update the node as pending for its synchronous nodes. + """ + idx = (node.task_idx, node.operation.type) + for task_idx, op_type in node.sync_idxs: + sync_node = graph[task_idx][op_type] + sync_node.pending_sync_idxs.add(idx) + + +def _push_candidate_node_if_ready( + actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]], + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], + node: _DAGOperationGraphNode, +) -> None: + """ + Push the node with a zero in-degree to the candidates if its operation is ready. + If it has synchronous nodes, its accelerator operation is not ready until all + the nodes are pending, then all the nodes will be pushed to the candidates. + """ + assert node.in_degree == 0, "Expected to have a zero in-degree" + # For the accelerator write node, update the in-degrees of the downstream + # accelerator read nodes and update them as pending. This is necessary because + # the data dependency edges between accelerator write and read nodes are only + # updated here. The accelerator P2P operation becomes ready after both the write + # and read nodes are marked as pending. + if node.is_accelerator_write: + for task_idx, op_type in node.out_edges: + read_node = graph[task_idx][op_type] + read_node.in_edges.pop((node.task_idx, node.operation.type)) + assert read_node.is_accelerator_read and len(read_node.in_edges) == 0 + _update_pending_sync_idxs(graph, read_node) + # For the accelerator operation node, update it as pending. + if len(node.sync_idxs) != 0: + _update_pending_sync_idxs(graph, node) + # The accelerator operation is ready when all the nodes have zero in-degrees. + # When the last node in the operation is updated as pending, push all the nodes + # to the candidates. + if node.is_ready: + if len(node.sync_idxs) == 0: + heapq.heappush( + actor_to_candidates[node.actor_handle._actor_id], + node, + ) + else: + for task_idx, op_type in node.sync_idxs: + sync_node = graph[task_idx][op_type] + heapq.heappush( + actor_to_candidates[sync_node.actor_handle._actor_id], + sync_node, + ) + + +def _select_next_nodes( + actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]], + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], +) -> Optional[List[_DAGOperationGraphNode]]: + """ + This function selects the next nodes for the topological sort to generate + execution schedule. If there are multiple candidate _DAGOperationGraphNodes, + select the node with the top priority. The priority is defined in + `_DAGOperationGraphNode.__lt__`. + + For the implementation details, we maintain a priority queue for each actor, + where the head of the priority queue is the node with the smallest `exec_task_idx`. + When a node has a zero in-degree, it is added to the corresponding actor's + priority queue. For a node other than an accelerator collective node, it is ready to be + executed if it has a zero in-degree. For an accelerator collective node, it is ready to + be executed when all the nodes in its collective operation have zero in-degrees. + + If a node is an accelerator collective node, it updates the `ready_collective_nodes` of + all the nodes in its collective operation. Unless all the nodes in its collective + group have zero in-degrees, this node is removed from the candidate list. + Eventually, exactly one accelerator collective node from its collective operation is + selected from the candidate list. + + If the selected node is an accelerator write node, select all the downstream accelerator + read nodes. If the selected node is an accelerator collective node, select all the accelerator + compute nodes in its collective operation. + + Args: + actor_to_candidates: A dictionary mapping an actor id to a list of + candidate nodes. The list is maintained as a priority queue, so + the head of the queue, i.e., `candidates[0]`, is the node with + the smallest `bind_index`. + graph: A dictionary mapping the index of a task to a dictionary of its + _DAGOperationGraphNodes for different operations. + + Returns: + A list of _DAGOperationGraphNodes to be placed into the corresponding + execution schedules. + """ + top_priority_node = None + for candidates in actor_to_candidates.values(): + if len(candidates) == 0: + continue + if top_priority_node is None or candidates[0] < top_priority_node: + top_priority_node = candidates[0] + + if top_priority_node is None: + return None + next_nodes = [top_priority_node] + + # Select all the synchronous nodes in the accelerator operation. + if len(top_priority_node.sync_idxs) != 0: + for task_idx, op_type in top_priority_node.sync_idxs: + node = graph[task_idx][op_type] + if node != top_priority_node: + next_nodes.append(node) + + # Remove the selected nodes from the candidates. + for node in next_nodes: + candidates = actor_to_candidates[node.actor_handle._actor_id] + candidates.remove(node) + heapq.heapify(candidates) + + # Remove the selected nodes from the candidates. + for node in next_nodes: + candidates = actor_to_candidates[node.actor_handle._actor_id] + # The accelerator read nodes are not added to the candidates. + if node in candidates: + candidates.remove(node) + heapq.heapify(candidates) + + return next_nodes + + +def _build_dag_node_operation_graph( + idx_to_task: Dict[int, "ray.dag.compiled_dag_node.CompiledTask"], + actor_to_operation_nodes: Dict[ + "ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]] + ], +) -> Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]]: + """ + Generate a DAG node operation graph by adding edges based on the + following rules: + + #1 Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which + belong to the same task. + #2 Add an edge from COMPUTE with bind_index i to COMPUTE with bind_index + i+1 if they belong to the same actor. + #3 Add an edge from WRITE of the writer task to READ of the reader task. + + This is the step one of building an execution schedule for each actor. + + Args: + idx_to_task: A dictionary that maps the `task_idx` to the `CompiledTask`. + `CompiledTask` contains information about a DAGNode and its downstream + nodes. + + actor_to_operation_nodes: A dictionary that maps an actor handle to + a list of lists of _DAGOperationGraphNode. For the same actor, the + index of the outer list corresponds to the index of the ExecutableTask + in the list of `executable_tasks` in `actor_to_executable_tasks`. In + the inner list, the order of operations is READ, COMPUTE, and WRITE. + + Returns: + A graph where each node is a _DAGOperationGraphNode. The key is `task_idx`, + the index to retrieve its task from `idx_to_task`, and the value is a + dictionary that maps the _DAGNodeOperationType (READ, COMPUTE, or WRITE) + to the corresponding _DAGOperationGraphNode + """ + assert idx_to_task + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]] = {} + + for _, operation_nodes_list in actor_to_operation_nodes.items(): + prev_compute_node = None + for operation_nodes in operation_nodes_list: + task_idx = operation_nodes[0].task_idx + read_node, compute_node, write_node = ( + operation_nodes[0], + operation_nodes[1], + operation_nodes[2], + ) + # Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which + # belong to the same task. + _add_edge(read_node, compute_node) + _add_edge(compute_node, write_node) + # Add an edge from COMPUTE with `bind_index` i to COMPUTE with + # `bind_index` i+1 if they belong to the same actor. + if prev_compute_node is not None: + _add_edge(prev_compute_node, compute_node, "", True) + prev_compute_node = compute_node + assert task_idx not in graph + graph[task_idx] = { + _DAGNodeOperationType.READ: read_node, + _DAGNodeOperationType.COMPUTE: compute_node, + _DAGNodeOperationType.WRITE: write_node, + } + + # Import `ray.dag` here to avoid circular import. + from ray.dag import ClassMethodNode, CollectiveOutputNode, MultiOutputNode + from ray.dag.collective_node import _CollectiveOperation + + # Add an edge from WRITE of the writer task to READ of the reader task. + # Set synchronous nodes for accelerator P2P operations. + for task_idx, task in idx_to_task.items(): + if not ( + isinstance(task.dag_node, ClassMethodNode) + or isinstance(task.dag_node, CollectiveOutputNode) + ): + # The graph is used to generate an execution schedule for each actor. + # The edge from the InputNode has no impact on the final execution + # schedule. + continue + if ( + isinstance(task.dag_node, ClassMethodNode) + and task.dag_node.is_class_method_output + ): + # Class method output node dependencies are handled at its upstream: + # i.e., class method node + continue + for downstream_task_idx in task.downstream_task_idxs: + downstream_dag_node = idx_to_task[downstream_task_idx].dag_node + if isinstance(downstream_dag_node, MultiOutputNode): + continue + write_node = graph[task_idx][_DAGNodeOperationType.WRITE] + if ( + isinstance(downstream_dag_node, ClassMethodNode) + and downstream_dag_node.is_class_method_output + ): + consumer_idxs = idx_to_task[downstream_task_idx].downstream_task_idxs + for consumer_idx in consumer_idxs: + if consumer_idx in graph: + read_node = graph[consumer_idx][_DAGNodeOperationType.READ] + _add_edge( + write_node, + read_node, + "accelerator" if write_node.requires_accelerator else "shm", + ) + if write_node.requires_accelerator: + idxs = { + (task_idx, _DAGNodeOperationType.WRITE), + (consumer_idx, _DAGNodeOperationType.READ), + } + for node in [write_node, read_node]: + node.sync_idxs.update(idxs) + continue + read_node = graph[downstream_task_idx][_DAGNodeOperationType.READ] + _add_edge( + write_node, + read_node, + "accelerator" if write_node.requires_accelerator else "shm", + ) + if write_node.requires_accelerator: + idxs = { + (task_idx, _DAGNodeOperationType.WRITE), + (downstream_task_idx, _DAGNodeOperationType.READ), + } + for node in [write_node, read_node]: + node.sync_idxs.update(idxs) + + # Set synchronous nodes for accelerator collective operations. + collective_op_to_idxs: Dict[ + _CollectiveOperation, Set[Tuple[int, _DAGNodeOperationType]] + ] = defaultdict(set) + for task_idx, task in idx_to_task.items(): + if ( + isinstance(task.dag_node, CollectiveOutputNode) + and not task.dag_node.is_class_method_output + ): + collective_op_to_idxs[task.dag_node.collective_op].add( + (task_idx, _DAGNodeOperationType.COMPUTE) + ) + for idxs in collective_op_to_idxs.values(): + for task_idx, op_type in idxs: + graph[task_idx][op_type].sync_idxs = idxs + + return graph + + +def _actor_viz_label(actor: "ray.actor.ActorHandle"): + """ + Returns the label of an actor in the visualization of the execution schedule. + + Args: + actor: The actor to be represented. + """ + class_name = actor._ray_actor_creation_function_descriptor.class_name + actor_id = actor._ray_actor_id.hex() + return f"Actor class name: {class_name}\nActor ID: {actor_id}" + + +def _node_viz_id_and_label( + node: _DAGOperationGraphNode, idx: int, optimized_index: int +): + """ + Returns the visualization id and label of a node. The visualization id is unique + across all nodes. + + Args: + node: The node to be represented. + idx: The index of the node in the execution schedule. + optimized_index: The index of the node in the optimized execution schedule. + """ + node_viz_label = node.viz_str() + f" {idx},{optimized_index}" + node_viz_id = f"{node._actor_id}_{node_viz_label}" + return node_viz_id, node_viz_label + + +def _visualize_execution_schedule( + actor_to_execution_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGOperationGraphNode] + ], + actor_to_overlapped_schedule: Optional[ + Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]] + ], + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], +): + """ + Visualize the execution schedule for each actor. + + The visualization will be saved as a PNG file named `compiled_graph_schedule.png`. + Details of the visualization: # noqa + + Node description format: + [] , + + Node description fields: + operation: is R(READ), C(COMPUTE), or W(WRITE) + orig_index: the index in the original execution schedule + overlap_index: the index in the overlap-communication optimized execution schedule + If this is different from orig_index, the node is highlighted in red color + + Node grouping: + The nodes belonging to the same actor are grouped in the same rectangle + The actor class name and the actor id are shown in the rectangle + + Edges: + black color (without label): data dependency + black color (annotated with "shm"): shared memory channel + blue color (annotated with "accelerator): accelerator channel + dashed edge: control dependency between compute operations + + Args: + actor_to_execution_schedule: A dictionary that maps an actor handle to + the execution schedule which is a list of operation nodes. + actor_to_overlapped_schedule: A dictionary that maps an actor handle to the + optimized execution schedule which is a list of operation nodes. + graph: A graph where each node is a _DAGOperationGraphNode. The key is + `task_idx`, the index to retrieve its task from `idx_to_task`, and + the value is a dictionary that maps the _DAGNodeOperationType (READ, + COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is + generated by `_build_dag_node_operation_graph`. + """ + try: + import graphviz + except ImportError: + raise ImportError( + "Please install graphviz to visualize the execution schedule. " + "You can install it by running `pip install graphviz`." + ) + + dot = graphviz.Digraph(comment="DAG") + # A dictionary that maps a node to its visualization id + node_to_viz_id: Dict[_DAGOperationGraphNode, str] = {} + + if actor_to_overlapped_schedule is None: + # TODO(rui): make the visualization more concise by only displaying + # the original schedule + actor_to_overlapped_schedule = actor_to_execution_schedule + for actor, execution_nodes in actor_to_execution_schedule.items(): + overlapped_schedule = actor_to_overlapped_schedule[actor] + node_to_optimized_index = { + node: i for i, node in enumerate(overlapped_schedule) + } + + actor_id = actor._ray_actor_id.hex() + with dot.subgraph(name=f"cluster_{actor_id}") as subgraph: + subgraph.attr(rank=actor_id, label=_actor_viz_label(actor)) + for i, node in enumerate(execution_nodes): + optimized_index = node_to_optimized_index.get(node) + node_viz_id, node_viz_label = _node_viz_id_and_label( + node, i, optimized_index + ) + color = "red" if optimized_index != i else "black" + subgraph.node(node_viz_id, node_viz_label, color=color) + node_to_viz_id[node] = node_viz_id + + for actor, execution_nodes in actor_to_execution_schedule.items(): + for i, node in enumerate(execution_nodes): + node_viz_id = node_to_viz_id[node] + for out_edge, viz_info in node.out_edges.items(): + label, control_dependency = viz_info + out_task_idx, out_op_type = out_edge + out_node = graph[out_task_idx][out_op_type] + out_node_viz_id = node_to_viz_id[out_node] + color = "blue" if label == "accelerator" else "black" + style = "dashed" if control_dependency else "solid" + dot.edge( + node_viz_id, out_node_viz_id, label=label, color=color, style=style + ) + + # Add legend + with dot.subgraph(name="cluster_legend") as legend: + legend.attr(label="Legend", labelloc="t", fontsize="20", bgcolor="lightgrey") + + # Single node and its explanation + legend.node("example_node", "[0] bwd C 10,10\n") + explanation = ( + '<' # noqa + '' + '' # noqa + "" + '' + '' # noqa + '' # noqa + '' # noqa + '' # noqa + "" + '' + '' # noqa + '' # noqa + "" + '' + '' # noqa + '' # noqa + '' # noqa + '' # noqa + "
Node description format:
[<task_index>] <method_name> <operation> <orig_index>, <overlap_index>
Node description fields:
operation: is R(READ), C(COMPUTE), or W(WRITE)
orig_index: the index in the original execution schedule
overlap_index: the index in the overlap-communication optimized execution schedule
If this is different from orig_index, the node is highlighted in red color
Node grouping:
The nodes belonging to the same actor are grouped in the same rectangle
The actor class name and the actor id are shown in the rectangle
Edges:
black color (without label): data dependency
black color (annotated with "shm"): shared memory channel
blue color (annotated with "accelerator): accelerator channel
dashed edge: control dependency between compute operations
>" + ) + + legend.node("example_explanation", explanation, shape="plaintext") + legend.edge("example_node", "example_explanation", style="invis") + + logger.info( + "Writing compiled graph schedule visualization " + "to compiled_graph_schedule.png" + ) + dot.render("compiled_graph_schedule", format="png", view=False) + + +def _generate_actor_to_execution_schedule( + graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]], +) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]: + """ + Generate an execution schedule for each actor. The schedule is a list of + operation nodes to be executed. The function uses a topological sort + algorithm to generate the schedule. + + Args: + graph: A graph where each node is a _DAGOperationGraphNode. The key is + `task_idx`, the index to retrieve its task from `idx_to_task`, and + the value is a dictionary that maps the _DAGNodeOperationType (READ, + COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is + generated by `_build_dag_node_operation_graph`. + + Returns: + actor_to_execution_schedule: A dictionary that maps an actor handle to + the execution schedule which is a list of operation nodes to be + executed. + """ + + # Mapping from the actor handle to the execution schedule which is a list + # of operations to be executed. + actor_to_execution_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGOperationGraphNode] + ] = defaultdict(list) + + # A dictionary mapping an actor id to a list of candidate nodes. The list + # is maintained as a priority queue, so the head of the queue, i.e., + # `candidates[0]`, is the node with the smallest `bind_index`. + actor_to_candidates: Dict[ + "ray._raylet.ActorID", List[_DAGOperationGraphNode] + ] = defaultdict(list) + for _, node_dict in graph.items(): + for _, node in node_dict.items(): + # A node with a zero in-degree edge means all of its dependencies + # have been satisfied, including both data and control dependencies. + # Therefore, it is a candidate for execution. + if node.in_degree == 0: + _push_candidate_node_if_ready(actor_to_candidates, graph, node) + + visited_nodes = set() + + # Use topological sort algorithm to generate the execution schedule. + while True: + # Select a list of nodes to be executed. There are three cases: + # 1. If a selected node is not an accelerator operation, only itself is returned. + # 2. If a selected node is an accelerator write operation, the corresponding accelerator + # read operations are also returned. + # 3. If a selected node is an accelerator collective operation, all the nodes in + # its collective operation are returned. + nodes = _select_next_nodes(actor_to_candidates, graph) + if nodes is None: + break + # Add the selected nodes to the execution schedule. + for node in nodes: + assert node not in visited_nodes + visited_nodes.add(node) + actor_to_execution_schedule[node.actor_handle].append(node) + # Update the in-degree of the downstream nodes. + for node in nodes: + for out_node_task_idx, out_node_type in node.out_edges: + out_node = graph[out_node_task_idx][out_node_type] + if out_node in visited_nodes: + # If the downstream node is already visited, it has been added + # to the execution schedule. They are the accelerator read nodes in + # case 2. + continue + out_node.in_edges.pop((node.task_idx, node.operation.type)) + if out_node.in_degree == 0: + _push_candidate_node_if_ready(actor_to_candidates, graph, out_node) + assert len(visited_nodes) == len(graph) * 3, "Expected all nodes to be visited" + + return actor_to_execution_schedule + + +def _generate_overlapped_execution_schedule( + actor_to_execution_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGOperationGraphNode] + ], +) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]: + """ + From an existing execution schedule, generate a new schedule by overlapping + computation and communication. + + Currently, the algorithm generates a new schedule for each actor as follows: + For each accelerator read operation (i.e., recv), scan backwards to find the nearest + compute node to swap with so that the accelerator read operation can be overlapped + with computation. + + Collective operations are not yet supported. + + Args: + actor_to_execution_schedule: A dictionary that maps an actor handle to + the existing execution schedule for the actor. The schedule is a list + is a list of operations to be executed. + + Returns: + A dictionary that maps an actor handle to the overlapped execution schedule + for the actor. + """ + + actor_to_overlapped_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGOperationGraphNode] + ] = copy.deepcopy(actor_to_execution_schedule) + for overlapped_schedule in actor_to_overlapped_schedule.values(): + for i in range(len(overlapped_schedule)): + if ( + overlapped_schedule[i].operation.type == _DAGNodeOperationType.READ + and overlapped_schedule[i].requires_accelerator + ): + # For each accelerator read operation (i.e., recv), scan backwards + # to find the nearest compute node to swap with so that + # the accelerator read operation can be overlapped with computation. + for j in range(i - 1, -1, -1): + if ( + overlapped_schedule[j].operation.type + == _DAGNodeOperationType.COMPUTE + ): + # Found a desired compute operation, make the swap + accelerator_read_op = overlapped_schedule[i] + prev_ops = overlapped_schedule[j:i] + overlapped_schedule[j + 1 : i + 1] = prev_ops + overlapped_schedule[j] = accelerator_read_op + break + if ( + overlapped_schedule[j].operation.type + == _DAGNodeOperationType.READ + or overlapped_schedule[j].operation.type + == _DAGNodeOperationType.WRITE + ) and overlapped_schedule[j].requires_accelerator: + # Found an accelerator read/write operation, skip the overlap + # optimization to keep relative order of accelerator operations + break + return actor_to_overlapped_schedule + + +def _extract_execution_schedule( + actor_to_execution_schedule: Dict[ + "ray.actor.ActorHandle", List[_DAGOperationGraphNode] + ], +) -> Dict["ray.actor.ActorHandle", List[_DAGNodeOperation]]: + """ + Extract _DAGNodeOperation from _DAGOperationGraphNode in the schedule + and discard unnecessary information. + """ + return { + actor: [node.operation for node in nodes] + for actor, nodes in actor_to_execution_schedule.items() + } diff --git a/lib/python3.12/site-packages/ray/dag/dag_operation_future.py b/lib/python3.12/site-packages/ray/dag/dag_operation_future.py new file mode 100644 index 0000000000000000000000000000000000000000..392c86286a9938139c6b0542b08e4a19ac7f4593 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/dag_operation_future.py @@ -0,0 +1,144 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, TypeVar + +from ray.experimental.channel.accelerator_context import AcceleratorContext +from ray.util.annotations import DeveloperAPI + +T = TypeVar("T") + + +@DeveloperAPI +class DAGOperationFuture(ABC, Generic[T]): + """ + A future representing the result of a DAG operation. + + This is an abstraction that is internal to each actor, + and is not exposed to the DAG caller. + """ + + @abstractmethod + def wait(self): + """ + Wait for the future and return the result of the operation. + """ + raise NotImplementedError + + +@DeveloperAPI +class ResolvedFuture(DAGOperationFuture): + """ + A future that is already resolved. Calling `wait()` on this will + immediately return the result without blocking. + """ + + def __init__(self, result): + """ + Initialize a resolved future. + + Args: + result: The result of the future. + """ + self._result = result + + def wait(self): + """ + Wait and immediately return the result. This operation will not block. + """ + return self._result + + +@DeveloperAPI +class GPUFuture(DAGOperationFuture[Any]): + """ + A future for a GPU event on a CUDA stream. + + This future wraps a buffer, and records an event on the given stream + when it is created. When the future is waited on, it makes the current + CUDA stream wait on the event, then returns the buffer. + + The buffer must be a GPU tensor produced by an earlier operation launched + on the given stream, or it could be CPU data. Then the future guarantees + that when the wait() returns, the buffer is ready on the current stream. + + The `wait()` does not block CPU. + """ + + # Caching GPU futures ensures CUDA events associated with futures are properly + # destroyed instead of relying on garbage collection. The CUDA event contained + # in a GPU future is destroyed right before removing the future from the cache. + # The dictionary key is the future ID, which is the task idx of the dag operation + # that produced the future. When a future is created, it is immediately added to + # the cache. When a future has been waited on, it is removed from the cache. + # When adding a future, if its ID is already a key in the cache, the old future + # is removed. This can happen when an exception is thrown in a previous execution + # of the dag, in which case the old future is never waited on. + # Upon dag teardown, all pending futures produced by the dag are removed. + gpu_futures: Dict[int, "GPUFuture"] = {} + + @staticmethod + def add_gpu_future(fut_id: int, fut: "GPUFuture") -> None: + """ + Cache the GPU future. + Args: + fut_id: GPU future ID. + fut: GPU future to be cached. + """ + if fut_id in GPUFuture.gpu_futures: + # The old future was not waited on because of an execution exception. + GPUFuture.gpu_futures.pop(fut_id).destroy_event() + GPUFuture.gpu_futures[fut_id] = fut + + @staticmethod + def remove_gpu_future(fut_id: int) -> None: + """ + Remove the cached GPU future and destroy its CUDA event. + Args: + fut_id: GPU future ID. + """ + if fut_id in GPUFuture.gpu_futures: + GPUFuture.gpu_futures.pop(fut_id).destroy_event() + + def __init__(self, buf: Any, fut_id: int, stream: Any = None): + """ + Initialize a GPU future on the given stream. + + Args: + buf: The buffer to return when the future is resolved. + fut_id: The future ID to cache the future. + stream: The torch stream to record the event on, this event is waited + on when the future is resolved. If None, the current stream is used. + """ + if stream is None: + stream = AcceleratorContext.get().current_stream() + + self._buf = buf + self._event = AcceleratorContext.get().create_event() + self._event.record(stream) + self._fut_id = fut_id + self._waited: bool = False + + # Cache the GPU future such that its CUDA event is properly destroyed. + GPUFuture.add_gpu_future(fut_id, self) + + def wait(self) -> Any: + """ + Wait for the future on the current CUDA stream and return the result from + the GPU operation. This operation does not block CPU. + """ + current_stream = AcceleratorContext.get().current_stream() + if not self._waited: + self._waited = True + current_stream.wait_event(self._event) + # Destroy the CUDA event after it is waited on. + GPUFuture.remove_gpu_future(self._fut_id) + + return self._buf + + def destroy_event(self) -> None: + """ + Destroy the CUDA event associated with this future. + """ + if self._event is None: + return + + self._event = None diff --git a/lib/python3.12/site-packages/ray/dag/experimental/__init__.py b/lib/python3.12/site-packages/ray/dag/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/dag/experimental/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/dag/experimental/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b262df1754ff29ba473c997cb48ecec34f7ff08 Binary files /dev/null and b/lib/python3.12/site-packages/ray/dag/experimental/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/dag/format_utils.py b/lib/python3.12/site-packages/ray/dag/format_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1428317da1c37a93800c66d3162ddffa58ad82e8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/format_utils.py @@ -0,0 +1,155 @@ +from ray.dag import DAGNode +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def get_dag_node_str( + dag_node: DAGNode, + body_line, +): + indent = _get_indentation() + other_args_to_resolve_lines = _get_other_args_to_resolve_lines( + dag_node._bound_other_args_to_resolve + ) + return ( + f"({dag_node.__class__.__name__}, {dag_node._stable_uuid})(\n" + f"{indent}body={body_line}\n" + f"{indent}args={_get_args_lines(dag_node._bound_args)}\n" + f"{indent}kwargs={_get_kwargs_lines(dag_node._bound_kwargs)}\n" + f"{indent}options={_get_options_lines(dag_node._bound_options)}\n" + f"{indent}other_args_to_resolve={other_args_to_resolve_lines}\n" + f")" + ) + + +def _get_indentation(num_spaces=4): + return " " * num_spaces + + +def _get_args_lines(bound_args): + """Pretty prints bounded args of a DAGNode, and recursively handle + DAGNode in list / dict containers. + """ + indent = _get_indentation() + lines = [] + for arg in bound_args: + if isinstance(arg, DAGNode): + node_repr_lines = str(arg).split("\n") + for node_repr_line in node_repr_lines: + lines.append(f"{indent}" + node_repr_line) + elif isinstance(arg, list): + for ele in arg: + node_repr_lines = str(ele).split("\n") + for node_repr_line in node_repr_lines: + lines.append(f"{indent}" + node_repr_line) + elif isinstance(arg, dict): + for _, val in arg.items(): + node_repr_lines = str(val).split("\n") + for node_repr_line in node_repr_lines: + lines.append(f"{indent}" + node_repr_line) + # TODO: (jiaodong) Handle nested containers and other obj types + else: + lines.append(f"{indent}" + str(arg) + ", ") + + if len(lines) == 0: + args_line = "[]" + else: + args_line = "[" + for args in lines: + args_line += f"\n{indent}{args}" + args_line += f"\n{indent}]" + + return args_line + + +def _get_kwargs_lines(bound_kwargs): + """Pretty prints bounded kwargs of a DAGNode, and recursively handle + DAGNode in list / dict containers. + """ + # TODO: (jiaodong) Nits, we're missing keys and indentation was a bit off. + if not bound_kwargs: + return "{}" + indent = _get_indentation() + kwargs_lines = [] + for key, val in bound_kwargs.items(): + if isinstance(val, DAGNode): + node_repr_lines = str(val).split("\n") + for index, node_repr_line in enumerate(node_repr_lines): + if index == 0: + kwargs_lines.append( + f"{indent}{key}:" + f"{indent}" + node_repr_line + ) + else: + kwargs_lines.append(f"{indent}{indent}" + node_repr_line) + + elif isinstance(val, list): + for ele in val: + node_repr_lines = str(ele).split("\n") + for node_repr_line in node_repr_lines: + kwargs_lines.append(f"{indent}" + node_repr_line) + elif isinstance(val, dict): + for _, inner_val in val.items(): + node_repr_lines = str(inner_val).split("\n") + for node_repr_line in node_repr_lines: + kwargs_lines.append(f"{indent}" + node_repr_line) + # TODO: (jiaodong) Handle nested containers and other obj types + else: + kwargs_lines.append(val) + + if len(kwargs_lines) > 0: + kwargs_line = "{" + for line in kwargs_lines: + kwargs_line += f"\n{indent}{line}" + kwargs_line += f"\n{indent}}}" + else: + kwargs_line = "{}" + + return kwargs_line + + +def _get_options_lines(bound_options): + """Pretty prints .options() in DAGNode. Only prints non-empty values.""" + if not bound_options: + return "{}" + indent = _get_indentation() + options_lines = [] + for key, val in bound_options.items(): + if val: + options_lines.append(f"{indent}{key}: " + str(val)) + + options_line = "{" + for line in options_lines: + options_line += f"\n{indent}{line}" + options_line += f"\n{indent}}}" + return options_line + + +def _get_other_args_to_resolve_lines(other_args_to_resolve): + if not other_args_to_resolve: + return "{}" + indent = _get_indentation() + other_args_to_resolve_lines = [] + for key, val in other_args_to_resolve.items(): + if isinstance(val, DAGNode): + node_repr_lines = str(val).split("\n") + for index, node_repr_line in enumerate(node_repr_lines): + if index == 0: + other_args_to_resolve_lines.append( + f"{indent}{key}:" + + f"{indent}" + + "\n" + + f"{indent}{indent}{indent}" + + node_repr_line + ) + else: + other_args_to_resolve_lines.append( + f"{indent}{indent}" + node_repr_line + ) + else: + other_args_to_resolve_lines.append(f"{indent}{key}: " + str(val)) + + other_args_to_resolve_line = "{" + for line in other_args_to_resolve_lines: + other_args_to_resolve_line += f"\n{indent}{line}" + other_args_to_resolve_line += f"\n{indent}}}" + return other_args_to_resolve_line diff --git a/lib/python3.12/site-packages/ray/dag/function_node.py b/lib/python3.12/site-packages/ray/dag/function_node.py new file mode 100644 index 0000000000000000000000000000000000000000..b48c63509f2c9507a2b03aaa4a3aec96b956b477 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/function_node.py @@ -0,0 +1,59 @@ +from typing import Any, Dict, List + +import ray +from ray.dag.dag_node import DAGNode +from ray.dag.format_utils import get_dag_node_str +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class FunctionNode(DAGNode): + """Represents a bound task node in a Ray task DAG.""" + + def __init__( + self, + func_body, + func_args, + func_kwargs, + func_options, + other_args_to_resolve=None, + ): + self._body = func_body + super().__init__( + func_args, + func_kwargs, + func_options, + other_args_to_resolve=other_args_to_resolve, + ) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return FunctionNode( + self._body, + new_args, + new_kwargs, + new_options, + other_args_to_resolve=new_other_args_to_resolve, + ) + + def _execute_impl(self, *args, **kwargs): + """Executor of FunctionNode by ray.remote(). + + Args and kwargs are to match base class signature, but not in the + implementation. All args and kwargs should be resolved and replaced + with value in bound_args and bound_kwargs via bottom-up recursion when + current node is executed. + """ + return ( + ray.remote(self._body) + .options(**self._bound_options) + .remote(*self._bound_args, **self._bound_kwargs) + ) + + def __str__(self) -> str: + return get_dag_node_str(self, str(self._body)) diff --git a/lib/python3.12/site-packages/ray/dag/input_node.py b/lib/python3.12/site-packages/ray/dag/input_node.py new file mode 100644 index 0000000000000000000000000000000000000000..0386f84cb999fab71c4d3241eb70084c90d8eef8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/input_node.py @@ -0,0 +1,321 @@ +from typing import Any, Dict, List, Optional, Union + +from ray.dag import DAGNode +from ray.dag.format_utils import get_dag_node_str +from ray.experimental.gradio_utils import type_to_string +from ray.util.annotations import DeveloperAPI + +IN_CONTEXT_MANAGER = "__in_context_manager__" + + +@DeveloperAPI +class InputNode(DAGNode): + r"""Ray dag node used in DAG building API to mark entrypoints of a DAG. + + Should only be function or class method. A DAG can have multiple + entrypoints, but only one instance of InputNode exists per DAG, shared + among all DAGNodes. + + Example: + + .. code-block:: + + m1.forward + / \ + dag_input ensemble -> dag_output + \ / + m2.forward + + In this pipeline, each user input is broadcasted to both m1.forward and + m2.forward as first stop of the DAG, and authored like + + .. code-block:: python + + import ray + + @ray.remote + class Model: + def __init__(self, val): + self.val = val + def forward(self, input): + return self.val * input + + @ray.remote + def combine(a, b): + return a + b + + with InputNode() as dag_input: + m1 = Model.bind(1) + m2 = Model.bind(2) + m1_output = m1.forward.bind(dag_input[0]) + m2_output = m2.forward.bind(dag_input.x) + ray_dag = combine.bind(m1_output, m2_output) + + # Pass mix of args and kwargs as input. + ray_dag.execute(1, x=2) # 1 sent to m1, 2 sent to m2 + + # Alternatively user can also pass single data object, list or dict + # and access them via list index, object attribute or dict key str. + ray_dag.execute(UserDataObject(m1=1, m2=2)) + # dag_input.m1, dag_input.m2 + ray_dag.execute([1, 2]) + # dag_input[0], dag_input[1] + ray_dag.execute({"m1": 1, "m2": 2}) + # dag_input["m1"], dag_input["m2"] + """ + + def __init__( + self, + *args, + input_type: Optional[Union[type, Dict[Union[int, str], type]]] = None, + _other_args_to_resolve=None, + **kwargs, + ): + """InputNode should only take attributes of validating and converting + input data rather than the input data itself. User input should be + provided via `ray_dag.execute(user_input)`. + + Args: + input_type: Describes the data type of inputs user will be giving. + - if given through singular InputNode: type of InputNode + - if given through InputAttributeNodes: map of key -> type + Used when deciding what Gradio block to represent the input nodes with. + _other_args_to_resolve: Internal only to keep InputNode's execution + context throughput pickling, replacement and serialization. + User should not use or pass this field. + """ + if len(args) != 0 or len(kwargs) != 0: + raise ValueError("InputNode should not take any args or kwargs.") + + self.input_attribute_nodes = {} + + self.input_type = input_type + if input_type is not None and isinstance(input_type, type): + if _other_args_to_resolve is None: + _other_args_to_resolve = {} + _other_args_to_resolve["result_type_string"] = type_to_string(input_type) + + super().__init__([], {}, {}, other_args_to_resolve=_other_args_to_resolve) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return InputNode(_other_args_to_resolve=new_other_args_to_resolve) + + def _execute_impl(self, *args, **kwargs): + """Executor of InputNode.""" + # Catch and assert singleton context at dag execution time. + assert self._in_context_manager(), ( + "InputNode is a singleton instance that should be only used in " + "context manager for dag building and execution. See the docstring " + "of class InputNode for examples." + ) + # If user only passed in one value, for simplicity we just return it. + if len(args) == 1 and len(kwargs) == 0: + return args[0] + + return DAGInputData(*args, **kwargs) + + def _in_context_manager(self) -> bool: + """Return if InputNode is created in context manager.""" + if ( + not self._bound_other_args_to_resolve + or IN_CONTEXT_MANAGER not in self._bound_other_args_to_resolve + ): + return False + else: + return self._bound_other_args_to_resolve[IN_CONTEXT_MANAGER] + + def set_context(self, key: str, val: Any): + """Set field in parent DAGNode attribute that can be resolved in both + pickle and JSON serialization + """ + self._bound_other_args_to_resolve[key] = val + + def __str__(self) -> str: + return get_dag_node_str(self, "__InputNode__") + + def __getattr__(self, key: str): + assert isinstance( + key, str + ), "Please only access dag input attributes with str key." + if key not in self.input_attribute_nodes: + self.input_attribute_nodes[key] = InputAttributeNode( + self, key, "__getattr__" + ) + return self.input_attribute_nodes[key] + + def __getitem__(self, key: Union[int, str]) -> Any: + assert isinstance(key, (str, int)), ( + "Please only use int index or str as first-level key to " + "access fields of dag input." + ) + + input_type = None + if self.input_type is not None and key in self.input_type: + input_type = type_to_string(self.input_type[key]) + + if key not in self.input_attribute_nodes: + self.input_attribute_nodes[key] = InputAttributeNode( + self, key, "__getitem__", input_type + ) + return self.input_attribute_nodes[key] + + def __enter__(self): + self.set_context(IN_CONTEXT_MANAGER, True) + return self + + def __exit__(self, *args): + pass + + def get_result_type(self) -> str: + """Get type of the output of this DAGNode. + + Generated by ray.experimental.gradio_utils.type_to_string(). + """ + if "result_type_string" in self._bound_other_args_to_resolve: + return self._bound_other_args_to_resolve["result_type_string"] + + +@DeveloperAPI +class InputAttributeNode(DAGNode): + """Represents partial access of user input based on an index (int), + object attribute or dict key (str). + + Examples: + + .. code-block:: python + + with InputNode() as dag_input: + a = dag_input[0] + b = dag_input.x + ray_dag = add.bind(a, b) + + # This makes a = 1 and b = 2 + ray_dag.execute(1, x=2) + + with InputNode() as dag_input: + a = dag_input[0] + b = dag_input[1] + ray_dag = add.bind(a, b) + + # This makes a = 2 and b = 3 + ray_dag.execute(2, 3) + + # Alternatively, you can input a single object + # and the inputs are automatically indexed from the object: + # This makes a = 2 and b = 3 + ray_dag.execute([2, 3]) + """ + + def __init__( + self, + dag_input_node: InputNode, + key: Union[int, str], + accessor_method: str, + input_type: str = None, + ): + self._dag_input_node = dag_input_node + self._key = key + self._accessor_method = accessor_method + super().__init__( + [], + {}, + {}, + { + "dag_input_node": dag_input_node, + "key": key, + "accessor_method": accessor_method, + # Type of the input tied to this node. Used by + # gradio_visualize_graph.GraphVisualizer to determine which Gradio + # component should be used for this node. + "result_type_string": input_type, + }, + ) + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ): + return InputAttributeNode( + new_other_args_to_resolve["dag_input_node"], + new_other_args_to_resolve["key"], + new_other_args_to_resolve["accessor_method"], + new_other_args_to_resolve["result_type_string"], + ) + + def _execute_impl(self, *args, **kwargs): + """Executor of InputAttributeNode. + + Args and kwargs are to match base class signature, but not in the + implementation. All args and kwargs should be resolved and replaced + with value in bound_args and bound_kwargs via bottom-up recursion when + current node is executed. + """ + + if isinstance(self._dag_input_node, DAGInputData): + return self._dag_input_node[self._key] + else: + # dag.execute() is called with only one arg, thus when an + # InputAttributeNode is executed, its dependent InputNode is + # resolved with original user input python object. + user_input_python_object = self._dag_input_node + if isinstance(self._key, str): + if self._accessor_method == "__getitem__": + return user_input_python_object[self._key] + elif self._accessor_method == "__getattr__": + return getattr(user_input_python_object, self._key) + elif isinstance(self._key, int): + return user_input_python_object[self._key] + else: + raise ValueError( + "Please only use int index or str as first-level key to " + "access fields of dag input." + ) + + def __str__(self) -> str: + return get_dag_node_str(self, f'["{self._key}"]') + + def get_result_type(self) -> str: + """Get type of the output of this DAGNode. + + Generated by ray.experimental.gradio_utils.type_to_string(). + """ + if "result_type_string" in self._bound_other_args_to_resolve: + return self._bound_other_args_to_resolve["result_type_string"] + + @property + def key(self) -> Union[int, str]: + return self._key + + +@DeveloperAPI +class DAGInputData: + """If user passed multiple args and kwargs directly to dag.execute(), we + generate this wrapper for all user inputs as one object, accessible via + list index or object attribute key. + """ + + def __init__(self, *args, **kwargs): + self._args = list(args) + self._kwargs = kwargs + + def __getitem__(self, key: Union[int, str]) -> Any: + if isinstance(key, int): + # Access list args by index. + return self._args[key] + elif isinstance(key, str): + # Access kwarg by key. + return self._kwargs[key] + else: + raise ValueError( + "Please only use int index or str as first-level key to " + "access fields of dag input." + ) diff --git a/lib/python3.12/site-packages/ray/dag/output_node.py b/lib/python3.12/site-packages/ray/dag/output_node.py new file mode 100644 index 0000000000000000000000000000000000000000..fc0ec1a1002676811cb4080ba3c4e74b8847e557 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/output_node.py @@ -0,0 +1,45 @@ +from typing import Any, Dict, List, Tuple, Union + +import ray +from ray.dag import DAGNode +from ray.dag.format_utils import get_dag_node_str +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class MultiOutputNode(DAGNode): + """Ray dag node used in DAG building API to mark the endpoint of DAG""" + + def __init__( + self, + args: Union[List[DAGNode], Tuple[DAGNode]], + other_args_to_resolve: Dict[str, Any] = None, + ): + if isinstance(args, tuple): + args = list(args) + if not isinstance(args, list): + raise ValueError(f"Invalid input type for `args`, {type(args)}.") + super().__init__( + args, + {}, + {}, + other_args_to_resolve=other_args_to_resolve or {}, + ) + + def _execute_impl( + self, *args, **kwargs + ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]: + return self._bound_args + + def _copy_impl( + self, + new_args: List[Any], + new_kwargs: Dict[str, Any], + new_options: Dict[str, Any], + new_other_args_to_resolve: Dict[str, Any], + ) -> "DAGNode": + """Return a copy of this node with the given new args.""" + return MultiOutputNode(new_args, new_other_args_to_resolve) + + def __str__(self) -> str: + return get_dag_node_str(self, "__MultiOutputNode__") diff --git a/lib/python3.12/site-packages/ray/dag/py_obj_scanner.py b/lib/python3.12/site-packages/ray/dag/py_obj_scanner.py new file mode 100644 index 0000000000000000000000000000000000000000..d86b982c10e10692a0624e6e86d87e1a45f60b46 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/py_obj_scanner.py @@ -0,0 +1,103 @@ +import io +import pickle # noqa: F401 +from typing import Any, Dict, Generic, List, Tuple, Type, TypeVar, Union + +import ray +from ray.dag.base import DAGNodeBase + +# Used in deserialization hooks to reference scanner instances. +_instances: Dict[int, "_PyObjScanner"] = {} + +# Generic types for the scanner to transform from and to. +SourceType = TypeVar("SourceType") +TransformedType = TypeVar("TransformedType") + + +def _get_node(instance_id: int, node_index: int) -> SourceType: + """Get the node instance. + + Note: This function should be static and globally importable, + otherwise the serialization overhead would be very significant. + """ + return _instances[instance_id]._replace_index(node_index) + + +class _PyObjScanner(ray.cloudpickle.CloudPickler, Generic[SourceType, TransformedType]): + """Utility to find and replace the `source_type` in Python objects. + + `source_type` can either be a single type or a tuple of multiple types. + + The caller must first call `find_nodes()`, then compute a replacement table and + pass it to `replace_nodes`. + + This uses cloudpickle under the hood, so all sub-objects that are not `source_type` + must be serializable. + + Args: + source_type: the type(s) of object to find and replace. Default to DAGNodeBase. + """ + + def __init__(self, source_type: Union[Type, Tuple] = DAGNodeBase): + self.source_type = source_type + # Buffer to keep intermediate serialized state. + self._buf = io.BytesIO() + # List of top-level SourceType found during the serialization pass. + self._found = None + # List of other objects found during the serialization pass. + # This is used to store references to objects so they won't be + # serialized by cloudpickle. + self._objects = [] + # Replacement table to consult during deserialization. + self._replace_table: Dict[SourceType, TransformedType] = None + _instances[id(self)] = self + super().__init__(self._buf) + + def reducer_override(self, obj): + """Hook for reducing objects. + + Objects of `self.source_type` are saved to `self._found` and a global map so + they can later be replaced. + + All other objects fall back to the default `CloudPickler` serialization. + """ + if isinstance(obj, self.source_type): + index = len(self._found) + self._found.append(obj) + return _get_node, (id(self), index) + + return super().reducer_override(obj) + + def find_nodes(self, obj: Any) -> List[SourceType]: + """ + Serialize `obj` and store all instances of `source_type` found in `_found`. + + Args: + obj: The object to scan for `source_type`. + Returns: + A list of all instances of `source_type` found in `obj`. + """ + assert ( + self._found is None + ), "find_nodes cannot be called twice on the same PyObjScanner instance." + self._found = [] + self._objects = [] + self.dump(obj) + return self._found + + def replace_nodes(self, table: Dict[SourceType, TransformedType]) -> Any: + """Replace previously found DAGNodes per the given table.""" + assert self._found is not None, "find_nodes must be called first" + self._replace_table = table + self._buf.seek(0) + return pickle.load(self._buf) + + def _replace_index(self, i: int) -> SourceType: + return self._replace_table[self._found[i]] + + def clear(self): + """Clear the scanner from the _instances""" + if id(self) in _instances: + del _instances[id(self)] + + def __del__(self): + self.clear() diff --git a/lib/python3.12/site-packages/ray/dag/utils.py b/lib/python3.12/site-packages/ray/dag/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe1f3adf3d7eb73333731cfa7b83bdc509acd25 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/utils.py @@ -0,0 +1,66 @@ +from typing import Dict + +from ray.dag import ( + ClassMethodNode, + ClassNode, + DAGNode, + FunctionNode, + InputAttributeNode, + InputNode, + MultiOutputNode, +) + + +class _DAGNodeNameGenerator(object): + """ + Generate unique suffix for each given Node in the DAG. + Apply monotonic increasing id suffix for duplicated names. + """ + + def __init__(self): + self.name_to_suffix: Dict[str, int] = dict() + + def get_node_name(self, node: DAGNode): + # InputNode should be unique. + if isinstance(node, InputNode): + return "INPUT_NODE" + if isinstance(node, MultiOutputNode): + return "MultiOutputNode" + # InputAttributeNode suffixes should match the user-defined key. + elif isinstance(node, InputAttributeNode): + return f"INPUT_ATTRIBUTE_NODE_{node._key}" + + # As class, method, and function nodes may have duplicated names, + # generate unique suffixes for such nodes. + if isinstance(node, ClassMethodNode): + node_name = node.get_options().get("name", None) or node._method_name + elif isinstance(node, (ClassNode, FunctionNode)): + node_name = node.get_options().get("name", None) or node._body.__name__ + # we use instance class name check here to avoid importing ServeNodes as + # serve components are not included in Ray Core. + elif type(node).__name__ in ("DeploymentNode", "DeploymentFunctionNode"): + node_name = node.get_deployment_name() + elif type(node).__name__ == "DeploymentFunctionExecutorNode": + node_name = node._deployment_function_handle.deployment_name + else: + raise ValueError( + "get_node_name() should only be called on DAGNode instances." + ) + + if node_name not in self.name_to_suffix: + self.name_to_suffix[node_name] = 0 + return node_name + else: + self.name_to_suffix[node_name] += 1 + suffix_num = self.name_to_suffix[node_name] + + return f"{node_name}_{suffix_num}" + + def reset(self): + self.name_to_suffix = dict() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.reset() diff --git a/lib/python3.12/site-packages/ray/dag/vis_utils.py b/lib/python3.12/site-packages/ray/dag/vis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1274a53cc20d0da412fe55a7d7224197c73af500 --- /dev/null +++ b/lib/python3.12/site-packages/ray/dag/vis_utils.py @@ -0,0 +1,114 @@ +import os +import tempfile + +from ray.dag import DAGNode +from ray.dag.utils import _DAGNodeNameGenerator +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def plot(dag: DAGNode, to_file=None): + if to_file is None: + tmp_file = tempfile.NamedTemporaryFile(suffix=".png") + to_file = tmp_file.name + extension = "png" + else: + _, extension = os.path.splitext(to_file) + if not extension: + extension = "png" + else: + extension = extension[1:] + + graph = _dag_to_dot(dag) + graph.write(to_file, format=extension) + + # Render the image directly if running inside a Jupyter notebook + try: + from IPython import display + + return display.Image(filename=to_file) + except ImportError: + pass + + # close temp file if needed + try: + tmp_file.close() + except NameError: + pass + + +def _check_pydot_and_graphviz(): + """Check if pydot and graphviz are installed. + + pydot and graphviz are required for plotting. We check this + during runtime rather than adding them to Ray dependencies. + + """ + try: + import pydot + except ImportError: + raise ImportError( + "pydot is required to plot DAG, install it with `pip install pydot`." + ) + try: + pydot.Dot.create(pydot.Dot()) + except (OSError, pydot.InvocationException): + raise ImportError( + "graphviz is required to plot DAG, " + "download it from https://graphviz.gitlab.io/download/" + ) + + +def _get_nodes_and_edges(dag: DAGNode): + """Get all unique nodes and edges in the DAG. + + A basic dfs with memorization to get all unique nodes + and edges in the DAG. + Unique nodes will be used to generate unique names, + while edges will be used to construct the graph. + """ + + edges = [] + nodes = [] + + def _dfs(node): + nodes.append(node) + for child_node in node._get_all_child_nodes(): + edges.append((child_node, node)) + return node + + dag.apply_recursive(_dfs) + return nodes, edges + + +def _dag_to_dot(dag: DAGNode): + """Create a Dot graph from dag. + + TODO(lchu): + 1. add more Dot configs in kwargs, + e.g. rankdir, alignment, etc. + 2. add more contents to graph, + e.g. args, kwargs and options of each node + + """ + # Step 0: check dependencies and init graph + _check_pydot_and_graphviz() + import pydot + + graph = pydot.Dot(rankdir="LR") + + # Step 1: generate unique name for each node in dag + nodes, edges = _get_nodes_and_edges(dag) + name_generator = _DAGNodeNameGenerator() + node_names = {} + for node in nodes: + node_names[node] = name_generator.get_node_name(node) + + # Step 2: create graph with all the edges + for edge in edges: + graph.add_edge(pydot.Edge(node_names[edge[0]], node_names[edge[1]])) + # if there is only one node + if len(nodes) == 1 and len(edges) == 0: + graph.add_node(pydot.Node(node_names[nodes[0]])) + + return graph diff --git a/lib/python3.12/site-packages/ray/experimental/__init__.py b/lib/python3.12/site-packages/ray/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..37cb09a1513b6c2de73c6640566dbf0f382d6692 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/__init__.py @@ -0,0 +1,11 @@ +from ray.experimental.dynamic_resources import set_resource +from ray.experimental.gpu_object_manager import GPUObjectManager, wait_tensor_freed +from ray.experimental.locations import get_local_object_locations, get_object_locations + +__all__ = [ + "get_object_locations", + "get_local_object_locations", + "set_resource", + "GPUObjectManager", + "wait_tensor_freed", +] diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__init__.py b/lib/python3.12/site-packages/ray/experimental/collective/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..42289cee16530a426c891f11239432f3b5d5c30d --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/__init__.py @@ -0,0 +1,23 @@ +from ray.experimental.collective.collective import ( + create_collective_group, + destroy_all_collective_groups, + destroy_collective_group, + get_collective_groups, +) +from ray.experimental.collective.operations import ( + allgather, + allreduce, + reducescatter, +) +from ray.experimental.collective.util import get_tensor_transport_manager + +__all__ = [ + "allgather", + "allreduce", + "reducescatter", + "get_collective_groups", + "create_collective_group", + "destroy_collective_group", + "destroy_all_collective_groups", + "get_tensor_transport_manager", +] diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5302c6817ddc0e7b355e5349b662c4eb75c6da3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2d080f7714e579221a9c1ffcc20e271d8190ff2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective_tensor_transport.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective_tensor_transport.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec91ee0aa9caa92e41656c1fed4111649dd8e51b Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/collective_tensor_transport.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/communicator.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/communicator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72d650942b7d27f64e2b014ada79f773345487c6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/communicator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/conftest.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/conftest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c69bc80e7cc503bac51b6dbb320073e30d45f055 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/conftest.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/nixl_tensor_transport.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/nixl_tensor_transport.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f40f203092fe984f963e17fa2de9afffc39e50d Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/nixl_tensor_transport.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/operations.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/operations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..338257871270963d6fd73ad0f6681e651eb4fac3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/operations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/tensor_transport_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/tensor_transport_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49e487f29b9766c8ee669755251bc1ee15da65ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/tensor_transport_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c99a82647a8631f899d08f086e123ba8d5cac8d8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/collective/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/collective.py b/lib/python3.12/site-packages/ray/experimental/collective/collective.py new file mode 100644 index 0000000000000000000000000000000000000000..3a31128088a1db4ec472c2c3931f5a50126922a9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/collective.py @@ -0,0 +1,233 @@ +import threading +import uuid +from typing import Dict, List, Optional, Union + +import ray +import ray.experimental.internal_kv as internal_kv +from ray.experimental.collective.communicator import CommunicatorHandle +from ray.experimental.collective.util import get_address_and_port +from ray.util.annotations import PublicAPI +from ray.util.collective.collective_group.torch_gloo_collective_group import ( + get_master_address_metadata_key, +) +from ray.util.collective.types import Backend + +_remote_communicator_manager: "Optional[RemoteCommunicatorManager]" = None +_remote_communicator_manager_lock = threading.Lock() + + +class RemoteCommunicatorManager: + """Singleton class to store the mapping between actors and communicators + that the actors are a part of. + """ + + def __init__(self): + # Handles to communicators that we created. Key is a user-provided + # name or UUID. + self._remote_communicators: Dict[str, CommunicatorHandle] = {} + + @staticmethod + def get() -> "RemoteCommunicatorManager": + global _remote_communicator_manager + with _remote_communicator_manager_lock: + if _remote_communicator_manager is None: + _remote_communicator_manager = RemoteCommunicatorManager() + return _remote_communicator_manager + + def add_remote_communicator(self, comm_handle: CommunicatorHandle): + self._remote_communicators[comm_handle.name] = comm_handle + + def remove_remote_communicator(self, name: str): + return self._remote_communicators.pop(name, None) + + def get_collective_groups( + self, + actors: Optional[List[ray.actor.ActorHandle]] = None, + backend: Optional[str] = None, + ): + """ + Get the collective groups that the given actors are a subset of. Filter by + backend if provided. + """ + actors = actors or [] + actors = set(actors) + + collectives = [] + # Find all collective groups that the given actors are a subset + # of, with the matching backend if provided. + for collective in self._remote_communicators.values(): + if actors.issubset(set(collective.actors)): + if backend is None or collective.backend == backend: + collectives.append(collective) + return collectives + + +def _do_init_collective_group( + self, + world_size: int, + rank: int, + backend: str = Backend.NCCL, + name: str = "default", +): + """Helper method that runs as a task on a remote actor to create a + collective group. + """ + ray.util.collective.init_collective_group( + world_size, rank, backend, group_name=name + ) + + +def _do_destroy_collective_group(self, name): + """Helper method that runs as a task on a remote actor to destroy a + collective group. + """ + ray.util.collective.destroy_collective_group(name) + + +@PublicAPI(stability="alpha") +def get_collective_groups( + actors: List[ray.actor.ActorHandle], backend: Optional[str] = None +) -> List[CommunicatorHandle]: + """ + Get the collective groups that the given actors are a subset of. Filter by + backend if provided. + + Args: + actors: List of actors. Return handles to all collective groups that + these actors are a subset of. + backend: An optional backend to filter by. See + ray.util.collective.types.Backend for valid backends. + + Returns: + A list of communicator handles that the actors are a subset of. + """ + manager = RemoteCommunicatorManager.get() + return manager.get_collective_groups(actors, backend) + + +@PublicAPI(stability="alpha") +def create_collective_group( + actors: List[ray.actor.ActorHandle], + backend: str, + name: Optional[str] = None, +) -> CommunicatorHandle: + """Create a collective group on the given list of actors. If this function + returns successfully, then the collective group has been initialized on all + actors, using the given order of actors as the ranks. + + Currently, an actor can only participate in one collective group per + backend at a time. To reuse an actor, destroy its collective group and + create a new one. + + Args: + actors: The actors to participate in the collective group. + backend: The backend to use. See ray.util.collective.types.Backend for + valid backends. + name: A name to use for the collective group. If None is provided, a + random name will be generated. + + Returns: + Handle to the communicator. + """ + manager = RemoteCommunicatorManager.get() + + if name is None: + name = str(uuid.uuid4()) + + # Validate the backend. + backend = Backend(backend) + + world_size = len(actors) + + for actor in actors: + if manager.get_collective_groups([actor], backend): + raise RuntimeError( + f"Actor {actor} already in group for backend {backend}. Actors can currently only participate in at most one group per backend." + ) + + actor_ids = [actor._ray_actor_id for actor in actors] + if len(set(actor_ids)) != len(actor_ids): + raise ValueError(f"All actors must be unique, got: {actors}") + + metadata_key = None + if backend == Backend.GLOO: + # Perform extra setup for torch.distributed. + # torch.distributed requires a master address and port. Find a suitable + # port on one of the actors. + master_addr, master_port = ray.get( + actors[0].__ray_call__.remote(lambda self: get_address_and_port()) + ) + + # Store the metadata on a named actor that all of the other + # actors can access. + metadata_key = get_master_address_metadata_key(name) + internal_kv._internal_kv_put(metadata_key, f"{master_addr}:{master_port}") + + try: + init_tasks = [ + actor.__ray_call__.remote( + _do_init_collective_group, world_size, rank, backend, name + ) + for rank, actor in enumerate(actors) + ] + ray.get(init_tasks) + finally: + # Clean up the metadata once collective group is initialized + # (or failed to initialize). + if metadata_key is not None: + internal_kv._internal_kv_del(metadata_key) + + # Group was successfully created. + # Register GLOO groups under TORCH_GLOO since GLOO uses torch.distributed. + registration_backend = Backend.TORCH_GLOO if backend == Backend.GLOO else backend + comm = CommunicatorHandle(actors, name, registration_backend) + manager.add_remote_communicator(comm) + return comm + + +@PublicAPI(stability="alpha") +def destroy_collective_group(group_or_name: Union[CommunicatorHandle, str]): + """ + Destroy a collective group. If this functions returns successfully, then + the actors that were in the collective can be reused to create a new + collective group. + + Args: + group_or_name: Either a communicator handle or the name of the group to + destroy. + """ + if isinstance(group_or_name, CommunicatorHandle): + name = group_or_name.name + elif isinstance(group_or_name, str): + name = group_or_name + else: + raise ValueError("Expected CommunicatorHandle or str (group name).") + + manager = RemoteCommunicatorManager.get() + group = manager.remove_remote_communicator(name) + if group is not None: + destroy_tasks = [ + actor.__ray_call__.options(concurrency_group="_ray_system").remote( + _do_destroy_collective_group, name + ) + for actor in group.actors + ] + try: + ray.get(destroy_tasks) + except ray.exceptions.ActorDiedError: + pass + else: + raise ValueError(f"No group with name {name} found.") + + +@PublicAPI(stability="alpha") +def destroy_all_collective_groups(): + """ + Destroy all collective groups. This will destroy all collective groups that + were previously created by this process. After this function returns, the + actors participating in those collective groups can be reused to create a + new collective group. + """ + manager = RemoteCommunicatorManager.get() + for collective in manager.get_collective_groups(): + destroy_collective_group(collective.name) diff --git a/lib/python3.12/site-packages/ray/experimental/collective/collective_tensor_transport.py b/lib/python3.12/site-packages/ray/experimental/collective/collective_tensor_transport.py new file mode 100644 index 0000000000000000000000000000000000000000..78f97d64c87afb75ba45594e4ebe6c17a6a01c54 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/collective_tensor_transport.py @@ -0,0 +1,199 @@ +from typing import TYPE_CHECKING, List, Optional + +import ray +from ray.experimental.collective.tensor_transport_manager import ( + TensorTransportManager, +) +from ray.util.collective.types import ( + Backend, + CollectiveCommunicatorMetadata, + CollectiveTransportMetadata, +) + +if TYPE_CHECKING: + import torch + + +class CollectiveTensorTransport(TensorTransportManager): + def __init__(self, tensor_transport_backend: Backend): + self._tensor_transport_backend = tensor_transport_backend + + @property + def tensor_transport_backend(self) -> Backend: + return self._tensor_transport_backend + + @staticmethod + def is_one_sided() -> bool: + return False + + @staticmethod + def can_abort_transport() -> bool: + return False + + def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool: + from ray.experimental.collective import get_collective_groups + + communicators = get_collective_groups( + [actor], backend=self.tensor_transport_backend + ) + return len(communicators) > 0 + + @staticmethod + def extract_tensor_transport_metadata( + obj_id: str, + gpu_object: List["torch.Tensor"], + ) -> CollectiveTransportMetadata: + tensor_meta = [] + device = None + if gpu_object: + device = gpu_object[0].device + for t in gpu_object: + if t.device.type != device.type: + raise ValueError( + "All tensors in an RDT object must have the same device type." + ) + tensor_meta.append((t.shape, t.dtype)) + return CollectiveTransportMetadata( + tensor_meta=tensor_meta, + tensor_device=device, + ) + + @staticmethod + def get_tensor_transport_metadata( + src_actor: "ray.actor.ActorHandle", + obj_id: str, + ) -> CollectiveTransportMetadata: + def __ray_get_tensor_transport_metadata__( + self: "ray.actor.ActorHandle", + obj_id: str, + ) -> CollectiveTransportMetadata: + + from ray._private.worker import global_worker + + gpu_object_store = global_worker.gpu_object_manager.gpu_object_store + # NOTE: We do not specify a timeout here because the user task that returns + # it could take arbitrarily long and we don't want to trigger a spurious + # timeout. + gpu_object = gpu_object_store.wait_and_get_object(obj_id) + return CollectiveTensorTransport.extract_tensor_transport_metadata( + obj_id, gpu_object + ) + + # Submit a Ray actor task to the source actor to get the tensor metadata. + # The metadata is a list of tuples, where each tuple contains the shape and dtype + # of a tensor in the GPU object store. This function returns an ObjectRef that + # points to the tensor metadata. + # NOTE(swang): We put this task on the background thread to avoid tasks + # executing on the main thread blocking this task. + + return src_actor.__ray_call__.options(concurrency_group="_ray_system").remote( + __ray_get_tensor_transport_metadata__, obj_id + ) + + @staticmethod + def get_communicator_metadata( + src_actor: "ray.actor.ActorHandle", + dst_actor: "ray.actor.ActorHandle", + backend: Optional[str] = None, + ) -> CollectiveCommunicatorMetadata: + + from ray.experimental.collective import get_collective_groups + + communicators = get_collective_groups( + [src_actor, dst_actor], + backend=backend, + ) + # TODO(kevin85421): Support multiple communicators. + if len(communicators) == 0: + raise ValueError( + f"No communicators found for actors {src_actor} and {dst_actor}. " + "Create a communicator with " + "`ray.experimental.collective.create_collective_group` " + "before calling actor tasks. with non-default tensor_transport." + ) + elif len(communicators) > 1: + raise ValueError( + f"There are {len(communicators)} possible communicators that contain actors {src_actor} and {dst_actor}. " + "Currently, RDT objects only support one communicator. Please make sure only " + "one communicator exists." + ) + communicator = communicators[0] + src_rank = communicator.get_rank(src_actor) + if src_rank == -1: + raise ValueError( + f"Sender actor {src_actor} not found in communicator. " + "Please make sure the sender and receiver are in the same communicator." + ) + dst_rank = communicator.get_rank(dst_actor) + if dst_rank == -1: + raise ValueError( + f"Receiver actor {dst_actor} not found in communicator. " + "Please make sure the sender and receiver are in the same communicator." + ) + + communicator_metadata = CollectiveCommunicatorMetadata( + communicator_name=communicator.name, + src_rank=src_rank, + dst_rank=dst_rank, + ) + return communicator_metadata + + @staticmethod + def recv_multiple_tensors( + tensors, + obj_id: str, + tensor_transport_metadata: CollectiveTransportMetadata, + communicator_metadata: CollectiveCommunicatorMetadata, + ): + from ray.util.collective import types + from ray.util.collective.collective import recv + + assert isinstance( + tensor_transport_metadata, types.CollectiveTransportMetadata + ), "metadata must be a CollectiveTransportMetadata object for non-NIXL transport" + assert isinstance( + communicator_metadata, types.CollectiveCommunicatorMetadata + ), "metadata must be a CollectiveCommunicatorMetadata object for non-NIXL transport" + + for tensor in tensors: + recv( + tensor, + communicator_metadata.src_rank, + communicator_metadata.communicator_name, + ) + + @staticmethod + def send_multiple_tensors( + tensors: List["torch.Tensor"], + tensor_transport_metadata: CollectiveTransportMetadata, + communicator_metadata: CollectiveCommunicatorMetadata, + ): + import ray.util.collective as collective + + device = tensors[0].device if tensors else None + + for tensor in tensors: + if tensor.device.type != device.type: + raise ValueError( + f"tensor device {tensor.device} does not match device {device}" + ) + collective.send( + tensor, + communicator_metadata.dst_rank, + communicator_metadata.communicator_name, + ) + + @staticmethod + def garbage_collect( + obj_id: str, tensor_transport_meta: CollectiveTransportMetadata + ): + pass + + @staticmethod + def abort_transport( + obj_id: str, + communicator_metadata: CollectiveCommunicatorMetadata, + ): + raise NotImplementedError( + "Collective transport does not support abort_transport for now." + ) diff --git a/lib/python3.12/site-packages/ray/experimental/collective/communicator.py b/lib/python3.12/site-packages/ray/experimental/collective/communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..2379bf220389c7218184b271e035bf640db0b5b7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/communicator.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass +from typing import List + +import ray +from ray.util.collective.types import Backend + + +@dataclass +class Communicator: + """ + A handle to a communicator that we are a member of. + """ + + # The name of the communicator. + name: str + # Our rank in the collective group. + rank: int + # A valid backend, as defined by + # ray.util.collective.types.Backend. + backend: str + + +class CommunicatorHandle: + """ + A communicator handle used by the driver to store handles to the + actors in the communicator. + """ + + def __init__(self, actors: List[ray.actor.ActorHandle], name: str, backend: str): + """ + Initializes the CommunicatorHandle with the given actor handles. + Assumes that the communicator has already been initialized on all actors. + + Args: + actors: A list of actor handles to be stored. + name: Name of the communicator. + backend: Communicator backend. See + ray.util.collective.types for valid values. + """ + self._actors = actors + self._name = name + self._backend = Backend(backend) + + def get_rank(self, actor: ray.actor.ActorHandle): + for i, a in enumerate(self._actors): + if a == actor: + return i + return -1 + + @property + def actors(self) -> List[ray.actor.ActorHandle]: + """ + Return all actor handles in this communicator. + """ + return self._actors[:] + + @property + def name(self) -> str: + return self._name + + @property + def backend(self) -> str: + return self._backend diff --git a/lib/python3.12/site-packages/ray/experimental/collective/conftest.py b/lib/python3.12/site-packages/ray/experimental/collective/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..4544931a9fd0bd8f775039eaa4b967cf85950752 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/conftest.py @@ -0,0 +1,243 @@ +import uuid +from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Type + +import torch + +import ray +from ray.experimental.channel.common import ChannelContext +from ray.experimental.channel.communicator import ( + Communicator, + ReduceOp, + TorchTensorAllocator, +) + + +class AbstractNcclGroup(Communicator): + """ + A dummy NCCL group for testing. + """ + + def __init__(self, actor_handles: List[ray.actor.ActorHandle]): + self._actor_handles = actor_handles + self._rank = None + + def initialize(self, rank: int) -> None: + self._rank = rank + + def get_rank(self, actor: ray.actor.ActorHandle) -> int: + return self._actor_handles.index(actor) + + def get_world_size(self) -> int: + return len(self._actor_handles) + + def get_self_rank(self) -> Optional[int]: + return self._rank + + def get_actor_handles(self) -> List["ray.actor.ActorHandle"]: + return self._actor_handles + + def send(self, value: "torch.Tensor", peer_rank: int) -> None: + raise NotImplementedError + + def recv( + self, + shape: Tuple[int], + dtype: "torch.dtype", + peer_rank: int, + allocator: Optional[TorchTensorAllocator] = None, + ) -> "torch.Tensor": + raise NotImplementedError + + def allgather( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + ) -> None: + raise NotImplementedError + + def allreduce( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ) -> None: + raise NotImplementedError + + def reducescatter( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ) -> None: + raise NotImplementedError + + @property + def recv_stream(self): + return None + + @property + def send_stream(self): + return None + + def destroy(self) -> None: + pass + + def get_transport_name(self) -> str: + return "accelerator" + + @classmethod + def generate_communicator_id(cls) -> str: + pass + + +class MockNcclGroupSet: + def __init__(self): + # Represents a mapping from a NCCL group ID to a set of actors and a custom + # NCCL group. + self.ids_to_actors_and_custom_comms: Dict[ + str, Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]] + ] = {} + + def __call__( + self, + actors: List["ray.actor.ActorHandle"], + custom_nccl_group: Optional[Communicator] = None, + use_communication_streams: bool = False, + accelerator_module_name: Optional[str] = None, + accelerator_communicator_cls: Optional[Type[Communicator]] = None, + ) -> str: + group_id = str(uuid.uuid4()) + self.ids_to_actors_and_custom_comms[group_id] = ( + frozenset(actors), + custom_nccl_group, + ) + + if custom_nccl_group is None: + ranks = list(range(len(actors))) + else: + ranks = [custom_nccl_group.get_rank(actor) for actor in actors] + init_tasks = [ + actor.__ray_call__.remote( + mock_do_init_nccl_group, + group_id, + rank, + actors, + custom_nccl_group, + ) + for rank, actor in zip(ranks, actors) + ] + ray.get(init_tasks, timeout=30) + + ctx = ChannelContext.get_current() + if custom_nccl_group is not None: + ctx.communicators[group_id] = custom_nccl_group + else: + ctx.communicators[group_id] = AbstractNcclGroup(actors) + + return group_id + + def mock_destroy_nccl_group(self, group_id: str) -> None: + ctx = ChannelContext.get_current() + if group_id not in ctx.communicators: + return + + actors, _ = self.ids_to_actors_and_custom_comms[group_id] + destroy_tasks = [ + actor.__ray_call__.remote( + mock_do_destroy_nccl_group, + group_id, + ) + for actor in actors + ] + ray.wait(destroy_tasks, timeout=30) + + if group_id in self.ids_to_actors_and_custom_comms: + del self.ids_to_actors_and_custom_comms[group_id] + ctx.communicators[group_id].destroy() + del ctx.communicators[group_id] + + def check_teardown(self, nccl_group_ids: List[str]) -> None: + ctx = ChannelContext.get_current() + for nccl_group_id in nccl_group_ids: + assert nccl_group_id not in self.ids_to_actors_and_custom_comms + assert nccl_group_id not in ctx.communicators + + +@ray.remote +class CPUTorchTensorWorker: + def __init__(self): + self.device = "cpu" + + def return_tensor( + self, size: int, dtype: Optional[torch.dtype] = None + ) -> torch.Tensor: + return torch.ones(size, dtype=dtype, device=self.device) + + def recv(self, tensor: torch.Tensor) -> Tuple[int, int]: + assert tensor.device == self.device + return tensor.shape, tensor[0] + + def recv_tensors(self, *tensors) -> Tuple[torch.Tensor, ...]: + return tuple(tensors) + + +def mock_do_init_nccl_group( + self, + group_id: str, + rank: int, + actors: List[ray.actor.ActorHandle], + custom_nccl_group: Optional[Communicator], +) -> None: + ctx = ChannelContext.get_current() + if custom_nccl_group is None: + nccl_group = AbstractNcclGroup(actors) + nccl_group.initialize(rank) + ctx.communicators[group_id] = nccl_group + else: + custom_nccl_group.initialize(rank) + ctx.communicators[group_id] = custom_nccl_group + + +def mock_do_destroy_nccl_group(self, group_id: str) -> None: + ctx = ChannelContext.get_current() + if group_id not in ctx.communicators: + return + ctx.communicators[group_id].destroy() + del ctx.communicators[group_id] + + +def check_nccl_group_init( + monkeypatch, + dag: "ray.dag.DAGNode", + actors_and_custom_comms: Set[ + Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]] + ], +) -> "ray.dag.CompiledDAG": + mock_nccl_group_set = MockNcclGroupSet() + monkeypatch.setattr( + "ray.dag.compiled_dag_node._init_communicator", + mock_nccl_group_set, + ) + + compiled_dag = dag.experimental_compile() + assert ( + set(mock_nccl_group_set.ids_to_actors_and_custom_comms.values()) + == actors_and_custom_comms + ) + + return compiled_dag, mock_nccl_group_set + + +def check_nccl_group_teardown( + monkeypatch, + compiled_dag: "ray.dag.CompiledDAG", + mock_nccl_group_set: MockNcclGroupSet, +): + monkeypatch.setattr( + "ray.dag.compiled_dag_node._destroy_communicator", + mock_nccl_group_set.mock_destroy_nccl_group, + ) + + created_communicator_ids = compiled_dag._actors_to_created_communicator_id.values() + compiled_dag.teardown() + mock_nccl_group_set.check_teardown(created_communicator_ids) diff --git a/lib/python3.12/site-packages/ray/experimental/collective/nixl_tensor_transport.py b/lib/python3.12/site-packages/ray/experimental/collective/nixl_tensor_transport.py new file mode 100644 index 0000000000000000000000000000000000000000..6ab5b9d03ee3657144ce829abd69671594b00950 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/nixl_tensor_transport.py @@ -0,0 +1,197 @@ +from typing import TYPE_CHECKING, List, Optional + +import ray +from ray.experimental.collective.tensor_transport_manager import ( + TensorTransportManager, +) +from ray.util.collective.types import ( + NIXL_GROUP_NAME, + Backend, + NixlCommunicatorMetadata, + NixlTransportMetadata, +) + +if TYPE_CHECKING: + import torch + + +class NixlTensorTransport(TensorTransportManager): + @property + def tensor_transport_backend(self) -> Backend: + return Backend.NIXL + + @staticmethod + def is_one_sided() -> bool: + return True + + @staticmethod + def can_abort_transport() -> bool: + return True + + def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool: + def __ray_actor_has_tensor_transport__( + self: "ray.actor.ActorHandle", + ) -> bool: + try: + from ray.util.collective.collective import get_group_handle + + nixl_backend = get_group_handle(NIXL_GROUP_NAME) + return nixl_backend is not None + except Exception: + return False + + return ray.get( + actor.__ray_call__.options(concurrency_group="_ray_system").remote( + __ray_actor_has_tensor_transport__ + ) + ) + + @staticmethod + def extract_tensor_transport_metadata( + obj_id: str, + gpu_object: List["torch.Tensor"], + ) -> NixlTransportMetadata: + from ray._private.worker import global_worker + from ray.util.collective.collective import get_group_handle + from ray.util.collective.collective_group.nixl_backend import NixlBackend + from ray.util.collective.types import NixlTransportMetadata + + gpu_object_store = global_worker.gpu_object_manager.gpu_object_store + nixl_backend: NixlBackend = get_group_handle(NIXL_GROUP_NAME) + device = None + tensor_meta = [] + duplicate_meta = gpu_object_store.record_and_get_meta_if_duplicate( + obj_id, gpu_object + ) + if duplicate_meta is not None: + return duplicate_meta + if gpu_object: + reg_descs, serialized_descs, agent_meta = nixl_backend.get_nixl_metadata( + gpu_object + ) + # We assume all tensors in one GPU object have the same device type. + device = gpu_object[0].device + for t in gpu_object: + if t.device.type != device.type: + raise ValueError( + "All tensors in an RDT object must have the same device type." + ) + tensor_meta.append((t.shape, t.dtype)) + else: + reg_descs, serialized_descs, agent_meta = None, None, None + ret = NixlTransportMetadata( + tensor_meta=tensor_meta, + tensor_device=device, + nixl_reg_descs=reg_descs, + nixl_serialized_descs=serialized_descs, + nixl_agent_meta=agent_meta, + ) + gpu_object_store.record_managed_meta_nixl(obj_id, ret) + return ret + + @staticmethod + def get_tensor_transport_metadata( + src_actor: "ray.actor.ActorHandle", + obj_id: str, + ) -> NixlTransportMetadata: + def __ray_get_tensor_transport_metadata__( + self: "ray.actor.ActorHandle", + obj_id: str, + ) -> NixlTransportMetadata: + + from ray._private.worker import global_worker + + gpu_object_manager = global_worker.gpu_object_manager + gpu_object_store = gpu_object_manager.gpu_object_store + # NOTE: We do not specify a timeout here because the user task that returns + # it could take arbitrarily long and we don't want to trigger a spurious + # timeout. + gpu_object = gpu_object_store.wait_and_get_object(obj_id) + return NixlTensorTransport.extract_tensor_transport_metadata( + obj_id, gpu_object + ) + + # Submit a Ray actor task to the source actor to get the tensor metadata. + # The metadata is a list of tuples, where each tuple contains the shape and dtype + # of a tensor in the GPU object store. This function returns an ObjectRef that + # points to the tensor metadata. + # NOTE(swang): We put this task on the background thread to avoid tasks + # executing on the main thread blocking this task. + + return src_actor.__ray_call__.options(concurrency_group="_ray_system").remote( + __ray_get_tensor_transport_metadata__, obj_id + ) + + @staticmethod + def get_communicator_metadata( + src_actor: "ray.actor.ActorHandle", + dst_actor: "ray.actor.ActorHandle", + backend: Optional[str] = None, + ) -> NixlCommunicatorMetadata: + + communicator_metadata = NixlCommunicatorMetadata( + communicator_name=NIXL_GROUP_NAME, + ) + + return communicator_metadata + + @staticmethod + def recv_multiple_tensors( + tensors, + obj_id: str, + tensor_transport_metadata: NixlTransportMetadata, + communicator_metadata: NixlCommunicatorMetadata, + ): + from ray.util.collective import types + from ray.util.collective.collective import get_group_handle + + if tensors: + g = get_group_handle(communicator_metadata.communicator_name) + + assert isinstance( + tensor_transport_metadata, types.NixlTransportMetadata + ), "metadata must be a NixlTransportMetadata object for NIXL transport" + assert isinstance( + communicator_metadata, types.NixlCommunicatorMetadata + ), "metadata must be a NixlCommunicatorMetadata object for NIXL transport" + + g.recv( + tensors, + obj_id, + tensor_transport_metadata.nixl_serialized_descs, + tensor_transport_metadata.nixl_agent_meta, + ) + + @staticmethod + def send_multiple_tensors( + tensors: List["torch.Tensor"], + communicator_metadata: NixlCommunicatorMetadata, + device: "torch.device", + ): + raise NotImplementedError( + "NIXL transport does not support send_multiple_tensors, since it is a one-sided transport." + ) + + @staticmethod + def garbage_collect(obj_id: str, tensor_transport_meta: NixlTransportMetadata): + from ray._private.worker import global_worker + from ray.util.collective.collective import get_group_handle + + gpu_object_store = global_worker.gpu_object_manager.gpu_object_store + count = gpu_object_store.remove_managed_meta_nixl(obj_id) + if count == 0: + descs = tensor_transport_meta.nixl_reg_descs + if descs is not None: + nixl_backend = get_group_handle(NIXL_GROUP_NAME) + nixl_backend.deregister_memory(descs) + + @staticmethod + def abort_transport( + obj_id: str, + communicator_metadata: NixlCommunicatorMetadata, + ): + from ray.util.collective.collective import get_group_handle + + g = get_group_handle(communicator_metadata.communicator_name) + if g: + g.abort(obj_id) diff --git a/lib/python3.12/site-packages/ray/experimental/collective/operations.py b/lib/python3.12/site-packages/ray/experimental/collective/operations.py new file mode 100644 index 0000000000000000000000000000000000000000..aa8ba37ecc1f1257ffa3d12e0e9b7c18a7d04801 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/operations.py @@ -0,0 +1,203 @@ +import logging +from typing import List, Optional, Union + +import ray +from ray.dag.collective_node import CollectiveOutputNode, _CollectiveOperation +from ray.dag.constants import ( + BIND_INDEX_KEY, + COLLECTIVE_OPERATION_KEY, + IS_CLASS_METHOD_OUTPUT_KEY, + PARENT_CLASS_NODE_KEY, +) +from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType +from ray.experimental.util.types import ( + AllGatherOp, + AllReduceOp, + ReduceOp, + ReduceScatterOp, + _CollectiveOp, +) +from ray.util.collective.types import ReduceOp as RayReduceOp + +logger = logging.getLogger(__name__) + + +def _bind( + inputs: Union[List["ray.dag.DAGNode"], List[List["ray.dag.DAGNode"]]], + op: _CollectiveOp, + transport: Optional[Union[str, Communicator]] = None, +): + """ + Bind inputs (input nodes or lists of input nodes) with a collective operation. + The collective operation is applied to each list of input nodes. The output nodes + will have the same shape as the input nodes. + + Example of binding a list of input node: + with InputNode() as inp: + res_comp1 = [actor.comp1.bind(inp) for actor in actors] + res_comp2 = [actor.comp2.bind(inp) for actor in actors] + res_ar = allreduce.bind([res_comp1, res_comp2]) + + Requirements: + 1. Each input node returns a torch tensor. + 2. Each input node within a list is from a different actor. + 3. If lists of input nodes are provided, the order of actors should + be the same for each nested list. + 4. If a custom transport is specified, its actor set matches the actor + set of the input nodes. + 5. If input nodes are provided, then all tensors have the same shape. + If lists of input nodes are provided, then all tensors in each + list have the same shape. + + Requirements 1-3 are checked in the `CollectiveGroup` constructor. + Requirement 4 is not checked yet. + + Args: + inputs: A list of DAG nodes or a list of lists of DAG nodes. Each leaf list + should contain one object per actor. + op: The collective operation. + transport: GPU communicator for the collective operation. If not + specified, the default ACCELERATOR is used. + + Returns: + A list of collective output nodes or a list of lists of collective output nodes, + with the same shape as the input nodes. Each output node has the same order and + belongs to the same actor as the corresponding input node. + """ + if isinstance(inputs[0], list) and not isinstance(op, AllReduceOp): + raise ValueError( + "Currently binding a nested list of dag nodes is only supported for allreduce" + ) + + # Convert list of DAGNode into nested list for type checking + if not isinstance(inputs[0], list): + inputs = [inputs] + + if transport is None: + transport = TorchTensorType.ACCELERATOR + collective_op = _CollectiveOperation(inputs, op, transport) + collective_output_nodes: List[CollectiveOutputNode] = [] + + if isinstance(op, AllGatherOp): + method_name = "allgather" + elif isinstance(op, AllReduceOp): + method_name = f"allreduce.{op.reduceOp}" + elif isinstance(op, ReduceScatterOp): + method_name = f"reducescatter.{op.reduceOp}" + else: + raise ValueError(f"Expected a collective operation, but got {op}") + + for i in range(len(inputs[0])): + input_node_list = [l[i] for l in inputs if l] + actor_handle: Optional["ray.actor.ActorHandle"] = input_node_list[ + 0 + ]._get_actor_handle() + assert actor_handle is not None + collective_output_node = CollectiveOutputNode( + method_name=method_name, + method_args=tuple(input_node_list), + method_kwargs=dict(), + method_options=dict(), + other_args_to_resolve={ + PARENT_CLASS_NODE_KEY: actor_handle, + BIND_INDEX_KEY: actor_handle._ray_dag_bind_index, + COLLECTIVE_OPERATION_KEY: collective_op, + }, + ) + actor_handle._ray_dag_bind_index += 1 + + if len(input_node_list) > 1: + output_nodes: List[CollectiveOutputNode] = [] + for i in range(len(input_node_list)): + output_node = CollectiveOutputNode( + f"return_idx_{i}", + (collective_output_node, i), + dict(), + dict(), + { + BIND_INDEX_KEY: collective_output_node._get_bind_index(), + IS_CLASS_METHOD_OUTPUT_KEY: True, + PARENT_CLASS_NODE_KEY: actor_handle, + }, + ) + output_nodes.append(output_node) + collective_output_nodes.append(output_nodes) + else: + collective_output_nodes.append(collective_output_node) + return collective_output_nodes + + +class AllGatherWrapper: + """Wrapper for NCCL all-gather.""" + + def bind( + self, + input_nodes: List["ray.dag.DAGNode"], + transport: Optional[Union[str, Communicator]] = None, + ) -> List[CollectiveOutputNode]: + return _bind(input_nodes, AllGatherOp(), transport) + + def __call__( + self, + tensor_list, + tensor, + group_name: str = "default", + ): + from ray.util.collective.collective import allgather + + return allgather(tensor_list, tensor, group_name) + + +class AllReduceWrapper: + """Wrapper for NCCL all-reduce.""" + + def bind( + self, + input_nodes: List["ray.dag.DAGNode"], + op: ReduceOp = ReduceOp.SUM, + transport: Optional[Union[str, Communicator]] = None, + ) -> List[CollectiveOutputNode]: + if not isinstance(op, ReduceOp): + raise ValueError(f"Unexpected operation: {op}") + + return _bind(input_nodes, AllReduceOp(reduceOp=op), transport) + + def __call__( + self, + tensor, + group_name: str = "default", + op: RayReduceOp = RayReduceOp.SUM, + ): + from ray.util.collective.collective import allreduce + + return allreduce(tensor, group_name, op) + + +class ReduceScatterWrapper: + """Wrapper for NCCL reduce-scatter.""" + + def bind( + self, + input_nodes: List["ray.dag.DAGNode"], + op: ReduceOp = ReduceOp.SUM, + transport: Optional[Union[str, Communicator]] = None, + ) -> List[CollectiveOutputNode]: + if not isinstance(op, ReduceOp): + raise ValueError(f"Unexpected operation: {op}") + + return _bind(input_nodes, ReduceScatterOp(reduceOp=op), transport) + + def __call__( + self, + tensor, + group_name: str = "default", + op: RayReduceOp = RayReduceOp.SUM, + ): + from ray.util.collective.collective import reducescatter + + return reducescatter(tensor, group_name, op) + + +allgather = AllGatherWrapper() +allreduce = AllReduceWrapper() +reducescatter = ReduceScatterWrapper() diff --git a/lib/python3.12/site-packages/ray/experimental/collective/tensor_transport_manager.py b/lib/python3.12/site-packages/ray/experimental/collective/tensor_transport_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..7f0210b7d9ca4ad5d3d5b6ec74e12296840b1ce2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/tensor_transport_manager.py @@ -0,0 +1,168 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, List, Optional + +import ray +from ray.util.collective.types import ( + Backend, + CommunicatorMetadata, + TensorTransportMetadata, +) + +if TYPE_CHECKING: + import torch + + +class TensorTransportManager(ABC): + @property + @abstractmethod + def tensor_transport_backend(self) -> Backend: + """The tensor transport backend, e.g., NCCL. + + Returns: + Backend: The backend of the tensor transport. + """ + + @staticmethod + @abstractmethod + def is_one_sided() -> bool: + """Whether the backend is one-sided. + + Returns: + bool: True if the backend is one-sided, False otherwise. + """ + + @staticmethod + @abstractmethod + def can_abort_transport() -> bool: + """ + Whether the backend can abort the transport. + If this returns False, then Ray will kill involved actors upon system errors to avoid hanging. + + Returns: + bool: True if the backend can abort the transport. + """ + + @abstractmethod + def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool: + """Whether the actor has the tensor transport available. + + Args: + actor: The actor to check. + + Returns: + bool: True if the actor has the tensor transport available, False otherwise. + """ + + @staticmethod + @abstractmethod + def get_tensor_transport_metadata( + src_actor: "ray.actor.ActorHandle", + obj_id: str, + ) -> TensorTransportMetadata: + """ + Get the tensor transport metadata for the GPU object. + This function retrieves metadata about tensors stored in the GPU object store, + including their shapes, dtypes, and any transport-specific metadata, e.g., NIXL descriptors. + + Args: + src_actor: The actor that runs this function. + obj_id: The ID of the GPU object to get metadata for + + Returns: + TensorTransportMetadata: A named tuple containing the tensor metadata. + """ + + @staticmethod + @abstractmethod + def extract_tensor_transport_metadata( + obj_id: str, + gpu_object: List["torch.Tensor"], + ) -> TensorTransportMetadata: + """ + Extract the tensor transport metadata from the GPU object. + + Args: + obj_id: The ID of the GPU object to extract the tensor transport metadata from. + gpu_object: The GPU object to extract the tensor transport metadata from. + + Returns: + TensorTransportMetadata: The tensor transport metadata. + """ + + @staticmethod + @abstractmethod + def get_communicator_metadata( + src_actor: "ray.actor.ActorHandle", + dst_actor: "ray.actor.ActorHandle", + backend: Optional[str] = None, + ) -> CommunicatorMetadata: + """ + Get the communicator metadata (e.g. communicator name, src/dst rank) for the send/recv operation. + This function is called before sending the GPU object. + + Args: + src_actor: The actor that runs this function. + dst_actor: The actor that runs this function. + backend: The backend to use for the collective operation. + + Returns: + CommunicatorMetadata: The communicator metadata. + """ + + @staticmethod + @abstractmethod + def recv_multiple_tensors( + tensors: List["torch.Tensor"], + obj_id: str, + tensor_transport_metadata: TensorTransportMetadata, + communicator_metadata: CommunicatorMetadata, + ): + """ + Receive multiple tensors from the source actor. + + Args: + tensors: The pre-allocated tensor space to receive the tensors. + obj_id: The object ID for related GPU object. + tensor_transport_metadata: The tensor transport metadata for the GPU object. + communicator_metadata: The communicator metadata for the send/recv operation. + + """ + + @staticmethod + @abstractmethod + def send_multiple_tensors( + tensors: List["torch.Tensor"], + communicator_metadata: CommunicatorMetadata, + ): + """ + Send multiple tensors to the destination actor. + + Args: + tensors: The tensors to send. + communicator_metadata: The communicator metadata for the send/recv operation. + """ + + @staticmethod + @abstractmethod + def garbage_collect(obj_id: str, tensor_transport_meta: TensorTransportMetadata): + """ + Garbage collect for the tensor transport after the GPU object is freed. + + Args: + obj_id: The ID of the GPU object to garbage collect. + tensor_transport_meta: The tensor transport metadata. + """ + + @staticmethod + @abstractmethod + def abort_transport( + obj_id: str, + communicator_metadata: CommunicatorMetadata, + ): + """ + Abort the transport. + + Args: + obj_id: The object ID for related GPU object. + communicator_metadata: The communicator metadata for the send/recv operation. + """ diff --git a/lib/python3.12/site-packages/ray/experimental/collective/util.py b/lib/python3.12/site-packages/ray/experimental/collective/util.py new file mode 100644 index 0000000000000000000000000000000000000000..241a95e890afeed8b8045b4c391b46d254c4203a --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/collective/util.py @@ -0,0 +1,68 @@ +import socket +from typing import TYPE_CHECKING, Tuple + +import ray +from ray._common.network_utils import find_free_port, is_ipv6 +from ray.experimental.collective.collective_tensor_transport import ( + CollectiveTensorTransport, +) +from ray.experimental.collective.nixl_tensor_transport import NixlTensorTransport +from ray.experimental.collective.tensor_transport_manager import TensorTransportManager +from ray.util.collective.types import Backend + +if TYPE_CHECKING: + import torch + +# Singleton instances for tensor transport managers +_nixl_tensor_transport_manager = None +_gloo_tensor_transport_manager = None +_nccl_tensor_transport_manager = None + + +def get_tensor_transport_manager( + tensor_transport: Backend, +) -> "TensorTransportManager": + """Get the tensor transport manager for the given tensor transport protocol. + + Args: + tensor_transport: The tensor transport protocol to use for the GPU object. + + Returns: + TensorTransportManager: The tensor transport manager for the given tensor transport protocol. + """ + if tensor_transport == Backend.NIXL: + global _nixl_tensor_transport_manager + if _nixl_tensor_transport_manager is None: + _nixl_tensor_transport_manager = NixlTensorTransport() + return _nixl_tensor_transport_manager + elif tensor_transport == Backend.TORCH_GLOO: + global _gloo_tensor_transport_manager + if _gloo_tensor_transport_manager is None: + _gloo_tensor_transport_manager = CollectiveTensorTransport(tensor_transport) + return _gloo_tensor_transport_manager + elif tensor_transport == Backend.NCCL: + global _nccl_tensor_transport_manager + if _nccl_tensor_transport_manager is None: + _nccl_tensor_transport_manager = CollectiveTensorTransport(tensor_transport) + return _nccl_tensor_transport_manager + else: + raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}") + + +def device_match_transport(device: "torch.device", tensor_transport: Backend) -> bool: + """Check if the device matches the transport.""" + if tensor_transport == Backend.NIXL: + return device.type == "cuda" or device.type == "cpu" + elif tensor_transport == Backend.TORCH_GLOO: + return device.type == "cpu" + elif tensor_transport == Backend.NCCL: + return device.type == "cuda" + else: + raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}") + + +def get_address_and_port() -> Tuple[str, int]: + """Returns the IP address and a free port on this node.""" + addr = ray.util.get_node_ip_address() + port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET) + return addr, port diff --git a/lib/python3.12/site-packages/ray/experimental/compiled_dag_ref.py b/lib/python3.12/site-packages/ray/experimental/compiled_dag_ref.py new file mode 100644 index 0000000000000000000000000000000000000000..d624457b8b53d419531813643480719274c27136 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/compiled_dag_ref.py @@ -0,0 +1,228 @@ +import asyncio +from typing import Any, List, Optional + +import ray +from ray.exceptions import ( + GetTimeoutError, + RayChannelError, + RayChannelTimeoutError, + RayTaskError, +) +from ray.util.annotations import PublicAPI + + +def _process_return_vals(return_vals: List[Any], return_single_output: bool): + """ + Process return values for return to the DAG caller. Any exceptions found in + return_vals will be raised. If return_single_output=True, it indicates that + the original DAG did not have a MultiOutputNode, so the DAG caller expects + a single return value instead of a list. + """ + # Check for exceptions. + if isinstance(return_vals, Exception): + raise return_vals + + for val in return_vals: + if isinstance(val, RayTaskError): + raise val.as_instanceof_cause() + + if return_single_output: + assert len(return_vals) == 1 + return return_vals[0] + + return return_vals + + +@PublicAPI(stability="alpha") +class CompiledDAGRef: + """ + A reference to a compiled DAG execution result. + + This is a subclass of ObjectRef and resembles ObjectRef. For example, + similar to ObjectRef, ray.get() can be called on it to retrieve the result. + However, there are several major differences: + 1. ray.get() can only be called once per CompiledDAGRef. + 2. ray.wait() is not supported. + 3. CompiledDAGRef cannot be copied, deep copied, or pickled. + 4. CompiledDAGRef cannot be passed as an argument to another task. + """ + + def __init__( + self, + dag: "ray.experimental.CompiledDAG", + execution_index: int, + channel_index: Optional[int] = None, + ): + """ + Args: + dag: The compiled DAG that generated this CompiledDAGRef. + execution_index: The index of the execution for the DAG. + A DAG can be executed multiple times, and execution index + indicates which execution this CompiledDAGRef corresponds to. + actor_execution_loop_refs: The actor execution loop refs that + are used to execute the DAG. This can be used internally to + check the task execution errors in case of exceptions. + channel_index: The index of the DAG's output channel to fetch + the result from. A DAG can have multiple output channels, and + channel index indicates which channel this CompiledDAGRef + corresponds to. If channel index is not provided, this CompiledDAGRef + wraps the results from all output channels. + + """ + self._dag = dag + self._execution_index = execution_index + self._channel_index = channel_index + # Whether ray.get() was called on this CompiledDAGRef. + self._ray_get_called = False + self._dag_output_channels = dag.dag_output_channels + + def __str__(self): + return ( + f"CompiledDAGRef({self._dag.get_id()}, " + f"execution_index={self._execution_index}, " + f"channel_index={self._channel_index})" + ) + + def __copy__(self): + raise ValueError("CompiledDAGRef cannot be copied.") + + def __deepcopy__(self, memo): + raise ValueError("CompiledDAGRef cannot be deep copied.") + + def __reduce__(self): + raise ValueError("CompiledDAGRef cannot be pickled.") + + def __del__(self): + # If the dag is already teardown, it should do nothing. + if self._dag.is_teardown: + return + + if self._ray_get_called: + # get() was already called, no further cleanup is needed. + return + + self._dag._delete_execution_results(self._execution_index, self._channel_index) + + def get(self, timeout: Optional[float] = None): + if self._ray_get_called: + raise ValueError( + "ray.get() can only be called once " + "on a CompiledDAGRef, and it was already called." + ) + + self._ray_get_called = True + try: + self._dag._execute_until( + self._execution_index, self._channel_index, timeout + ) + return_vals = self._dag._get_execution_results( + self._execution_index, self._channel_index + ) + except RayChannelTimeoutError: + raise + except RayChannelError as channel_error: + # If we get a channel error, we'd like to call ray.get() on + # the actor execution loop refs to check if this is a result + # of task execution error which could not be passed down + # (e.g., when a pure NCCL channel is used, it is only + # able to send tensors, but not the wrapped exceptions). + # In this case, we'd like to raise the task execution error + # (which is the actual cause of the channel error) instead + # of the channel error itself. + # TODO(rui): determine which error to raise if multiple + # actor task refs have errors. + actor_execution_loop_refs = list(self._dag.worker_task_refs.values()) + try: + ray.get(actor_execution_loop_refs, timeout=10) + except GetTimeoutError as timeout_error: + raise Exception( + "Timed out when getting the actor execution loop exception. " + "This should not happen, please file a GitHub issue." + ) from timeout_error + except Exception as execution_error: + # Use 'from None' to suppress the context of the original + # channel error, which is not useful to the user. + raise execution_error from None + else: + raise channel_error + except Exception: + raise + return _process_return_vals(return_vals, True) + + +@PublicAPI(stability="alpha") +class CompiledDAGFuture: + """ + A reference to a compiled DAG execution result, when executed with asyncio. + This differs from CompiledDAGRef in that `await` must be called on the + future to get the result, instead of `ray.get()`. + + This resembles async usage of ObjectRefs. For example, similar to + ObjectRef, `await` can be called directly on the CompiledDAGFuture to + retrieve the result. However, there are several major differences: + 1. `await` can only be called once per CompiledDAGFuture. + 2. ray.wait() is not supported. + 3. CompiledDAGFuture cannot be copied, deep copied, or pickled. + 4. CompiledDAGFuture cannot be passed as an argument to another task. + """ + + def __init__( + self, + dag: "ray.experimental.CompiledDAG", + execution_index: int, + fut: "asyncio.Future", + channel_index: Optional[int] = None, + ): + self._dag = dag + self._execution_index = execution_index + self._fut = fut + self._channel_index = channel_index + + def __str__(self): + return ( + f"CompiledDAGFuture({self._dag.get_id()}, " + f"execution_index={self._execution_index}, " + f"channel_index={self._channel_index})" + ) + + def __copy__(self): + raise ValueError("CompiledDAGFuture cannot be copied.") + + def __deepcopy__(self, memo): + raise ValueError("CompiledDAGFuture cannot be deep copied.") + + def __reduce__(self): + raise ValueError("CompiledDAGFuture cannot be pickled.") + + def __await__(self): + if self._fut is None: + raise ValueError( + "CompiledDAGFuture can only be awaited upon once, and it has " + "already been awaited upon." + ) + + # NOTE(swang): If the object is zero-copy deserialized, then it will + # stay in scope as long as this future is in scope. Therefore, we + # delete self._fut here before we return the result to the user. + fut = self._fut + self._fut = None + + if not self._dag._has_execution_results(self._execution_index): + result = yield from fut.__await__() + self._dag._max_finished_execution_index += 1 + self._dag._cache_execution_results(self._execution_index, result) + + return_vals = self._dag._get_execution_results( + self._execution_index, self._channel_index + ) + return _process_return_vals(return_vals, True) + + def __del__(self): + if self._dag.is_teardown: + return + + if self._fut is None: + # await() was already called, no further cleanup is needed. + return + + self._dag._delete_execution_results(self._execution_index, self._channel_index) diff --git a/lib/python3.12/site-packages/ray/experimental/dynamic_resources.py b/lib/python3.12/site-packages/ray/experimental/dynamic_resources.py new file mode 100644 index 0000000000000000000000000000000000000000..bb7bd3948095e911d4a8d1bbb82e13db305e0c4a --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/dynamic_resources.py @@ -0,0 +1,7 @@ +def set_resource(resource_name, capacity, node_id=None): + raise DeprecationWarning( + "Dynamic custom resources are deprecated. Consider using placement " + "groups instead (docs.ray.io/en/master/placement-group.html). You " + "can also specify resources at Ray start time with the 'resources' " + "field in the cluster autoscaler." + ) diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__init__.py b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..13be5939544503b51ffcb2b975eb5d787611eb77 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__init__.py @@ -0,0 +1,6 @@ +from ray.experimental.gpu_object_manager.gpu_object_manager import ( + GPUObjectManager, + wait_tensor_freed, +) + +__all__ = ["GPUObjectManager", "wait_tensor_freed"] diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf60b78f4ab1dc3c73f0bbcf4cd1d26c7b438b6b Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d8f2ce1e5631a364aedc1749d3fde5b6d8f6df8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_store.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_store.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56570ff056a69db06f65cef9abc89fc8ab285013 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/__pycache__/gpu_object_store.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_manager.py b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..788ee9ece376a65ddc4b5c5e80c380cd1373e23d --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_manager.py @@ -0,0 +1,655 @@ +import logging +import threading +import time +import warnings +from queue import Queue +from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Set, Tuple + +import ray +from ray._private import ray_constants +from ray._private.custom_types import TensorTransportEnum +from ray._raylet import ObjectRef + +if TYPE_CHECKING: + import torch + + from ray.experimental.gpu_object_manager.gpu_object_store import ( + GPUObjectStore, + ) + from ray.util.collective.types import CommunicatorMetadata, TensorTransportMetadata + +logger = logging.getLogger(__name__) + +# GPUObjectMeta is a named tuple containing the source actor, tensor transport +# backend, tensor metadata, and other information that needs to be recorded. +# - The tensor transport backend is the backend used to transport the tensors. +# Currently, the supported backends are "nccl" and "torch_gloo". +# - The tensor metadata is a list of tuples, each containing the shape and dtype +# of a tensor in the GPU object store. +class GPUObjectMeta(NamedTuple): + src_actor: "ray.actor.ActorHandle" + # Must be a valid backend name as defined in + # `ray.util.collective.types.Backend`. + tensor_transport_backend: str + tensor_transport_meta: "TensorTransportMetadata" + # sent_dest_actors tracks the set of actor IDs that this object has been sent to. + sent_dest_actors: Set[str] + # sent_to_src_actor_and_others_warned indicates whether the object has already triggered a warning about being sent back to the source actor and other actors simultaneously. + sent_to_src_actor_and_others_warned: bool + + +# This is used to periodically check in on the RDT transfer through the refs from +# __ray_send__ and __ray_recv__ and abort operations in case of failures / timeouts. +class TransferMetadata(NamedTuple): + src_actor: "ray.actor.ActorHandle" + dst_actor: "ray.actor.ActorHandle" + send_ref: Optional[ObjectRef] + recv_ref: ObjectRef + communicator_meta: "CommunicatorMetadata" + backend: str + obj_id: str + timeout: float + + +# TODO(swang): Uncomment and add an API docs page and example usage. +# @PublicAPI(stability="alpha") +def wait_tensor_freed(tensor: "torch.Tensor", timeout: Optional[float] = None): + """ + Wait for the tensor to be freed. + + This function is useful for cases where an actor keeps a reference to a + tensor after returning the tensor from a task annotated with + `@ray.method(tensor_transport=...)`. In this case, Ray will store a + *reference* to the tensor, so any in-place modifications made by the actor + that returned the tensor could be seen by other actors. See + :ref:`Ray Direct Transport (RDT) ` for more details. + + Call this function for RDT objects to ensure that all corresponding + `ray.ObjectRefs` have gone out of scope and therefore the tensor is safe to + write to again. + + Args: + tensor: The tensor to wait to be freed. This should be a tensor that was + previously returned by a task annotated with + `@ray.method(tensor_transport=...)` or stored via + `ray.put(_tensor_transport="...")`. + timeout: The timeout in seconds to wait for all references to the tensor + to go out of scope. Set to None to wait indefinitely. Note that if + None is used, this function could hang if the `ray.ObjectRefs` that + refer to this tensor never go out of scope. + """ + gpu_object_manager = ray.worker.global_worker.gpu_object_manager + gpu_object_manager.gpu_object_store.wait_tensor_freed(tensor, timeout) + + +class GPUObjectManager: + def __init__(self): + # A dictionary that maps from owned object's ID to GPUObjectMeta. + # This dictionary is hosted on the "driver" process of the actors that + # store and send/receive GPU objects. + self.managed_gpu_object_metadata: Dict[str, GPUObjectMeta] = {} + + # Per-actor local storage for GPU objects. We create the GPU object + # store lazily, if a user specifies a non-default tensor_transport, to + # avoid circular import and because it imports third-party dependencies + # like PyTorch. + self._gpu_object_store: Optional["GPUObjectStore"] = None + # Lock to ensure we only create the GPU object store once. + self.gpu_object_store_lock = threading.Lock() + + # Thread safe queue of transport refs that the monitor thread needs to start monitoring + self._unmonitored_transfers: Queue[TransferMetadata] = Queue() + # Background thread to poll on the transfer operation. + self._monitor_failures_thread = None + # Event to signal the monitor_failures thread to shutdown + self._monitor_failures_shutdown_event = threading.Event() + + @property + def gpu_object_store(self) -> "ray.experimental.GPUObjectStore": + with self.gpu_object_store_lock: + if self._gpu_object_store is None: + from ray.experimental.gpu_object_manager.gpu_object_store import ( + GPUObjectStore, + ) + + self._gpu_object_store = GPUObjectStore() + return self._gpu_object_store + + def shutdown(self): + """ + Interrupt and join the monitor_failures thread. + """ + if self._monitor_failures_thread: + self._monitor_failures_shutdown_event.set() + self._monitor_failures_thread.join() + self._monitor_failures_shutdown_event.clear() + self._monitor_failures_thread = None + + def _monitor_failures(self): + """ + Monitor the refs from send and recv tasks and abort the transfers + if they error out or timeout to prevent hanging. + """ + not_done = [] + done = [] + ref_info_map = {} + while not self._monitor_failures_shutdown_event.is_set(): + while not self._unmonitored_transfers.empty(): + ref_info = self._unmonitored_transfers.get() + if ref_info.send_ref: + not_done.append(ref_info.send_ref) + ref_info_map[ref_info.send_ref.hex()] = ref_info + not_done.append(ref_info.recv_ref) + ref_info_map[ref_info.recv_ref.hex()] = ref_info + if len(not_done) > 0: + done, not_done = ray.wait(not_done, num_returns=1, timeout=1) + if len(done) > 0: + try: + ray.get(done[0]) + ref_info_map.pop(done[0].hex(), None) + except Exception as e: + self._abort_transport(done[0], ref_info_map, e) + + while len(not_done) > 0: + if not_done[0].hex() not in ref_info_map: + # The associated transfer was already aborted. + not_done.pop(0) + elif ref_info_map[not_done[0].hex()].timeout < time.time(): + self._abort_transport( + not_done[0], + ref_info_map, + TimeoutError( + f"RDT transfer failed after {ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS}s. " + "You can increase the timeout by setting RAY_rdt_fetch_fail_timeout_milliseconds" + ), + ) + else: + # wait returns lists in the same order they were passed in, so if + # the timeout of first hasn't been reached, neither have the others. + break + if len(not_done) == 0: + # If we emptied out _unmonitored_transfers on this iteration, wait for a bit. + self._monitor_failures_shutdown_event.wait(1) + + def _abort_transport( + self, + failed_ref: ObjectRef, + ref_info_map: Dict[str, TransferMetadata], + exception: Exception, + ): + """ + Cleans up the ref_info_map, kill the src and dst actors, and destroy the + collective group if necessary. + """ + from ray.experimental.collective import ( + destroy_collective_group, + get_tensor_transport_manager, + ) + from ray.experimental.gpu_object_manager.gpu_object_store import ( + __ray_abort_transport__, + ) + from ray.util.collective.types import CollectiveCommunicatorMetadata + + ref_info = ref_info_map.pop(failed_ref.hex(), None) + if ref_info is None: + return + + if ref_info.send_ref: + ref_info_map.pop(ref_info.send_ref.hex(), None) + ref_info_map.pop(ref_info.recv_ref.hex(), None) + + tensor_transport_manager = get_tensor_transport_manager(ref_info.backend) + if tensor_transport_manager.can_abort_transport(): + if not tensor_transport_manager.is_one_sided(): + # This is dead code until we implement a NCCL abort since NIXL + # is the only abortable transport for now and is one-sided. + ref_info.src_actor.__ray_call__.options( + concurrency_group="_ray_system_error" + ).remote( + __ray_abort_transport__, + ref_info.obj_id, + ref_info.communicator_meta, + ) + ref_info.dst_actor.__ray_call__.options( + concurrency_group="_ray_system_error" + ).remote( + __ray_abort_transport__, + ref_info.obj_id, + ref_info.communicator_meta, + ) + logger.info( + "RDT transfer with src actor %s and dst actor %s failed due to %s.", + ref_info.src_actor, + ref_info.dst_actor, + exception, + ) + else: + # TODO(#51276): Kill all actors in the collective group when we support more collective operations + ray.kill(ref_info.src_actor) + ray.kill(ref_info.dst_actor) + logger.error( + "RDT transfer with src actor %s and dst actor %s failed. Killing the actors. " + "Transfer failed with exception: %s", + ref_info.src_actor, + ref_info.dst_actor, + exception, + ) + + # isinstance does an implicit cast and makes communicator_name inaccessible + # so we have to get communicator_name before the cast. + collective_group_name = ref_info.communicator_meta.communicator_name + if isinstance(ref_info.communicator_meta, CollectiveCommunicatorMetadata): + try: + destroy_collective_group(collective_group_name) + logger.error( + "Destroyed collective group %s due to a hanging/failed RDT transfer", + collective_group_name, + ) + except ValueError: + # Collective group was already destroyed + pass + + def is_managed_object(self, obj_id: str) -> bool: + """ + Check if the GPU object is managed by this process. + + Args: + obj_id: The object ID of the GPU object. + + Returns: + True if the current process is the driver process coordinating the data transfer + of this GPU object. + """ + return obj_id in self.managed_gpu_object_metadata + + def add_gpu_object_metadata( + self, obj_ref: ObjectRef, gpu_object_meta: GPUObjectMeta + ): + """ + Add the GPU object metadata to the GPU object manager. + + Args: + obj_ref: The ObjectRef of the GPU object. + gpu_object_meta: The GPU object metadata. + """ + obj_id = obj_ref.hex() + self.managed_gpu_object_metadata[obj_id] = gpu_object_meta + + def add_gpu_object_ref( + self, + obj_ref: ObjectRef, + src_actor: "ray.actor.ActorHandle", + tensor_transport: TensorTransportEnum, + tensor_transport_meta: Optional["TensorTransportMetadata"] = None, + ): + """Add a GPU object reference to the GPU object manager. This should be + called whenever the current process calls a task that is annotated with + `@ray.method(tensor_transport=...)`. + + Args: + obj_ref: The ObjectRef of the task output. + src_actor: The actor that executes the task and that creates the GPU object. + tensor_transport: The tensor transport protocol to use for the GPU object. + tensor_transport_meta: The tensor transport metadata that is pre-computed. + """ + from ray.experimental.collective import get_tensor_transport_manager + from ray.experimental.gpu_object_manager.gpu_object_store import ( + _tensor_transport_to_collective_backend, + ) + + tensor_transport_backend = _tensor_transport_to_collective_backend( + tensor_transport + ) + obj_id = obj_ref.hex() + tensor_transport_manager = get_tensor_transport_manager( + tensor_transport_backend + ) + if not tensor_transport_meta: + tensor_meta = tensor_transport_manager.get_tensor_transport_metadata( + src_actor, obj_id + ) + else: + tensor_meta = tensor_transport_meta + self.managed_gpu_object_metadata[obj_id] = GPUObjectMeta( + src_actor=src_actor, + tensor_transport_backend=tensor_transport_backend, + tensor_transport_meta=tensor_meta, + sent_dest_actors=set(), + sent_to_src_actor_and_others_warned=False, + ) + + def _get_gpu_object_metadata(self, obj_ref: ObjectRef) -> GPUObjectMeta: + obj_id = obj_ref.hex() + return self.managed_gpu_object_metadata[obj_id] + + def _fetch_object( + self, + obj_id: str, + tensor_transport: TensorTransportEnum = TensorTransportEnum.OBJECT_STORE, + ): + """ + Fetches the GPU object from the source actor's GPU object store via the object store + instead of out-of-band tensor transfer and stores the tensors in the local GPU object store. + + This is useful when the current process does not support the designated out-of-band tensor transport. + For example, if the tensor transport is NCCL but the driver does not have a GPU, we use this call to + fulfill a `ray.get` call. + + Args: + obj_id: The object ID of the GPU object. + tensor_transport: The tensor transport to use to fetch the GPU object. + + Returns: + None + """ + from ray.experimental.collective import get_tensor_transport_manager + from ray.experimental.gpu_object_manager.gpu_object_store import ( + __ray_fetch_gpu_object__, + ) + + if tensor_transport not in [ + TensorTransportEnum.OBJECT_STORE, + TensorTransportEnum.NIXL, + ]: + raise ValueError( + f"Currently ray.get() only supports OBJECT_STORE and NIXL tensor transport, got {tensor_transport}, please specify the correct tensor transport in ray.get()." + ) + + if self.gpu_object_store.has_object(obj_id): + return + gpu_object_meta = self.managed_gpu_object_metadata[obj_id] + src_actor = gpu_object_meta.src_actor + tensor_transport_backend = gpu_object_meta.tensor_transport_backend + tensor_transport_manager = get_tensor_transport_manager( + tensor_transport_backend + ) + if tensor_transport == TensorTransportEnum.OBJECT_STORE: + tensors = ray.get( + src_actor.__ray_call__.options(concurrency_group="_ray_system").remote( + __ray_fetch_gpu_object__, obj_id + ) + ) + self.gpu_object_store.add_object(obj_id, tensors, is_primary=False) + else: + if isinstance(gpu_object_meta.tensor_transport_meta, ObjectRef): + # If the tensor transport meta is an ObjectRef, gpu object manager + # needs to fetch the tensor transport meta from the src actor first. + fetched_meta = ray.get(gpu_object_meta.tensor_transport_meta) + + gpu_object_meta = gpu_object_meta._replace( + tensor_transport_meta=fetched_meta + ) + # Update the managed GPU object metadata so that the next time + # it doesn't need to fetch the tensor transport meta again. + self.managed_gpu_object_metadata[obj_id] = gpu_object_meta + + from ray.experimental.gpu_object_manager.gpu_object_store import ( + __ray_recv__, + ) + + communicator_meta = tensor_transport_manager.get_communicator_metadata( + None, None, tensor_transport_backend + ) + __ray_recv__( + None, obj_id, [gpu_object_meta.tensor_transport_meta], communicator_meta + ) + + def trigger_out_of_band_tensor_transfer( + self, dst_actor: "ray.actor.ActorHandle", task_args: Tuple[Any, ...] + ): + """ + Triggers tensor communication operations between actors. When a managed ObjectRef is passed + to another actor task, CPU data will still be passed through the object store, but the in-actor + tensors will be passed out-of-band. + + This function triggers the out-of-band tensor transfer by submitting Ray actor + tasks `__ray_send__` to the sender actor and `__ray_recv__` to the receiver actor to initiate + tensor communication using protocols like NCCL or GLOO. + + Before the receiver actor executes the actor task, the deserializer combines the + CPU data with the tensors from the sender actor to reconstruct the original task output + generated by the sender actor. + + Args: + dst_actor: The target actor to receive tensors + task_args: List of arguments for the target actor task that may contain ObjectRefs. + """ + + gpu_object_refs = set() + for arg in task_args: + # If an ObjectRef is managed, it means the actual value is a list of tensors stored + # on a remote actor. Therefore, this function will trigger a tensor communication + # operation between the sender and receiver actors. + if not isinstance(arg, ObjectRef): + continue + if self.is_managed_object(arg.hex()): + gpu_object_refs.add(arg) + if gpu_object_refs: + from ray.experimental.collective import get_tensor_transport_manager + from ray.experimental.gpu_object_manager.gpu_object_store import ( + __ray_recv__, + __ray_send__, + ) + + # Count the number of readers for each GPU object. + for obj_ref in gpu_object_refs: + # Import get_collective_groups here to avoid dependency on + # collective libraries for default Ray installation. + + gpu_object_meta = self._get_gpu_object_metadata(obj_ref) + + src_actor = gpu_object_meta.src_actor + tensor_transport_meta = gpu_object_meta.tensor_transport_meta + + obj_id = obj_ref.hex() + + # Update the set of destination actors for this object + # The set inside NamedTuple is mutable, so we can modify it directly + gpu_object_meta.sent_dest_actors.add(dst_actor._actor_id) + # Check if a warning should be triggered for this object: + # 1. object has not triggered a warning yet. + # 2. object is sent back to its source actor. + # 3. object is also sent to at least one other actor + if ( + not gpu_object_meta.sent_to_src_actor_and_others_warned + and src_actor._actor_id in gpu_object_meta.sent_dest_actors + and len(gpu_object_meta.sent_dest_actors) > 1 + ): + warnings.warn( + f"GPU ObjectRef({obj_id}) is being passed back to the actor that created it {src_actor}. " + "Note that GPU objects are mutable. If the tensor is modified, Ray's internal copy will also be updated, and subsequent passes to other actors " + "will receive the updated version instead of the original.", + UserWarning, + ) + # Mark the object as warned by creating a new NamedTuple instance + self.managed_gpu_object_metadata[obj_id] = gpu_object_meta._replace( + sent_to_src_actor_and_others_warned=True + ) + + if src_actor._actor_id == dst_actor._actor_id: + # If the source and destination actors are the same, the tensors can + # be transferred intra-process, so we skip the out-of-band tensor + # transfer. + continue + + tensor_transport_manager = get_tensor_transport_manager( + gpu_object_meta.tensor_transport_backend + ) + communicator_meta = tensor_transport_manager.get_communicator_metadata( + src_actor, + dst_actor, + gpu_object_meta.tensor_transport_backend, + ) + + send_ref = None + if not tensor_transport_manager.is_one_sided(): + # Send tensors stored in the `src_actor`'s GPU object store to the + # destination rank `dst_rank`. + # NOTE: We put this task on the background thread to avoid tasks + # executing on the main thread blocking the data transfer. + send_ref = src_actor.__ray_call__.options( + concurrency_group="_ray_system" + ).remote( + __ray_send__, + obj_id, + tensor_transport_meta, + communicator_meta, + ) + + # Receive tensors from the source rank and store them in the + # `dst_actor`'s GPU object store. + # NOTE: Putting this task on the background thread is technically only + # needed for the sender task, but we put the receiver task on the same + # background thread to ensure that all communication operations are + # executed in a global order. + recv_ref = dst_actor.__ray_call__.options( + concurrency_group="_ray_system" + ).remote( + __ray_recv__, + obj_id, + [tensor_transport_meta], + communicator_meta, + ) + + self._unmonitored_transfers.put( + TransferMetadata( + src_actor=src_actor, + dst_actor=dst_actor, + send_ref=send_ref, + recv_ref=recv_ref, + communicator_meta=communicator_meta, + backend=gpu_object_meta.tensor_transport_backend, + obj_id=obj_id, + timeout=time.time() + ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS, + ) + ) + if self._monitor_failures_thread is None: + self._monitor_failures_thread = threading.Thread( + target=self._monitor_failures, daemon=True + ) + self._monitor_failures_thread.start() + + def get_gpu_object( + self, + object_id: str, + tensor_transport: TensorTransportEnum = TensorTransportEnum.OBJECT_STORE, + ) -> List["torch.Tensor"]: + """ + Get the GPU object for a given object ID. + + Args: + object_id: The object ID of the GPU object. + tensor_transport: The tensor transport to use to fetch the GPU object. + + Returns: + The GPU object. + """ + gpu_object_store = self.gpu_object_store + if self.is_managed_object(object_id): + self._fetch_object(object_id, tensor_transport) + + # If the GPU object is the primary copy, it means the transfer is intra-actor. + # In this case, we should not remove the GPU object after it is consumed once, + # because the GPU object reference may be used again. + # Instead, we should wait for the GC callback to clean it up. + pop_object = not gpu_object_store.is_primary_copy(object_id) + if pop_object: + gpu_object = gpu_object_store.wait_and_pop_object( + object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS + ) + else: + gpu_object = gpu_object_store.wait_and_get_object( + object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS + ) + return gpu_object + + def free_object_primary_copy(self, object_id: str): + """ + Free the RDT object on the primary copy holder and free metadata + on the owner. + """ + from ray.experimental.gpu_object_manager.gpu_object_store import ( + __ray_free__, + ) + + try: + gpu_object_meta = self.managed_gpu_object_metadata[object_id] + src_actor = gpu_object_meta.src_actor + tensor_transport_backend = gpu_object_meta.tensor_transport_backend + tensor_transport_meta = gpu_object_meta.tensor_transport_meta + src_actor.__ray_call__.options(concurrency_group="_ray_system").remote( + __ray_free__, object_id, tensor_transport_backend, tensor_transport_meta + ) + # NOTE: This may have to change if we support lineage reconstruction for RDT + # TODO(#57962): Metadata is currently not removed on borrowers that borrow through + # the NIXL ray.put / ray.get + self.managed_gpu_object_metadata.pop(object_id, None) + except Exception as e: + logger.error( + "Something went wrong while freeing the RDT object!", exc_info=e + ) + + def actor_has_tensor_transport( + self, actor: "ray.actor.ActorHandle", tensor_transport: TensorTransportEnum + ): + """ + Check if the actor has a communicator for the given tensor transport backend. + + Args: + actor: The actor to check. + tensor_transport: The tensor transport backend to check. + + Returns: + True if the actor has a communicator for the given tensor transport backend, False otherwise. + """ + # Import get_collective_groups here to avoid dependency on + # collective libraries for default Ray installation. + from ray.experimental.collective import get_tensor_transport_manager + from ray.experimental.gpu_object_manager.gpu_object_store import ( + _tensor_transport_to_collective_backend, + ) + + tensor_transport_backend = _tensor_transport_to_collective_backend( + tensor_transport + ) + tensor_transport_manager = get_tensor_transport_manager( + tensor_transport_backend + ) + return tensor_transport_manager.actor_has_tensor_transport(actor) + + def put_object( + self, + obj_ref: ObjectRef, + tensor_transport: TensorTransportEnum, + tensors: List["torch.Tensor"], + ): + """ + Put the GPU object into the GPU object manager. + + Args: + obj_ref: The object ref of the GPU object. + tensor_transport: The tensor transport backend to use. + tensors: The tensors to put into the GPU object manager. + + """ + from ray.experimental.collective import get_tensor_transport_manager + from ray.experimental.gpu_object_manager.gpu_object_store import ( + _tensor_transport_to_collective_backend, + ) + + tensor_transport_backend = _tensor_transport_to_collective_backend( + tensor_transport + ) + transport_manager = get_tensor_transport_manager(tensor_transport_backend) + tensor_transport_meta = transport_manager.extract_tensor_transport_metadata( + obj_ref.hex(), tensors + ) + + src_actor = ray.get_runtime_context().current_actor + self.gpu_object_store.add_object(obj_ref.hex(), tensors, is_primary=True) + self.add_gpu_object_ref( + obj_ref, + src_actor, + tensor_transport, + tensor_transport_meta=tensor_transport_meta, + ) diff --git a/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_store.py b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_store.py new file mode 100644 index 0000000000000000000000000000000000000000..280b3452fe9a2ebe0c7964f78ccc9fc43e89a686 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/gpu_object_manager/gpu_object_store.py @@ -0,0 +1,427 @@ +import threading +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Set, Union + +import ray +import ray.util.collective as collective +from ray._private.custom_types import TensorTransportEnum +from ray._raylet import ObjectRef +from ray.experimental.collective import get_tensor_transport_manager +from ray.experimental.collective.util import device_match_transport +from ray.util.collective.types import ( + Backend, + CommunicatorMetadata, + TensorTransportMetadata, +) + +try: + import torch +except ImportError: + raise ImportError( + "`tensor_transport` requires PyTorch. " + "Please install torch with 'pip install torch' to use this feature." + ) + +TENSOR_TRANSPORT_TO_COLLECTIVE_BACKEND = { + TensorTransportEnum.NCCL: Backend.NCCL, + TensorTransportEnum.GLOO: Backend.TORCH_GLOO, + TensorTransportEnum.NIXL: Backend.NIXL, +} + + +def _tensor_transport_to_collective_backend( + tensor_transport: TensorTransportEnum, +) -> Backend: + try: + return TENSOR_TRANSPORT_TO_COLLECTIVE_BACKEND[tensor_transport] + except KeyError: + raise ValueError( + f"Invalid tensor transport {tensor_transport.name}, must be one of {list(TENSOR_TRANSPORT_TO_COLLECTIVE_BACKEND.keys())}." + ) + + +def __ray_send__( + self, + obj_id: str, + tensor_transport_meta: TensorTransportMetadata, + communicator_meta: CommunicatorMetadata, +): + """Helper function that runs on the src actor to send tensors to the dst actor.""" + from ray._private.worker import global_worker + + gpu_object_store = global_worker.gpu_object_manager._gpu_object_store + assert gpu_object_store.has_object( + obj_id + ), f"obj_id={obj_id} not found in GPU object store" + + tensors = gpu_object_store.get_object(obj_id) + + backend = collective.get_group_handle(communicator_meta.communicator_name).backend() + + tensor_transport_manager = get_tensor_transport_manager(backend) + if tensors and not device_match_transport(tensors[0].device, backend): + raise ValueError( + f"Tensor transport backend {backend} does not support tensor transfer on device {tensors[0].device}." + ) + tensor_transport_manager.send_multiple_tensors( + tensors, + tensor_transport_meta, + communicator_meta, + ) + + +def __ray_recv__( + self, + obj_id: str, + tensor_transport_meta: List[Union[ObjectRef, TensorTransportMetadata]], + communicator_meta: CommunicatorMetadata, +): + """Helper function that runs on the dst actor to receive tensors from the src actor.""" + from ray._private.worker import global_worker + + gpu_object_store = global_worker.gpu_object_manager.gpu_object_store + try: + tensor_transport_meta: TensorTransportMetadata = ( + ray.get(tensor_transport_meta[0]) + if isinstance(tensor_transport_meta[0], ObjectRef) + else tensor_transport_meta[0] + ) + device = tensor_transport_meta.tensor_device + tensor_meta = tensor_transport_meta.tensor_meta + + backend = collective.get_group_handle( + communicator_meta.communicator_name + ).backend() + + if tensor_meta and not device_match_transport(device, backend): + raise ValueError( + f"Tensor transport backend {backend} does not support tensor transfer on device {device}." + ) + + tensors = [] + for meta in tensor_meta: + shape, dtype = meta + tensor = torch.empty(shape, dtype=dtype, device=device) + tensors.append(tensor) + + tensor_transport_manager = get_tensor_transport_manager(backend) + tensor_transport_manager.recv_multiple_tensors( + tensors, + obj_id, + tensor_transport_meta, + communicator_meta, + ) + gpu_object_store.add_object(obj_id, tensors, is_primary=False) + except Exception as e: + # Store the error as a gpu object if the recv fails, + # so waiters will raise the error. + gpu_object_store.add_object(obj_id, e, is_primary=False) + + +def __ray_abort_transport__(self, obj_id: str, communicator_meta: CommunicatorMetadata): + """Helper function that can run on an actor doing a send or recv to abort the transport.""" + backend = collective.get_group_handle(communicator_meta.communicator_name).backend() + tensor_transport_manager = get_tensor_transport_manager(backend) + tensor_transport_manager.abort_transport(obj_id, communicator_meta) + + +def __ray_free__( + self, + obj_id: str, + tensor_transport_backend: Backend, + tensor_transport_meta: TensorTransportMetadata, +): + try: + from ray._private.worker import global_worker + from ray.experimental.collective import get_tensor_transport_manager + + tensor_transport_manager = get_tensor_transport_manager( + tensor_transport_backend + ) + tensor_transport_manager.garbage_collect(obj_id, tensor_transport_meta) + + gpu_object_manager = global_worker.gpu_object_manager + gpu_object_store = gpu_object_manager.gpu_object_store + gpu_object_store.pop_object(obj_id) + except AssertionError: + # This could fail if this is a retry and it's already been freed. + pass + + +def __ray_fetch_gpu_object__(self, obj_id: str): + """Helper function that runs on the src actor to fetch tensors from the GPU object store via the object store.""" + from ray._private.worker import global_worker + + gpu_object_store = global_worker.gpu_object_manager.gpu_object_store + assert gpu_object_store.has_object( + obj_id + ), f"obj_id={obj_id} not found in GPU object store" + gpu_object = gpu_object_store.get_object(obj_id) + return gpu_object + + +@dataclass +class _GPUObject: + # A list of tensors representing the GPU object. + data: List["torch.Tensor"] + # Whether the GPU object is the primary copy. + is_primary: bool + # If a recv failed, we store the error here. + error: Optional[Exception] = None + + +class GPUObjectStore: + """ + This class is thread-safe. The GPU object store is meant to be read and + written by the following threads: + 1. The main thread, which is executing user code. This thread may get, put, + and pop objects. + 2. The background _ray_system thread, which executes data transfers. This + thread may get and put objects. + 3. The background CoreWorker server thread, which executes garbage + collection callbacks that pop objects that are no longer in use. + """ + + def __init__(self): + # A dictionary that maps from an object ID to a queue of tensor lists. + # + # Note: Currently, `_gpu_object_store` is only supported for Ray Actors. + self._gpu_object_store: Dict[str, deque[_GPUObject]] = defaultdict(deque) + # Mapping from tensor data pointer to the IDs of objects that contain it. + self._tensor_to_object_ids: Dict[int, Set[str]] = defaultdict[int, Set[str]]( + set + ) + # Synchronization for GPU object store. + self._lock = threading.RLock() + # Signal when an object becomes present in the object store. + self._object_present_cv = threading.Condition(self._lock) + # Signal when an object is freed from the object store. + self._object_freed_cv = threading.Condition(self._lock) + + # These are only used for NIXL. Will be removed in the future. + # Mapping from object ID to the NIXL managed meta. + self._managed_meta_nixl: Dict[str, Any] = {} + # Mapping from NIXL managed meta to the number of objects that contain it. + self._managed_meta_counts_nixl: Dict[Any, int] = defaultdict[Any, int](int) + + def has_object(self, obj_id: str) -> bool: + with self._lock: + existed = obj_id in self._gpu_object_store + if existed: + return len(self._gpu_object_store[obj_id]) > 0 + return existed + + def has_tensor(self, tensor: "torch.Tensor") -> bool: + with self._lock: + return tensor.data_ptr() in self._tensor_to_object_ids + + def get_object(self, obj_id: str) -> Optional[List["torch.Tensor"]]: + with self._lock: + if self._gpu_object_store[obj_id][0].error: + raise self._gpu_object_store[obj_id][0].error + return self._gpu_object_store[obj_id][0].data + + def add_object( + self, + obj_id: str, + gpu_object: Union[List["torch.Tensor"], Exception], + is_primary: bool, + ): + """ + Add a GPU object to the GPU object store. + + Args: + obj_id: The object ID of the GPU object. + gpu_object: A list of tensors representing the GPU object. + is_primary: Whether the GPU object is the primary copy. + """ + with self._object_present_cv: + if isinstance(gpu_object, Exception): + self._gpu_object_store[obj_id].append( + _GPUObject([], is_primary, error=gpu_object) + ) + else: + for tensor in gpu_object: + self._tensor_to_object_ids[tensor.data_ptr()].add(obj_id) + # Append to the queue instead of overwriting + self._gpu_object_store[obj_id].append( + _GPUObject( + gpu_object, + is_primary, + ) + ) + self._object_present_cv.notify_all() + + def is_primary_copy(self, obj_id: str) -> bool: + with self._lock: + return ( + self.has_object(obj_id) and self._gpu_object_store[obj_id][0].is_primary + ) + + def wait_and_get_object( + self, obj_id: str, timeout: Optional[float] = None + ) -> List["torch.Tensor"]: + """Atomically waits for the GPU object to be present in the GPU object + store, then gets it. If the object is not present after the optional + timeout, raise a TimeoutError. + + Args: + obj_id: The object ID to wait for. + timeout: The maximum time in seconds to wait for the object to be + present in the GPU object store. If not specified, wait indefinitely. + + Returns: + The tensors in the GPU object. + """ + with self._lock: + self._wait_object(obj_id, timeout) + return self.get_object(obj_id) + + def get_duplicate_objects( + self, + src_obj_id: str, + src_gpu_object: List["torch.Tensor"], + ) -> Optional[str]: + """Get another object ID of the GPU object that duplicates the given GPU object.""" + with self._lock: + if len(src_gpu_object) == 0: + return None + obj_id_set = set() + for tensor in src_gpu_object: + for obj_id in self._tensor_to_object_ids[tensor.data_ptr()]: + obj_id_set.add(obj_id) + + for dst_obj_id in obj_id_set: + if dst_obj_id != src_obj_id: + dst_gpu_object = self._gpu_object_store[dst_obj_id][0].data + is_same_tensors = len(src_gpu_object) == len( + dst_gpu_object + ) and all( + t1.data_ptr() == t2.data_ptr() + for t1, t2 in zip(src_gpu_object, dst_gpu_object) + ) + if not is_same_tensors: + raise ValueError( + f"Some of the tensors in this object are still in scope as part of another RDT object. " + f"Ensure that ObjectRef({src_obj_id}) is out of scope before creating this object." + ) + return dst_obj_id + return None + + def record_managed_meta_nixl(self, obj_id: str, meta: Any): + """Record the NIXL managed meta for the given object ID.""" + with self._lock: + self._managed_meta_nixl[obj_id] = meta + self._managed_meta_counts_nixl[meta] += 1 + + def record_and_get_meta_if_duplicate( + self, src_obj_id: str, src_gpu_object: List["torch.Tensor"] + ) -> Optional[str]: + """Record the NIXL managed meta for the given object ID if it is a duplicate of another object, and return the meta if it is.""" + with self._lock: + duplicate_obj_id = self.get_duplicate_objects(src_obj_id, src_gpu_object) + if duplicate_obj_id is not None: + meta = self._managed_meta_nixl[duplicate_obj_id] + self._managed_meta_counts_nixl[meta] += 1 + self._managed_meta_nixl[src_obj_id] = meta + return meta + return None + + def remove_managed_meta_nixl(self, obj_id: str): + """Remove the NIXL managed meta for the given object ID and return the count of the managed meta after removal.""" + with self._lock: + meta = self._managed_meta_nixl.pop(obj_id) + self._managed_meta_counts_nixl[meta] -= 1 + count = self._managed_meta_counts_nixl[meta] + if count <= 0: + self._managed_meta_counts_nixl.pop(meta) + return count + + def wait_and_pop_object( + self, obj_id: str, timeout: Optional[float] = None + ) -> List["torch.Tensor"]: + """Atomically waits for the GPU object to be present in the GPU object + store, then pops it. If the object is not present after the optional + timeout, raise a TimeoutError. + + Args: + obj_id: The object ID to wait for. + timeout: The maximum time in seconds to wait for the object to be + present in the GPU object store. If not specified, wait + indefinitely. + + Returns: + The GPU object. + """ + with self._lock: + self._wait_object(obj_id, timeout) + return self.pop_object(obj_id) + + def _wait_object(self, obj_id: str, timeout: Optional[float] = None) -> None: + """Helper method to wait for the GPU object to be present in the GPU object store. + If the object is not present after the optional timeout, raise a + TimeoutError. + + Args: + obj_id: The object ID to wait for. + timeout: The maximum time in seconds to wait for the object to be + present in the GPU object store. If not specified, wait + indefinitely. + """ + with self._object_present_cv: + if not self._object_present_cv.wait_for( + lambda: self.has_object(obj_id), + timeout=timeout, + ): + raise TimeoutError( + f"ObjectRef({obj_id}) not found in RDT object store after {timeout}s, transfer may have failed. Please report this issue on GitHub: https://github.com/ray-project/ray/issues/new/choose" + ) + + def pop_object(self, obj_id: str) -> List["torch.Tensor"]: + with self._lock: + assert self.has_object( + obj_id + ), f"obj_id={obj_id} not found in GPU object store" + queue = self._gpu_object_store.get(obj_id) + gpu_object = queue.popleft() + if len(queue) == 0: + del self._gpu_object_store[obj_id] + if gpu_object.error: + raise gpu_object.error + for tensor in gpu_object.data: + self._tensor_to_object_ids[tensor.data_ptr()].remove(obj_id) + if len(self._tensor_to_object_ids[tensor.data_ptr()]) == 0: + self._tensor_to_object_ids.pop(tensor.data_ptr()) + self._object_freed_cv.notify_all() + return gpu_object.data + + def wait_tensor_freed( + self, tensor: "torch.Tensor", timeout: Optional[float] = None + ) -> None: + """ + Wait for the object to be freed from the GPU object store. + """ + with self._object_freed_cv: + if not self._object_freed_cv.wait_for( + lambda: tensor.data_ptr() not in self._tensor_to_object_ids, + timeout=timeout, + ): + raise TimeoutError( + f"Tensor {tensor} not freed from RDT object store after {timeout}s. The tensor will not be freed until all ObjectRefs containing the tensor have gone out of scope." + ) + + def get_num_objects(self) -> int: + """ + Return the number of objects in the GPU object store. + """ + with self._lock: + # Count total objects across all queues + return sum(len(queue) for queue in self._gpu_object_store.values()) + + def get_num_managed_meta_nixl(self) -> int: + """ + Return the number of NIXL managed meta in the GPU object store. + """ + with self._lock: + return len(self._managed_meta_nixl) diff --git a/lib/python3.12/site-packages/ray/experimental/gradio_utils.py b/lib/python3.12/site-packages/ray/experimental/gradio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..279948c5408b6b0e1a51483dc8575f78bf94d2d3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/gradio_utils.py @@ -0,0 +1,12 @@ +def type_to_string(_type: type) -> str: + """Gets the string representation of a type. + + THe original type can be derived from the returned string representation through + pydoc.locate(). + """ + if _type.__module__ == "typing": + return f"{_type.__module__}.{_type._name}" + elif _type.__module__ == "builtins": + return _type.__name__ + else: + return f"{_type.__module__}.{_type.__name__}" diff --git a/lib/python3.12/site-packages/ray/experimental/internal_kv.py b/lib/python3.12/site-packages/ray/experimental/internal_kv.py new file mode 100644 index 0000000000000000000000000000000000000000..862ff3bacc89364ffa8e2ec247a8ff8ec163fbfe --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/internal_kv.py @@ -0,0 +1,121 @@ +from typing import List, Optional, Union + +from ray._private.client_mode_hook import client_mode_hook +from ray._raylet import GcsClient + +_initialized = False +global_gcs_client = None + + +def _internal_kv_reset(): + global global_gcs_client, _initialized + global_gcs_client = None + _initialized = False + + +def internal_kv_get_gcs_client(): + return global_gcs_client + + +def _initialize_internal_kv(gcs_client: GcsClient): + """Initialize the internal KV for use in other function calls.""" + global global_gcs_client, _initialized + assert gcs_client is not None + global_gcs_client = gcs_client + _initialized = True + + +@client_mode_hook +def _internal_kv_initialized(): + return global_gcs_client is not None + + +@client_mode_hook +def _internal_kv_get( + key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None +) -> bytes: + """Fetch the value of a binary key.""" + + if isinstance(key, str): + key = key.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + assert isinstance(key, bytes) + return global_gcs_client.internal_kv_get(key, namespace) + + +@client_mode_hook +def _internal_kv_exists( + key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None +) -> bool: + """Check key exists or not.""" + + if isinstance(key, str): + key = key.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + assert isinstance(key, bytes) + return global_gcs_client.internal_kv_exists(key, namespace) + + +@client_mode_hook +def _pin_runtime_env_uri(uri: str, *, expiration_s: int) -> None: + """Pin a runtime_env URI for expiration_s.""" + return global_gcs_client.pin_runtime_env_uri(uri, expiration_s) + + +@client_mode_hook +def _internal_kv_put( + key: Union[str, bytes], + value: Union[str, bytes], + overwrite: bool = True, + *, + namespace: Optional[Union[str, bytes]] = None +) -> bool: + """Globally associates a value with a given binary key. + + This only has an effect if the key does not already have a value. + + Returns: + already_exists: whether the value already exists. + """ + + if isinstance(key, str): + key = key.encode() + if isinstance(value, str): + value = value.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + assert ( + isinstance(key, bytes) + and isinstance(value, bytes) + and isinstance(overwrite, bool) + ) + return global_gcs_client.internal_kv_put(key, value, overwrite, namespace) == 0 + + +@client_mode_hook +def _internal_kv_del( + key: Union[str, bytes], + *, + del_by_prefix: bool = False, + namespace: Optional[Union[str, bytes]] = None +) -> int: + if isinstance(key, str): + key = key.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + assert isinstance(key, bytes) + return global_gcs_client.internal_kv_del(key, del_by_prefix, namespace) + + +@client_mode_hook +def _internal_kv_list( + prefix: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None +) -> List[bytes]: + """List all keys in the internal KV store that start with the prefix.""" + if isinstance(prefix, str): + prefix = prefix.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + return global_gcs_client.internal_kv_keys(prefix, namespace) diff --git a/lib/python3.12/site-packages/ray/experimental/locations.py b/lib/python3.12/site-packages/ray/experimental/locations.py new file mode 100644 index 0000000000000000000000000000000000000000..01d4018904808d1a90468fecc92f2bfc840efe97 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/locations.py @@ -0,0 +1,76 @@ +from typing import Any, Dict, List + +import ray +from ray._raylet import ObjectRef + + +def get_object_locations( + obj_refs: List[ObjectRef], timeout_ms: int = -1 +) -> Dict[ObjectRef, Dict[str, Any]]: + """Lookup the locations for a list of objects. + + It returns a dict maps from an object to its location. The dict excludes + those objects whose location lookup failed. + + Args: + object_refs (List[ObjectRef]): List of object refs. + timeout_ms: The maximum amount of time in micro seconds to wait + before returning. Wait infinitely if it's negative. + + Returns: + A dict maps from an object to its location. The dict excludes those + objects whose location lookup failed. + + The location is stored as a dict with following attributes: + + - node_ids (List[str]): The hex IDs of the nodes that have a + copy of this object. Objects less than 100KB will be in memory + store not plasma store and therefore will have nodes_id = []. + + - object_size (int): The size of data + metadata in bytes. Can be None if the + size is unknown yet (e.g. task not completed). + + Raises: + RuntimeError: if the processes were not started by ray.init(). + ray.exceptions.GetTimeoutError: if it couldn't finish the + request in time. + """ + if not ray.is_initialized(): + raise RuntimeError("Ray hasn't been initialized.") + return ray._private.worker.global_worker.core_worker.get_object_locations( + obj_refs, timeout_ms + ) + + +def get_local_object_locations( + obj_refs: List[ObjectRef], +) -> Dict[ObjectRef, Dict[str, Any]]: + """Lookup the locations for a list of objects *from the local core worker*. No RPCs + are made in this method. + + It returns a dict maps from an object to its location. The dict excludes + those objects whose location lookup failed. + + Args: + object_refs (List[ObjectRef]): List of object refs. + + Returns: + A dict maps from an object to its location. The dict excludes those + objects whose location lookup failed. + + The location is stored as a dict with following attributes: + + - node_ids (List[str]): The hex IDs of the nodes that have a + copy of this object. Objects less than 100KB will be in memory + store not plasma store and therefore will have nodes_id = []. + + - object_size (int): The size of data + metadata in bytes. Can be None if the + size is unknown yet (e.g. task not completed). + + Raises: + RuntimeError: if the processes were not started by ray.init(). + """ + if not ray.is_initialized(): + raise RuntimeError("Ray hasn't been initialized.") + core_worker = ray._private.worker.global_worker.core_worker + return core_worker.get_local_object_locations(obj_refs) diff --git a/lib/python3.12/site-packages/ray/experimental/queue.py b/lib/python3.12/site-packages/ray/experimental/queue.py new file mode 100644 index 0000000000000000000000000000000000000000..c255e987993e02b2c09dbd154c9965e0b4b869c3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/queue.py @@ -0,0 +1,13 @@ +import warnings + +from ray.util.queue import Empty, Full, Queue + +warnings.warn( + DeprecationWarning( + "ray.experimental.queue has been moved to ray.util.queue. " + "Please update your import path." + ), + stacklevel=2, +) + +__all__ = ["Empty", "Full", "Queue"] diff --git a/lib/python3.12/site-packages/ray/experimental/shuffle.py b/lib/python3.12/site-packages/ray/experimental/shuffle.py new file mode 100644 index 0000000000000000000000000000000000000000..64ce1792d9d1b42447c177c0ce15f6891b9b2386 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/shuffle.py @@ -0,0 +1,357 @@ +"""A simple distributed shuffle implementation in Ray. + +This utility provides a `simple_shuffle` function that can be used to +redistribute M input partitions into N output partitions. It does this with +a single wave of shuffle map tasks followed by a single wave of shuffle reduce +tasks. Each shuffle map task generates O(N) output objects, and each shuffle +reduce task consumes O(M) input objects, for a total of O(N*M) objects. + +To try an example 10GB shuffle, run: + + $ python -m ray.experimental.shuffle \ + --num-partitions=50 --partition-size=200e6 \ + --object-store-memory=1e9 + +This will print out some statistics on the shuffle execution such as: + + --- Aggregate object store stats across all nodes --- + Plasma memory usage 0 MiB, 0 objects, 0.0% full + Spilled 9487 MiB, 2487 objects, avg write throughput 1023 MiB/s + Restored 9487 MiB, 2487 objects, avg read throughput 1358 MiB/s + Objects consumed by Ray tasks: 9537 MiB. + + Shuffled 9536 MiB in 16.579771757125854 seconds +""" +import time +from typing import Any, Callable, Iterable, List, Tuple, Union + +import ray +from ray import ObjectRef +from ray.cluster_utils import Cluster + +# TODO(ekl) why doesn't TypeVar() deserialize properly in Ray? +# The type produced by the input reader function. +InType = Any +# The type produced by the output writer function. +OutType = Any +# Integer identifying the partition number. +PartitionID = int + + +class ObjectStoreWriter: + """This class is used to stream shuffle map outputs to the object store. + + It can be subclassed to optimize writing (e.g., batching together small + records into larger objects). This will be performance critical if your + input records are small (the example shuffle uses very large records, so + the naive strategy works well). + """ + + def __init__(self): + self.results = [] + + def add(self, item: InType) -> None: + """Queue a single item to be written to the object store. + + This base implementation immediately writes each given item to the + object store as a standalone object. + """ + self.results.append(ray.put(item)) + + def finish(self) -> List[ObjectRef]: + """Return list of object refs representing written items.""" + return self.results + + +class ObjectStoreWriterNonStreaming(ObjectStoreWriter): + def __init__(self): + self.results = [] + + def add(self, item: InType) -> None: + self.results.append(item) + + def finish(self) -> List[Any]: + return self.results + + +def round_robin_partitioner( + input_stream: Iterable[InType], num_partitions: int +) -> Iterable[Tuple[PartitionID, InType]]: + """Round robin partitions items from the input reader. + + You can write custom partitioning functions for your use case. + + Args: + input_stream: Iterator over items from the input reader. + num_partitions: Number of output partitions. + + Yields: + Tuples of (partition id, input item). + """ + i = 0 + for item in input_stream: + yield (i, item) + i += 1 + i %= num_partitions + + +@ray.remote +class _StatusTracker: + def __init__(self): + self.num_map = 0 + self.num_reduce = 0 + self.map_refs = [] + self.reduce_refs = [] + + def register_objectrefs(self, map_refs, reduce_refs): + self.map_refs = map_refs + self.reduce_refs = reduce_refs + + def get_progress(self): + if self.map_refs: + ready, self.map_refs = ray.wait( + self.map_refs, + timeout=1, + num_returns=len(self.map_refs), + fetch_local=False, + ) + self.num_map += len(ready) + elif self.reduce_refs: + ready, self.reduce_refs = ray.wait( + self.reduce_refs, + timeout=1, + num_returns=len(self.reduce_refs), + fetch_local=False, + ) + self.num_reduce += len(ready) + return self.num_map, self.num_reduce + + +def render_progress_bar(tracker, input_num_partitions, output_num_partitions): + from tqdm import tqdm + + num_map = 0 + num_reduce = 0 + map_bar = tqdm(total=input_num_partitions, position=0) + map_bar.set_description("Map Progress.") + reduce_bar = tqdm(total=output_num_partitions, position=1) + reduce_bar.set_description("Reduce Progress.") + + while num_map < input_num_partitions or num_reduce < output_num_partitions: + new_num_map, new_num_reduce = ray.get(tracker.get_progress.remote()) + map_bar.update(new_num_map - num_map) + reduce_bar.update(new_num_reduce - num_reduce) + num_map = new_num_map + num_reduce = new_num_reduce + time.sleep(0.1) + map_bar.close() + reduce_bar.close() + + +def simple_shuffle( + *, + input_reader: Callable[[PartitionID], Iterable[InType]], + input_num_partitions: int, + output_num_partitions: int, + output_writer: Callable[[PartitionID, List[Union[ObjectRef, Any]]], OutType], + partitioner: Callable[ + [Iterable[InType], int], Iterable[PartitionID] + ] = round_robin_partitioner, + object_store_writer: ObjectStoreWriter = ObjectStoreWriter, + tracker: _StatusTracker = None, + streaming: bool = True, +) -> List[OutType]: + """Simple distributed shuffle in Ray. + + Args: + input_reader: Function that generates the input items for a + partition (e.g., data records). + input_num_partitions: The number of input partitions. + output_num_partitions: The desired number of output partitions. + output_writer: Function that consumes a iterator of items for a + given output partition. It returns a single value that will be + collected across all output partitions. + partitioner: Partitioning function to use. Defaults to round-robin + partitioning of input items. + object_store_writer: Class used to write input items to the + object store in an efficient way. Defaults to a naive + implementation that writes each input record as one object. + tracker: Tracker actor that is used to display the progress bar. + streaming: Whether or not if the shuffle will be streaming. + + Returns: + List of outputs from the output writers. + """ + + @ray.remote(num_returns=output_num_partitions) + def shuffle_map(i: PartitionID) -> List[List[Union[Any, ObjectRef]]]: + writers = [object_store_writer() for _ in range(output_num_partitions)] + for out_i, item in partitioner(input_reader(i), output_num_partitions): + writers[out_i].add(item) + return [c.finish() for c in writers] + + @ray.remote + def shuffle_reduce( + i: PartitionID, *mapper_outputs: List[List[Union[Any, ObjectRef]]] + ) -> OutType: + input_objects = [] + assert len(mapper_outputs) == input_num_partitions + for obj_refs in mapper_outputs: + for obj_ref in obj_refs: + input_objects.append(obj_ref) + return output_writer(i, input_objects) + + shuffle_map_out = [shuffle_map.remote(i) for i in range(input_num_partitions)] + + shuffle_reduce_out = [ + shuffle_reduce.remote( + j, *[shuffle_map_out[i][j] for i in range(input_num_partitions)] + ) + for j in range(output_num_partitions) + ] + + if tracker: + tracker.register_objectrefs.remote( + [map_out[0] for map_out in shuffle_map_out], shuffle_reduce_out + ) + render_progress_bar(tracker, input_num_partitions, output_num_partitions) + + return ray.get(shuffle_reduce_out) + + +def build_cluster(num_nodes, num_cpus, object_store_memory): + cluster = Cluster() + for _ in range(num_nodes): + cluster.add_node(num_cpus=num_cpus, object_store_memory=object_store_memory) + cluster.wait_for_nodes() + return cluster + + +def run( + ray_address=None, + object_store_memory=1e9, + num_partitions=5, + partition_size=200e6, + num_nodes=None, + num_cpus=8, + no_streaming=False, + use_wait=False, + tracker=None, +): + import time + + import numpy as np + + is_multi_node = num_nodes + if ray_address: + print("Connecting to a existing cluster...") + ray.init(address=ray_address, ignore_reinit_error=True) + elif is_multi_node: + print("Emulating a cluster...") + print(f"Num nodes: {num_nodes}") + print(f"Num CPU per node: {num_cpus}") + print(f"Object store memory per node: {object_store_memory}") + cluster = build_cluster(num_nodes, num_cpus, object_store_memory) + ray.init(address=cluster.address) + else: + print("Start a new cluster...") + ray.init(num_cpus=num_cpus, object_store_memory=object_store_memory) + + partition_size = int(partition_size) + num_partitions = num_partitions + rows_per_partition = partition_size // (8 * 2) + if tracker is None: + tracker = _StatusTracker.remote() + use_wait = use_wait + + def input_reader(i: PartitionID) -> Iterable[InType]: + for _ in range(num_partitions): + yield np.ones((rows_per_partition // num_partitions, 2), dtype=np.int64) + + def output_writer(i: PartitionID, shuffle_inputs: List[ObjectRef]) -> OutType: + total = 0 + if not use_wait: + for obj_ref in shuffle_inputs: + arr = ray.get(obj_ref) + total += arr.size * arr.itemsize + else: + while shuffle_inputs: + [ready], shuffle_inputs = ray.wait(shuffle_inputs, num_returns=1) + arr = ray.get(ready) + total += arr.size * arr.itemsize + + return total + + def output_writer_non_streaming( + i: PartitionID, shuffle_inputs: List[Any] + ) -> OutType: + total = 0 + for arr in shuffle_inputs: + total += arr.size * arr.itemsize + return total + + if no_streaming: + output_writer_callable = output_writer_non_streaming + object_store_writer = ObjectStoreWriterNonStreaming + else: + object_store_writer = ObjectStoreWriter + output_writer_callable = output_writer + + start = time.time() + output_sizes = simple_shuffle( + input_reader=input_reader, + input_num_partitions=num_partitions, + output_num_partitions=num_partitions, + output_writer=output_writer_callable, + object_store_writer=object_store_writer, + tracker=tracker, + ) + delta = time.time() - start + + time.sleep(0.5) + print() + + summary = None + for i in range(5): + try: + summary = ray._private.internal_api.memory_summary(stats_only=True) + except Exception: + time.sleep(1) + pass + if summary: + break + print(summary) + print() + print( + "Shuffled", int(sum(output_sizes) / (1024 * 1024)), "MiB in", delta, "seconds" + ) + + +def main(): + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--ray-address", type=str, default=None) + parser.add_argument("--object-store-memory", type=float, default=1e9) + parser.add_argument("--num-partitions", type=int, default=5) + parser.add_argument("--partition-size", type=float, default=200e6) + parser.add_argument("--num-nodes", type=int, default=None) + parser.add_argument("--num-cpus", type=int, default=8) + parser.add_argument("--no-streaming", action="store_true", default=False) + parser.add_argument("--use-wait", action="store_true", default=False) + args = parser.parse_args() + + run( + ray_address=args.ray_address, + object_store_memory=args.object_store_memory, + num_partitions=args.num_partitions, + partition_size=args.partition_size, + num_nodes=args.num_nodes, + num_cpus=args.num_cpus, + no_streaming=args.no_streaming, + use_wait=args.use_wait, + ) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/experimental/tf_utils.py b/lib/python3.12/site-packages/ray/experimental/tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..128dc9c8db866ce1b12df8342e32ae899ac2f428 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/tf_utils.py @@ -0,0 +1,4 @@ +raise ImportError( + "ray.experimental.tf_utils has been removed. " + "Use: from ray.rllib.utils import tf_utils." +) diff --git a/lib/python3.12/site-packages/ray/experimental/tqdm_ray.py b/lib/python3.12/site-packages/ray/experimental/tqdm_ray.py new file mode 100644 index 0000000000000000000000000000000000000000..c5d7324a8cbdab99954385b302d285a810ae446c --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/tqdm_ray.py @@ -0,0 +1,417 @@ +import builtins +import copy +import json +import logging +import os +import sys +import threading +import time +import uuid +from typing import Any, Dict, Iterable, Optional + +import colorama + +import ray +from ray._private.ray_constants import env_bool +from ray.util.debug import log_once + +try: + import tqdm.auto as real_tqdm +except ImportError: + real_tqdm = None + +logger = logging.getLogger(__name__) + +# Describes the state of a single progress bar. +ProgressBarState = Dict[str, Any] + +# Magic token used to identify Ray TQDM log lines. +RAY_TQDM_MAGIC = "__ray_tqdm_magic_token__" + +# Global manager singleton. +_manager: Optional["_BarManager"] = None +_mgr_lock = threading.Lock() +_print = builtins.print + + +def safe_print(*args, **kwargs): + """Use this as an alternative to `print` that will not corrupt tqdm output. + + By default, the builtin print will be patched to this function when tqdm_ray is + used. To disable this, set RAY_TQDM_PATCH_PRINT=0. + """ + + # Ignore prints to StringIO objects, etc. + if kwargs.get("file") not in [sys.stdout, sys.stderr, None]: + return _print(*args, **kwargs) + + try: + instance().hide_bars() + _print(*args, **kwargs) + finally: + instance().unhide_bars() + + +class tqdm: + """Experimental: Ray distributed tqdm implementation. + + This class lets you use tqdm from any Ray remote task or actor, and have the + progress centrally reported from the driver. This avoids issues with overlapping + / conflicting progress bars, as the driver centrally manages tqdm positions. + + Supports a limited subset of tqdm args. + """ + + DEFAULT_FLUSH_INTERVAL_SECONDS = 1.0 + + def __init__( + self, + iterable: Optional[Iterable] = None, + desc: Optional[str] = None, + total: Optional[int] = None, + unit: Optional[str] = None, + position: Optional[int] = None, + flush_interval_s: Optional[float] = None, + ): + import ray._private.services as services + + if total is None and iterable is not None: + try: + total = len(iterable) + except (TypeError, AttributeError): + total = None + + self._iterable = iterable + self._desc = desc or "" + self._total = total + self._unit = unit or "it" + self._ip = services.get_node_ip_address() + self._pid = os.getpid() + self._pos = position or 0 + self._uuid = uuid.uuid4().hex + self._x = 0 + self._closed = False + self._flush_interval_s = ( + flush_interval_s + if flush_interval_s is not None + else self.DEFAULT_FLUSH_INTERVAL_SECONDS + ) + self._last_flush_time = 0.0 + + def set_description(self, desc): + """Implements tqdm.tqdm.set_description.""" + self._desc = desc + self._dump_state() + + def update(self, n=1): + """Implements tqdm.tqdm.update.""" + self._x += n + self._dump_state() + + def close(self): + """Implements tqdm.tqdm.close.""" + self._closed = True + # Don't bother if ray is shutdown (in __del__ hook). + if ray is not None: + self._dump_state(force_flush=True) + + def refresh(self): + """Implements tqdm.tqdm.refresh.""" + self._dump_state() + + @property + def total(self) -> Optional[int]: + return self._total + + @total.setter + def total(self, total: int): + self._total = total + + def _dump_state(self, force_flush=False) -> None: + now = time.time() + if not force_flush and now - self._last_flush_time < self._flush_interval_s: + return + self._last_flush_time = now + if ray._private.worker.global_worker.mode == ray.WORKER_MODE: + # Include newline in payload to avoid split prints. + # TODO(ekl) we should move this to events.json to avoid log corruption. + print(json.dumps(self._get_state()) + "\n", end="") + else: + instance().process_state_update(copy.deepcopy(self._get_state())) + + def _get_state(self) -> ProgressBarState: + return { + "__magic_token__": RAY_TQDM_MAGIC, + "x": self._x, + "pos": self._pos, + "desc": self._desc, + "total": self._total, + "unit": self._unit, + "ip": self._ip, + "pid": self._pid, + "uuid": self._uuid, + "closed": self._closed, + } + + def __iter__(self): + if self._iterable is None: + raise ValueError("No iterable provided") + for x in iter(self._iterable): + self.update(1) + yield x + + +class _Bar: + """Manages a single virtual progress bar on the driver. + + The actual position of individual bars is calculated as (pos_offset + position), + where `pos_offset` is the position offset determined by the BarManager. + """ + + def __init__(self, state: ProgressBarState, pos_offset: int): + """Initialize a bar. + + Args: + state: The initial progress bar state. + pos_offset: The position offset determined by the BarManager. + """ + self.state = state + self.pos_offset = pos_offset + self.bar = real_tqdm.tqdm( + desc=state["desc"] + " " + str(state["pos"]), + total=state["total"], + unit=state["unit"], + position=pos_offset + state["pos"], + dynamic_ncols=True, + unit_scale=True, + ) + if state["x"]: + self.bar.update(state["x"]) + + def update(self, state: ProgressBarState) -> None: + """Apply the updated worker progress bar state.""" + if state["desc"] != self.state["desc"]: + self.bar.set_description(state["desc"]) + if state["total"] != self.state["total"]: + self.bar.total = state["total"] + self.bar.refresh() + delta = state["x"] - self.state["x"] + if delta: + self.bar.update(delta) + self.bar.refresh() + self.state = state + + def close(self): + """The progress bar has been closed.""" + self.bar.close() + + def update_offset(self, pos_offset: int) -> None: + """Update the position offset assigned by the BarManager.""" + if pos_offset != self.pos_offset: + self.pos_offset = pos_offset + self.bar.clear() + self.bar.pos = -(pos_offset + self.state["pos"]) + self.bar.refresh() + + +class _BarGroup: + """Manages a group of virtual progress bar produced by a single worker. + + All the progress bars in the group have the same `pos_offset` determined by the + BarManager for the process. + """ + + def __init__(self, ip, pid, pos_offset): + self.ip = ip + self.pid = pid + self.pos_offset = pos_offset + self.bars_by_uuid: Dict[str, _Bar] = {} + + def has_bar(self, bar_uuid) -> bool: + """Return whether this bar exists.""" + return bar_uuid in self.bars_by_uuid + + def allocate_bar(self, state: ProgressBarState) -> None: + """Add a new bar to this group.""" + self.bars_by_uuid[state["uuid"]] = _Bar(state, self.pos_offset) + + def update_bar(self, state: ProgressBarState) -> None: + """Update the state of a managed bar in this group.""" + bar = self.bars_by_uuid[state["uuid"]] + bar.update(state) + + def close_bar(self, state: ProgressBarState) -> None: + """Remove a bar from this group.""" + bar = self.bars_by_uuid[state["uuid"]] + # Note: Hide and then unhide bars to prevent flashing of the + # last bar when we are closing multiple bars sequentially. + instance().hide_bars() + bar.close() + del self.bars_by_uuid[state["uuid"]] + instance().unhide_bars() + + def slots_required(self): + """Return the number of pos slots we need to accommodate bars in this group.""" + if not self.bars_by_uuid: + return 0 + return 1 + max(bar.state["pos"] for bar in self.bars_by_uuid.values()) + + def update_offset(self, offset: int) -> None: + """Update the position offset assigned by the BarManager.""" + if offset != self.pos_offset: + self.pos_offset = offset + for bar in self.bars_by_uuid.values(): + bar.update_offset(offset) + + def hide_bars(self) -> None: + """Temporarily hide visible bars to avoid conflict with other log messages.""" + for bar in self.bars_by_uuid.values(): + bar.bar.clear() + + def unhide_bars(self) -> None: + """Opposite of hide_bars().""" + for bar in self.bars_by_uuid.values(): + bar.bar.refresh() + + +class _BarManager: + """Central tqdm manager run on the driver. + + This class holds a collection of BarGroups and updates their `pos_offset` as + needed to ensure individual progress bars do not collide in position, kind of + like a virtual memory manager. + """ + + def __init__(self): + import ray._private.services as services + + self.ip = services.get_node_ip_address() + self.pid = os.getpid() + self.bar_groups = {} + self.in_hidden_state = False + self.num_hides = 0 + self.lock = threading.RLock() + # Avoid colorizing Jupyter output, since the tqdm bar is rendered in + # ipywidgets instead of in the console. + self.should_colorize = not ray.widgets.util.in_notebook() + + def process_state_update(self, state: ProgressBarState) -> None: + """Apply the remote progress bar state update. + + This creates a new bar locally if it doesn't already exist. When a bar is + created or destroyed, we also recalculate and update the `pos_offset` of each + BarGroup on the screen. + """ + with self.lock: + self._process_state_update_locked(state) + + def _process_state_update_locked(self, state: ProgressBarState) -> None: + if not real_tqdm: + if log_once("no_tqdm"): + logger.warning("tqdm is not installed. Progress bars will be disabled.") + return + if state["ip"] == self.ip: + if state["pid"] == self.pid: + prefix = "" + else: + prefix = "(pid={}) ".format(state.get("pid")) + if self.should_colorize: + prefix = "{}{}{}{}".format( + colorama.Style.DIM, + colorama.Fore.CYAN, + prefix, + colorama.Style.RESET_ALL, + ) + else: + prefix = "(pid={}, ip={}) ".format( + state.get("pid"), + state.get("ip"), + ) + if self.should_colorize: + prefix = "{}{}{}{}".format( + colorama.Style.DIM, + colorama.Fore.CYAN, + prefix, + colorama.Style.RESET_ALL, + ) + state["desc"] = prefix + state["desc"] + process = self._get_or_allocate_bar_group(state) + if process.has_bar(state["uuid"]): + # Always call `update_bar` to sync any last remaining updates + # prior to closing. Otherwise, the displayed progress bars + # can be left incomplete, even after execution finishes. + # Fixes https://github.com/ray-project/ray/issues/44983 + process.update_bar(state) + + if state["closed"]: + process.close_bar(state) + self._update_offsets() + else: + process.allocate_bar(state) + self._update_offsets() + + def hide_bars(self) -> None: + """Temporarily hide visible bars to avoid conflict with other log messages.""" + with self.lock: + if not self.in_hidden_state: + self.in_hidden_state = True + self.num_hides += 1 + for group in self.bar_groups.values(): + group.hide_bars() + + def unhide_bars(self) -> None: + """Opposite of hide_bars().""" + with self.lock: + if self.in_hidden_state: + self.in_hidden_state = False + for group in self.bar_groups.values(): + group.unhide_bars() + + def _get_or_allocate_bar_group(self, state: ProgressBarState): + ptuple = (state["ip"], state["pid"]) + if ptuple not in self.bar_groups: + offset = sum(p.slots_required() for p in self.bar_groups.values()) + self.bar_groups[ptuple] = _BarGroup(state["ip"], state["pid"], offset) + return self.bar_groups[ptuple] + + def _update_offsets(self): + offset = 0 + for proc in self.bar_groups.values(): + proc.update_offset(offset) + offset += proc.slots_required() + + +def instance() -> _BarManager: + """Get or create a BarManager for this process.""" + global _manager + + with _mgr_lock: + if _manager is None: + _manager = _BarManager() + if env_bool("RAY_TQDM_PATCH_PRINT", True): + import builtins + + builtins.print = safe_print + return _manager + + +if __name__ == "__main__": + + @ray.remote + def processing(delay): + def sleep(x): + print("Intermediate result", x) + time.sleep(delay) + return x + + ray.data.range(1000, override_num_blocks=100).map( + sleep, compute=ray.data.ActorPoolStrategy(size=1) + ).count() + + ray.get( + [ + processing.remote(0.03), + processing.remote(0.01), + processing.remote(0.05), + ] + ) diff --git a/lib/python3.12/site-packages/ray/experimental/util/__init__.py b/lib/python3.12/site-packages/ray/experimental/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/experimental/util/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/util/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdcff5b83796c795b2fef1f16a9c5f3f5346ee18 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/util/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/util/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/util/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44a71102c6e5327fd754b049a905835049cc5c4e Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/util/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/util/types.py b/lib/python3.12/site-packages/ray/experimental/util/types.py new file mode 100644 index 0000000000000000000000000000000000000000..92ac6488139878f48e819efd6fab301d28d9a20b --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/util/types.py @@ -0,0 +1,46 @@ +from dataclasses import dataclass +from enum import Enum + +from ray.util.annotations import PublicAPI + + +class _CollectiveOp: + pass + + +@PublicAPI +class ReduceOp(Enum): + SUM = 0 + PRODUCT = 1 + MAX = 2 + MIN = 3 + AVG = 4 + + +@PublicAPI +@dataclass +class AllGatherOp(_CollectiveOp): + pass + + +@PublicAPI +@dataclass +class AllReduceOp(_CollectiveOp): + reduceOp: ReduceOp = ReduceOp.SUM + + +@PublicAPI +@dataclass +class ReduceScatterOp(_CollectiveOp): + reduceOp: ReduceOp = ReduceOp.SUM + + +@PublicAPI(stability="alpha") +class Device(Enum): + DEFAULT = "default" + CPU = "cpu" + GPU = "gpu" + CUDA = "cuda" + + def __str__(self): + return self.value diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/__init__.py b/lib/python3.12/site-packages/ray/rllib/callbacks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..788d9c0222bc509f6f7bd2602abc1ba4d20d61d2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/callbacks.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/callbacks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b6cf5286cdbef2a573c416543b7bfd5c4a959a6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/callbacks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1db60d3f3ae406128a4646a38f94dc21d942da6f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/callbacks/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/callbacks.py b/lib/python3.12/site-packages/ray/rllib/callbacks/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..fb4107872ad73dec973d24d29179265b3146b5fa --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/callbacks/callbacks.py @@ -0,0 +1,741 @@ +import gc +import os +import platform +import tracemalloc +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import gymnasium as gym + +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.env.env_context import EnvContext +from ray.rllib.evaluation.episode_v2 import EpisodeV2 +from ray.rllib.offline.offline_evaluation_runner_group import ( + OfflineEvaluationRunnerGroup, +) +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic, + PublicAPI, + override, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import AgentID, EnvType, EpisodeType, PolicyID +from ray.tune.callback import _CallbackMeta + +# Import psutil after ray so the packaged version is used. +import psutil + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm import Algorithm + from ray.rllib.env.env_runner import EnvRunner + from ray.rllib.env.env_runner_group import EnvRunnerGroup + + +@PublicAPI +class RLlibCallback(metaclass=_CallbackMeta): + """Abstract base class for RLlib callbacks (similar to Keras callbacks). + + These callbacks can be used for custom metrics and custom postprocessing. + + By default, all of these callbacks are no-ops. To configure custom training + callbacks, subclass RLlibCallback and then set + {"callbacks": YourCallbacksClass} in the algo config. + """ + + @OverrideToImplementCustomLogic + def on_algorithm_init( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + **kwargs, + ) -> None: + """Callback run when a new Algorithm instance has finished setup. + + This method gets called at the end of Algorithm.setup() after all + the initialization is done, and before actually training starts. + + Args: + algorithm: Reference to the Algorithm instance. + metrics_logger: The MetricsLogger object inside the `Algorithm`. Can be + used to log custom metrics after algo initialization. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_train_result( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + result: dict, + **kwargs, + ) -> None: + """Called at the end of Algorithm.train(). + + Args: + algorithm: Current Algorithm instance. + metrics_logger: The MetricsLogger object inside the Algorithm. Can be + used to log custom metrics after traing results are available. + result: Dict of results returned from Algorithm.train() call. + You can mutate this object to add additional metrics. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_evaluate_start( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + **kwargs, + ) -> None: + """Callback before evaluation starts. + + This method gets called at the beginning of Algorithm.evaluate(). + + Args: + algorithm: Reference to the algorithm instance. + metrics_logger: The MetricsLogger object inside the `Algorithm`. Can be + used to log custom metrics before running the next round of evaluation. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_evaluate_offline_start( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + **kwargs, + ) -> None: + """Callback before offline evaluation starts. + + This method gets called at the beginning of Algorithm.evaluate_offline(). + + Args: + algorithm: Reference to the algorithm instance. + metrics_logger: The MetricsLogger object inside the `Algorithm`. Can be + used to log custom metrics before running the next round of offline + evaluation. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_evaluate_end( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + evaluation_metrics: dict, + **kwargs, + ) -> None: + """Runs when the evaluation is done. + + Runs at the end of Algorithm.evaluate(). + + Args: + algorithm: Reference to the algorithm instance. + metrics_logger: The MetricsLogger object inside the `Algorithm`. Can be + used to log custom metrics after the most recent evaluation round. + evaluation_metrics: Results dict to be returned from algorithm.evaluate(). + You can mutate this object to add additional metrics. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_evaluate_offline_end( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + evaluation_metrics: dict, + **kwargs, + ) -> None: + """Runs when the offline evaluation is done. + + Runs at the end of Algorithm.evaluate_offline(). + + Args: + algorithm: Reference to the algorithm instance. + metrics_logger: The MetricsLogger object inside the `Algorithm`. Can be + used to log custom metrics after the most recent offline evaluation + round. + evaluation_metrics: Results dict to be returned from + Algorithm.evaluate_offline(). You can mutate this object to add + additional metrics. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_env_runners_recreated( + self, + *, + algorithm: "Algorithm", + env_runner_group: "EnvRunnerGroup", + env_runner_indices: List[int], + is_evaluation: bool, + **kwargs, + ) -> None: + """Callback run after one or more EnvRunner actors have been recreated. + + You can access and change the EnvRunners in question through the following code + snippet inside your custom override of this method: + + .. testcode:: + from ray.rllib.callbacks.callbacks import RLlibCallback + + class MyCallbacks(RLlibCallback): + def on_env_runners_recreated( + self, + *, + algorithm, + env_runner_group, + env_runner_indices, + is_evaluation, + **kwargs, + ): + # Define what you would like to do on the recreated EnvRunner: + def func(env_runner): + # Here, we just set some arbitrary property to 1. + if is_evaluation: + env_runner._custom_property_for_evaluation = 1 + else: + env_runner._custom_property_for_training = 1 + + # Use the `foreach_env_runner` method of the worker set and + # only loop through those worker IDs that have been restarted. + # Note that we set `local_worker=False` to NOT include it (local + # workers are never recreated; if they fail, the entire Algorithm + # fails). + env_runner_group.foreach_env_runner( + func, + remote_worker_ids=env_runner_indices, + local_env_runner=False, + ) + + Args: + algorithm: Reference to the Algorithm instance. + env_runner_group: The EnvRunnerGroup object in which the workers in question + reside. You can use a `env_runner_group.foreach_env_runner( + remote_worker_ids=..., local_env_runner=False)` method call to execute + custom code on the recreated (remote) workers. Note that the local + worker is never recreated as a failure of this would also crash the + Algorithm. + env_runner_indices: The list of (remote) worker IDs that have been + recreated. + is_evaluation: Whether `worker_set` is the evaluation EnvRunnerGroup + (located in `Algorithm.eval_env_runner_group`) or not. + """ + pass + + @OverrideToImplementCustomLogic + def on_offline_eval_runners_recreated( + self, + *, + algorithm: "Algorithm", + offline_eval_runner_group: "OfflineEvaluationRunnerGroup", + offline_eval_runner_indices: List[int], + **kwargs, + ) -> None: + """Callback run after one or more OfflineEvaluationRunner actors have been recreated. + + You can access and change the OfflineEvaluationRunners in question through the following code + snippet inside your custom override of this method: + + .. testcode:: + from ray.rllib.callbacks.callbacks import RLlibCallback + + class MyCallbacks(RLlibCallback): + def on_offline_eval_runners_recreated( + self, + *, + algorithm, + offline_eval_runner_group, + offline_eval_runner_indices, + **kwargs, + ): + # Define what you would like to do on the recreated EnvRunner: + def func(offline_eval_runner): + # Here, we just set some arbitrary property to 1. + if is_evaluation: + offline_eval_runner._custom_property_for_evaluation = 1 + else: + offline_eval_runner._custom_property_for_training = 1 + + # Use the `foreach_runner` method of the worker set and + # only loop through those worker IDs that have been restarted. + # Note that `local_runner=False` as long as there are remote + # runners. + offline_eval_runner_group.foreach_runner( + func, + remote_runner_ids=offline_eval_runner_indices, + local_runner=False, + ) + + Args: + algorithm: Reference to the Algorithm instance. + offline_eval_runner_group: The OfflineEvaluationRunnerGroup object in which + the workers in question reside. You can use a `runner_group.foreach_runner( + remote_worker_ids=..., local_runner=False)` method call to execute + custom code on the recreated (remote) workers. + offline_eval_runner_indices: The list of (remote) worker IDs that have been + recreated. + """ + pass + + @OverrideToImplementCustomLogic + def on_checkpoint_loaded( + self, + *, + algorithm: "Algorithm", + **kwargs, + ) -> None: + """Callback run when an Algorithm has loaded a new state from a checkpoint. + + This method gets called at the end of `Algorithm.load_checkpoint()`. + + Args: + algorithm: Reference to the Algorithm instance. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_environment_created( + self, + *, + env_runner: "EnvRunner", + metrics_logger: Optional[MetricsLogger] = None, + env: gym.Env, + env_context: EnvContext, + **kwargs, + ) -> None: + """Callback run when a new environment object has been created. + + Note: This only applies to the new API stack. The env used is usually a + gym.Env (or more specifically a gym.vector.Env). + + Args: + env_runner: Reference to the current EnvRunner instance. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics after environment creation. + env: The environment object that has been created on `env_runner`. This is + usually a gym.Env (or a gym.vector.Env) object. + env_context: The `EnvContext` object that has been passed to the + `gym.make()` call as kwargs (and to the gym.Env as `config`). It should + have all the config key/value pairs in it as well as the + EnvContext-typical properties: `worker_index`, `num_workers`, and + `remote`. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_episode_created( + self, + *, + # TODO (sven): Deprecate Episode/EpisodeV2 with new API stack. + episode: Union[EpisodeType, EpisodeV2], + # TODO (sven): Deprecate this arg new API stack (in favor of `env_runner`). + worker: Optional["EnvRunner"] = None, + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + # TODO (sven): Deprecate this arg new API stack (in favor of `env`). + base_env: Optional[BaseEnv] = None, + env: Optional[gym.Env] = None, + # TODO (sven): Deprecate this arg new API stack (in favor of `rl_module`). + policies: Optional[Dict[PolicyID, Policy]] = None, + rl_module: Optional[RLModule] = None, + env_index: int, + **kwargs, + ) -> None: + """Callback run when a new episode is created (but has not started yet!). + + This method gets called after a new SingleAgentEpisode or MultiAgentEpisode + instance has been created. This happens before the respective sub-environment's + `reset()` is called by RLlib. + + 1) SingleAgentEpisode/MultiAgentEpisode created: This callback is called. + 2) Respective sub-environment (gym.Env) is `reset()`. + 3) Callback `on_episode_start` is called. + 4) Stepping through sub-environment/episode commences. + + Args: + episode: The newly created SingleAgentEpisode or MultiAgentEpisode. + This is the episode that is about to be started with an upcoming + `env.reset()`. Only after this reset call, the `on_episode_start` + callback will be called. + env_runner: Reference to the current EnvRunner. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics after Episode creation. + env: The gym.Env running the episode. + rl_module: The RLModule used to compute actions for stepping the env. In + single-agent mode, this is a simple RLModule, in multi-agent mode, this + is a MultiRLModule. + env_index: The index of the sub-environment that is about to be reset. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_episode_start( + self, + *, + episode: Union[EpisodeType, EpisodeV2], + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + rl_module: Optional[RLModule] = None, + # TODO (sven): Deprecate these args. + worker: Optional["EnvRunner"] = None, + base_env: Optional[BaseEnv] = None, + policies: Optional[Dict[PolicyID, Policy]] = None, + **kwargs, + ) -> None: + """Callback run right after an Episode has been started. + + This method gets called after a SingleAgentEpisode or MultiAgentEpisode instance + has been reset with a call to `env.reset()` by the EnvRunner. + + 1) Single-/MultiAgentEpisode created: `on_episode_created()` is called. + 2) Respective sub-environment (gym.Env) is `reset()`. + 3) Single-/MultiAgentEpisode starts: This callback is called. + 4) Stepping through sub-environment/episode commences. + + Args: + episode: The just started (after `env.reset()`) SingleAgentEpisode or + MultiAgentEpisode object. + env_runner: Reference to the EnvRunner running the env and episode. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics during env/episode stepping. + env: The gym.Env or gym.vector.Env object running the started episode. + env_index: The index of the sub-environment that is about to be reset + (within the vector of sub-environments of the BaseEnv). + rl_module: The RLModule used to compute actions for stepping the env. In + single-agent mode, this is a simple RLModule, in multi-agent mode, this + is a MultiRLModule. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_episode_step( + self, + *, + episode: Union[EpisodeType, EpisodeV2], + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + rl_module: Optional[RLModule] = None, + # TODO (sven): Deprecate these args. + worker: Optional["EnvRunner"] = None, + base_env: Optional[BaseEnv] = None, + policies: Optional[Dict[PolicyID, Policy]] = None, + **kwargs, + ) -> None: + """Called on each episode step (after the action(s) has/have been logged). + + Note that on the new API stack, this callback is also called after the final + step of an episode, meaning when terminated/truncated are returned as True + from the `env.step()` call, but is still provided with the non-numpy'ized + episode object (meaning the data has NOT been converted to numpy arrays yet). + + The exact time of the call of this callback is after `env.step([action])` and + also after the results of this step (observation, reward, terminated, truncated, + infos) have been logged to the given `episode` object. + + Args: + episode: The just stepped SingleAgentEpisode or MultiAgentEpisode object + (after `env.step()` and after returned obs, rewards, etc.. have been + logged to the episode object). + env_runner: Reference to the EnvRunner running the env and episode. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics during env/episode stepping. + env: The gym.Env or gym.vector.Env object running the started episode. + env_index: The index of the sub-environment that has just been stepped. + rl_module: The RLModule used to compute actions for stepping the env. In + single-agent mode, this is a simple RLModule, in multi-agent mode, this + is a MultiRLModule. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_episode_end( + self, + *, + episode: Union[EpisodeType, EpisodeV2], + prev_episode_chunks: Optional[List[EpisodeType]] = None, + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + rl_module: Optional[RLModule] = None, + # TODO (sven): Deprecate these args. + worker: Optional["EnvRunner"] = None, + base_env: Optional[BaseEnv] = None, + policies: Optional[Dict[PolicyID, Policy]] = None, + **kwargs, + ) -> None: + """Called when an episode is done (after terminated/truncated have been logged). + + The exact time of the call of this callback is after `env.step([action])` and + also after the results of this step (observation, reward, terminated, truncated, + infos) have been logged to the given `episode` object, where either terminated + or truncated were True: + + - The env is stepped: `final_obs, rewards, ... = env.step([action])` + + - The step results are logged `episode.add_env_step(final_obs, rewards)` + + - Callback `on_episode_step` is fired. + + - Another env-to-module connector call is made (even though we won't need any + RLModule forward pass anymore). We make this additional call to ensure that in + case users use the connector pipeline to process observations (and write them + back into the episode), the episode object has all observations - even the + terminal one - properly processed. + + - ---> This callback `on_episode_end()` is fired. <--- + + - The episode is numpy'ized (i.e. lists of obs/rewards/actions/etc.. are + converted into numpy arrays). + + Args: + episode: The terminated/truncated SingleAgent- or MultiAgentEpisode object + (after `env.step()` that returned terminated=True OR truncated=True and + after the returned obs, rewards, etc.. have been logged to the episode + object). Note that this method is still called before(!) the episode + object is numpy'ized, meaning all its timestep data is still present in + lists of individual timestep data. + prev_episode_chunks: A complete list of all previous episode chunks + with the same ID as `episode` that have been sampled on this EnvRunner. + In order to compile metrics across the complete episode, users should + loop through the list: `[episode] + previous_episode_chunks` and + accumulate the required information. + env_runner: Reference to the EnvRunner running the env and episode. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics during env/episode stepping. + env: The gym.Env or gym.vector.Env object running the started episode. + env_index: The index of the sub-environment that has just been terminated + or truncated. + rl_module: The RLModule used to compute actions for stepping the env. In + single-agent mode, this is a simple RLModule, in multi-agent mode, this + is a MultiRLModule. + kwargs: Forward compatibility placeholder. + """ + pass + + @OverrideToImplementCustomLogic + def on_sample_end( + self, + *, + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + samples: Union[SampleBatch, List[EpisodeType]], + # TODO (sven): Deprecate these args. + worker: Optional["EnvRunner"] = None, + **kwargs, + ) -> None: + """Called at the end of `EnvRunner.sample()`. + + Args: + env_runner: Reference to the current EnvRunner object. + metrics_logger: The MetricsLogger object inside the `env_runner`. Can be + used to log custom metrics during env/episode stepping. + samples: Lists of SingleAgentEpisode or MultiAgentEpisode instances to be + returned. You can mutate the episodes to modify the returned training + data. + kwargs: Forward compatibility placeholder. + """ + pass + + @OldAPIStack + def on_sub_environment_created( + self, + *, + worker: "EnvRunner", + sub_environment: EnvType, + env_context: EnvContext, + env_index: Optional[int] = None, + **kwargs, + ) -> None: + """Callback run when a new sub-environment has been created. + + This method gets called after each sub-environment (usually a + gym.Env) has been created, validated (RLlib built-in validation + + possible custom validation function implemented by overriding + `Algorithm.validate_env()`), wrapped (e.g. video-wrapper), and seeded. + + Args: + worker: Reference to the current EnvRunner. + sub_environment: The sub-environment instance that has been + created. This is usually a gym.Env object. + env_context: The `EnvContext` object that has been passed to + the env's constructor. + env_index: The index of the sub-environment that has been created + (within the vector of sub-environments of the gym.vector.Env). + kwargs: Forward compatibility placeholder. + """ + pass + + @OldAPIStack + def on_postprocess_trajectory( + self, + *, + worker: "EnvRunner", + episode, + agent_id: AgentID, + policy_id: PolicyID, + policies: Dict[PolicyID, Policy], + postprocessed_batch: SampleBatch, + original_batches: Dict[AgentID, Tuple[Policy, SampleBatch]], + **kwargs, + ) -> None: + """Called immediately after a policy's postprocess_fn is called. + + You can use this callback to do additional postprocessing for a policy, + including looking at the trajectory data of other agents in multi-agent + settings. + + Args: + worker: Reference to the current rollout worker. + episode: Episode object. + agent_id: Id of the current agent. + policy_id: Id of the current policy for the agent. + policies: Dict mapping policy IDs to policy objects. In single + agent mode there will only be a single "default_policy". + postprocessed_batch: The postprocessed sample batch + for this agent. You can mutate this object to apply your own + trajectory postprocessing. + original_batches: Dict mapping agent IDs to their unpostprocessed + trajectory data. You should not mutate this object. + kwargs: Forward compatibility placeholder. + """ + pass + + @OldAPIStack + def on_create_policy(self, *, policy_id: PolicyID, policy: Policy) -> None: + """Callback run whenever a new policy is added to an algorithm. + + Args: + policy_id: ID of the newly created policy. + policy: The policy just created. + """ + pass + + @OldAPIStack + def on_learn_on_batch( + self, *, policy: Policy, train_batch: SampleBatch, result: dict, **kwargs + ) -> None: + """Called at the beginning of Policy.learn_on_batch(). + + Note: This is called before 0-padding via + `pad_batch_to_sequences_of_same_size`. + + Also note, SampleBatch.INFOS column will not be available on + train_batch within this callback if framework is tf1, due to + the fact that tf1 static graph would mistake it as part of the + input dict if present. + It is available though, for tf2 and torch frameworks. + + Args: + policy: Reference to the current Policy object. + train_batch: SampleBatch to be trained on. You can + mutate this object to modify the samples generated. + result: A results dict to add custom metrics to. + kwargs: Forward compatibility placeholder. + """ + pass + + # Deprecated, use `on_env_runners_recreated`, instead. + def on_workers_recreated( + self, + *, + algorithm, + worker_set, + worker_ids, + is_evaluation, + **kwargs, + ) -> None: + pass + + +class MemoryTrackingCallbacks(RLlibCallback): + """MemoryTrackingCallbacks can be used to trace and track memory usage + in rollout workers. + + The Memory Tracking Callbacks uses tracemalloc and psutil to track + python allocations during rollouts, + in training or evaluation. + + The tracking data is logged to the custom_metrics of an episode and + can therefore be viewed in tensorboard + (or in WandB etc..) + + Add MemoryTrackingCallbacks callback to the tune config + e.g. { ...'callbacks': MemoryTrackingCallbacks ...} + + Note: + This class is meant for debugging and should not be used + in production code as tracemalloc incurs + a significant slowdown in execution speed. + """ + + def __init__(self): + super().__init__() + + # Will track the top 10 lines where memory is allocated + tracemalloc.start(10) + + @override(RLlibCallback) + def on_episode_end( + self, + *, + episode: Union[EpisodeType, EpisodeV2], + env_runner: Optional["EnvRunner"] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + rl_module: Optional[RLModule] = None, + # TODO (sven): Deprecate these args. + worker: Optional["EnvRunner"] = None, + base_env: Optional[BaseEnv] = None, + policies: Optional[Dict[PolicyID, Policy]] = None, + **kwargs, + ) -> None: + gc.collect() + snapshot = tracemalloc.take_snapshot() + top_stats = snapshot.statistics("lineno") + + for stat in top_stats[:10]: + count = stat.count + # Convert total size from Bytes to KiB. + size = stat.size / 1024 + + trace = str(stat.traceback) + + episode.custom_metrics[f"tracemalloc/{trace}/size"] = size + episode.custom_metrics[f"tracemalloc/{trace}/count"] = count + + process = psutil.Process(os.getpid()) + worker_rss = process.memory_info().rss + worker_vms = process.memory_info().vms + if platform.system() == "Linux": + # This is only available on Linux + worker_data = process.memory_info().data + episode.custom_metrics["tracemalloc/worker/data"] = worker_data + episode.custom_metrics["tracemalloc/worker/rss"] = worker_rss + episode.custom_metrics["tracemalloc/worker/vms"] = worker_vms diff --git a/lib/python3.12/site-packages/ray/rllib/callbacks/utils.py b/lib/python3.12/site-packages/ray/rllib/callbacks/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b6bdf33f3049b51b7b284f084f1044292283845b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/callbacks/utils.py @@ -0,0 +1,156 @@ +from typing import Any, Callable, Dict, List, Optional + +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import OldAPIStack + + +def make_callback( + callback_name: str, + callbacks_objects: Optional[List[RLlibCallback]] = None, + callbacks_functions: Optional[List[Callable]] = None, + *, + args: List[Any] = None, + kwargs: Dict[str, Any] = None, +) -> None: + """Calls an RLlibCallback method or a registered callback callable. + + Args: + callback_name: The name of the callback method or key, for example: + "on_episode_start" or "on_train_result". + callbacks_objects: The RLlibCallback object or list of RLlibCallback objects + to call the `callback_name` method on (in the order they appear in the + list). + callbacks_functions: The callable or list of callables to call + (in the order they appear in the list). + args: Call args to pass to the method/callable calls. + kwargs: Call kwargs to pass to the method/callable calls. + """ + # Loop through all available RLlibCallback objects. + callbacks_objects = force_list(callbacks_objects) + for callback_obj in callbacks_objects: + getattr(callback_obj, callback_name)(*(args or ()), **(kwargs or {})) + + # Loop through all available RLlibCallback objects. + callbacks_functions = force_list(callbacks_functions) + for callback_fn in callbacks_functions: + callback_fn(*(args or ()), **(kwargs or {})) + + +@OldAPIStack +def _make_multi_callbacks(callback_class_list): + class _MultiCallbacks(RLlibCallback): + IS_CALLBACK_CONTAINER = True + + def __init__(self): + super().__init__() + self._callback_list = [ + callback_class() for callback_class in callback_class_list + ] + + def on_algorithm_init(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_algorithm_init(**kwargs) + + def on_workers_recreated(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_workers_recreated(**kwargs) + + # Only on new API stack. + def on_env_runners_recreated(self, **kwargs) -> None: + pass + + def on_offline_eval_runners_recreated(self, **kwargs) -> None: + pass + + def on_checkpoint_loaded(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_checkpoint_loaded(**kwargs) + + def on_create_policy(self, *, policy_id, policy) -> None: + for callback in self._callback_list: + callback.on_create_policy(policy_id=policy_id, policy=policy) + + def on_environment_created(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_environment_created(**kwargs) + + def on_sub_environment_created(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_sub_environment_created(**kwargs) + + def on_episode_created(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_episode_created(**kwargs) + + def on_episode_start(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_episode_start(**kwargs) + + def on_episode_step(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_episode_step(**kwargs) + + def on_episode_end(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_episode_end(**kwargs) + + def on_evaluate_start(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_evaluate_start(**kwargs) + + def on_evaluate_end(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_evaluate_end(**kwargs) + + # TODO (simon, sven): Fix the test such that we can simply remove + # these. + def on_evaluate_offline_start(self, **kwargs): + for callback in self._callback_list: + callback.on_evaluate_offline_start(**kwargs) + + def on_evaluate_offline_end(self, **kwargs): + for callback in self._callback_list: + callback.on_evaluate_offline_end(**kwargs) + + def on_postprocess_trajectory( + self, + *, + worker, + episode, + agent_id, + policy_id, + policies, + postprocessed_batch, + original_batches, + **kwargs, + ) -> None: + for callback in self._callback_list: + callback.on_postprocess_trajectory( + worker=worker, + episode=episode, + agent_id=agent_id, + policy_id=policy_id, + policies=policies, + postprocessed_batch=postprocessed_batch, + original_batches=original_batches, + **kwargs, + ) + + def on_sample_end(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_sample_end(**kwargs) + + def on_learn_on_batch( + self, *, policy, train_batch, result: dict, **kwargs + ) -> None: + for callback in self._callback_list: + callback.on_learn_on_batch( + policy=policy, train_batch=train_batch, result=result, **kwargs + ) + + def on_train_result(self, **kwargs) -> None: + for callback in self._callback_list: + callback.on_train_result(**kwargs) + + return _MultiCallbacks diff --git a/lib/python3.12/site-packages/ray/rllib/core/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ab64de5ec9b684992aa60f2de78f01c47e661bb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/__pycache__/columns.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/__pycache__/columns.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f26c881db5b31587b8b9c92c0c7faddf66e0e876 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/__pycache__/columns.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/distribution/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4b18dd297751df46f68cdfdf91b6427a2f402e8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/distribution.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/distribution.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a63233501329dc1a99bd0721dbdc4010672e79c0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/distribution/__pycache__/distribution.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/distribution.py b/lib/python3.12/site-packages/ray/rllib/core/distribution/distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..a5812058713eb670140a99f2219adfbf1aa30e8b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/distribution/distribution.py @@ -0,0 +1,250 @@ +"""This is the next version of action distribution base class.""" +import abc +from typing import Tuple + +import gymnasium as gym + +from ray.rllib.utils.annotations import ExperimentalAPI, override +from ray.rllib.utils.typing import TensorType, Union + + +@ExperimentalAPI +class Distribution(abc.ABC): + """The base class for distribution over a random variable. + + Examples: + + .. testcode:: + + import torch + from ray.rllib.core.models.configs import MLPHeadConfig + from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical + ) + + model = MLPHeadConfig(input_dims=[1]).build(framework="torch") + + # Create an action distribution from model logits + action_logits = model(torch.Tensor([[1]])) + action_dist = TorchCategorical.from_logits(action_logits) + action = action_dist.sample() + + # Create another distribution from a dummy Tensor + action_dist2 = TorchCategorical.from_logits(torch.Tensor([0])) + + # Compute some common metrics + logp = action_dist.logp(action) + kl = action_dist.kl(action_dist2) + entropy = action_dist.entropy() + """ + + @abc.abstractmethod + def sample( + self, + *, + sample_shape: Tuple[int, ...] = None, + return_logp: bool = False, + **kwargs, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + """Draw a sample from the distribution. + + Args: + sample_shape: The shape of the sample to draw. + return_logp: Whether to return the logp of the sampled values. + **kwargs: Forward compatibility placeholder. + + Returns: + The sampled values. If return_logp is True, returns a tuple of the + sampled values and its logp. + """ + + @abc.abstractmethod + def rsample( + self, + *, + sample_shape: Tuple[int, ...] = None, + return_logp: bool = False, + **kwargs, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + """Draw a re-parameterized sample from the action distribution. + + If this method is implemented, we can take gradients of samples w.r.t. the + distribution parameters. + + Args: + sample_shape: The shape of the sample to draw. + return_logp: Whether to return the logp of the sampled values. + **kwargs: Forward compatibility placeholder. + + Returns: + The sampled values. If return_logp is True, returns a tuple of the + sampled values and its logp. + """ + + @abc.abstractmethod + def logp(self, value: TensorType, **kwargs) -> TensorType: + """The log-likelihood of the distribution computed at `value` + + Args: + value: The value to compute the log-likelihood at. + **kwargs: Forward compatibility placeholder. + + Returns: + The log-likelihood of the value. + """ + + @abc.abstractmethod + def kl(self, other: "Distribution", **kwargs) -> TensorType: + """The KL-divergence between two distributions. + + Args: + other: The other distribution. + **kwargs: Forward compatibility placeholder. + + Returns: + The KL-divergence between the two distributions. + """ + + @abc.abstractmethod + def entropy(self, **kwargs) -> TensorType: + """The entropy of the distribution. + + Args: + **kwargs: Forward compatibility placeholder. + + Returns: + The entropy of the distribution. + """ + + @staticmethod + @abc.abstractmethod + def required_input_dim(space: gym.Space, **kwargs) -> int: + """Returns the required length of an input parameter tensor. + + Args: + space: The space this distribution will be used for, + whose shape attributes will be used to determine the required shape of + the input parameter tensor. + **kwargs: Forward compatibility placeholder. + + Returns: + size of the required input vector (minus leading batch dimension). + """ + + @classmethod + def from_logits(cls, logits: TensorType, **kwargs) -> "Distribution": + """Creates a Distribution from logits. + + The caller does not need to have knowledge of the distribution class in order + to create it and sample from it. The passed batched logits vectors might be + split up and are passed to the distribution class' constructor as kwargs. + + Args: + logits: The logits to create the distribution from. + **kwargs: Forward compatibility placeholder. + + Returns: + The created distribution. + + .. testcode:: + + import numpy as np + from ray.rllib.core.distribution.distribution import Distribution + + class Uniform(Distribution): + def __init__(self, lower, upper): + self.lower = lower + self.upper = upper + + def sample(self): + return self.lower + (self.upper - self.lower) * np.random.rand() + + def logp(self, x): + ... + + def kl(self, other): + ... + + def entropy(self): + ... + + @staticmethod + def required_input_dim(space): + ... + + def rsample(self): + ... + + @classmethod + def from_logits(cls, logits, **kwargs): + return Uniform(logits[:, 0], logits[:, 1]) + + logits = np.array([[0.0, 1.0], [2.0, 3.0]]) + my_dist = Uniform.from_logits(logits) + sample = my_dist.sample() + """ + raise NotImplementedError + + @classmethod + def get_partial_dist_cls( + parent_cls: "Distribution", **partial_kwargs + ) -> "Distribution": + """Returns a partial child of TorchMultiActionDistribution. + + This is useful if inputs needed to instantiate the Distribution from logits + are available, but the logits are not. + """ + + class DistributionPartial(parent_cls): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @staticmethod + def _merge_kwargs(**kwargs): + """Checks if keys in kwargs don't clash with partial_kwargs.""" + overlap = set(kwargs) & set(partial_kwargs) + if overlap: + raise ValueError( + f"Cannot override the following kwargs: {overlap}.\n" + f"This is because they were already set at the time this " + f"partial class was defined." + ) + merged_kwargs = {**partial_kwargs, **kwargs} + return merged_kwargs + + @classmethod + @override(parent_cls) + def required_input_dim(cls, space: gym.Space, **kwargs) -> int: + merged_kwargs = cls._merge_kwargs(**kwargs) + assert space == merged_kwargs["space"] + return parent_cls.required_input_dim(**merged_kwargs) + + @classmethod + @override(parent_cls) + def from_logits( + cls, + logits: TensorType, + **kwargs, + ) -> "DistributionPartial": + merged_kwargs = cls._merge_kwargs(**kwargs) + distribution = parent_cls.from_logits(logits, **merged_kwargs) + # Replace the class of the returned distribution with this partial + # This makes it so that we can use type() on this distribution and + # get back the partial class. + distribution.__class__ = cls + return distribution + + # Substitute name of this partial class to match the original class. + DistributionPartial.__name__ = f"{parent_cls}Partial" + + return DistributionPartial + + def to_deterministic(self) -> "Distribution": + """Returns a deterministic equivalent for this distribution. + + Specifically, the deterministic equivalent for a Categorical distribution is a + Deterministic distribution that selects the action with maximum logit value. + Generally, the choice of the deterministic replacement is informed by + established conventions. + """ + return self diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48a184c4d30c0f4602c23398cb007e7a8d6532f4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/torch_distribution.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/torch_distribution.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..474af2da3a17f091e5a2db6835bfc3c046cc96ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/__pycache__/torch_distribution.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/torch_distribution.py b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/torch_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..d0b94828a9c020734593f252a632ebb7a2cb7991 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/distribution/torch/torch_distribution.py @@ -0,0 +1,707 @@ +"""The main difference between this and the old ActionDistribution is that this one +has more explicit input args. So that the input format does not have to be guessed from +the code. This matches the design pattern of torch distribution which developers may +already be familiar with. +""" +import abc +from typing import Dict, Iterable, List, Optional + +import gymnasium as gym +import numpy as np +import tree + +from ray.rllib.core.distribution.distribution import Distribution +from ray.rllib.utils.annotations import DeveloperAPI, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import MAX_LOG_NN_OUTPUT, MIN_LOG_NN_OUTPUT, SMALL_NUMBER +from ray.rllib.utils.typing import TensorType, Tuple, Union + +torch, nn = try_import_torch() + + +@DeveloperAPI +class TorchDistribution(Distribution, abc.ABC): + """Wrapper class for torch.distributions.""" + + def __init__(self, *args, **kwargs): + super().__init__() + self._dist = self._get_torch_distribution(*args, **kwargs) + + @abc.abstractmethod + def _get_torch_distribution( + self, *args, **kwargs + ) -> "torch.distributions.Distribution": + """Returns the torch.distributions.Distribution object to use.""" + + @override(Distribution) + def logp(self, value: TensorType, **kwargs) -> TensorType: + return self._dist.log_prob(value, **kwargs) + + @override(Distribution) + def entropy(self) -> TensorType: + return self._dist.entropy() + + @override(Distribution) + def kl(self, other: "Distribution") -> TensorType: + return torch.distributions.kl.kl_divergence(self._dist, other._dist) + + @override(Distribution) + def sample( + self, + *, + sample_shape=None, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + sample = self._dist.sample( + sample_shape if sample_shape is not None else torch.Size() + ) + return sample + + @override(Distribution) + def rsample( + self, + *, + sample_shape=None, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + rsample = self._dist.rsample( + sample_shape if sample_shape is not None else torch.Size() + ) + return rsample + + @classmethod + @override(Distribution) + def from_logits(cls, logits: TensorType, **kwargs) -> "TorchDistribution": + return cls(logits=logits, **kwargs) + + +@DeveloperAPI +class TorchCategorical(TorchDistribution): + r"""Wrapper class for PyTorch Categorical distribution. + + Creates a categorical distribution parameterized by either :attr:`probs` or + :attr:`logits` (but not both). + + Samples are integers from :math:`\{0, \ldots, K-1\}` where `K` is + ``probs.size(-1)``. + + If `probs` is 1-dimensional with length-`K`, each element is the relative + probability of sampling the class at that index. + + If `probs` is N-dimensional, the first N-1 dimensions are treated as a batch of + relative probability vectors. + + .. testcode:: + :skipif: True + + m = TorchCategorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ])) + m.sample(sample_shape=(2,)) # equal probability of 0, 1, 2, 3 + + .. testoutput:: + + tensor([3, 4]) + + Args: + logits: Event log probabilities (unnormalized) + probs: The probabilities of each event. + temperature: In case of using logits, this parameter can be used to determine + the sharpness of the distribution. i.e. + ``probs = softmax(logits / temperature)``. The temperature must be strictly + positive. A low value (e.g. 1e-10) will result in argmax sampling while a + larger value will result in uniform sampling. + """ + + @override(TorchDistribution) + def __init__( + self, + logits: "torch.Tensor" = None, + probs: "torch.Tensor" = None, + ) -> None: + # We assert this here because to_deterministic makes this assumption. + assert (probs is None) != ( + logits is None + ), "Exactly one out of `probs` and `logits` must be set!" + + self.probs = probs + self.logits = logits + super().__init__(logits=logits, probs=probs) + + # Build this distribution only if really needed (in `self.rsample()`). It's + # quite expensive according to cProfile. + self._one_hot = None + + @override(TorchDistribution) + def _get_torch_distribution( + self, + logits: "torch.Tensor" = None, + probs: "torch.Tensor" = None, + ) -> "torch.distributions.Distribution": + return torch.distributions.categorical.Categorical( + logits=logits, probs=probs, validate_args=False + ) + + @staticmethod + @override(Distribution) + def required_input_dim(space: gym.Space, **kwargs) -> int: + assert isinstance(space, gym.spaces.Discrete) + return int(space.n) + + @override(Distribution) + def rsample(self, sample_shape=()): + if self._one_hot is None: + self._one_hot = torch.distributions.one_hot_categorical.OneHotCategorical( + logits=self.logits, probs=self.probs, validate_args=False + ) + one_hot_sample = self._one_hot.sample(sample_shape) + return (one_hot_sample - self.probs).detach() + self.probs + + def to_deterministic(self) -> "TorchDeterministic": + if self.probs is not None: + probs_or_logits = self.probs + else: + probs_or_logits = self.logits + + return TorchDeterministic(loc=torch.argmax(probs_or_logits, dim=-1)) + + +@DeveloperAPI +class TorchDiagGaussian(TorchDistribution): + """Wrapper class for PyTorch Normal distribution. + + Creates a normal distribution parameterized by :attr:`loc` and :attr:`scale`. In + case of multi-dimensional distribution, the variance is assumed to be diagonal. + + .. testcode:: + :skipif: True + + loc, scale = torch.tensor([0.0, 0.0]), torch.tensor([1.0, 1.0]) + m = TorchDiagGaussian(loc=loc, scale=scale) + m.sample(sample_shape=(2,)) # 2d normal dist with loc=0 and scale=1 + + .. testoutput:: + + tensor([[ 0.1046, -0.6120], [ 0.234, 0.556]]) + + .. testcode:: + :skipif: True + + # scale is None + m = TorchDiagGaussian(loc=torch.tensor([0.0, 1.0])) + m.sample(sample_shape=(2,)) # normally distributed with loc=0 and scale=1 + + .. testoutput:: + + tensor([0.1046, 0.6120]) + + + Args: + loc: mean of the distribution (often referred to as mu). If scale is None, the + second half of the `loc` will be used as the log of scale. + scale: standard deviation of the distribution (often referred to as sigma). + Has to be positive. + """ + + @override(TorchDistribution) + def __init__( + self, + loc: Union[float, "torch.Tensor"], + scale: Optional[Union[float, "torch.Tensor"]], + ): + self.loc = loc + super().__init__(loc=loc, scale=scale) + + def _get_torch_distribution(self, loc, scale) -> "torch.distributions.Distribution": + return torch.distributions.normal.Normal(loc, scale, validate_args=False) + + @override(TorchDistribution) + def logp(self, value: TensorType) -> TensorType: + return super().logp(value).sum(-1) + + @override(TorchDistribution) + def entropy(self) -> TensorType: + return super().entropy().sum(-1) + + @override(TorchDistribution) + def kl(self, other: "TorchDistribution") -> TensorType: + return super().kl(other).sum(-1) + + @staticmethod + @override(Distribution) + def required_input_dim(space: gym.Space, **kwargs) -> int: + assert isinstance(space, gym.spaces.Box) + return int(np.prod(space.shape, dtype=np.int32) * 2) + + @classmethod + @override(Distribution) + def from_logits(cls, logits: TensorType, **kwargs) -> "TorchDiagGaussian": + loc, log_std = logits.chunk(2, dim=-1) + scale = log_std.exp() + return cls(loc=loc, scale=scale) + + def to_deterministic(self) -> "TorchDeterministic": + return TorchDeterministic(loc=self.loc) + + +@DeveloperAPI +class TorchSquashedGaussian(TorchDistribution): + @override(TorchDistribution) + def __init__( + self, + loc: Union[float, "torch.Tensor"], + scale: Optional[Union[float, "torch.Tensor"]] = 1.0, + low: float = -1.0, + high: float = 1.0, + ): + self.loc = loc + self.low = low + self.high = high + + super().__init__(loc=loc, scale=scale) + + def _get_torch_distribution(self, loc, scale) -> "torch.distributions.Distribution": + return torch.distributions.normal.Normal(loc, scale, validate_args=False) + + @override(TorchDistribution) + def sample( + self, *, sample_shape=None + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + # Sample from the Normal distribution. + sample = super().sample( + sample_shape=sample_shape if sample_shape is not None else torch.Size() + ) + # Return the squashed sample. + return self._squash(sample) + + @override(TorchDistribution) + def rsample( + self, *, sample_shape=None + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + # Sample from the Normal distribution. + sample = super().rsample( + sample_shape=sample_shape if sample_shape is not None else torch.Size() + ) + # Return the squashed sample. + return self._squash(sample) + + @override(TorchDistribution) + def logp(self, value: TensorType, **kwargs) -> TensorType: + # Unsquash value. + value = self._unsquash(value) + # Get log-probabilities from Normal distribution. + logp = super().logp(value, **kwargs) + # Clip the log probabilities as a safeguard and sum. + logp = torch.clamp(logp, -100, 100).sum(-1) + # Return the log probabilities for squashed Normal. + value = torch.tanh(value) + return logp - torch.log(1 - value**2 + SMALL_NUMBER).sum(-1) + + @override(TorchDistribution) + def entropy(self) -> TensorType: + raise ValueError("ENtropy not defined for `TorchSquashedGaussian`.") + + @override(TorchDistribution) + def kl(self, other: Distribution) -> TensorType: + raise ValueError("KL not defined for `TorchSquashedGaussian`.") + + def _squash(self, sample: TensorType) -> TensorType: + # Rescale the sample to interval given by the bounds (including the bounds). + sample = ((torch.tanh(sample) + 1.0) / 2.0) * (self.high - self.low) + self.low + # Return a clipped sample to comply with the bounds. + return torch.clamp(sample, self.low, self.high) + + def _unsquash(self, sample: TensorType) -> TensorType: + # Rescale to [-1.0, 1.0]. + sample = (sample - self.low) / (self.high - self.low) * 2.0 - 1.0 + # Stabilize input to atanh function. + sample = torch.clamp(sample, -1.0 + SMALL_NUMBER, 1.0 - SMALL_NUMBER) + return torch.atanh(sample) + + @staticmethod + @override(Distribution) + def required_input_dim(space: gym.Space, **kwargs) -> int: + assert isinstance(space, gym.spaces.Box), space + return int(np.prod(space.shape, dtype=np.int32) * 2) + + @classmethod + @override(TorchDistribution) + def from_logits( + cls, logits: TensorType, low: float = -1.0, high: float = 1.0, **kwargs + ) -> "TorchSquashedGaussian": + loc, log_std = logits.chunk(2, dim=-1) + # Clip the `scale` values (coming from the `RLModule.forward()`) to + # reasonable values. + log_std = torch.clamp(log_std, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT) + scale = log_std.exp() + + # Assert that `low` is smaller than `high`. + assert np.all(np.less(low, high)) + # Return class instance. + return cls(loc=loc, scale=scale, low=low, high=high, **kwargs) + + def to_deterministic(self) -> Distribution: + return TorchDeterministic(loc=self.loc) + + +@DeveloperAPI +class TorchDeterministic(Distribution): + """The distribution that returns the input values directly. + + This is similar to DiagGaussian with standard deviation zero (thus only + requiring the "mean" values as NN output). + + Note: entropy is always zero, ang logp and kl are not implemented. + + .. testcode:: + :skipif: True + + m = TorchDeterministic(loc=torch.tensor([0.0, 0.0])) + m.sample(sample_shape=(2,)) + + .. testoutput:: + + tensor([[ 0.0, 0.0], [ 0.0, 0.0]]) + + Args: + loc: the determinsitic value to return + """ + + @override(Distribution) + def __init__(self, loc: "torch.Tensor") -> None: + super().__init__() + self.loc = loc + + @override(Distribution) + def sample( + self, + *, + sample_shape=None, + **kwargs, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + device = self.loc.device + dtype = self.loc.dtype + shape = ( + sample_shape if sample_shape is not None else torch.Size() + ) + self.loc.shape + return torch.ones(shape, device=device, dtype=dtype) * self.loc + + def rsample( + self, + *, + sample_shape: Tuple[int, ...] = None, + **kwargs, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + raise NotImplementedError + + @override(Distribution) + def logp(self, value: TensorType, **kwargs) -> TensorType: + return torch.zeros_like(self.loc) + + @override(Distribution) + def entropy(self, **kwargs) -> TensorType: + raise RuntimeError(f"`entropy()` not supported for {self.__class__.__name__}.") + + @override(Distribution) + def kl(self, other: "Distribution", **kwargs) -> TensorType: + raise RuntimeError(f"`kl()` not supported for {self.__class__.__name__}.") + + @staticmethod + @override(Distribution) + def required_input_dim(space: gym.Space, **kwargs) -> int: + assert isinstance(space, gym.spaces.Box) + return int(np.prod(space.shape, dtype=np.int32)) + + def to_deterministic(self) -> "TorchDeterministic": + return self + + +@DeveloperAPI +class TorchMultiCategorical(Distribution): + """MultiCategorical distribution for MultiDiscrete action spaces.""" + + @override(Distribution) + def __init__( + self, + categoricals: List[TorchCategorical], + ): + super().__init__() + self._cats = categoricals + + @override(Distribution) + def sample(self) -> TensorType: + arr = [cat.sample() for cat in self._cats] + sample_ = torch.stack(arr, dim=-1) + return sample_ + + @override(Distribution) + def rsample(self, sample_shape=()): + arr = [cat.rsample() for cat in self._cats] + sample_ = torch.stack(arr, dim=-1) + return sample_ + + @override(Distribution) + def logp(self, value: "torch.Tensor") -> TensorType: + value = torch.unbind(value, dim=-1) + logps = torch.stack([cat.logp(act) for cat, act in zip(self._cats, value)]) + return torch.sum(logps, dim=0) + + @override(Distribution) + def entropy(self) -> TensorType: + return torch.sum( + torch.stack([cat.entropy() for cat in self._cats], dim=-1), dim=-1 + ) + + @override(Distribution) + def kl(self, other: Distribution) -> TensorType: + kls = torch.stack( + [cat.kl(oth_cat) for cat, oth_cat in zip(self._cats, other._cats)], + dim=-1, + ) + return torch.sum(kls, dim=-1) + + @staticmethod + @override(Distribution) + def required_input_dim(space: gym.Space, **kwargs) -> int: + assert isinstance(space, gym.spaces.MultiDiscrete) + return int(np.sum(space.nvec)) + + @classmethod + @override(Distribution) + def from_logits( + cls, + logits: "torch.Tensor", + input_lens: List[int], + temperatures: List[float] = None, + **kwargs, + ) -> "TorchMultiCategorical": + """Creates this Distribution from logits (and additional arguments). + + If you wish to create this distribution from logits only, please refer to + `Distribution.get_partial_dist_cls()`. + + Args: + logits: The tensor containing logits to be separated by logit_lens. + child_distribution_cls_struct: A struct of Distribution classes that can + be instantiated from the given logits. + input_lens: A list of integers that indicate the length of the logits + vectors to be passed into each child distribution. + temperatures: A list of floats representing the temperature to use for + each Categorical distribution. If not provided, 1.0 is used for all. + **kwargs: Forward compatibility kwargs. + """ + if not temperatures: + # If temperatures are not provided, use 1.0 for all actions. + temperatures = [1.0] * len(input_lens) + + assert ( + sum(input_lens) == logits.shape[-1] + ), "input_lens must sum to logits.shape[-1]" + assert len(input_lens) == len( + temperatures + ), "input_lens and temperatures must be same length" + + categoricals = [ + TorchCategorical(logits=logits) + for logits in torch.split(logits, input_lens, dim=-1) + ] + + return cls(categoricals=categoricals) + + def to_deterministic(self) -> "TorchDeterministic": + """Converts `TorchMultiCategorical` into `TorchDeterministic`.""" + logits_list = [cat.logits for cat in self._cats] + # Check, if the module is recurrent. + is_recurrent = logits_list[0].dim() == 3 # (B, T, K_i) + + # Determine max number of categories across all categorical distributions + max_K = max(logits.shape[-1] for logits in logits_list) + + padded_logits = [] + for logits in logits_list: + # Pad last dimension (category dim) to max_K + pad_width = max_K - logits.shape[-1] + # If the distributions have different number of categories, pad. + if pad_width > 0: + # Pad only last dimension + pad_dims = (0, pad_width) + logits = nn.functional.pad(logits, pad_dims, value=-float("inf")) + padded_logits.append(logits) + + # Stack along new dim=0 (categorical dimension). + # Shape: (num_components, B, T, max_K) or (num_components, B, max_K) + stacked = torch.stack(padded_logits, dim=0) + + # Move categorical dim (0) to last if needed, and take argmax. + if is_recurrent: + # Current shape is (num_components, B, T, K) and we want to have + # (B, T, num_components) via argmax over last dimension. So take + # argmax over last dim (K), then permute. + argmax = torch.argmax(stacked, dim=-1) # shape: (num_components, B, T) + loc = argmax.permute(1, 2, 0) # (B, T, num_components) + else: + # stacked: (num_components, B, K) + # → argmax over last dim (K), shape: (num_components, B) + # → transpose to (B, num_components) + argmax = torch.argmax(stacked, dim=-1) # (num_components, B) + loc = argmax.transpose(0, 1) # (B, num_components) + + return TorchDeterministic(loc=loc) + + +@DeveloperAPI +class TorchMultiDistribution(Distribution): + """Action distribution that operates on multiple, possibly nested actions.""" + + def __init__( + self, + child_distribution_struct: Union[Tuple, List, Dict], + ): + """Initializes a TorchMultiDistribution object. + + Args: + child_distribution_struct: A complex struct that contains the child + distribution instances that make up this multi-distribution. + """ + super().__init__() + self._original_struct = child_distribution_struct + self._flat_child_distributions = tree.flatten(child_distribution_struct) + + @override(Distribution) + def rsample( + self, + *, + sample_shape: Tuple[int, ...] = None, + **kwargs, + ) -> Union[TensorType, Tuple[TensorType, TensorType]]: + rsamples = [] + for dist in self._flat_child_distributions: + rsample = dist.rsample(sample_shape=sample_shape, **kwargs) + rsamples.append(rsample) + + rsamples = tree.unflatten_as(self._original_struct, rsamples) + + return rsamples + + @override(Distribution) + def logp(self, value: TensorType) -> TensorType: + # Different places in RLlib use this method with different inputs. + # We therefore need to handle a flattened and concatenated input, as well as + # a nested one. + # TODO(Artur): Deprecate tensor inputs, only allow nested structures. + if isinstance(value, torch.Tensor): + split_indices = [] + for dist in self._flat_child_distributions: + if isinstance(dist, TorchCategorical): + split_indices.append(1) + elif isinstance(dist, TorchMultiCategorical): + split_indices.append(len(dist._cats)) + else: + sample = dist.sample() + # Cover Box(shape=()) case. + if len(sample.shape) == 1: + split_indices.append(1) + else: + split_indices.append(sample.size()[1]) + split_value = list(torch.split(value, split_indices, dim=1)) + else: + split_value = tree.flatten(value) + + def map_(val, dist): + # Remove extra dimension if present. + if ( + isinstance(dist, TorchCategorical) + and val.shape[-1] == 1 + and len(val.shape) > 1 + ): + val = torch.squeeze(val, dim=-1) + return dist.logp(val) + + flat_logps = tree.map_structure( + map_, split_value, self._flat_child_distributions + ) + + return sum(flat_logps) + + @override(Distribution) + def kl(self, other: Distribution) -> TensorType: + kl_list = [ + d.kl(o) + for d, o in zip( + self._flat_child_distributions, other._flat_child_distributions + ) + ] + return sum(kl_list) + + @override(Distribution) + def entropy(self): + entropy_list = [d.entropy() for d in self._flat_child_distributions] + return sum(entropy_list) + + @override(Distribution) + def sample(self): + child_distributions_struct = tree.unflatten_as( + self._original_struct, self._flat_child_distributions + ) + return tree.map_structure(lambda s: s.sample(), child_distributions_struct) + + @staticmethod + @override(Distribution) + def required_input_dim( + space: gym.Space, input_lens: List[int], as_list: bool = False, **kwargs + ) -> int: + if as_list: + return input_lens + else: + return sum(input_lens) + + @classmethod + @override(Distribution) + def from_logits( + cls, + logits: "torch.Tensor", + child_distribution_cls_struct: Union[Dict, Iterable], + input_lens: Union[Dict, List[int]], + **kwargs, + ) -> "TorchMultiDistribution": + """Creates this Distribution from logits (and additional arguments). + + If you wish to create this distribution from logits only, please refer to + `Distribution.get_partial_dist_cls()`. + + Args: + logits: The tensor containing logits to be separated by `input_lens`. + child_distribution_cls_struct: A struct of Distribution classes that can + be instantiated from the given logits. + child_distribution_cls_struct: A struct of Distribution classes that can + be instantiated from the given logits. + input_lens: A list or dict of integers that indicate the length of each + logit. If this is given as a dict, the structure should match the + structure of child_distribution_cls_struct. + **kwargs: Forward compatibility kwargs. + + Returns: + A TorchMultiDistribution object. + """ + logit_lens = tree.flatten(input_lens) + child_distribution_cls_list = tree.flatten(child_distribution_cls_struct) + split_logits = torch.split(logits, logit_lens, dim=-1) + + child_distribution_list = tree.map_structure( + lambda dist, input_: dist.from_logits(input_), + child_distribution_cls_list, + list(split_logits), + ) + + child_distribution_struct = tree.unflatten_as( + child_distribution_cls_struct, child_distribution_list + ) + + return cls( + child_distribution_struct=child_distribution_struct, + ) + + def to_deterministic(self) -> "TorchMultiDistribution": + flat_deterministic_dists = [ + dist.to_deterministic() for dist in self._flat_child_distributions + ] + deterministic_dists = tree.unflatten_as( + self._original_struct, flat_deterministic_dists + ) + return TorchMultiDistribution(deterministic_dists) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/learner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8fc450012cde28455c0ac9962fd2a6f716c68385 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/__init__.py @@ -0,0 +1,7 @@ +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.learner_group import LearnerGroup + +__all__ = [ + "Learner", + "LearnerGroup", +] diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f6c3becdf2534e8f24b3f4799a77ea6b9fe5449 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c51b6285e61727ef59f5bdfa62d310d03da2e5c3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..551355e103d239cd36aad53c8a7eccf1ab10a531 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/differentiable_learner_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bbec8bf70a5b4e8d045448aeb944b0c3b365208 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d75abc122c7db27209ab2fe0ef322e1309078c03 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/learner_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/training_data.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/training_data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a892c737f7de6f085ce05ab2cc8e6d58e85526d5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/training_data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35f7e7eb99e395a16325b7e835898d4c25e47259 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner.py b/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..b2520f722b3082892e83ac1a2fc2388dbe3822a1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner.py @@ -0,0 +1,821 @@ +import abc +import logging +from typing import ( + TYPE_CHECKING, + Any, + Collection, + Dict, + Iterable, + Optional, + Tuple, + Union, +) + +import numpy + +from ray.rllib.connectors.learner.learner_connector_pipeline import ( + LearnerConnectorPipeline, +) +from ray.rllib.core import ALL_MODULES, COMPONENT_METRICS_LOGGER +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils import unflatten_dict +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.metrics import ( + DATASET_NUM_ITERS_TRAINED, + DATASET_NUM_ITERS_TRAINED_LIFETIME, + MODULE_TRAIN_BATCH_SIZE_MEAN, + NUM_ENV_STEPS_TRAINED, + NUM_ENV_STEPS_TRAINED_LIFETIME, + NUM_MODULE_STEPS_TRAINED, + NUM_MODULE_STEPS_TRAINED_LIFETIME, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.minibatch_utils import ( + MiniBatchCyclicIterator, + MiniBatchDummyIterator, + MiniBatchRayDataIterator, +) +from ray.rllib.utils.typing import ( + DeviceType, + ModuleID, + NamedParamDict, + ResultDict, + StateDict, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.core.learner.differentiable_learner_config import ( + DifferentiableLearnerConfig, + ) + +logger = logging.getLogger(__name__) + + +class DifferentiableLearner(Checkpointable): + """A differentiable `Learner` class enabling functional parameter updates. + + This class is part of RLlib's Meta Learning API and provides a differentiable + `Learner`, allowing higher-order updates within the meta-learning loop. + + Unlike standard learners, this class operates on a provided `RLModule` reference + instead of creating its own. Updated cloned module parameters are returned to the + caller, facilitating further update processes. + + The `update`, `_update`, `compute_gradients`, and `apply_gradients` methods + require, in addition to other arguments, a dictionary of cloned `RLModule` + parameters for functional updates. + """ + + # The framework an instance uses. + framework_str: str + + # The key for the total loss. + TOTAL_LOSS_KEY: str = "total_loss" + + def __init__( + self, + *, + config: "AlgorithmConfig", + learner_config: "DifferentiableLearnerConfig", + module: Optional[RLModule] = None, + **kwargs, + ) -> None: + + # An `AlgorithmConfig` used to access certain global configurations. + self.config: "AlgorithmConfig" = config.copy(copy_frozen=False) + # The specific configuration for the differentiable learner instance. + self.learner_config: "DifferentiableLearnerConfig" = learner_config + # The reference to the caller's module. + self._module: Optional[MultiRLModule] = module + # The reference to the caller's device. + self._device: Optional[DeviceType] = None + # A counter for functional weight updates. + self._weights_seq_no: int = 0 + + # Whether self.build has already been called. + self._is_built: bool = False + + # The learner connector pipeline. + self._learner_connector: Optional[LearnerConnectorPipeline] = None + + # The Learner's own MetricsLogger to be used to log RLlib's built-in metrics or + # custom user-defined ones (e.g. custom loss values). When returning from an + # `update_from_...()` method call, the Learner will do a `self.metrics.reduce()` + # and return the resulting (reduced) dict. + self.metrics: MetricsLogger = MetricsLogger() + + # In case of offline learning and multiple learners, each learner receives a + # repeatable iterator that iterates over a split of the streamed data. + self.iterator: MiniBatchRayDataIterator = None + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def build(self, device: Optional[DeviceType] = None) -> None: + + if self._is_built: + logger.debug("DifferentiableLearner already built. Skipping built.") + + # If a device was passed, set the `DifferentiableLearner`'s device. + if device: + self._device = device + + # TODO (simon): Move the `build_learner_connector` to the + # `DifferentiableLearnerConfig`. + self._learner_connector = self.learner_config.build_learner_connector( + input_observation_space=None, + input_action_space=None, + device=self._device, + ) + + # This instance is now ready for use. + self._is_built = True + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def compute_gradients( + self, + loss_per_module: Dict[ModuleID, TensorType], + params: Dict[ModuleID, NamedParamDict], + **kwargs, + ) -> Dict[ModuleID, NamedParamDict]: + """Computes functional gradients based on the given losses. + + Note that this method requires computing gradients functionally, + without relying on an optimizer. If an optimizer is needed, a + differentiable optimizer from a third-party package must be used. + + Args: + loss_per_module: Dict mapping module IDs to their individual total loss + terms, computed by the individual `compute_loss_for_module()` calls. + The overall total loss (sum of loss terms over all modules) is stored + under `loss_per_module[ALL_MODULES]`. + params: A dictionary containing cloned parameters for each module id. + **kwargs: Forward compatibility kwargs. + + Returns: + The gradients in the same (dict) format as `params`. + """ + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def apply_gradients( + self, + gradients: Dict[ModuleID, NamedParamDict], + params: Dict[ModuleID, NamedParamDict], + ) -> Dict[ModuleID, NamedParamDict]: + """Applies given gradients functionally. + + Note that this method requires functional parameter updates, + meaning modifications must not be performed in-place (e.g., via an + optimizer or directly within the `MultiRLModule`). + + Args: + gradients: A dictionary containing named gradients for each module id. + params: A dictionary containing named parameters for each module id. + + Returns: + The updated parameters in the same (dict) format as `params`. + """ + + @OverrideToImplementCustomLogic + def should_module_be_updated(self, module_id, multi_agent_batch=None): + """Returns whether a module should be updated or not based on `self.config`. + + Args: + module_id: The ModuleID that we want to query on whether this module + should be updated or not. + multi_agent_batch: An optional MultiAgentBatch to possibly provide further + information on the decision on whether the RLModule should be updated + or not. + """ + should_module_be_updated_fn = self.config.policies_to_train + # If None, return True (by default, all modules should be updated). + if should_module_be_updated_fn is None: + return True + # If collection given, return whether `module_id` is in that container. + elif not callable(should_module_be_updated_fn): + return module_id in set(should_module_be_updated_fn) + + return should_module_be_updated_fn(module_id, multi_agent_batch) + + @OverrideToImplementCustomLogic + def compute_losses( + self, *, fwd_out: Dict[str, Any], batch: Dict[str, Any] + ) -> Dict[str, Any]: + """Computes the loss(es) for the module being optimized. + + This method must be overridden by MultiRLModule-specific Learners in order to + define the specific loss computation logic. If the algorithm is single-agent, + only `compute_loss_for_module()` should be overridden instead. If the algorithm + uses independent multi-agent learning (default behavior for RLlib's multi-agent + setups), also only `compute_loss_for_module()` should be overridden, but it will + be called for each individual RLModule inside the MultiRLModule. + For the functional update to work, no `forward` call should be made + within this method, especially not a non-functional one. Instead, use + the model outputs provided by `fwd_out`. + + Args: + fwd_out: Output from a functional call to the `forward_train()` method of the + underlying MultiRLModule (`self.module`) during training + (`self.update()`). + batch: The train batch that was used to compute `fwd_out`. + + Returns: + A dictionary mapping module IDs to individual loss terms. + """ + loss_per_module = {} + for module_id in fwd_out: + loss = self.compute_loss_for_module( + module_id=module_id, + # TODO (simon): Check, if this should be provided per + # `DifferentiableLearnerConfig`. + config=self.config.get_config_for_module(module_id), + batch=batch[module_id], + fwd_out=fwd_out[module_id], + ) + loss_per_module[module_id] = loss + + return loss_per_module + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + """Computes the loss for a single module. + + Think of this as computing loss for a single agent. For multi-agent use-cases + that require more complicated computation for loss, consider overriding the + `compute_losses` method instead. + + Args: + module_id: The id of the module. + config: The AlgorithmConfig specific to the given `module_id`. + batch: The train batch for this particular module. + fwd_out: The output of the forward pass for this particular module. + + Returns: + A single total loss tensor. If you have more than one optimizer on the + provided `module_id` and would like to compute gradients separately using + these different optimizers, simply add up the individual loss terms for + each optimizer and return the sum. Also, for recording/logging any + individual loss terms, you can use the `Learner.metrics.log_value( + key=..., value=...)` or `DifferentiableLearner.metrics.log_dict()` APIs. + See: :py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for + more information. + """ + + def update( + self, + params: Dict[ModuleID, NamedParamDict], + training_data: TrainingData, + *, + _no_metrics_reduce: bool = False, + **kwargs, + ) -> Tuple[Dict[ModuleID, NamedParamDict], ResultDict]: + """Make a functional update on provided parameters. + + You can use this method to take more than one backward pass on the batch. All + configuration parameters for the iteration loop are set within the + `learner_config`. Note, the same configuration will be used for all module ids + in `MultiRLModule`. + + Args: + params: A parameter dictionary holding named parameters for each module id. + These parameters must be a clone of the module's original parameters to + perform a functional update on them. + training_data: A `TrainingData` instance containing the data or data iterators + to be used in updating the given parameters in `param`. + + Returns: + The functionally updated parameters in the (dict) format they were passed in + and a `ResultDict` object produced by a call to `self.metrics.reduce()`. The + returned dict may be arbitrarily nested and must have `Stats` objects at + all its leafs, allowing components further downstream (i.e. a user of this + Learner) to further reduce these results (for example over n parallel + Learners). + """ + self._check_is_built() + + # TODO (simon): Implement a `before_gradient_based_update`, if necessary. + # Call `before_gradient_based_update` to allow for non-gradient based + # preparations-, logging-, and update logic to happen. + # self.before_gradient_based_update(timesteps=timesteps or {}) + + training_data.validate() + training_data.solve_refs() + assert training_data.batches is None, "`training_data.batches` must be None!" + + self._weights_seq_no += 1 + + batch_iter = self._create_iterator_if_necessary( + training_data=training_data, + num_total_minibatches=self.learner_config.num_total_minibatches, + num_epochs=self.learner_config.num_epochs, + minibatch_size=self.learner_config.minibatch_size, + shuffle_batch_per_epoch=self.learner_config.shuffle_batch_per_epoch, + ) + + self.metrics.activate_tensor_mode() + + # Perform the actual looping through the minibatches or the given data iterator. + for iteration, tensor_minibatch in enumerate(batch_iter): + # Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs + # found in this batch. If not, throw an error. + unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set( + self.module.keys() + ) + if unknown_module_ids: + raise ValueError( + f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that " + f"are not in this Learner!" + ) + + # Make the actual in-graph/traced `_update` call. This should return + # all tensor values (no numpy). + fwd_out, loss_per_module, params, _ = self._update( + # TODO (simon): Maybe filter ParamDict by module keys in batch. + tensor_minibatch.policy_batches, + params, + ) + + # TODO (sven): Maybe move this into loop above to get metrics more accuratcely + # cover the minibatch/epoch logic. + # Log all timesteps (env, agent, modules) based on given episodes/batch. + self._log_steps_trained_metrics(tensor_minibatch) + + self._set_slicing_by_batch_id(tensor_minibatch, value=False) + + if self.iterator: + # Record the number of batches pulled from the dataset. + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED), + iteration + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED_LIFETIME), + iteration + 1, + reduce="sum", + ) + # Log all individual RLModules' loss terms and its registered optimizers' + # current learning rates. + # Note: We do this only once for the last of the minibatch updates, b/c the + # window is only 1 anyways. + for mid, loss in loss_per_module.items(): + self.metrics.log_value( + key=(mid, self.TOTAL_LOSS_KEY), + value=loss, + window=1, + ) + + # Call `after_gradient_based_update` to allow for non-gradient based + # cleanups-, logging-, and update logic to happen. + # TODO (simon): Check, if this should stay here, when running multiple + # gradient steps inside the iterator loop above (could be a complete epoch) + # the target networks might need to be updated earlier. + # self.after_gradient_based_update(timesteps=timesteps or {}) + self.metrics.deactivate_tensor_mode() + + # Reduce results across all minibatch update steps. + if not _no_metrics_reduce: + return params, loss_per_module, self.metrics.reduce() + else: + return params, loss_per_module, {} + + def _create_iterator_if_necessary( + self, + *, + training_data: TrainingData, + num_total_minibatches: int = 0, + num_epochs: int = 1, + minibatch_size: Optional[int] = None, + shuffle_batch_per_epoch: bool = False, + **kwargs, + ) -> Iterable: + """Provides a batch iterator.""" + + # Data iterator provided. + if training_data.data_iterators: + num_iters = kwargs.pop("num_iters", None) + if num_iters is None: + raise ValueError( + "Learner.update(data_iterators=..) requires `num_iters` kwarg!" + ) + + def _collate_fn(_batch: Dict[str, numpy.ndarray]) -> MultiAgentBatch: + _batch = unflatten_dict(_batch) + _batch = MultiAgentBatch( + { + module_id: SampleBatch(module_data) + for module_id, module_data in _batch.items() + }, + env_steps=sum( + len(next(iter(module_data.values()))) + for module_data in _batch.values() + ), + ) + _batch = self._convert_batch_type(_batch, to_device=False) + return self._set_slicing_by_batch_id(_batch, value=True) + + def _finalize_fn(batch: MultiAgentBatch) -> MultiAgentBatch: + return self._convert_batch_type(batch, to_device=True, use_stream=True) + + if not self.iterator: + # This iterator holds a `ray.data.DataIterator` and manages it state. + self.iterator = MiniBatchRayDataIterator( + iterator=training_data.data_iterators[0], + collate_fn=_collate_fn, + finalize_fn=_finalize_fn, + minibatch_size=minibatch_size, + num_iters=num_iters, + **kwargs, + ) + + batch_iter = self.iterator + else: + batch = self._make_batch_if_necessary(training_data=training_data) + assert batch is not None + + # TODO: Move this into LearnerConnector pipeline? + # Filter out those RLModules from the final train batch that should not be + # updated. + for module_id in list(batch.policy_batches.keys()): + if not self.should_module_be_updated(module_id, batch): + del batch.policy_batches[module_id] + if not batch.policy_batches: + return {} + + batch = self._set_slicing_by_batch_id(batch, value=True) + + if minibatch_size: + batch_iter_cls = MiniBatchCyclicIterator + elif num_epochs > 1: + # `minibatch_size` was not set but `num_epochs` > 1. + minibatch_size = batch.count + # Note that there is no need to shuffle here, b/c we don't have + # minibatches. + batch_iter_cls = MiniBatchCyclicIterator + else: + # `minibatch_size` and `num_epochs` are not set by the user. + batch_iter_cls = MiniBatchDummyIterator + + batch_iter = batch_iter_cls( + batch, + num_epochs=num_epochs, + minibatch_size=minibatch_size, + shuffle_batch_per_epoch=shuffle_batch_per_epoch and (num_epochs > 1), + num_total_minibatches=num_total_minibatches, + ) + return batch_iter + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def _update( + self, + batch: Dict[str, Any], + params: Dict[ModuleID, NamedParamDict], + **kwargs, + ) -> Tuple[Any, Dict[ModuleID, NamedParamDict], Any, Any]: + """Contains all logic for an in-graph/traceable functional update step. + + Framework specific subclasses must implement this method. This should include + functional calls to the RLModule's `forward_train`, `compute_loss`, + `compute_gradients`, `postprocess_gradients`, and `apply_gradients` methods + and return a tuple with all the individual results as well as the functionally + updated parameter dictionary. + + Args: + batch: The train batch already converted to a Dict mapping str to (possibly + nested) tensors. + kwargs: Forward compatibility kwargs. + + Returns: + A tuple consisting of: + 1) the output of a functional forward call to the RLModule using + `params`, + 2) the functionally updated parameters in the (dict) format passed in, + 3) the `loss_per_module` dictionary mapping module IDs to individual + loss tensors, + 4) a metrics dict mapping module IDs to metrics key/value pairs. + + """ + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + """Gets the state of the `DifferentiableLearner` instance. + + Note, because the `MultiRLModule` held by this class is only a reference + it is not contained in the class' state. + """ + self._check_is_built() + + state = { + "should_module_be_updated": self.config.policies_to_train, + WEIGHTS_SEQ_NO: self._weights_seq_no, + } + + if self._check_component(COMPONENT_METRICS_LOGGER, components, not_components): + # TODO (sven): Make `MetricsLogger` a Checkpointable. + state[COMPONENT_METRICS_LOGGER] = self.metrics.get_state() + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + """Sets the state of the `DifferentiableLearner` instance.""" + self._check_is_built() + + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update our trainable Modules information/function via our config. + # If not provided in state (None), all Modules will be trained by default. + if "should_module_be_updated" in state: + self.config.multi_agent(policies_to_train=state["should_module_be_updated"]) + + # TODO (sven): Make `MetricsLogger` a Checkpointable. + if COMPONENT_METRICS_LOGGER in state: + self.metrics.set_state(state[COMPONENT_METRICS_LOGGER]) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args, + { + "config": self.config, + }, # **kwargs + ) + + @override(Checkpointable) + def get_checkpointable_components(self): + if not self._check_is_built(error=False): + self.build() + return [] + + # TODO (simon): Duplicate in Learner. Move to base class "Learnable". + def _make_batch_if_necessary(self, training_data): + batch = training_data.batch + + # Call the learner connector on the given `episodes` (if we have one). + if training_data.episodes is not None: + # If we want to learn from Episodes, we must have a LearnerConnector + # pipeline to translate into a train batch first. + if self._learner_connector is None: + raise ValueError( + f"If episodes provided for training, Learner ({self}) must have a " + "LearnerConnector pipeline (but pipeline is None)!" + ) + + # Call the learner connector pipeline. + shared_data = {} + batch = self._learner_connector( + rl_module=self.module, + batch=training_data.batch if training_data.batch is not None else {}, + episodes=training_data.episodes, + shared_data=shared_data, + metrics=self.metrics, + ) + # Convert to a batch. + # TODO (sven): Try to not require MultiAgentBatch anymore. + batch = MultiAgentBatch( + { + module_id: ( + SampleBatch(module_data, _zero_padded=True) + if shared_data.get(f"_zero_padded_for_mid={module_id}") + else SampleBatch(module_data) + ) + for module_id, module_data in batch.items() + }, + env_steps=sum(len(e) for e in training_data.episodes), + ) + # Single-agent SampleBatch: Have to convert to MultiAgentBatch. + elif isinstance(training_data.batch, SampleBatch): + if len(self.module) != 1: + raise ValueError( + f"SampleBatch provided, but RLModule ({self.module}) has more than " + f"one sub-RLModule! Need to provide MultiAgentBatch instead." + ) + + batch = MultiAgentBatch( + {next(iter(self.module.keys())): training_data.batch}, + env_steps=len(training_data.batch), + ) + # If we already have an `MultiAgentBatch` but with `numpy` array, convert to + # tensors. + elif ( + isinstance(training_data.batch, MultiAgentBatch) + and training_data.batch.policy_batches + and ( + isinstance( + next(iter(training_data.batch.policy_batches.values()))["obs"], + numpy.ndarray, + ) + or next(iter(training_data.batch.policy_batches.values()))["obs"].device + != self._device + ) + ): + batch = self._convert_batch_type(training_data.batch) + + return batch + + # TODO (simon): Duplicate in Learner. Move to base class "Learnable". + def _set_slicing_by_batch_id( + self, batch: MultiAgentBatch, *, value: bool + ) -> MultiAgentBatch: + """Enables slicing by batch id in the given batch. + + If the input batch contains batches of sequences we need to make sure when + slicing happens it is sliced via batch id and not timestamp. Calling this + method enables the same flag on each SampleBatch within the input + MultiAgentBatch. + + Args: + batch: The MultiAgentBatch to enable slicing by batch id on. + value: The value to set the flag to. + + Returns: + The input MultiAgentBatch with the indexing flag is enabled / disabled on. + """ + + for pid, policy_batch in batch.policy_batches.items(): + # We assume that arriving batches for recurrent modules OR batches that + # have a SEQ_LENS column are already zero-padded to the max sequence length + # and have tensors of shape [B, T, ...]. Therefore, we slice sequence + # lengths in B. See SampleBatch for more information. + if ( + self.module[pid].is_stateful() + or policy_batch.get("seq_lens") is not None + ): + if value: + policy_batch.enable_slicing_by_batch_id() + else: + policy_batch.disable_slicing_by_batch_id() + + return batch + + # TODO (simon): Duplicate in Learner. Move to base class "Learnable". + def _check_is_built(self, error: bool = True) -> bool: + if self.module is None: + if error: + raise ValueError( + "Learner.build() must be called after constructing a " + "Learner and before calling any methods on it." + ) + return False + return True + + @abc.abstractmethod + def _get_tensor_variable( + self, + value: Any, + dtype: Any = None, + trainable: bool = False, + ) -> TensorType: + """Returns a framework-specific tensor variable with the initial given value. + + This is a framework specific method that should be implemented by the + framework specific sub-classes. + + Args: + value: The initial value for the tensor variable variable. + + Returns: + The framework specific tensor variable of the given initial value, + dtype and trainable/requires_grad property. + """ + + # TODO (simon): Duplicate in Learner. Move to base class "Learnable". + def _reset(self): + self.metrics = MetricsLogger() + self._is_built = False + + # TODO (simon): Duplicate in Learner. Move to base class "Learnable". + def _log_steps_trained_metrics(self, batch: MultiAgentBatch): + """Logs this iteration's steps trained, based on given `batch`.""" + # Collect all module steps and add them for `ALL_MODULES` to avoid + # biasing the throughput by looping through modules. + total_module_steps = 0 + # Loop through all modules. + for mid, module_batch in batch.policy_batches.items(): + # Log weights seq no for this batch. + self.metrics.log_value( + (mid, WEIGHTS_SEQ_NO), + self._weights_seq_no, + window=1, + ) + + module_batch_size = len(module_batch) + # Log average batch size (for each module). + self.metrics.log_value( + key=(mid, MODULE_TRAIN_BATCH_SIZE_MEAN), + value=module_batch_size, + ) + # Log module steps (for each module). + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_TRAINED), + value=module_batch_size, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_TRAINED_LIFETIME), + value=module_batch_size, + reduce="sum", + with_throughput=True, + ) + total_module_steps += module_batch_size + + # Log module steps (sum of all modules). + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED), + value=total_module_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED_LIFETIME), + value=total_module_steps, + reduce="sum", + with_throughput=True, + ) + # Log env steps (all modules). + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_TRAINED), + batch.env_steps(), + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME), + batch.env_steps(), + reduce="sum", + with_throughput=True, + ) + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def _make_functional_call( + self, + module: MultiRLModule, + params: Dict[ModuleID, NamedParamDict], + batch: MultiAgentBatch, + **kwargs, + ) -> Dict[ModuleID, Dict[str, TensorType]]: + """Makes a functional call to a module. + + Functional calls enable the Learner to (a) use the same module as its + `MetaLearner` and (b) to generate and apply gradients without modifying + the `RLModule` parameters directly. + + Args: + module: The `MultiRLModule` to be used for the functional call. Note, this + module's `forward` method must call the `foward_train`. + params: A dictionary containing containing for each `RLModule`'s id its + named parameter dictionary. For functional calls to work, these + parameters need to be cloned. + batch: A `MultiAgentBatch` instance to be used in the functional call. + + Returns: + A dictionary with the output of the module's forward pass. + """ + + @property + def module(self) -> MultiRLModule: + """The MultiRLModule reference that is being used in updates.""" + return self._module + + @module.setter + def module(self, module: MultiRLModule) -> None: + """Sets the `MultiRLModule`. + + Args: + module: The reference to the `MultiRLModule` of the class that holds the + instance of this `DifferentiableLearner` instance. + """ + self._module = module diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner_config.py b/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner_config.py new file mode 100644 index 0000000000000000000000000000000000000000..a9629cefe71559b47489a7f618ecb42f1b78e5e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/differentiable_learner_config.py @@ -0,0 +1,148 @@ +from dataclasses import dataclass, fields +from typing import Callable, List, Optional, Union + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.learner.differentiable_learner import DifferentiableLearner +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.typing import DeviceType, ModuleID + + +@dataclass +class DifferentiableLearnerConfig: + """Configures a `DifferentiableLearner`.""" + + # TODO (simon): We implement only for `PyTorch`, so maybe we use here directly + # TorchDifferentiableLearner` and check for this? + # The `DifferentiableLearner` class. Must be derived from `DifferentiableLearner`. + learner_class: Callable + + learner_connector: Optional[ + Callable[["RLModule"], Union["ConnectorV2", List["ConnectorV2"]]] + ] = None + + add_default_connectors_to_learner_pipeline: bool = True + + is_multi_agent: bool = False + + policies_to_update: List[ModuleID] = None + + # The learning rate to use for the nested update. Note, in the default case this + # learning rate is only used to update parameters in a functional form, i.e. the + # `RLModule`'s stateful parameters are only updated in the `MetaLearner`. Different + # logic can be implemented in customized `DifferentiableLearner`s. + lr: float = 3e-5 + + # TODO (simon): Add further hps like clip_grad, ... + # The total number of minibatches to be formed from the batch per learner, e.g. + # setting `train_batch_size_per_learner=10` and `num_total_minibatches` to 2 + # runs 2 SGD minibatch updates with a batch of 5 per training iteration. + num_total_minibatches: int = 0 + + # The number of epochs per training iteration. + num_epochs: int = 1 + + # The minibatch size per SGD minibatch update, e.g. with a `train_batch_size_per_learner=10` + # and a `minibatch_size=2` the training step runs 5 SGD minibatch updates with minibatches + # of 2. + minibatch_size: int = None + + # If the batch should be shuffled between epochs. + shuffle_batch_per_epoch: bool = False + + def __post_init__(self): + """Additional initialization processes.""" + + # Ensure we have a `DifferentiableLearner` class. + if not issubclass(self.learner_class, DifferentiableLearner): + raise ValueError( + "`learner_class` must be a subclass of `DifferentiableLearner " + f"but is {self.learner_class}." + ) + + def build_learner_connector( + self, + input_observation_space: Optional[gym.spaces.Space], + input_action_space: Optional[gym.spaces.Space], + device: Optional[DeviceType] = None, + ): + from ray.rllib.connectors.learner import ( + AddColumnsFromEpisodesToTrainBatch, + AddObservationsFromEpisodesToBatch, + AddStatesFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AgentToModuleMapping, + BatchIndividualItems, + LearnerConnectorPipeline, + NumpyToTensor, + ) + + custom_connectors = [] + # Create a learner connector pipeline (including RLlib's default + # learner connector piece) and return it. + if self.learner_connector is not None: + val_ = self.learner_connector( + input_observation_space, + input_action_space, + # device, # TODO (sven): Also pass device into custom builder. + ) + + from ray.rllib.connectors.connector_v2 import ConnectorV2 + + # ConnectorV2 (piece or pipeline). + if isinstance(val_, ConnectorV2): + custom_connectors = [val_] + # Sequence of individual ConnectorV2 pieces. + elif isinstance(val_, (list, tuple)): + custom_connectors = list(val_) + # Unsupported return value. + else: + raise ValueError( + "`AlgorithmConfig.training(learner_connector=..)` must return " + "a ConnectorV2 object or a list thereof (to be added to a " + f"pipeline)! Your function returned {val_}." + ) + + pipeline = LearnerConnectorPipeline( + connectors=custom_connectors, + input_observation_space=input_observation_space, + input_action_space=input_action_space, + ) + if self.add_default_connectors_to_learner_pipeline: + # Append OBS handling. + pipeline.append( + AddObservationsFromEpisodesToBatch(as_learner_connector=True) + ) + # Append all other columns handling. + pipeline.append(AddColumnsFromEpisodesToTrainBatch()) + # Append time-rank handler. + pipeline.append(AddTimeDimToBatchAndZeroPad(as_learner_connector=True)) + # Append STATE_IN/STATE_OUT handler. + pipeline.append(AddStatesFromEpisodesToBatch(as_learner_connector=True)) + # If multi-agent -> Map from AgentID-based data to ModuleID based data. + if self.is_multi_agent: + pipeline.append( + AgentToModuleMapping( + rl_module_specs=( + self.rl_module_spec.rl_module_specs + if isinstance(self.rl_module_spec, MultiRLModuleSpec) + else set(self.policies) + ), + agent_to_module_mapping_fn=self.policy_mapping_fn, + ) + ) + # Batch all data. + pipeline.append(BatchIndividualItems(multi_agent=self.is_multi_agent)) + # Convert to Tensors. + pipeline.append(NumpyToTensor(as_learner_connector=True, device=device)) + return pipeline + + def update_from_kwargs(self, **kwargs): + """Sets all slots with values defined in `kwargs`.""" + # Get all field names (i.e., slot names). + field_names = {f.name for f in fields(self)} + for key, value in kwargs.items(): + if key in field_names: + setattr(self, key, value) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/learner.py b/lib/python3.12/site-packages/ray/rllib/core/learner/learner.py new file mode 100644 index 0000000000000000000000000000000000000000..c1548703b1fba722b80bdc4cd5598c9682620857 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/learner.py @@ -0,0 +1,1765 @@ +import abc +import copy +import logging +import platform +from collections import defaultdict +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Dict, + Hashable, + Iterable, + List, + Optional, + Sequence, + Tuple, + Union, +) + +import numpy +import tree + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.connectors.learner.learner_connector_pipeline import ( + LearnerConnectorPipeline, +) +from ray.rllib.core import ( + COMPONENT_METRICS_LOGGER, + COMPONENT_OPTIMIZER, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI +from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, +) +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.debug import update_global_seed_if_necessary +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + ALL_MODULES, + DATASET_NUM_ITERS_TRAINED, + DATASET_NUM_ITERS_TRAINED_LIFETIME, + MODULE_TRAIN_BATCH_SIZE_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_TRAINED, + NUM_ENV_STEPS_TRAINED_LIFETIME, + NUM_MODULE_STEPS_TRAINED, + NUM_MODULE_STEPS_TRAINED_LIFETIME, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.metrics.ray_metrics import ( + DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + TimerAndPrometheusLogger, +) +from ray.rllib.utils.minibatch_utils import ( + MiniBatchCyclicIterator, + MiniBatchDummyIterator, + MiniBatchRayDataIterator, +) +from ray.rllib.utils.schedules.scheduler import Scheduler +from ray.rllib.utils.typing import ( + EpisodeType, + LearningRateOrSchedule, + ModuleID, + Optimizer, + Param, + ParamDict, + ParamRef, + ResultDict, + ShouldModuleBeUpdatedFn, + StateDict, + TensorType, +) +from ray.util.annotations import PublicAPI +from ray.util.metrics import Counter, Histogram + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + +DEFAULT_OPTIMIZER = "default_optimizer" + +# COMMON LEARNER LOSS_KEYS +POLICY_LOSS_KEY = "policy_loss" +VF_LOSS_KEY = "vf_loss" +ENTROPY_KEY = "entropy" + +# Additional update keys +LR_KEY = "learning_rate" + + +@PublicAPI(stability="alpha") +class Learner(Checkpointable): + """Base class for Learners. + + This class will be used to train RLModules. It is responsible for defining the loss + function, and updating the neural network weights that it owns. It also provides a + way to add/remove modules to/from RLModules in a multi-agent scenario, in the + middle of training (This is useful for league based training). + + Deep learning framework-specific implementations of this class fill in the + details for distributed training, and for computing and applying + gradients. User should not need to subclass this class, but instead inherit from + the deep learning framework (for example torch) specific subclasses to implement + their algorithm-specific update logic. + + Args: + config: The AlgorithmConfig object from which to derive most of the settings + needed to build the Learner. + module_spec: The module specification for the RLModule that is being trained. + If the module is a single agent module, after building the module it will + be converted to a multi-agent module with a default key. Can be none if the + module is provided directly via the `module` argument. Refer to + ray.rllib.core.rl_module.RLModuleSpec + or ray.rllib.core.rl_module.MultiRLModuleSpec for more info. + module: If learner is being used stand-alone, the RLModule can be optionally + passed in directly instead of the through the `module_spec`. + + Note: We use PPO and torch as an example here because many of the showcased + components need implementations to come together. However, the same + pattern is generally applicable. + + .. testcode:: + + import gymnasium as gym + + from ray.rllib.algorithms.ppo.ppo import PPOConfig + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import ( + PPOTorchRLModule + ) + from ray.rllib.core import COMPONENT_RL_MODULE, DEFAULT_MODULE_ID + from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + + env = gym.make("CartPole-v1") + + # Create a PPO config object first. + config = ( + PPOConfig() + .framework("torch") + .training(model={"fcnet_hiddens": [128, 128]}) + ) + + # Create a learner instance directly from our config. All we need as + # extra information here is the env to be able to extract space information + # (needed to construct the RLModule inside the Learner). + learner = config.build_learner(env=env) + + # Take one gradient update on the module and report the results. + # results = learner.update(...) + + # Add a new module, perhaps for league based training. + learner.add_module( + module_id="new_player", + module_spec=RLModuleSpec( + module_class=PPOTorchRLModule, + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[64, 64]), + catalog_class=PPOCatalog, + ) + ) + + # Take another gradient update with both previous and new modules. + # results = learner.update(...) + + # Remove a module. + learner.remove_module("new_player") + + # Will train previous modules only. + # results = learner.update(...) + + # Get the state of the learner. + state = learner.get_state() + + # Set the state of the learner. + learner.set_state(state) + + # Get the weights of the underlying MultiRLModule. + weights = learner.get_state(components=COMPONENT_RL_MODULE) + + # Set the weights of the underlying MultiRLModule. + learner.set_state({COMPONENT_RL_MODULE: weights}) + + + Extension pattern: + + .. testcode:: + + from ray.rllib.core.learner.torch.torch_learner import TorchLearner + + class MyLearner(TorchLearner): + + def compute_losses(self, fwd_out, batch): + # Compute the losses per module based on `batch` and output of the + # forward pass (`fwd_out`). To access the (algorithm) config for a + # specific RLModule, do: + # `self.config.get_config_for_module([moduleID])`. + return {DEFAULT_MODULE_ID: module_loss} + """ + + framework: str = None + TOTAL_LOSS_KEY: str = "total_loss" + + def __init__( + self, + *, + config: "AlgorithmConfig", + module_spec: Optional[Union[RLModuleSpec, MultiRLModuleSpec]] = None, + module: Optional[RLModule] = None, + ): + # TODO (sven): Figure out how to do this + self.config = config.copy(copy_frozen=False) + self._module_spec: Optional[MultiRLModuleSpec] = module_spec + self._module_obj: Optional[MultiRLModule] = module + + # Make node and device of this Learner available. + self._node = platform.node() + self._device = None + + # Set a seed, if necessary. + if self.config.seed is not None: + update_global_seed_if_necessary(self.framework, self.config.seed) + + # Whether self.build has already been called. + self._is_built = False + + # Attributes to be set separately (not by user's custom `build()` code) + # by a LearnerGroup. + self._learner_index = 0 + self._placement_group = None + + # These are the attributes that are set during build. + + # The actual MultiRLModule used by this Learner. + self._module: Optional[MultiRLModule] = None + self._weights_seq_no = 0 + # Our Learner connector pipeline. + self._learner_connector: Optional[LearnerConnectorPipeline] = None + # These are set for properly applying optimizers and adding or removing modules. + self._optimizer_parameters: Dict[Optimizer, List[ParamRef]] = {} + self._named_optimizers: Dict[str, Optimizer] = {} + self._params: ParamDict = {} + # Dict mapping ModuleID to a list of optimizer names. Note that the optimizer + # name includes the ModuleID as a prefix: optimizer_name=`[ModuleID]_[.. rest]`. + self._module_optimizers: Dict[ModuleID, List[str]] = defaultdict(list) + self._optimizer_name_to_module: Dict[str, ModuleID] = {} + + # Only manage optimizer's learning rate if user has NOT overridden + # the `configure_optimizers_for_module` method. Otherwise, leave responsibility + # to handle lr-updates entirely in user's hands. + self._optimizer_lr_schedules: Dict[Optimizer, Scheduler] = {} + + # The Learner's own MetricsLogger to be used to log RLlib's built-in metrics or + # custom user-defined ones (e.g. custom loss values). When returning from an + # `update_from_...()` method call, the Learner will do a `self.metrics.reduce()` + # and return the resulting (reduced) dict. + self.metrics = MetricsLogger() + + # In case of offline learning and multiple learners, each learner receives a + # repeatable iterator that iterates over a split of the streamed data. + self.iterator: MiniBatchRayDataIterator = None + + # Ray metrics + self._metrics_all_modules_num_env_steps_trained = Counter( + name="rllib_learner_all_modules_num_env_steps_trained_counter", + description="Number of env steps trained (sum over all modules).", + tag_keys=("rllib",), + ) + self._metrics_all_modules_num_env_steps_trained.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_all_modules_num_module_steps_trained = Counter( + name="rllib_learner_all_modules_num_module_steps_trained_counter", + description="Number of module steps trained (sum over all modules).", + tag_keys=("rllib",), + ) + self._metrics_all_modules_num_module_steps_trained.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_learner_inner_update = Histogram( + name="rllib_learner_update_inner_update_time", + description="Duration of the Learner's inner update.", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_learner_inner_update.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + # TODO (sven): Do we really need this API? It seems like LearnerGroup constructs + # all Learner workers and then immediately builds them any ways? Unless there is + # a reason related to Train worker group setup. + @OverrideToImplementCustomLogic_CallToSuperRecommended + def build(self) -> None: + """Builds the Learner. + + This method should be called before the learner is used. It is responsible for + setting up the LearnerConnectorPipeline, the RLModule, optimizer(s), and + (optionally) the optimizers' learning rate schedulers. + """ + if self._is_built: + logger.debug("Learner already built. Skipping build.") + return + + # Build learner connector pipeline used on this Learner worker. + self._learner_connector = None + # If the Algorithm uses aggregation actors to run episodes through the learner + # connector, its Learners don't need a connector pipelines and instead learn + # directly from pre-loaded batches already on the GPU. + if self.config.num_aggregator_actors_per_learner == 0: + # TODO (sven): Figure out which space to provide here. For now, + # it doesn't matter, as the default connector piece doesn't use + # this information anyway. + # module_spec = self._module_spec.as_multi_rl_module_spec() + self._learner_connector = self.config.build_learner_connector( + input_observation_space=None, + input_action_space=None, + device=self._device, + ) + + # Build the module to be trained by this learner. + self._module = self._make_module() + + # Configure, construct, and register all optimizers needed to train + # `self.module`. + self.configure_optimizers() + + # Log the number of trainable/non-trainable parameters. + self._log_trainable_parameters() + self._is_built = True + + @property + def distributed(self) -> bool: + """Whether the learner is running in distributed mode.""" + return self.config.num_learners > 1 + + @property + def module(self) -> MultiRLModule: + """The MultiRLModule that is being trained.""" + return self._module + + @property + def node(self) -> Any: + return self._node + + @property + def device(self) -> Any: + return self._device + + def register_optimizer( + self, + *, + module_id: ModuleID = ALL_MODULES, + optimizer_name: str = DEFAULT_OPTIMIZER, + optimizer: Optimizer, + params: Sequence[Param], + lr_or_lr_schedule: Optional[LearningRateOrSchedule] = None, + ) -> None: + """Registers an optimizer with a ModuleID, name, param list and lr-scheduler. + + Use this method in your custom implementations of either + `self.configure_optimizers()` or `self.configure_optimzers_for_module()` (you + should only override one of these!). If you register a learning rate Scheduler + setting together with an optimizer, RLlib will automatically keep this + optimizer's learning rate updated throughout the training process. + Alternatively, you can construct your optimizers directly with a learning rate + and manage learning rate scheduling or updating yourself. + + Args: + module_id: The `module_id` under which to register the optimizer. If not + provided, will assume ALL_MODULES. + optimizer_name: The name (str) of the optimizer. If not provided, will + assume DEFAULT_OPTIMIZER. + optimizer: The already instantiated optimizer object to register. + params: A list of parameters (framework-specific variables) that will be + trained/updated + lr_or_lr_schedule: An optional fixed learning rate or learning rate schedule + setup. If provided, RLlib will automatically keep the optimizer's + learning rate updated. + """ + # Validate optimizer instance and its param list. + self._check_registered_optimizer(optimizer, params) + + full_registration_name = module_id + "_" + optimizer_name + + # Store the given optimizer under the given `module_id`. + self._module_optimizers[module_id].append(full_registration_name) + self._optimizer_name_to_module[full_registration_name] = module_id + + # Store the optimizer instance under its full `module_id`_`optimizer_name` + # key. + self._named_optimizers[full_registration_name] = optimizer + + # Store all given parameters under the given optimizer. + self._optimizer_parameters[optimizer] = [] + for param in params: + param_ref = self.get_param_ref(param) + self._optimizer_parameters[optimizer].append(param_ref) + self._params[param_ref] = param + + # Optionally, store a scheduler object along with this optimizer. If such a + # setting is provided, RLlib will handle updating the optimizer's learning rate + # over time. + if lr_or_lr_schedule is not None: + # Validate the given setting. + Scheduler.validate( + fixed_value_or_schedule=lr_or_lr_schedule, + setting_name="lr_or_lr_schedule", + description="learning rate or schedule", + ) + # Create the scheduler object for this optimizer. + scheduler = Scheduler( + fixed_value_or_schedule=lr_or_lr_schedule, + framework=self.framework, + device=self._device, + ) + self._optimizer_lr_schedules[optimizer] = scheduler + # Set the optimizer to the current (first) learning rate. + self._set_optimizer_lr( + optimizer=optimizer, + lr=scheduler.get_current_value(), + ) + + @OverrideToImplementCustomLogic + def configure_optimizers(self) -> None: + """Configures, creates, and registers the optimizers for this Learner. + + Optimizers are responsible for updating the model's parameters during training, + based on the computed gradients. + + Normally, you should not override this method for your custom algorithms + (which require certain optimizers), but rather override the + `self.configure_optimizers_for_module(module_id=..)` method and register those + optimizers in there that you need for the given `module_id`. + + You can register an optimizer for any RLModule within `self.module` (or for + the ALL_MODULES ID) by calling `self.register_optimizer()` and passing the + module_id, optimizer_name (only in case you would like to register more than + one optimizer for a given module), the optimizer instane itself, a list + of all the optimizer's parameters (to be updated by the optimizer), and + an optional learning rate or learning rate schedule setting. + + This method is called once during building (`self.build()`). + """ + # The default implementation simply calls `self.configure_optimizers_for_module` + # on each RLModule within `self.module`. + for module_id in self.module.keys(): + if self.rl_module_is_compatible(self.module[module_id]): + config = self.config.get_config_for_module(module_id) + self.configure_optimizers_for_module(module_id=module_id, config=config) + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def configure_optimizers_for_module( + self, module_id: ModuleID, config: "AlgorithmConfig" = None + ) -> None: + """Configures an optimizer for the given module_id. + + This method is called for each RLModule in the MultiRLModule being + trained by the Learner, as well as any new module added during training via + `self.add_module()`. It should configure and construct one or more optimizers + and register them via calls to `self.register_optimizer()` along with the + `module_id`, an optional optimizer name (str), a list of the optimizer's + framework specific parameters (variables), and an optional learning rate value + or -schedule. + + Args: + module_id: The module_id of the RLModule that is being configured. + config: The AlgorithmConfig specific to the given `module_id`. + """ + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def compute_gradients( + self, loss_per_module: Dict[ModuleID, TensorType], **kwargs + ) -> ParamDict: + """Computes the gradients based on the given losses. + + Args: + loss_per_module: Dict mapping module IDs to their individual total loss + terms, computed by the individual `compute_loss_for_module()` calls. + The overall total loss (sum of loss terms over all modules) is stored + under `loss_per_module[ALL_MODULES]`. + **kwargs: Forward compatibility kwargs. + + Returns: + The gradients in the same (flat) format as self._params. Note that all + top-level structures, such as module IDs, will not be present anymore in + the returned dict. It will merely map parameter tensor references to their + respective gradient tensors. + """ + + @OverrideToImplementCustomLogic + def postprocess_gradients(self, gradients_dict: ParamDict) -> ParamDict: + """Applies potential postprocessing operations on the gradients. + + This method is called after gradients have been computed and modifies them + before they are applied to the respective module(s) by the optimizer(s). + This might include grad clipping by value, norm, or global-norm, or other + algorithm specific gradient postprocessing steps. + + This default implementation calls `self.postprocess_gradients_for_module()` + on each of the sub-modules in our MultiRLModule: `self.module` and + returns the accumulated gradients dicts. + + Args: + gradients_dict: A dictionary of gradients in the same (flat) format as + self._params. Note that top-level structures, such as module IDs, + will not be present anymore in this dict. It will merely map gradient + tensor references to gradient tensors. + + Returns: + A dictionary with the updated gradients and the exact same (flat) structure + as the incoming `gradients_dict` arg. + """ + + # The flat gradients dict (mapping param refs to params), returned by this + # method. + postprocessed_gradients = {} + + for module_id in self.module.keys(): + # Send a gradients dict for only this `module_id` to the + # `self.postprocess_gradients_for_module()` method. + module_grads_dict = {} + for optimizer_name, optimizer in self.get_optimizers_for_module(module_id): + optim_grads = self.filter_param_dict_for_optimizer( + gradients_dict, optimizer + ) + for ref, grad in optim_grads.items(): + assert ref not in module_grads_dict + module_grads_dict[ref] = grad + + module_grads_dict = self.postprocess_gradients_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id), + module_gradients_dict=module_grads_dict, + ) + assert isinstance(module_grads_dict, dict) + + # Update our return dict. + postprocessed_gradients.update(module_grads_dict) + + return postprocessed_gradients + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def postprocess_gradients_for_module( + self, + *, + module_id: ModuleID, + config: Optional["AlgorithmConfig"] = None, + module_gradients_dict: ParamDict, + ) -> ParamDict: + """Applies postprocessing operations on the gradients of the given module. + + Args: + module_id: The module ID for which we will postprocess computed gradients. + Note that `module_gradients_dict` already only carries those gradient + tensors that belong to this `module_id`. Other `module_id`'s gradients + are not available in this call. + config: The AlgorithmConfig specific to the given `module_id`. + module_gradients_dict: A dictionary of gradients in the same (flat) format + as self._params, mapping gradient refs to gradient tensors, which are to + be postprocessed. You may alter these tensors in place or create new + ones and return these in a new dict. + + Returns: + A dictionary with the updated gradients and the exact same (flat) structure + as the incoming `module_gradients_dict` arg. + """ + postprocessed_grads = {} + + if config.grad_clip is None and not config.log_gradients: + postprocessed_grads.update(module_gradients_dict) + return postprocessed_grads + + for optimizer_name, optimizer in self.get_optimizers_for_module(module_id): + grad_dict_to_clip = self.filter_param_dict_for_optimizer( + param_dict=module_gradients_dict, + optimizer=optimizer, + ) + if config.grad_clip: + # Perform gradient clipping, if configured. + global_norm = self._get_clip_function()( + grad_dict_to_clip, + grad_clip=config.grad_clip, + grad_clip_by=config.grad_clip_by, + ) + if config.grad_clip_by == "global_norm" or config.log_gradients: + # If we want to log gradients, but do not use the global norm + # for clipping compute it here. + if config.log_gradients and config.grad_clip_by != "global_norm": + # Compute the global norm of gradients. + global_norm = self._get_global_norm_function()( + # Note, `tf.linalg.global_norm` needs a list of tensors. + list(grad_dict_to_clip.values()), + ) + self.metrics.log_value( + key=(module_id, f"gradients_{optimizer_name}_global_norm"), + value=global_norm, + window=1, + ) + postprocessed_grads.update(grad_dict_to_clip) + # In the other case check, if we want to log gradients only. + elif config.log_gradients: + # Compute the global norm of gradients and log it. + global_norm = self._get_global_norm_function()( + # Note, `tf.linalg.global_norm` needs a list of tensors. + list(grad_dict_to_clip.values()), + ) + self.metrics.log_value( + key=(module_id, f"gradients_{optimizer_name}_global_norm"), + value=global_norm, + window=1, + ) + + return postprocessed_grads + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def apply_gradients(self, gradients_dict: ParamDict) -> None: + """Applies the gradients to the MultiRLModule parameters. + + Args: + gradients_dict: A dictionary of gradients in the same (flat) format as + self._params. Note that top-level structures, such as module IDs, + will not be present anymore in this dict. It will merely map gradient + tensor references to gradient tensors. + """ + + def get_optimizer( + self, + module_id: ModuleID = DEFAULT_MODULE_ID, + optimizer_name: str = DEFAULT_OPTIMIZER, + ) -> Optimizer: + """Returns the optimizer object, configured under the given module_id and name. + + If only one optimizer was registered under `module_id` (or ALL_MODULES) + via the `self.register_optimizer` method, `optimizer_name` is assumed to be + DEFAULT_OPTIMIZER. + + Args: + module_id: The ModuleID for which to return the configured optimizer. + If not provided, will assume DEFAULT_MODULE_ID. + optimizer_name: The name of the optimizer (registered under `module_id` via + `self.register_optimizer()`) to return. If not provided, will assume + DEFAULT_OPTIMIZER. + + Returns: + The optimizer object, configured under the given `module_id` and + `optimizer_name`. + """ + # `optimizer_name` could possibly be the full optimizer name (including the + # module_id under which it is registered). + if optimizer_name in self._named_optimizers: + return self._named_optimizers[optimizer_name] + + # Normally, `optimizer_name` is just the optimizer's name, not including the + # `module_id`. + full_registration_name = module_id + "_" + optimizer_name + if full_registration_name in self._named_optimizers: + return self._named_optimizers[full_registration_name] + + # No optimizer found. + raise KeyError( + f"Optimizer not found! module_id={module_id} " + f"optimizer_name={optimizer_name}" + ) + + def get_optimizers_for_module( + self, module_id: ModuleID = ALL_MODULES + ) -> List[Tuple[str, Optimizer]]: + """Returns a list of (optimizer_name, optimizer instance)-tuples for module_id. + + Args: + module_id: The ModuleID for which to return the configured + (optimizer name, optimizer)-pairs. If not provided, will return + optimizers registered under ALL_MODULES. + + Returns: + A list of tuples of the format: ([optimizer_name], [optimizer object]), + where optimizer_name is the name under which the optimizer was registered + in `self.register_optimizer`. If only a single optimizer was + configured for `module_id`, [optimizer_name] will be DEFAULT_OPTIMIZER. + """ + named_optimizers = [] + for full_registration_name in self._module_optimizers[module_id]: + optimizer = self._named_optimizers[full_registration_name] + # TODO (sven): How can we avoid registering optimziers under this + # constructed `[module_id]_[optim_name]` format? + optim_name = full_registration_name[len(module_id) + 1 :] + named_optimizers.append((optim_name, optimizer)) + return named_optimizers + + def filter_param_dict_for_optimizer( + self, param_dict: ParamDict, optimizer: Optimizer + ) -> ParamDict: + """Reduces the given ParamDict to contain only parameters for given optimizer. + + Args: + param_dict: The ParamDict to reduce/filter down to the given `optimizer`. + The returned dict will be a subset of `param_dict` only containing keys + (param refs) that were registered together with `optimizer` (and thus + that `optimizer` is responsible for applying gradients to). + optimizer: The optimizer object to whose parameter refs the given + `param_dict` should be reduced. + + Returns: + A new ParamDict only containing param ref keys that belong to `optimizer`. + """ + # Return a sub-dict only containing those param_ref keys (and their values) + # that belong to the `optimizer`. + return { + ref: param_dict[ref] + for ref in self._optimizer_parameters[optimizer] + if ref in param_dict and param_dict[ref] is not None + } + + @abc.abstractmethod + def get_param_ref(self, param: Param) -> Hashable: + """Returns a hashable reference to a trainable parameter. + + This should be overridden in framework specific specialization. For example in + torch it will return the parameter itself, while in tf it returns the .ref() of + the variable. The purpose is to retrieve a unique reference to the parameters. + + Args: + param: The parameter to get the reference to. + + Returns: + A reference to the parameter. + """ + + @abc.abstractmethod + def get_parameters(self, module: RLModule) -> Sequence[Param]: + """Returns the list of parameters of a module. + + This should be overridden in framework specific learner. For example in torch it + will return .parameters(), while in tf it returns .trainable_variables. + + Args: + module: The module to get the parameters from. + + Returns: + The parameters of the module. + """ + + @abc.abstractmethod + def _convert_batch_type( + self, + batch: MultiAgentBatch, + to_device: bool = False, + pin_memory: bool = False, + use_stream: bool = False, + ) -> MultiAgentBatch: + """Converts the elements of a MultiAgentBatch to Tensors on the correct device. + + Args: + batch: The MultiAgentBatch object to convert. + + Returns: + The resulting MultiAgentBatch with framework-specific tensor values placed + on the correct device. + """ + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def add_module( + self, + *, + module_id: ModuleID, + module_spec: RLModuleSpec, + config_overrides: Optional[Dict] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + """Adds a module to the underlying MultiRLModule. + + Changes this Learner's config in order to make this architectural change + permanent wrt. to checkpointing. + + Args: + module_id: The ModuleID of the module to be added. + module_spec: The ModuleSpec of the module to be added. + config_overrides: The `AlgorithmConfig` overrides that should apply to + the new Module, if any. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + + Returns: + The new MultiRLModuleSpec (after the RLModule has been added). + """ + validate_module_id(module_id, error=True) + self._check_is_built() + + # Force-set inference-only = False. + module_spec = copy.deepcopy(module_spec) + module_spec.inference_only = False + + # Build the new RLModule and add it to self.module. + module = module_spec.build() + self.module.add_module(module_id, module) + + # Change our config (AlgorithmConfig) to contain the new Module. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + self.config.policies[module_id] = PolicySpec() + if config_overrides is not None: + self.config.multi_agent( + algorithm_config_overrides_per_module={module_id: config_overrides} + ) + self.config.rl_module(rl_module_spec=MultiRLModuleSpec.from_module(self.module)) + self._module_spec = self.config.rl_module_spec + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + + # Allow the user to configure one or more optimizers for this new module. + self.configure_optimizers_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id), + ) + + return self.config.rl_module_spec + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def remove_module( + self, + module_id: ModuleID, + *, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + """Removes a module from the Learner. + + Args: + module_id: The ModuleID of the module to be removed. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + + Returns: + The new MultiRLModuleSpec (after the RLModule has been removed). + """ + self._check_is_built() + module = self.module[module_id] + + # Delete the removed module's parameters and optimizers. + if self.rl_module_is_compatible(module): + parameters = self.get_parameters(module) + for param in parameters: + param_ref = self.get_param_ref(param) + if param_ref in self._params: + del self._params[param_ref] + for optimizer_name, optimizer in self.get_optimizers_for_module(module_id): + del self._optimizer_parameters[optimizer] + name = module_id + "_" + optimizer_name + del self._named_optimizers[name] + if optimizer in self._optimizer_lr_schedules: + del self._optimizer_lr_schedules[optimizer] + del self._module_optimizers[module_id] + + # Remove the module from the MultiRLModule. + self.module.remove_module(module_id) + + # Change self.config to reflect the new architecture. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + del self.config.policies[module_id] + self.config.algorithm_config_overrides_per_module.pop(module_id, None) + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + self.config.rl_module(rl_module_spec=MultiRLModuleSpec.from_module(self.module)) + + # Remove all stats from the module from our metrics logger, so we don't report + # results from this module again. + if module_id in self.metrics.stats: + del self.metrics.stats[module_id] + + return self.config.rl_module_spec + + @OverrideToImplementCustomLogic + def should_module_be_updated(self, module_id, multi_agent_batch=None): + """Returns whether a module should be updated or not based on `self.config`. + + Args: + module_id: The ModuleID that we want to query on whether this module + should be updated or not. + multi_agent_batch: An optional MultiAgentBatch to possibly provide further + information on the decision on whether the RLModule should be updated + or not. + """ + should_module_be_updated_fn = self.config.policies_to_train + # If None, return True (by default, all modules should be updated). + if should_module_be_updated_fn is None: + return True + # If collection given, return whether `module_id` is in that container. + elif not callable(should_module_be_updated_fn): + return module_id in set(should_module_be_updated_fn) + + return should_module_be_updated_fn(module_id, multi_agent_batch) + + @OverrideToImplementCustomLogic + def compute_losses( + self, *, fwd_out: Dict[str, Any], batch: Dict[str, Any] + ) -> Dict[str, Any]: + """Computes the loss(es) for the module being optimized. + + This method must be overridden by MultiRLModule-specific Learners in order to + define the specific loss computation logic. If the algorithm is single-agent, + only `compute_loss_for_module()` should be overridden instead. If the algorithm + uses independent multi-agent learning (default behavior for RLlib's multi-agent + setups), also only `compute_loss_for_module()` should be overridden, but it will + be called for each individual RLModule inside the MultiRLModule. + It is recommended to not compute any forward passes within this method, and to + use the `forward_train()` outputs of the RLModule(s) to compute the required + loss tensors. + See here for a custom loss function example script: + https://github.com/ray-project/ray/blob/master/rllib/examples/learners/ppo_with_custom_loss_fn.py # noqa + + Args: + fwd_out: Output from a call to the `forward_train()` method of the + underlying MultiRLModule (`self.module`) during training + (`self.update()`). + batch: The train batch that was used to compute `fwd_out`. + + Returns: + A dictionary mapping module IDs to individual loss terms. + """ + loss_per_module = {} + for module_id in fwd_out: + module_batch = batch[module_id] + module_fwd_out = fwd_out[module_id] + + module = self.module[module_id].unwrapped() + if isinstance(module, SelfSupervisedLossAPI): + loss = module.compute_self_supervised_loss( + learner=self, + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + ) + else: + loss = self.compute_loss_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + ) + loss_per_module[module_id] = loss + + return loss_per_module + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + """Computes the loss for a single module. + + Think of this as computing loss for a single agent. For multi-agent use-cases + that require more complicated computation for loss, consider overriding the + `compute_losses` method instead. + + Args: + module_id: The id of the module. + config: The AlgorithmConfig specific to the given `module_id`. + batch: The train batch for this particular module. + fwd_out: The output of the forward pass for this particular module. + + Returns: + A single total loss tensor. If you have more than one optimizer on the + provided `module_id` and would like to compute gradients separately using + these different optimizers, simply add up the individual loss terms for + each optimizer and return the sum. Also, for recording/logging any + individual loss terms, you can use the `Learner.metrics.log_value( + key=..., value=...)` or `Learner.metrics.log_dict()` APIs. See: + :py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for more + information. + """ + + @OverrideToImplementCustomLogic + def update( + self, + batch: Optional[MultiAgentBatch] = None, + batches: Optional[List[MultiAgentBatch]] = None, + batch_refs: Optional[List[ray.ObjectRef]] = None, + episodes: Optional[List[EpisodeType]] = None, + episodes_refs: Optional[List[ray.ObjectRef]] = None, + data_iterators: Optional[List[ray.data.DataIterator]] = None, + training_data: Optional[TrainingData] = None, + *, + # TODO (sven): Make this a more formal structure with its own type. + timesteps: Optional[Dict[str, Any]] = None, + num_total_minibatches: int = 0, + num_epochs: int = 1, + minibatch_size: Optional[int] = None, + shuffle_batch_per_epoch: bool = False, + _no_metrics_reduce: bool = False, + **kwargs, + ) -> ResultDict: + """Run `num_epochs` epochs over the given train batch. + + You can use this method to take more than one backward pass on the batch. + The same `minibatch_size` and `num_epochs` will be used for all module ids in + MultiRLModule. + + Args: + batch: A batch of training data to update from. + timesteps: Timesteps dict, which must have the key + `NUM_ENV_STEPS_SAMPLED_LIFETIME`. + # TODO (sven): Make this a more formal structure with its own type. + num_epochs: The number of complete passes over the entire train batch. Each + pass might be further split into n minibatches (if `minibatch_size` + provided). + minibatch_size: The size of minibatches to use to further split the train + `batch` into sub-batches. The `batch` is then iterated over n times + where n is `len(batch) // minibatch_size`. + shuffle_batch_per_epoch: Whether to shuffle the train batch once per epoch. + If the train batch has a time rank (axis=1), shuffling will only take + place along the batch axis to not disturb any intact (episode) + trajectories. Also, shuffling is always skipped if `minibatch_size` is + None, meaning the entire train batch is processed each epoch, making it + unnecessary to shuffle. + + Returns: + A `ResultDict` object produced by a call to `self.metrics.reduce()`. The + returned dict may be arbitrarily nested and must have `Stats` objects at + all its leafs, allowing components further downstream (i.e. a user of this + Learner) to further reduce these results (for example over n parallel + Learners). + """ + self._check_is_built() + + # Call `before_gradient_based_update` to allow for non-gradient based + # preparations-, logging-, and update logic to happen. + self.before_gradient_based_update(timesteps=timesteps or {}) + + if training_data is None: + training_data = TrainingData( + batch=batch, + batches=batches, + batch_refs=batch_refs, + episodes=episodes, + episodes_refs=episodes_refs, + data_iterators=data_iterators, + ) + training_data.validate() + training_data.solve_refs() + assert training_data.batches is None, "`training_data.batches` must be None!" + + self._weights_seq_no += 1 + + batch_iter = self._create_iterator_if_necessary( + training_data=training_data, + num_total_minibatches=num_total_minibatches, + num_epochs=num_epochs, + minibatch_size=minibatch_size, + shuffle_batch_per_epoch=shuffle_batch_per_epoch, + **kwargs, + ) + + self.metrics.activate_tensor_mode() + + # Perform the actual looping through the minibatches or the given data iterator. + for iteration, tensor_minibatch in enumerate(batch_iter): + # Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs + # found in this batch. If not, throw an error. + unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set( + self.module.keys() + ) + if unknown_module_ids: + raise ValueError( + f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that " + f"are not in this Learner!" + ) + + # Make the actual in-graph/traced `_update` call. This should return + # all tensor values (no numpy). + with TimerAndPrometheusLogger(self._metrics_learner_inner_update): + fwd_out, loss_per_module, _ = self._update( + tensor_minibatch.policy_batches + ) + + # Ray metrics + self._log_metrics(batch=tensor_minibatch) + + # TODO (sven): Maybe move this into loop above to get metrics more accuratcely + # cover the minibatch/epoch logic. + # Log all timesteps (env, agent, modules) based on given episodes/batch. + self._log_steps_trained_metrics(tensor_minibatch) + + self._set_slicing_by_batch_id(tensor_minibatch, value=False) + + if self.iterator: + # Record the number of batches pulled from the dataset. + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED), + iteration + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED_LIFETIME), + iteration + 1, + reduce="sum", + ) + # Log all individual RLModules' loss terms and its registered optimizers' + # current learning rates. + # Note: We do this only once for the last of the minibatch updates, b/c the + # window is only 1 anyways. + for mid, loss in loss_per_module.items(): + self.metrics.log_value( + key=(mid, self.TOTAL_LOSS_KEY), + value=loss, + window=1, + ) + + # Call `after_gradient_based_update` to allow for non-gradient based + # cleanups-, logging-, and update logic to happen. + # TODO (simon): Check, if this should stay here, when running multiple + # gradient steps inside the iterator loop above (could be a complete epoch) + # the target networks might need to be updated earlier. + self.after_gradient_based_update(timesteps=timesteps or {}) + + self.metrics.deactivate_tensor_mode() + + # Reduce results across all minibatch update steps. + if not _no_metrics_reduce: + return self.metrics.reduce() + + def _create_iterator_if_necessary( + self, + *, + training_data: TrainingData, + num_total_minibatches: int = 0, + num_epochs: int = 1, + minibatch_size: Optional[int] = None, + shuffle_batch_per_epoch: bool = False, + **kwargs, + ) -> Iterable: + # Data iterator provided. + if training_data.data_iterators: + num_iters = kwargs.pop("num_iters", None) + if num_iters is None: + raise ValueError( + "Learner.update(data_iterators=..) requires `num_iters` kwarg!" + ) + + if not self.iterator: + # This iterator holds a `ray.data.DataIterator` and manages it state. + self.iterator = MiniBatchRayDataIterator( + iterator=training_data.data_iterators[0], + device=self.device, + minibatch_size=minibatch_size, + num_iters=num_iters, + **kwargs, + ) + + batch_iter = self.iterator + else: + batch = self._make_batch_if_necessary(training_data=training_data) + assert batch is not None + + # TODO: Move this into LearnerConnector pipeline? + # Filter out those RLModules from the final train batch that should not be + # updated. + for module_id in list(batch.policy_batches.keys()): + if not self.should_module_be_updated(module_id, batch): + del batch.policy_batches[module_id] + if not batch.policy_batches: + return {} + + batch = self._set_slicing_by_batch_id(batch, value=True) + + if minibatch_size: + batch_iter_cls = MiniBatchCyclicIterator + elif num_epochs > 1: + # `minibatch_size` was not set but `num_epochs` > 1. + minibatch_size = batch.count + # Note that there is no need to shuffle here, b/c we don't have + # minibatches. + batch_iter_cls = MiniBatchCyclicIterator + else: + # `minibatch_size` and `num_epochs` are not set by the user. + batch_iter_cls = MiniBatchDummyIterator + + batch_iter = batch_iter_cls( + batch, + num_epochs=num_epochs, + minibatch_size=minibatch_size, + shuffle_batch_per_epoch=shuffle_batch_per_epoch and (num_epochs > 1), + num_total_minibatches=num_total_minibatches, + ) + return batch_iter + + @OverrideToImplementCustomLogic + @abc.abstractmethod + def _update( + self, + batch: Dict[str, Any], + **kwargs, + ) -> Tuple[Any, Any, Any]: + """Contains all logic for an in-graph/traceable update step. + + Framework specific subclasses must implement this method. This should include + calls to the RLModule's `forward_train`, `compute_loss`, compute_gradients`, + `postprocess_gradients`, and `apply_gradients` methods and return a tuple + with all the individual results. + + Args: + batch: The train batch already converted to a Dict mapping str to (possibly + nested) tensors. + kwargs: Forward compatibility kwargs. + + Returns: + A tuple consisting of: + 1) The `forward_train()` output of the RLModule, + 2) the loss_per_module dictionary mapping module IDs to individual loss + tensors + 3) a metrics dict mapping module IDs to metrics key/value pairs. + + """ + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + self._check_is_built() + + state = { + "should_module_be_updated": self.config.policies_to_train, + WEIGHTS_SEQ_NO: self._weights_seq_no, + } + + if self._check_component(COMPONENT_RL_MODULE, components, not_components): + state[COMPONENT_RL_MODULE] = self.module.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=self._get_subcomponents( + COMPONENT_RL_MODULE, not_components + ), + **kwargs, + ) + if self._check_component(COMPONENT_OPTIMIZER, components, not_components): + state[COMPONENT_OPTIMIZER] = self._get_optimizer_state() + + if self._check_component(COMPONENT_METRICS_LOGGER, components, not_components): + # TODO (sven): Make `MetricsLogger` a Checkpointable. + state[COMPONENT_METRICS_LOGGER] = self.metrics.get_state() + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + self._check_is_built() + + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + if COMPONENT_RL_MODULE in state: + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + self.module.set_state(state[COMPONENT_RL_MODULE]) + + if COMPONENT_OPTIMIZER in state: + self._set_optimizer_state(state[COMPONENT_OPTIMIZER]) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update our trainable Modules information/function via our config. + # If not provided in state (None), all Modules will be trained by default. + if "should_module_be_updated" in state: + self.config.multi_agent(policies_to_train=state["should_module_be_updated"]) + + # TODO (sven): Make `MetricsLogger` a Checkpointable. + if COMPONENT_METRICS_LOGGER in state: + self.metrics.set_state(state[COMPONENT_METRICS_LOGGER]) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args, + { + "config": self.config, + "module_spec": self._module_spec, + "module": self._module_obj, + }, # **kwargs + ) + + @override(Checkpointable) + def get_checkpointable_components(self): + if not self._check_is_built(error=False): + self.build() + return [ + (COMPONENT_RL_MODULE, self.module), + ] + + def _make_batch_if_necessary(self, training_data): + batch = training_data.batch + + # Call the learner connector on the given `episodes` (if we have one). + if training_data.episodes is not None: + # If we want to learn from Episodes, we must have a LearnerConnector + # pipeline to translate into a train batch first. + if self._learner_connector is None: + raise ValueError( + f"If episodes provided for training, Learner ({self}) must have a " + "LearnerConnector pipeline (but pipeline is None)!" + ) + + # Call the learner connector pipeline. + shared_data = {} + batch = self._learner_connector( + rl_module=self.module, + batch=training_data.batch if training_data.batch is not None else {}, + episodes=training_data.episodes, + shared_data=shared_data, + metrics=self.metrics, + ) + # Convert to a batch. + # TODO (sven): Try to not require MultiAgentBatch anymore. + batch = MultiAgentBatch( + { + module_id: ( + SampleBatch(module_data, _zero_padded=True) + if shared_data.get(f"_zero_padded_for_mid={module_id}") + else SampleBatch(module_data) + ) + for module_id, module_data in batch.items() + }, + env_steps=sum(len(e) for e in training_data.episodes), + ) + # Single-agent SampleBatch: Have to convert to MultiAgentBatch. + elif isinstance(training_data.batch, SampleBatch): + if len(self.module) != 1: + raise ValueError( + f"SampleBatch provided, but RLModule ({self.module}) has more than " + f"one sub-RLModule! Need to provide MultiAgentBatch instead." + ) + + batch = MultiAgentBatch( + {next(iter(self.module.keys())): training_data.batch}, + env_steps=len(training_data.batch), + ) + # If we already have an `MultiAgentBatch` but with `numpy` array, convert to + # tensors. + elif ( + isinstance(training_data.batch, MultiAgentBatch) + and training_data.batch.policy_batches + and ( + any( + tree.map_structure( + lambda a: isinstance(a, numpy.ndarray), + tree.flatten(training_data.batch.policy_batches), + ) + ) + or any( + tree.map_structure( + lambda a: isinstance(a, torch.Tensor) + and a.device != self._device, + tree.flatten(training_data.batch.policy_batches), + ) + ) + ) + ): + batch = self._convert_batch_type(training_data.batch) + + return batch + + def _get_optimizer_state(self) -> StateDict: + """Returns the state of all optimizers currently registered in this Learner. + + Returns: + The current state of all optimizers currently registered in this Learner. + """ + raise NotImplementedError + + def _set_optimizer_state(self, state: StateDict) -> None: + """Sets the state of all optimizers currently registered in this Learner. + + Args: + state: The state of the optimizers. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def before_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None: + """Called before gradient-based updates are completed. + + Should be overridden to implement custom preparation-, logging-, or + non-gradient-based Learner/RLModule update logic before(!) gradient-based + updates are performed. + + Args: + timesteps: Timesteps dict, which must have the key + `NUM_ENV_STEPS_SAMPLED_LIFETIME`. + # TODO (sven): Make this a more formal structure with its own type. + """ + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None: + """Called after gradient-based updates are completed. + + Should be overridden to implement custom cleanup-, logging-, or non-gradient- + based Learner/RLModule update logic after(!) gradient-based updates have been + completed. + + Args: + timesteps: Timesteps dict, which must have the key + `NUM_ENV_STEPS_SAMPLED_LIFETIME`. + # TODO (sven): Make this a more formal structure with its own type. + """ + # Only update this optimizer's lr, if a scheduler has been registered + # along with it. + for module_id, optimizer_names in self._module_optimizers.items(): + for optimizer_name in optimizer_names: + optimizer = self._named_optimizers[optimizer_name] + # Update and log learning rate of this optimizer. + lr_schedule = self._optimizer_lr_schedules.get(optimizer) + if lr_schedule is not None: + new_lr = lr_schedule.update( + timestep=timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0) + ) + self._set_optimizer_lr(optimizer, lr=new_lr) + self.metrics.log_value( + # Cut out the module ID from the beginning since it's already part + # of the key sequence: (ModuleID, "[optim name]_lr"). + key=(module_id, f"{optimizer_name[len(module_id) + 1:]}_{LR_KEY}"), + value=self._get_optimizer_lr(optimizer), + window=1, + ) + + def _set_slicing_by_batch_id( + self, batch: MultiAgentBatch, *, value: bool + ) -> MultiAgentBatch: + """Enables slicing by batch id in the given batch. + + If the input batch contains batches of sequences we need to make sure when + slicing happens it is sliced via batch id and not timestamp. Calling this + method enables the same flag on each SampleBatch within the input + MultiAgentBatch. + + Args: + batch: The MultiAgentBatch to enable slicing by batch id on. + value: The value to set the flag to. + + Returns: + The input MultiAgentBatch with the indexing flag is enabled / disabled on. + """ + + for pid, policy_batch in batch.policy_batches.items(): + # We assume that arriving batches for recurrent modules OR batches that + # have a SEQ_LENS column are already zero-padded to the max sequence length + # and have tensors of shape [B, T, ...]. Therefore, we slice sequence + # lengths in B. See SampleBatch for more information. + if ( + self.module[pid].is_stateful() + or policy_batch.get("seq_lens") is not None + ): + if value: + policy_batch.enable_slicing_by_batch_id() + else: + policy_batch.disable_slicing_by_batch_id() + + return batch + + def _make_module(self) -> MultiRLModule: + """Construct the multi-agent RL module for the learner. + + This method uses `self._module_specs` or `self._module_obj` to construct the + module. If the module_class is a single agent RL module it will be wrapped to a + multi-agent RL module. Override this method if there are other things that + need to happen for instantiation of the module. + + Returns: + A constructed MultiRLModule. + """ + # Module was provided directly through constructor -> Use as-is. + if self._module_obj is not None: + module = self._module_obj + self._module_spec = MultiRLModuleSpec.from_module(module) + # RLModuleSpec was provided directly through constructor -> Use it to build the + # RLModule. + elif self._module_spec is not None: + module = self._module_spec.build() + # Try using our config object. Note that this would only work if the config + # object has all the necessary space information already in it. + else: + module = self.config.get_multi_rl_module_spec().build() + + # If not already, convert to MultiRLModule. + module = module.as_multi_rl_module() + + return module + + def rl_module_is_compatible(self, module: RLModule) -> bool: + """Check whether the given `module` is compatible with this Learner. + + The default implementation checks the Learner-required APIs and whether the + given `module` implements all of them (if not, returns False). + + Args: + module: The RLModule to check. + + Returns: + True if the module is compatible with this Learner. + """ + return all(isinstance(module, api) for api in self.rl_module_required_apis()) + + @classmethod + def rl_module_required_apis(cls) -> list[type]: + """Returns the required APIs for an RLModule to be compatible with this Learner. + + The returned values may or may not be used inside the `rl_module_is_compatible` + method. + + Args: + module: The RLModule to check. + + Returns: + A list of RLModule API classes that an RLModule must implement in order + to be compatible with this Learner. + """ + return [] + + def _check_registered_optimizer( + self, + optimizer: Optimizer, + params: Sequence[Param], + ) -> None: + """Checks that the given optimizer and parameters are valid for the framework. + + Args: + optimizer: The optimizer object to check. + params: The list of parameters to check. + """ + if not isinstance(params, list): + raise ValueError( + f"`params` ({params}) must be a list of framework-specific parameters " + "(variables)!" + ) + + def _log_trainable_parameters(self) -> None: + """Logs the number of trainable and non-trainable parameters to self.metrics. + + Use MetricsLogger (self.metrics) tuple-keys: + (ALL_MODULES, NUM_TRAINABLE_PARAMETERS) and + (ALL_MODULES, NUM_NON_TRAINABLE_PARAMETERS) with EMA. + """ + pass + + def _check_is_built(self, error: bool = True) -> bool: + if self.module is None: + if error: + raise ValueError( + "Learner.build() must be called after constructing a " + "Learner and before calling any methods on it." + ) + return False + return True + + def _reset(self): + self._params = {} + self._optimizer_parameters = {} + self._named_optimizers = {} + self._module_optimizers = defaultdict(list) + self._optimizer_lr_schedules = {} + self.metrics = MetricsLogger() + self._is_built = False + + def apply(self, func, *_args, **_kwargs): + return func(self, *_args, **_kwargs) + + @abc.abstractmethod + def _get_tensor_variable( + self, + value: Any, + dtype: Any = None, + trainable: bool = False, + ) -> TensorType: + """Returns a framework-specific tensor variable with the initial given value. + + This is a framework specific method that should be implemented by the + framework specific sub-classes. + + Args: + value: The initial value for the tensor variable variable. + + Returns: + The framework specific tensor variable of the given initial value, + dtype and trainable/requires_grad property. + """ + + @staticmethod + @abc.abstractmethod + def _get_optimizer_lr(optimizer: Optimizer) -> float: + """Returns the current learning rate of the given local optimizer. + + Args: + optimizer: The local optimizer to get the current learning rate for. + + Returns: + The learning rate value (float) of the given optimizer. + """ + + @staticmethod + @abc.abstractmethod + def _set_optimizer_lr(optimizer: Optimizer, lr: float) -> None: + """Updates the learning rate of the given local optimizer. + + Args: + optimizer: The local optimizer to update the learning rate for. + lr: The new learning rate. + """ + + @staticmethod + @abc.abstractmethod + def _get_clip_function() -> Callable: + """Returns the gradient clipping function to use.""" + + @staticmethod + @abc.abstractmethod + def _get_global_norm_function() -> Callable: + """Returns the global norm function to use, given the framework.""" + + def _log_steps_trained_metrics(self, batch: MultiAgentBatch): + """Logs this iteration's steps trained, based on given `batch`.""" + # Collect all module steps and add them for `ALL_MODULES` to avoid + # biasing the throughput by looping through modules. + total_module_steps = 0 + # Loop through all modules. + for mid, module_batch in batch.policy_batches.items(): + # Log weights seq no for this batch. + self.metrics.log_value( + (mid, WEIGHTS_SEQ_NO), + self._weights_seq_no, + window=1, + ) + + module_batch_size = len(module_batch) + # Log average batch size (for each module). + self.metrics.log_value( + key=(mid, MODULE_TRAIN_BATCH_SIZE_MEAN), + value=module_batch_size, + ) + # Log module steps (for each module). + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_TRAINED), + value=module_batch_size, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_TRAINED_LIFETIME), + value=module_batch_size, + reduce="sum", + with_throughput=True, + ) + total_module_steps += module_batch_size + + # Log module steps (sum of all modules). + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED), + value=total_module_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED_LIFETIME), + value=total_module_steps, + reduce="sum", + with_throughput=True, + ) + # Log env steps (all modules). + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_TRAINED), + batch.env_steps(), + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME), + batch.env_steps(), + reduce="sum", + with_throughput=True, + ) + + def _set_learner_index_and_placement_group(self, *, learner_index, placement_group): + self._learner_index = learner_index + self._placement_group = placement_group + + @Deprecated(new="Learner.update(batch=.., ..)", error=False) + def update_from_batch(self, batch, **kwargs): + if isinstance(batch, list): + if isinstance(batch[0], ray.ObjectRef): + return self.update(batches_refs=batch, **kwargs) + else: + return self.update(batches=batch, **kwargs) + else: + return self.update(batch=batch, **kwargs) + + @Deprecated(new="Learner.update(episodes=.., ..)", error=False) + def update_from_episodes(self, episodes, **kwargs): + assert isinstance(episodes, list) + if isinstance(episodes[0], ray.ObjectRef): + return self.update(episodes_refs=episodes, **kwargs) + else: + return self.update(episodes=episodes, **kwargs) + + @Deprecated(new="Learner.compute_losses(...)", error=True) + def compute_loss(self, *args, **kwargs): + pass + + def _log_metrics(self, batch: MultiAgentBatch) -> None: + _env_steps = int(batch.env_steps()) + if _env_steps > 0: + self._metrics_all_modules_num_env_steps_trained.inc(value=_env_steps) + total_module_steps = sum( + len(module_batch) for module_batch in batch.policy_batches.values() + ) + self._metrics_all_modules_num_module_steps_trained.inc( + value=total_module_steps + ) + else: + logger.warning( + f"RLlib {self.__class__.__name__}: Skipping Prometheus logging for metrics: " + f"{self._metrics_all_modules_num_env_steps_trained.info['name']} and " + f"{self._metrics_all_modules_num_module_steps_trained.info['name']}. " + f"Received MultiAgentBatch.env_steps()={_env_steps}, but the number of steps must be greater than 0." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/learner_group.py b/lib/python3.12/site-packages/ray/rllib/core/learner/learner_group.py new file mode 100644 index 0000000000000000000000000000000000000000..5a06f79071f33680e984353bae96a1742615021f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/learner_group.py @@ -0,0 +1,790 @@ +import copy +import itertools +import pathlib +from functools import partial +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Dict, + List, + Optional, + Set, + Type, + Union, +) + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.actor_manager import ( + FaultTolerantActorManager, + RemoteCallResults, + ResultOrError, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.metrics.ray_metrics import ( + DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + TimerAndPrometheusLogger, +) +from ray.rllib.utils.typing import ( + EpisodeType, + ModuleID, + RLModuleSpecType, + ShouldModuleBeUpdatedFn, + StateDict, + T, +) +from ray.train._internal.backend_executor import BackendExecutor +from ray.util.annotations import PublicAPI +from ray.util.metrics import Histogram + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.util.placement_group import PlacementGroup + + +def _get_backend_config(learner_class: Type[Learner]) -> str: + if learner_class.framework == "torch": + from ray.train.torch.config import TorchConfig, _TorchBackend + + # Override `_TorchBackend` share_cuda_visible_devices=True setting. + # We need this to be False to make sure Learner actors only see their + # own GPU. There is no need in RLlib's LearnerGroups for 2 different Learner + # actors to communicate with each other through their GPUs. + class _RLlibTorchBackend(_TorchBackend): + share_cuda_visible_devices = False + + class RLlibTorchConfig(TorchConfig): + @property + def backend_cls(self): + return _RLlibTorchBackend + + backend_config = RLlibTorchConfig() + + else: + raise ValueError( + "`learner_class.framework` must be 'torch' (but is " + f"{learner_class.framework}!" + ) + + return backend_config + + +class RLlibBackendExecutor(BackendExecutor): + # Override `BackendExecutor` placement group creation logic. We need to pass our own + # to make sure the one of the Algorithm (Trainable) is used for all the + # Algorithm's actors. + def _create_placement_group(self): + pass + + # TODO (sven): Change this once there is a better (public) API for this in the + # superclass. + def set_placement_group(self, placement_group): + if placement_group is not None: + self._placement_group = placement_group + + +@PublicAPI(stability="alpha") +class LearnerGroup(Checkpointable): + """Coordinator of n (possibly remote) Learner workers. + + Each Learner worker has a copy of the RLModule, the loss function(s), and + one or more optimizers. + """ + + def __init__( + self, + *, + config: "AlgorithmConfig", + # TODO (sven): Rename into `rl_module_spec`. + module_spec: Optional[RLModuleSpecType] = None, + placement_group: Optional["PlacementGroup"] = None, + ): + """Initializes a LearnerGroup instance. + + Args: + config: The AlgorithmConfig object to use to configure this LearnerGroup. + Call the `learners(num_learners=...)` method on your config to + specify the number of learner workers to use. + Call the same method with arguments `num_cpus_per_learner` and/or + `num_gpus_per_learner` to configure the compute used by each + Learner worker in this LearnerGroup. + Call the `training(learner_class=...)` method on your config to specify, + which exact Learner class to use. + Call the `rl_module(rl_module_spec=...)` method on your config to set up + the specifics for your RLModule to be used in each Learner. + module_spec: If not already specified in `config`, a separate overriding + RLModuleSpec may be provided via this argument. + placement_group: An optional `PlacementGroup` instance to set the + `RLlibBackendExecutor`'s `self._placement_group` attribute to. + If run within an Algorithm (tune.Trainable), the placement group of tune + trial actor is passed through here. + """ + self.config = config.copy(copy_frozen=False) + self._module_spec = module_spec + + learner_class = self.config.learner_class + module_spec = module_spec or self.config.get_multi_rl_module_spec() + + self._learner = None + self._workers = None + # If a user calls self.shutdown() on their own then this flag is set to true. + # When del is called the backend executor isn't shutdown twice if this flag is + # true. the backend executor would otherwise log a warning to the console from + # ray train. + self._is_shut_down = False + + # How many timesteps had to be dropped due to a full input queue? + self._ts_dropped = 0 + + # A single local Learner. + if not self.is_remote: + self._learner = learner_class(config=config, module_spec=module_spec) + self._learner.build() + self._worker_manager = None + + # Ray metrics + self._metrics_local_learner_training_data_solve_refs = Histogram( + name="rllib_learner_local_training_data_solve_refs_time", + description="Time spent in resolve training data refs for local learner.", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_local_learner_training_data_solve_refs.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + # N remote Learner workers. + else: + backend_config = _get_backend_config(learner_class) + + num_cpus_per_learner = ( + self.config.num_cpus_per_learner + if self.config.num_cpus_per_learner != "auto" + else 1 + if self.config.num_gpus_per_learner == 0 + else 0 + ) + num_gpus_per_learner = max(0, self.config.num_gpus_per_learner) + resources_per_learner = { + "CPU": num_cpus_per_learner, + "GPU": num_gpus_per_learner, + } + + backend_executor = RLlibBackendExecutor( + backend_config=backend_config, + num_workers=self.config.num_learners, + resources_per_worker=resources_per_learner, + max_retries=0, + ) + # Set the placement group - if any - of the BackendExecutor. + backend_executor.set_placement_group(placement_group) + backend_executor.start( + train_cls=learner_class, + train_cls_kwargs={ + "config": config, + "module_spec": module_spec, + }, + ) + self._backend_executor = backend_executor + + self._workers = [w.actor for w in backend_executor.worker_group.workers] + + ray.get( + [ + worker._set_learner_index_and_placement_group.remote( + learner_index=idx, + placement_group=placement_group, + ) + for idx, worker in enumerate(self._workers) + ] + ) + + # Run the neural network building code on remote workers. + ray.get([w.build.remote() for w in self._workers]) + + self._worker_manager = FaultTolerantActorManager( + self._workers, + max_remote_requests_in_flight_per_actor=( + self.config.max_requests_in_flight_per_learner + ), + ) + + # Ray metrics + self._metrics_update_time = Histogram( + name="rllib_learner_group_update_time", + description="Time spent in LearnerGroup.update()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_update_time.set_default_tags({"rllib": self.__class__.__name__}) + + # TODO (sven): Replace this with call to `self.metrics.peek()`? + # Currently LearnerGroup does not have a metrics object. + def get_stats(self) -> Dict[str, Any]: + """Returns the current stats for the input queue for this learner group.""" + return { + "learner_group_ts_dropped_lifetime": self._ts_dropped, + "actor_manager_num_outstanding_async_reqs": ( + 0 + if self.is_local + else self._worker_manager.num_outstanding_async_reqs() + ), + } + + @property + def is_remote(self) -> bool: + return self.config.num_learners > 0 + + @property + def is_local(self) -> bool: + return not self.is_remote + + def update( + self, + *, + batch: Optional[MultiAgentBatch] = None, + batches: Optional[List[MultiAgentBatch]] = None, + batch_refs: Optional[List[ray.ObjectRef]] = None, + episodes: Optional[List[EpisodeType]] = None, + episodes_refs: Optional[List[ray.ObjectRef]] = None, + data_iterators: Optional[List[ray.data.DataIterator]] = None, + training_data: Optional[TrainingData] = None, + timesteps: Optional[Dict[str, Any]] = None, + async_update: bool = False, + return_state: bool = False, + # User kwargs passed onto the Learners. + **kwargs, + ) -> List[Dict[str, Any]]: + """Performs gradient based updates on Learners, based on given training data. + + Args: + batch: A data batch to use for the update. If there are more + than one Learner workers, the batch is split amongst these and one + shard is sent to each Learner. + batch_refs: + episodes: A list of Episodes to process and perform the update + for. If there are more than one Learner workers, the list of episodes + is split amongst these and one list shard is sent to each Learner. + episodes_refs: + timesteps: + async_update: Whether the update request(s) to the Learner workers should be + sent asynchronously. If True, will return NOT the results from the + update on the given data, but all results from prior asynchronous update + requests that have not been returned thus far. + return_state: Whether to include one of the Learner worker's state from + after the update step in the returned results dict (under the + `_rl_module_state_after_update` key). Note that after an update, all + Learner workers' states should be identical, so we use the first + Learner's state here. Useful for avoiding an extra `get_weights()` call, + e.g. for synchronizing EnvRunner weights. + num_epochs: The number of complete passes over the entire train batch. Each + pass might be further split into n minibatches (if `minibatch_size` + provided). + minibatch_size: The size of minibatches to use to further split the train + `batch` into sub-batches. The `batch` is then iterated over n times + where n is `len(batch) // minibatch_size`. + shuffle_batch_per_epoch: Whether to shuffle the train batch once per epoch. + If the train batch has a time rank (axis=1), shuffling will only take + place along the batch axis to not disturb any intact (episode) + trajectories. Also, shuffling is always skipped if `minibatch_size` is + None, meaning the entire train batch is processed each epoch, making it + unnecessary to shuffle. + **kwargs: + + Returns: + If `async_update` is False, a dictionary with the reduced results of the + updates from the Learner(s) or a list of dictionaries of results from the + updates from the Learner(s). + If `async_update` is True, a list of list of dictionaries of results, where + the outer list corresponds to separate previous calls to this method, and + the inner list corresponds to the results from each Learner(s). Or if the + results are reduced, a list of dictionaries of the reduced results from each + call to async_update that is ready. + """ + with TimerAndPrometheusLogger(self._metrics_update_time): + # Create and validate TrainingData object, if not already provided. + if training_data is None: + training_data = TrainingData( + batch=batch, + batches=batches, + batch_refs=batch_refs, + episodes=episodes, + episodes_refs=episodes_refs, + data_iterators=data_iterators, + ) + training_data.validate() + + # Local Learner instance. + if self.is_local: + if async_update: + raise ValueError( + "Can't call `update(async_update=True)` when running with " + "`num_learners=0`! Set `config.num_learners > 0` to allow async " + "updates." + ) + # Solve all ray refs locally already here. + + # Ray metrics + with TimerAndPrometheusLogger( + self._metrics_local_learner_training_data_solve_refs + ): + training_data.solve_refs() + if return_state: + kwargs["return_state"] = return_state + # Return the single Learner's update results. + return [ + self._learner.update( + training_data=training_data, + timesteps=timesteps, + **kwargs, + ) + ] + + # Remote Learner actors' kwargs. + remote_call_kwargs = [ + dict( + training_data=td_shard, + timesteps=timesteps, + # If `return_state=True`, only return it from the first Learner + # actor. + return_state=(return_state and i == 0), + **kw, + **kwargs, + ) + for i, (td_shard, kw) in enumerate( + training_data.shard( + num_shards=len(self), + len_lookback_buffer=self.config.episode_lookback_horizon, + **kwargs, + ) + ) + ] + + # Async updates. + if async_update: + # Retrieve all ready results (kicked off by prior calls to this method). + results = self._worker_manager.fetch_ready_async_reqs( + timeout_seconds=0.0 + ) + # Send out new request(s), if there is still capacity on the actors + # (each actor is allowed only some number of max in-flight requests + # at the same time). + num_sent_requests = self._worker_manager.foreach_actor_async( + "update", + kwargs=remote_call_kwargs, + ) + + # Some requests were dropped, record lost ts/data. + if num_sent_requests != len(self): + factor = 1 - (num_sent_requests / len(self)) + # TODO (sven): Move this logic into a TrainingData API as well + # (`TrainingData.env_steps()`). + if training_data.batch_refs is not None: + dropped = ( + len(training_data.batch_refs) + * self.config.train_batch_size_per_learner + ) + elif training_data.batch is not None: + dropped = len(training_data.batch) + # List of Ray ObjectRefs (each object ref is a list of episodes of + # total len=`rollout_fragment_length * num_envs_per_env_runner`) + elif training_data.episodes_refs is not None: + dropped = ( + len(training_data.episodes_refs) + * self.config.get_rollout_fragment_length() + * self.config.num_envs_per_env_runner + ) + else: + assert training_data.episodes is not None + dropped = sum(len(e) for e in training_data.episodes) + + self._ts_dropped += factor * dropped + # Sync updates. + else: + results = self._worker_manager.foreach_actor( + "update", + kwargs=remote_call_kwargs, + ) + + results = self._get_results(results) + + return results + + def add_module( + self, + *, + module_id: ModuleID, + module_spec: RLModuleSpec, + config_overrides: Optional[Dict] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + """Adds a module to the underlying MultiRLModule. + + Changes this Learner's config in order to make this architectural change + permanent wrt. to checkpointing. + + Args: + module_id: The ModuleID of the module to be added. + module_spec: The ModuleSpec of the module to be added. + config_overrides: The `AlgorithmConfig` overrides that should apply to + the new Module, if any. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + + Returns: + The new MultiRLModuleSpec (after the change has been performed). + """ + validate_module_id(module_id, error=True) + + # Force-set inference-only = False. + module_spec = copy.deepcopy(module_spec) + module_spec.inference_only = False + + results = self.foreach_learner( + func=lambda _learner: _learner.add_module( + module_id=module_id, + module_spec=module_spec, + config_overrides=config_overrides, + new_should_module_be_updated=new_should_module_be_updated, + ), + ) + marl_spec = self._get_results(results)[0] + + # Change our config (AlgorithmConfig) to contain the new Module. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + self.config.policies[module_id] = PolicySpec() + if config_overrides is not None: + self.config.multi_agent( + algorithm_config_overrides_per_module={module_id: config_overrides} + ) + self.config.rl_module(rl_module_spec=marl_spec) + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + + return marl_spec + + def remove_module( + self, + module_id: ModuleID, + *, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + """Removes a module from the Learner. + + Args: + module_id: The ModuleID of the module to be removed. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + + Returns: + The new MultiRLModuleSpec (after the change has been performed). + """ + results = self.foreach_learner( + func=lambda _learner: _learner.remove_module( + module_id=module_id, + new_should_module_be_updated=new_should_module_be_updated, + ), + ) + marl_spec = self._get_results(results)[0] + + # Change self.config to reflect the new architecture. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + del self.config.policies[module_id] + self.config.algorithm_config_overrides_per_module.pop(module_id, None) + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + self.config.rl_module(rl_module_spec=marl_spec) + + return marl_spec + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + state = {} + + if self._check_component(COMPONENT_LEARNER, components, not_components): + if self.is_local: + state[COMPONENT_LEARNER] = self._learner.get_state( + components=self._get_subcomponents(COMPONENT_LEARNER, components), + not_components=self._get_subcomponents( + COMPONENT_LEARNER, not_components + ), + **kwargs, + ) + else: + worker = self._worker_manager.healthy_actor_ids()[0] + assert len(self) == self._worker_manager.num_healthy_actors() + _comps = self._get_subcomponents(COMPONENT_LEARNER, components) + _not_comps = self._get_subcomponents(COMPONENT_LEARNER, not_components) + results = self._worker_manager.foreach_actor( + lambda w: w.get_state(_comps, not_components=_not_comps, **kwargs), + remote_actor_ids=[worker], + ) + state[COMPONENT_LEARNER] = self._get_results(results)[0] + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + if COMPONENT_LEARNER in state: + if self.is_local: + self._learner.set_state(state[COMPONENT_LEARNER]) + else: + state_ref = ray.put(state[COMPONENT_LEARNER]) + self.foreach_learner( + lambda _learner, _ref=state_ref: _learner.set_state(ray.get(_ref)) + ) + + def get_weights( + self, module_ids: Optional[Collection[ModuleID]] = None + ) -> StateDict: + """Convenience method instead of self.get_state(components=...). + + Args: + module_ids: An optional collection of ModuleIDs for which to return weights. + If None (default), return weights of all RLModules. + + Returns: + The results of + `self.get_state(components='learner/rl_module')['learner']['rl_module']`. + """ + # Return the entire RLModule state (all possible single-agent RLModules). + if module_ids is None: + components = COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + # Return a subset of the single-agent RLModules. + else: + components = [ + "".join(tup) + for tup in itertools.product( + [COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/"], + list(module_ids), + ) + ] + state = self.get_state(components)[COMPONENT_LEARNER][COMPONENT_RL_MODULE] + return state + + def set_weights(self, weights) -> None: + """Convenience method instead of self.set_state({'learner': {'rl_module': ..}}). + + Args: + weights: The weights dict of the MultiRLModule of a Learner inside this + LearnerGroup. + """ + self.set_state({COMPONENT_LEARNER: {COMPONENT_RL_MODULE: weights}}) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + { + "config": self.config, + "module_spec": self._module_spec, + }, # **kwargs + ) + + @override(Checkpointable) + def get_checkpointable_components(self): + # Return the entire ActorManager, if remote. Otherwise, return the + # local worker. Also, don't give the component (Learner) a name ("") + # as it's the only component in this LearnerGroup to be saved. + return [ + ( + COMPONENT_LEARNER, + self._learner if self.is_local else self._worker_manager, + ) + ] + + def foreach_learner( + self, + func: Callable[[Learner, Optional[Any]], T], + *, + healthy_only: bool = True, + remote_actor_ids: List[int] = None, + timeout_seconds: Optional[float] = None, + return_obj_refs: bool = False, + mark_healthy: bool = False, + **kwargs, + ) -> RemoteCallResults: + r"""Calls the given function on each Learner L with the args: (L, \*\*kwargs). + + Args: + func: The function to call on each Learner L with args: (L, \*\*kwargs). + healthy_only: If True, applies `func` only to Learner actors currently + tagged "healthy", otherwise to all actors. If `healthy_only=False` and + `mark_healthy=True`, will send `func` to all actors and mark those + actors "healthy" that respond to the request within `timeout_seconds` + and are currently tagged as "unhealthy". + remote_actor_ids: Apply func on a selected set of remote actors. Use None + (default) for all actors. + timeout_seconds: Time to wait (in seconds) for results. Set this to 0.0 for + fire-and-forget. Set this to None (default) to wait infinitely (i.e. for + synchronous execution). + return_obj_refs: whether to return ObjectRef instead of actual results. + Note, for fault tolerance reasons, these returned ObjectRefs should + never be resolved with ray.get() outside of the context of this manager. + mark_healthy: Whether to mark all those actors healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that actors are NOT set unhealthy, if they simply time out + (only if they return a RayActorError). + Also not that this setting is ignored if `healthy_only=True` (b/c this + setting only affects actors that are currently tagged as unhealthy). + + Returns: + A list of size len(Learners) with the return values of all calls to `func`. + """ + if self.is_local: + results = RemoteCallResults() + results.add_result( + None, + ResultOrError(result=func(self._learner, **kwargs)), + None, + ) + return results + + return self._worker_manager.foreach_actor( + func=partial(func, **kwargs), + healthy_only=healthy_only, + remote_actor_ids=remote_actor_ids, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + def __len__(self): + return 0 if self.is_local else len(self._workers) + + def shutdown(self): + """Shuts down the LearnerGroup.""" + if self.is_remote and hasattr(self, "_backend_executor"): + self._backend_executor.shutdown(graceful_termination=True) + self._is_shut_down = True + + def __del__(self): + if not self._is_shut_down: + self.shutdown() + + def _get_results(self, results): + processed_results = [] + for result in results: + result_or_error = result.get() + if result.ok: + processed_results.append(result_or_error) + else: + raise result_or_error + return processed_results + + @Deprecated(new="LearnerGroup.update(batch=.., **kwargs)", error=False) + def update_from_batch(self, batch, **kwargs): + return self.update(batch=batch, **kwargs) + + @Deprecated(new="LearnerGroup.update(episodes=.., **kwargs)", error=False) + def update_from_episodes(self, episodes, **kwargs): + return self.update(episodes=episodes, **kwargs) + + @Deprecated(new="LearnerGroup.update_from_batch(async=True)", error=True) + def async_update(self, *args, **kwargs): + pass + + @Deprecated(new="LearnerGroup.load_from_path(path=..., component=...)", error=False) + def load_module_state( + self, + *, + multi_rl_module_ckpt_dir: Optional[str] = None, + modules_to_load: Optional[Set[str]] = None, + rl_module_ckpt_dirs: Optional[Dict[ModuleID, str]] = None, + ) -> None: + """Load the checkpoints of the modules being trained by this LearnerGroup. + + `load_module_state` can be used 3 ways: + 1. Load a checkpoint for the MultiRLModule being trained by this + LearnerGroup. Limit the modules that are loaded from the checkpoint + by specifying the `modules_to_load` argument. + 2. Load the checkpoint(s) for single agent RLModules that + are in the MultiRLModule being trained by this LearnerGroup. + 3. Load a checkpoint for the MultiRLModule being trained by this + LearnerGroup and load the checkpoint(s) for single agent RLModules + that are in the MultiRLModule. The checkpoints for the single + agent RLModules take precedence over the module states in the + MultiRLModule checkpoint. + + NOTE: At lease one of multi_rl_module_ckpt_dir or rl_module_ckpt_dirs is + must be specified. modules_to_load can only be specified if + multi_rl_module_ckpt_dir is specified. + + Args: + multi_rl_module_ckpt_dir: The path to the checkpoint for the + MultiRLModule. + modules_to_load: A set of module ids to load from the checkpoint. + rl_module_ckpt_dirs: A mapping from module ids to the path to a + checkpoint for a single agent RLModule. + """ + if not (multi_rl_module_ckpt_dir or rl_module_ckpt_dirs): + raise ValueError( + "At least one of `multi_rl_module_ckpt_dir` or " + "`rl_module_ckpt_dirs` must be provided!" + ) + if multi_rl_module_ckpt_dir: + multi_rl_module_ckpt_dir = pathlib.Path(multi_rl_module_ckpt_dir) + if rl_module_ckpt_dirs: + for module_id, path in rl_module_ckpt_dirs.items(): + rl_module_ckpt_dirs[module_id] = pathlib.Path(path) + + # MultiRLModule checkpoint is provided. + if multi_rl_module_ckpt_dir: + # Restore the entire MultiRLModule state. + if modules_to_load is None: + self.restore_from_path( + multi_rl_module_ckpt_dir, + component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE, + ) + # Restore individual module IDs. + else: + for module_id in modules_to_load: + self.restore_from_path( + multi_rl_module_ckpt_dir / module_id, + component=( + COMPONENT_LEARNER + + "/" + + COMPONENT_RL_MODULE + + "/" + + module_id + ), + ) + if rl_module_ckpt_dirs: + for module_id, path in rl_module_ckpt_dirs.items(): + self.restore_from_path( + path, + component=( + COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/" + module_id + ), + ) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e655d98604617fe16fd2e5e6cc2ca20f3dbaf4f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_differentiable_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_differentiable_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d9a904faf591b4321e544269939d8d2e16133eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_differentiable_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5194910ff8b715a89d8ea025d01e51535421ac91 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_meta_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_meta_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef2af16b93cace71cc7bb3920b0f458d2620a948 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/__pycache__/torch_meta_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_differentiable_learner.py b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_differentiable_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..baa6a6bc18d3edc941a248df83c4ce276ed5f355 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_differentiable_learner.py @@ -0,0 +1,374 @@ +import contextlib +import logging +from typing import Any, Dict, Optional, Tuple + +from ray.rllib.algorithms.algorithm_config import ( + TorchCompileWhatToCompile, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.differentiable_learner import DifferentiableLearner +from ray.rllib.core.rl_module.torch.torch_rl_module import ( + TorchCompileConfig, +) +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import DeviceType, ModuleID, NamedParamDict, TensorType + +logger = logging.getLogger("__name__") + +torch, nn = try_import_torch() + + +# TODO (simon): Maybe create again a base class `TorchLearnable`. +class TorchDifferentiableLearner(DifferentiableLearner): + """A `DifferentiableLearner` class leveraging PyTorch for functional updates. + + This class utilizes PyTorch 2.0's `func` module to perform functional + updates on the provided parameters. + """ + + # Set the framework to `"torch"`. + framework: str = "torch" + + def __init__(self, **kwargs): + + # First initialize the `DifferentiableLearner` base class to set + # the configurations and `MultiRLModule`. + super().__init__(**kwargs) + + # Whether to compile the RL Module of this learner. This implies that the. + # forward_train method of the RL Module will be compiled. Further more, + # other forward methods of the RL Module will be compiled on demand. + # This is assumed to not happen, since other forwrad methods are not expected + # to be used during training. + self._torch_compile_forward_train = False + self._torch_compile_cfg = None + # Whether to compile the `_uncompiled_update` method of this learner. This + # implies that everything within `_uncompiled_update` will be compiled, + # not only the forward_train method of the RL Module. + # Note that this is experimental. + # Note that this requires recompiling the forward methods once we add/remove + # RL Modules. + self._torch_compile_complete_update = False + if self.config.torch_compile_learner: + if ( + self.config.torch_compile_learner_what_to_compile + == TorchCompileWhatToCompile.COMPLETE_UPDATE + ): + self._torch_compile_complete_update = True + self._compiled_update_initialized = False + else: + self._torch_compile_forward_train = True + + self._torch_compile_cfg = TorchCompileConfig( + torch_dynamo_backend=self.config.torch_compile_learner_dynamo_backend, + torch_dynamo_mode=self.config.torch_compile_learner_dynamo_mode, + ) + + # TODO (simon): See, if we can include these without a torch optimizer. + # self._lr_schedulers = {} + # self._lr_scheduler_classes = None + # if self.config._torch_lr_scheduler_classes: + # self._lr_scheduler_classes = self.config._torch_lr_scheduler_classes + + def _uncompiled_update( + self, + batch: Dict, + params: Dict[ModuleID, NamedParamDict], + **kwargs, + ) -> Tuple[Any, Any, Dict[ModuleID, NamedParamDict], Any]: + """Performs a single functional update using a batch of data. + + This update operates on parameters passed via a functional call to the + `MultiRLModule` and leverages PyTorch 2.0's `autograd` module. Parameters + are not modified in-place within `self._module`; instead, updates are + applied to the cloned parameters provided. + + Args: + batch: A dictionary (or `MultiAgentBatch`) containing training data for + all modules in the `MultiRLModule` (that should be trained). + params: A dictionary of named parameters for each module id. + + Returns: + A tuple consisting of: + 1) the output of a functional forward call to the RLModule using + `params`, + 2) the `loss_per_module` dictionary mapping module IDs to individual + loss tensors, + 3) the functionally updated parameters in the (dict) format passed in, + 4) a metrics dict mapping module IDs to metrics key/value pairs. + + """ + # TODO (sven): Causes weird cuda error when WandB is used. + # Diagnosis thus far: + # - All peek values during metrics.reduce are non-tensors. + # - However, in impala.py::training_step(), a tensor does arrive after learner + # group.update(), so somehow, there is still a race condition + # possible (learner, which performs the reduce() and learner thread, which + # performs the logging of tensors into metrics logger). + self._compute_off_policyness(batch) + + # Make a functional forward pass with the provided parameters. + fwd_out = self._make_functional_call(params, batch) + loss_per_module = self.compute_losses(fwd_out=fwd_out, batch=batch) + # Compute gradients for the provided parameters. + gradients = self.compute_gradients(loss_per_module, params) + + with contextlib.ExitStack() as stack: + if self.config.num_learners > 1: + for mod in self.module.values(): + # Skip non-torch modules, b/c they may not have the `no_sync` API. + if isinstance(mod, torch.nn.Module): + stack.enter_context(mod.no_sync()) + # TODO (simon): See, if we need here postprocessing of gradients. + # postprocessed_gradients = self.postprocess_gradients(gradients) + # Make a stateless (of `params`) update of the `RLModule` parameters. + params = self.apply_gradients(gradients, params) + + # Deactivate tensor-mode on our MetricsLogger and collect the (tensor) + # results. + return fwd_out, loss_per_module, params, {} + + # TODO (simon): Maybe make type for gradients. + @override(DifferentiableLearner) + def compute_gradients( + self, + loss_per_module: Dict[ModuleID, TensorType], + params: Dict[ModuleID, NamedParamDict], + **kwargs, + ) -> Dict[ModuleID, NamedParamDict]: + """Computes functionally gradients based on the given losses. + + This method uses `torch.autograd.grad` to make the backward pass on the + `MultiRLModule` which enables a functional backward pass. If a PyTorch + optimizer is needed a differentiable one must be used (e.g. `torchopt`). + + Args: + loss_per_module: Dict mapping module IDs to their individual total loss + terms, computed by the individual `compute_loss_for_module()` calls. + The overall total loss (sum of loss terms over all modules) is stored + under `loss_per_module[ALL_MODULES]` + params: A dictionary containing named parameters for each module id. + **kwargs: Forward compatibility kwargs. + + Returns: + The (named) gradients in the same (dict) format as `params`. + """ + + # TODO (simon): Add grad scalers later. + total_loss = sum(loss_per_module.values()) + + # Use `torch`'s `autograd` to compute gradients and create a graph, so we can + # compute higher order gradients. Allow specified inputs not being used in outputs + # as probably not all modules/parameters of a `MultiRLModule` are used in the loss. + # Note, parameters are named parameters as this is needed by the + # `torch.func.functional_call` + # TODO (simon): Make sure this works for `MultiRLModule`s. This here can have + # all parameter tensors in a list. But the `functional_call` above needs named + # parameters for each module. Implement this via `foreach_module`. + grads = torch.autograd.grad( + total_loss, + sum((list(param.values()) for mid, param in params.items()), []), + create_graph=True, + retain_graph=True, + allow_unused=True, + ) + + # Map all gradients to their keys. + named_grads = { + module_id: { + name: grad for (name, _), grad in zip(module_params.items(), grads) + } + for module_id, module_params in params.items() + } + + return named_grads + + @override(DifferentiableLearner) + def apply_gradients( + self, + gradients: Dict[ModuleID, NamedParamDict], + params: Dict[ModuleID, NamedParamDict], + ) -> Dict[ModuleID, NamedParamDict]: + """Applies the given gradients in a functional manner. + + This method requires functional parameter updates, meaning modifications + must not be performed in-place (e.g., using an optimizer or directly within + the `MultiRLModule`). + + Args: + gradients: A dictionary containing named gradients for each module id. + params: A dictionary containing named parameters for each module id. + + Returns: + The updated parameters in the same (dict) format as `params`. + """ + policies_to_update = self.learner_config.policies_to_update or list( + gradients.keys() + ) + # Note, because this is a functional update we cannot apply in-place + # modifications of parameters. + updated_params = {} + for module_id, module_grads in gradients.items(): + if module_id not in policies_to_update: + updated_params[module_id] = params[module_id] + continue + updated_params[module_id] = {} + for name, grad in module_grads.items(): + # If updates should not be skipped turn `nan` and `inf` gradients to zero. + if ( + not self.config.torch_skip_nan_gradients + and not torch.isfinite(grad).all() + ): + # Warn the user about `nan` gradients. + logger.warning(f"Gradients {name} contain `nan/inf` values.") + # If updates should be skipped, do not step the optimizer and return. + if not self.config.torch_skip_nan_gradients: + logger.warning( + "Setting `nan/inf` gradients to zero. If updates with " + "`nan/inf` gradients should not be set to zero and instead " + "the update be skipped entirely set `torch_skip_nan_gradients` " + "to `True`." + ) + # If necessary turn `nan` gradients to zero. Note, this can corrupt the + # internal state of the optimizer, if many `nan` gradients occur. + grad = torch.nan_to_num(grad) + + if self.config.torch_skip_nan_gradients or torch.isfinite(grad).all(): + # Update each parameter, by a simple gradient descent step. + updated_params[module_id][name] = ( + params[module_id][name] - self.learner_config.lr * grad + ) + elif grad is None or not torch.isfinite(grad).all(): + logger.warning( + "Skipping this update. If updates with `nan/inf` gradients " + "should not be skipped entirely and instead `nan/inf` " + "gradients set to `zero` set `torch_skip_nan_gradients` to " + "`False`." + ) + return updated_params + + def _make_functional_call( + self, params: Dict[ModuleID, NamedParamDict], batch: MultiAgentBatch + ) -> Dict[ModuleID, NamedParamDict]: + """Makes a functional call for each module in the `MultiRLModule`.""" + return self._module.foreach_module( + lambda mid, m: torch.func.functional_call(m, params[mid], batch[mid]), + return_dict=True, + ) + + @override(DifferentiableLearner) + def _get_tensor_variable( + self, value, dtype=None, trainable=False + ) -> "torch.Tensor": + tensor = torch.tensor( + value, + requires_grad=trainable, + # TODO (simon): Make GPU-trainable. + # device=self._device, + dtype=( + dtype + or ( + torch.float32 + if isinstance(value, float) + else torch.int32 + if isinstance(value, int) + else None + ) + ), + ) + return nn.Parameter(tensor) if trainable else tensor + + def _convert_batch_type( + self, + batch: MultiAgentBatch, + to_device: bool = True, + pin_memory: bool = False, + use_stream: bool = False, + ) -> MultiAgentBatch: + batch = convert_to_torch_tensor( + batch.policy_batches, + device=self._device if to_device else None, + pin_memory=pin_memory, + use_stream=use_stream, + ) + # TODO (sven): This computation of `env_steps` is not accurate! + length = max(len(b) for b in batch.values()) + batch = MultiAgentBatch(batch, env_steps=length) + return batch + + def _compute_off_policyness(self, batch): + # Log off-policy'ness of this batch wrt the current weights. + off_policyness = { + (mid, DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY): ( + (self._weights_seq_no - module_batch[WEIGHTS_SEQ_NO]).float() + ) + for mid, module_batch in batch.items() + if WEIGHTS_SEQ_NO in module_batch + } + for key in off_policyness.keys(): + mid = key[0] + if Columns.LOSS_MASK not in batch[mid]: + off_policyness[key] = torch.mean(off_policyness[key]) + else: + mask = batch[mid][Columns.LOSS_MASK] + num_valid = torch.sum(mask) + off_policyness[key] = torch.sum(off_policyness[key][mask]) / num_valid + self.metrics.log_dict(off_policyness, window=1) + + @override(DifferentiableLearner) + def build(self, device: Optional[DeviceType] = None) -> None: + """Builds the TorchDifferentiableLearner. + + This method is specific to TorchDifferentiableLearner. Before running super() it will + initialize the device properly based on `self.config`, so that `_make_module()` + can place the created module on the correct device. After running super() it + wraps the module in a TorchDDPRLModule if `config.num_learners > 0`. + Note, in inherited classes it is advisable to call the parent's `build()` + after setting up all variables because `configure_optimizer_for_module` is + called in this `Learner.build()`. + """ + # TODO (simon): Allow different `DifferentiableLearner` instances in a + # `MetaLearner` to run on different GPUs. + + super().build(device=device) + + if self._torch_compile_complete_update: + torch._dynamo.reset() + self._compiled_update_initialized = False + self._possibly_compiled_update = torch.compile( + self._uncompiled_update, + backend=self._torch_compile_cfg.torch_dynamo_backend, + mode=self._torch_compile_cfg.torch_dynamo_mode, + **self._torch_compile_cfg.kwargs, + ) + # Otherwise, we use the possibly compiled `forward_train` in + # the module, compiled in the `MetaLearner`. + else: + # Nothing, to do. + self._possibly_compiled_update = self._uncompiled_update + + @override(DifferentiableLearner) + def _update( + self, batch: Dict[str, Any], params: Dict[ModuleID, NamedParamDict] + ) -> Tuple[Any, Dict[ModuleID, NamedParamDict], Any, Any]: + # The first time we call _update after building the learner or + # adding/removing models, we update with the uncompiled update method. + # This makes it so that any variables that may be created during the first + # update step are already there when compiling. More specifically, + # this avoids errors that occur around using defaultdicts with + # torch.compile(). + if ( + self._torch_compile_complete_update + and not self._compiled_update_initialized + ): + self._compiled_update_initialized = True + return self._uncompiled_update(batch, params) + else: + return self._possibly_compiled_update(batch, params) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_learner.py b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..b1ad8fdbf9eb20a1e06fb14845e96303eeb6b3cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_learner.py @@ -0,0 +1,685 @@ +import contextlib +import logging +from collections import defaultdict +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Hashable, + Optional, + Sequence, + Tuple, +) + +from ray.rllib.algorithms.algorithm_config import ( + AlgorithmConfig, + TorchCompileWhatToCompile, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import LR_KEY, Learner +from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, +) +from ray.rllib.core.rl_module.rl_module import ( + RLModule, + RLModuleSpec, +) +from ray.rllib.core.rl_module.torch.torch_rl_module import ( + TorchCompileConfig, + TorchDDPRLModule, + TorchRLModule, +) +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.framework import get_device, try_import_torch +from ray.rllib.utils.metrics import ( + ALL_MODULES, + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_NON_TRAINABLE_PARAMETERS, + NUM_TRAINABLE_PARAMETERS, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import ( + ModuleID, + Optimizer, + Param, + ParamDict, + ShouldModuleBeUpdatedFn, + StateDict, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +torch, nn = try_import_torch() +logger = logging.getLogger(__name__) + + +class TorchLearner(Learner): + + framework: str = "torch" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # Whether to compile the RL Module of this learner. This implies that the. + # forward_train method of the RL Module will be compiled. Further more, + # other forward methods of the RL Module will be compiled on demand. + # This is assumed to not happen, since other forwrad methods are not expected + # to be used during training. + self._torch_compile_forward_train = False + self._torch_compile_cfg = None + # Whether to compile the `_uncompiled_update` method of this learner. This + # implies that everything within `_uncompiled_update` will be compiled, + # not only the forward_train method of the RL Module. + # Note that this is experimental. + # Note that this requires recompiling the forward methods once we add/remove + # RL Modules. + self._torch_compile_complete_update = False + if self.config.torch_compile_learner: + if ( + self.config.torch_compile_learner_what_to_compile + == TorchCompileWhatToCompile.COMPLETE_UPDATE + ): + self._torch_compile_complete_update = True + self._compiled_update_initialized = False + else: + self._torch_compile_forward_train = True + + self._torch_compile_cfg = TorchCompileConfig( + torch_dynamo_backend=self.config.torch_compile_learner_dynamo_backend, + torch_dynamo_mode=self.config.torch_compile_learner_dynamo_mode, + ) + + # Loss scalers for mixed precision training. Map optimizer names to + # associated torch GradScaler objects. + self._grad_scalers = None + if self.config._torch_grad_scaler_class: + self._grad_scalers = defaultdict( + lambda: self.config._torch_grad_scaler_class() + ) + self._lr_schedulers = {} + self._lr_scheduler_classes = None + if self.config._torch_lr_scheduler_classes: + self._lr_scheduler_classes = self.config._torch_lr_scheduler_classes + + @OverrideToImplementCustomLogic + @override(Learner) + def configure_optimizers_for_module( + self, + module_id: ModuleID, + config: "AlgorithmConfig" = None, + ) -> None: + module = self._module[module_id] + + # For this default implementation, the learning rate is handled by the + # attached lr Scheduler (controlled by self.config.lr, which can be a + # fixed value or a schedule setting). + params = self.get_parameters(module) + optimizer = torch.optim.Adam(params) + + # Register the created optimizer (under the default optimizer name). + self.register_optimizer( + module_id=module_id, + optimizer=optimizer, + params=params, + lr_or_lr_schedule=config.lr, + ) + + def _uncompiled_update( + self, + batch: Dict, + **kwargs, + ): + """Performs a single update given a batch of data.""" + + # TODO (sven): Causes weird cuda error when WandB is used. + # Diagnosis thus far: + # - All peek values during metrics.reduce are non-tensors. + # - However, in impala.py::training_step(), a tensor does arrive after learner + # group.update(), so somehow, there is still a race condition + # possible (learner, which performs the reduce() and learner thread, which + # performs the logging of tensors into metrics logger). + self._compute_off_policyness(batch) + + fwd_out = self.module.forward_train(batch) + loss_per_module = self.compute_losses(fwd_out=fwd_out, batch=batch) + gradients = self.compute_gradients(loss_per_module) + + with contextlib.ExitStack() as stack: + if self.config.num_learners > 1: + for mod in self.module.values(): + # Skip non-torch modules, b/c they may not have the `no_sync` API. + if isinstance(mod, torch.nn.Module): + stack.enter_context(mod.no_sync()) + postprocessed_gradients = self.postprocess_gradients(gradients) + self.apply_gradients(postprocessed_gradients) + + return fwd_out, loss_per_module, {} + + @override(Learner) + def compute_gradients( + self, loss_per_module: Dict[ModuleID, TensorType], **kwargs + ) -> ParamDict: + for optim in self._optimizer_parameters: + # `set_to_none=True` is a faster way to zero out the gradients. + optim.zero_grad(set_to_none=True) + + if self._grad_scalers is not None: + total_loss = sum( + self._grad_scalers[mid].scale(loss) + for mid, loss in loss_per_module.items() + ) + else: + total_loss = sum(loss_per_module.values()) + + # If we don't have any loss computations, `sum` returns 0. + if isinstance(total_loss, int): + assert total_loss == 0 + return {} + + total_loss.backward() + grads = {pid: p.grad for pid, p in self._params.items()} + + return grads + + @override(Learner) + def apply_gradients(self, gradients_dict: ParamDict) -> None: + # Set the gradient of the parameters. + for pid, grad in gradients_dict.items(): + # If updates should not be skipped turn `nan` and `inf` gradients to zero. + if ( + not self.config.torch_skip_nan_gradients + and not torch.isfinite(grad).all() + ): + # Warn the user about `nan` gradients. + logger.warning(f"Gradients {pid} contain `nan/inf` values.") + # If updates should be skipped, do not step the optimizer and return. + if not self.config.torch_skip_nan_gradients: + logger.warning( + "Setting `nan/inf` gradients to zero. If updates with " + "`nan/inf` gradients should not be set to zero and instead " + "the update be skipped entirely set `torch_skip_nan_gradients` " + "to `True`." + ) + # If necessary turn `nan` gradients to zero. Note this can corrupt the + # internal state of the optimizer, if many `nan` gradients occur. + self._params[pid].grad = torch.nan_to_num(grad) + # Otherwise, use the gradient as is. + else: + self._params[pid].grad = grad + + # For each optimizer call its step function. + for module_id, optimizer_names in self._module_optimizers.items(): + for optimizer_name in optimizer_names: + optim = self.get_optimizer(module_id, optimizer_name) + # If we have learning rate schedulers for a module add them, if + # necessary. + if self._lr_scheduler_classes is not None: + if ( + module_id not in self._lr_schedulers + or optimizer_name not in self._lr_schedulers[module_id] + ): + # Set for each module and optimizer a scheduler. + self._lr_schedulers[module_id] = {optimizer_name: []} + # If the classes are in a dictionary each module might have + # a different set of schedulers. + if isinstance(self._lr_scheduler_classes, dict): + scheduler_classes = self._lr_scheduler_classes[module_id] + # Else, each module has the same learning rate schedulers. + else: + scheduler_classes = self._lr_scheduler_classes + # Initialize and add the schedulers. + for scheduler_class in scheduler_classes: + self._lr_schedulers[module_id][optimizer_name].append( + scheduler_class(optim) + ) + + # Step through the scaler (unscales gradients, if applicable). + if self._grad_scalers is not None: + scaler = self._grad_scalers[module_id] + scaler.step(optim) + self.metrics.log_value( + (module_id, "_torch_grad_scaler_current_scale"), + scaler.get_scale(), + window=1, # snapshot in time, no EMA/mean. + ) + # Update the scaler. + scaler.update() + # `step` the optimizer (default), but only if all gradients are finite. + elif self.config.torch_skip_nan_gradients or all( + param.grad is None or torch.isfinite(param.grad).all() + for group in optim.param_groups + for param in group["params"] + ): + optim.step() + # If gradients are not all finite warn the user that the update will be + # skipped. + elif not all( + torch.isfinite(param.grad).all() + for group in optim.param_groups + for param in group["params"] + ): + logger.warning( + "Skipping this update. If updates with `nan/inf` gradients " + "should not be skipped entirely and instead `nan/inf` " + "gradients set to `zero` set `torch_skip_nan_gradients` to " + "`False`." + ) + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(Learner) + def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None: + """Called after gradient-based updates are completed. + + Should be overridden to implement custom cleanup-, logging-, or non-gradient- + based Learner/RLModule update logic after(!) gradient-based updates have been + completed. + + Note, for `framework="torch"` users can register + `torch.optim.lr_scheduler.LRScheduler` via + `AlgorithmConfig._torch_lr_scheduler_classes`. These schedulers need to be + stepped here after gradient updates and reported. + + Args: + timesteps: Timesteps dict, which must have the key + `NUM_ENV_STEPS_SAMPLED_LIFETIME`. + # TODO (sven): Make this a more formal structure with its own type. + """ + + # If we have no `torch.optim.lr_scheduler.LRScheduler` registered call the + # `super()`'s method to update RLlib's learning rate schedules. + if not self._lr_schedulers: + return super().after_gradient_based_update(timesteps=timesteps) + + # Only update this optimizer's lr, if a scheduler has been registered + # along with it. + for module_id, optimizer_names in self._module_optimizers.items(): + for optimizer_name in optimizer_names: + # If learning rate schedulers are provided step them here. Note, + # stepping them in `TorchLearner.apply_gradients` updates the + # learning rates during minibatch updates; we want to update + # between whole batch updates. + if ( + module_id in self._lr_schedulers + and optimizer_name in self._lr_schedulers[module_id] + ): + for scheduler in self._lr_schedulers[module_id][optimizer_name]: + scheduler.step() + optimizer = self.get_optimizer(module_id, optimizer_name) + self.metrics.log_value( + # Cut out the module ID from the beginning since it's already + # part of the key sequence: (ModuleID, "[optim name]_lr"). + key=( + module_id, + f"{optimizer_name[len(module_id) + 1:]}_{LR_KEY}", + ), + value=self._get_optimizer_lr(optimizer), + window=1, + ) + + @override(Learner) + def _get_optimizer_state(self) -> StateDict: + ret = {} + for name, optim in self._named_optimizers.items(): + ret[name] = { + "module_id": self._optimizer_name_to_module[name], + "state": convert_to_numpy(optim.state_dict()), + } + return ret + + @override(Learner) + def _set_optimizer_state(self, state: StateDict) -> None: + for name, state_dict in state.items(): + # Ignore updating optimizers matching to submodules not present in this + # Learner's MultiRLModule. + module_id = state_dict["module_id"] + if name not in self._named_optimizers and module_id in self.module: + self.configure_optimizers_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id=module_id), + ) + if name in self._named_optimizers: + self._named_optimizers[name].load_state_dict( + convert_to_torch_tensor(state_dict["state"], device=self._device) + ) + + @override(Learner) + def get_param_ref(self, param: Param) -> Hashable: + return param + + @override(Learner) + def get_parameters(self, module: RLModule) -> Sequence[Param]: + return list(module.parameters()) + + @override(Learner) + def _convert_batch_type( + self, + batch: MultiAgentBatch, + to_device: bool = True, + pin_memory: bool = False, + use_stream: bool = False, + ) -> MultiAgentBatch: + batch = convert_to_torch_tensor( + batch.policy_batches, + device=self._device if to_device else None, + pin_memory=pin_memory, + use_stream=use_stream, + ) + # TODO (sven): This computation of `env_steps` is not accurate! + length = max(len(b) for b in batch.values()) + batch = MultiAgentBatch(batch, env_steps=length) + return batch + + @override(Learner) + def add_module( + self, + *, + module_id: ModuleID, + # TODO (sven): Rename to `rl_module_spec`. + module_spec: RLModuleSpec, + config_overrides: Optional[Dict] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + # Call super's add_module method. + marl_spec = super().add_module( + module_id=module_id, + module_spec=module_spec, + config_overrides=config_overrides, + new_should_module_be_updated=new_should_module_be_updated, + ) + + # we need to ddpify the module that was just added to the pool + module = self._module[module_id] + + if self._torch_compile_forward_train: + module.compile(self._torch_compile_cfg) + elif self._torch_compile_complete_update: + # When compiling the update, we need to reset and recompile + # _uncompiled_update every time we add/remove a module anew. + torch._dynamo.reset() + self._compiled_update_initialized = False + self._possibly_compiled_update = torch.compile( + self._uncompiled_update, + backend=self._torch_compile_cfg.torch_dynamo_backend, + mode=self._torch_compile_cfg.torch_dynamo_mode, + **self._torch_compile_cfg.kwargs, + ) + + if isinstance(module, TorchRLModule): + self._module[module_id].to(self._device) + if self.distributed: + if ( + self._torch_compile_complete_update + or self._torch_compile_forward_train + ): + raise ValueError( + "Using torch distributed and torch compile " + "together tested for now. Please disable " + "torch compile." + ) + self._module.add_module( + module_id, + TorchDDPRLModule(module, **self.config.torch_ddp_kwargs), + override=True, + ) + + self._log_trainable_parameters() + + return marl_spec + + @override(Learner) + def remove_module(self, module_id: ModuleID, **kwargs) -> MultiRLModuleSpec: + marl_spec = super().remove_module(module_id, **kwargs) + + if self._torch_compile_complete_update: + # When compiling the update, we need to reset and recompile + # _uncompiled_update every time we add/remove a module anew. + torch._dynamo.reset() + self._compiled_update_initialized = False + self._possibly_compiled_update = torch.compile( + self._uncompiled_update, + backend=self._torch_compile_cfg.torch_dynamo_backend, + mode=self._torch_compile_cfg.torch_dynamo_mode, + **self._torch_compile_cfg.kwargs, + ) + + self._log_trainable_parameters() + + return marl_spec + + @override(Learner) + def build(self) -> None: + """Builds the TorchLearner. + + This method is specific to TorchLearner. Before running super() it will + initialize the device properly based on `self.config`, so that `_make_module()` + can place the created module on the correct device. After running super() it + wraps the module in a TorchDDPRLModule if `config.num_learners > 0`. + Note, in inherited classes it is advisable to call the parent's `build()` + after setting up all variables because `configure_optimizer_for_module` is + called in this `Learner.build()`. + """ + self._device = get_device(self.config, self.config.num_gpus_per_learner) + + super().build() + + if self._torch_compile_complete_update: + torch._dynamo.reset() + self._compiled_update_initialized = False + self._possibly_compiled_update = torch.compile( + self._uncompiled_update, + backend=self._torch_compile_cfg.torch_dynamo_backend, + mode=self._torch_compile_cfg.torch_dynamo_mode, + **self._torch_compile_cfg.kwargs, + ) + else: + if self._torch_compile_forward_train: + if isinstance(self._module, TorchRLModule): + self._module.compile(self._torch_compile_cfg) + elif isinstance(self._module, MultiRLModule): + for module in self._module._rl_modules.values(): + # Compile only TorchRLModules, e.g. we don't want to compile + # a RandomRLModule. + if isinstance(self._module, TorchRLModule): + module.compile(self._torch_compile_cfg) + else: + raise ValueError( + "Torch compile is only supported for TorchRLModule and " + "MultiRLModule." + ) + + self._possibly_compiled_update = self._uncompiled_update + + self._make_modules_ddp_if_necessary() + + @override(Learner) + def _update(self, batch: Dict[str, Any]) -> Tuple[Any, Any, Any]: + # The first time we call _update after building the learner or + # adding/removing models, we update with the uncompiled update method. + # This makes it so that any variables that may be created during the first + # update step are already there when compiling. More specifically, + # this avoids errors that occur around using defaultdicts with + # torch.compile(). + if ( + self._torch_compile_complete_update + and not self._compiled_update_initialized + ): + self._compiled_update_initialized = True + return self._uncompiled_update(batch) + else: + return self._possibly_compiled_update(batch) + + @OverrideToImplementCustomLogic + def _make_modules_ddp_if_necessary(self) -> None: + """Default logic for (maybe) making all Modules within self._module DDP.""" + + # If the module is a MultiRLModule and nn.Module we can simply assume + # all the submodules are registered. Otherwise, we need to loop through + # each submodule and move it to the correct device. + # TODO (Kourosh): This can result in missing modules if the user does not + # register them in the MultiRLModule. We should find a better way to + # handle this. + if self.config.num_learners > 1: + # Single agent module: Convert to `TorchDDPRLModule`. + if isinstance(self._module, TorchRLModule): + self._module = TorchDDPRLModule( + self._module, **self.config.torch_ddp_kwargs + ) + # Multi agent module: Convert each submodule to `TorchDDPRLModule`. + else: + assert isinstance(self._module, MultiRLModule) + for key in self._module.keys(): + sub_module = self._module[key] + if isinstance(sub_module, TorchRLModule): + # Wrap and override the module ID key in self._module. + self._module.add_module( + key, + TorchDDPRLModule( + sub_module, **self.config.torch_ddp_kwargs + ), + override=True, + ) + + def rl_module_is_compatible(self, module: RLModule) -> bool: + return isinstance(module, nn.Module) + + @override(Learner) + def _check_registered_optimizer( + self, + optimizer: Optimizer, + params: Sequence[Param], + ) -> None: + super()._check_registered_optimizer(optimizer, params) + if not isinstance(optimizer, torch.optim.Optimizer): + raise ValueError( + f"The optimizer ({optimizer}) is not a torch.optim.Optimizer! " + "Only use torch.optim.Optimizer subclasses for TorchLearner." + ) + for param in params: + if not isinstance(param, torch.Tensor): + raise ValueError( + f"One of the parameters ({param}) in the registered optimizer " + "is not a torch.Tensor!" + ) + + @override(Learner) + def _make_module(self) -> MultiRLModule: + module = super()._make_module() + self._map_module_to_device(module) + return module + + def _map_module_to_device(self, module: MultiRLModule) -> None: + """Moves the module to the correct device.""" + if isinstance(module, torch.nn.Module): + module.to(self._device) + else: + for key in module.keys(): + if isinstance(module[key], torch.nn.Module): + module[key].to(self._device) + + @override(Learner) + def _log_trainable_parameters(self) -> None: + # Log number of non-trainable and trainable parameters of our RLModule. + num_trainable_params = defaultdict(int) + num_non_trainable_params = defaultdict(int) + for mid, rlm in self.module._rl_modules.items(): + if isinstance(rlm, TorchRLModule): + for p in rlm.parameters(): + n = p.numel() + if p.requires_grad: + num_trainable_params[(mid, NUM_TRAINABLE_PARAMETERS)] += n + else: + num_non_trainable_params[ + (mid, NUM_NON_TRAINABLE_PARAMETERS) + ] += n + + self.metrics.log_dict( + { + **{ + (ALL_MODULES, NUM_TRAINABLE_PARAMETERS): sum( + num_trainable_params.values() + ), + (ALL_MODULES, NUM_NON_TRAINABLE_PARAMETERS): sum( + num_non_trainable_params.values() + ), + }, + **num_trainable_params, + **num_non_trainable_params, + } + ) + + def _compute_off_policyness(self, batch): + # Log off-policy'ness of this batch wrt the current weights. + off_policyness = { + (mid, DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY): ( + (self._weights_seq_no - module_batch[WEIGHTS_SEQ_NO]).float() + ) + for mid, module_batch in batch.items() + if WEIGHTS_SEQ_NO in module_batch + } + for key in off_policyness.keys(): + mid = key[0] + if Columns.LOSS_MASK not in batch[mid]: + off_policyness[key] = torch.mean(off_policyness[key]) + else: + mask = batch[mid][Columns.LOSS_MASK] + num_valid = torch.sum(mask) + off_policyness[key] = torch.sum(off_policyness[key][mask]) / num_valid + self.metrics.log_dict(off_policyness, window=1) + + @override(Learner) + def _get_tensor_variable( + self, value, dtype=None, trainable=False + ) -> "torch.Tensor": + tensor = torch.tensor( + value, + requires_grad=trainable, + device=self._device, + dtype=( + dtype + or ( + torch.float32 + if isinstance(value, float) + else torch.int32 + if isinstance(value, int) + else None + ) + ), + ) + return nn.Parameter(tensor) if trainable else tensor + + @staticmethod + @override(Learner) + def _get_optimizer_lr(optimizer: "torch.optim.Optimizer") -> float: + for g in optimizer.param_groups: + return g["lr"] + + @staticmethod + @override(Learner) + def _set_optimizer_lr(optimizer: "torch.optim.Optimizer", lr: float) -> None: + for g in optimizer.param_groups: + g["lr"] = lr + + @staticmethod + @override(Learner) + def _get_clip_function() -> Callable: + from ray.rllib.utils.torch_utils import clip_gradients + + return clip_gradients + + @staticmethod + @override(Learner) + def _get_global_norm_function() -> Callable: + from ray.rllib.utils.torch_utils import compute_global_norm + + return compute_global_norm diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_meta_learner.py b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_meta_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..0187942af02c538a02923ff8192f73830e0aca0f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/torch/torch_meta_learner.py @@ -0,0 +1,454 @@ +import contextlib +import logging +from itertools import cycle +from typing import Any, Dict, List, Optional, Tuple + +import ray +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.core import ALL_MODULES +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.torch.torch_differentiable_learner import ( + TorchDifferentiableLearner, +) +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + DATASET_NUM_ITERS_TRAINED, + DATASET_NUM_ITERS_TRAINED_LIFETIME, + DIFFERENTIABLE_LEARNER_RESULTS, +) +from ray.rllib.utils.metrics.utils import to_snake_case +from ray.rllib.utils.typing import ( + EpisodeType, + ModuleID, + NamedParamDict, + ResultDict, + TensorType, +) + +logger = logging.getLogger("__name__") + +torch, nn = try_import_torch() + + +class TorchMetaLearner(TorchLearner): + """A `TorchLearner` designed for meta-learning with functional updates. + + This `TorchLearner` manages one or more `DifferentiableLearner` instances, + which perform functional updates on the `MultiRLModule`. These updates enable + the computation of higher-order gradients, making this class suitable for + meta-learning applications. + + The `update` method executes one or more update loops on its + `DifferentiableLearner` instances, leveraging the functionally updated + parameters for its own learning process. + """ + + # A list of `TorchDifferentiableLearner`s that run inner functional update + # loops. + others: List[TorchDifferentiableLearner] + + def __init__( + self, + **kwargs, + ): + # First, initialize the `TorchLearner`. + super().__init__(**kwargs) + + # Initialize all configured `TorchDifferentiableLearner`s. + self.others = [ + other_config.learner_class( + config=self.config, + learner_config=other_config, + module=self.module, + ) + for other_config in self.config.differentiable_learner_configs + ] + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(TorchLearner) + def build(self) -> None: + """Builds the `TorchMetaLearner`. + + This method should be called before the meta-learner is used. It is + responsible for setting up the `TorchDifferentiableLearner`s. + """ + super().build() + + # Build all `DifferentiableLearner`s. + for other in self.others: + other.build(device=self._device) + + # Ensure 'others' have a module reference. + for other in self.others: + if not other.module: + other.module = self.module + + @OverrideToImplementCustomLogic + @override(Learner) + def update( + self, + batch: Optional[MultiAgentBatch] = None, + batches: Optional[List[MultiAgentBatch]] = None, + batch_refs: Optional[List[ray.ObjectRef]] = None, + episodes: Optional[List[EpisodeType]] = None, + episodes_refs: Optional[List[ray.ObjectRef]] = None, + data_iterators: Optional[List[ray.data.DataIterator]] = None, + training_data: Optional[TrainingData] = None, + *, + # TODO (simon): Maybe defined here other_training_data to provide + # TrainingData for differentiable Learners. List to nest them. + # other_training_data: List[Union[List[Any], TrainingData]] + # TODO (sven): Make this a more formal structure with its own type. + timesteps: Optional[Dict[str, Any]] = None, + num_total_minibatches: int = 0, + num_epochs: int = 1, + minibatch_size: Optional[int] = None, + shuffle_batch_per_epoch: bool = False, + _no_metrics_reduce: bool = False, + others_training_data: Optional[List[TrainingData]] = None, + **kwargs, + ) -> ResultDict: + """Performs a meta-update on the `MultiRLModule`. + + This method allows multiple backward passes on a batch by iteratively + applying functional updates from the `TorchDifferentiableLearner` instances + in `others`. The resulting differentiable parameters are then used in this + learner's own functional forward pass on the `MultiRLModule`, enabling the + computation of higher-order gradients during the backward pass. + """ + self._check_is_built() + + # Call `before_gradient_based_update` to allow for non-gradient based + # preparations-, logging-, and update logic to happen. + self.before_gradient_based_update(timesteps=timesteps or {}) + + if training_data is None: + training_data = TrainingData( + batch=batch, + batches=batches, + batch_refs=batch_refs, + episodes=episodes, + episodes_refs=episodes_refs, + data_iterators=data_iterators, + ) + # Validate training data. + # TODO (simon): Pass in a `TrainingData` list for all differentiable learners. + # to be used if necessary. Some algorithms need it. + training_data.validate() + training_data.solve_refs() + assert training_data.batches is None, "`training_data.batches` must be None!" + + # Increase the _weigths_seq_no in each `update` run. + self._weights_seq_no += 1 + + # Create a batch iterator. + # TODO (simon): Create this method in the `Learner`. + batch_iter = self._create_iterator_if_necessary( + training_data=training_data, + num_total_minibatches=num_total_minibatches, + num_epochs=num_epochs, + minibatch_size=minibatch_size, + shuffle_batch_per_epoch=shuffle_batch_per_epoch, + **kwargs, + ) + + # If no training data for `DifferentiableLearner`s have been passed in, cycle + # over the main training_data for each `DifferentiableLearner`. + if not others_training_data: + others_training_data = cycle([training_data]) + + self.metrics.activate_tensor_mode() + + # Perform the actual looping through the minibatches or the given data iterator. + for iteration, tensor_minibatch in enumerate(batch_iter): + # Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs + # found in this batch. If not, throw an error. + unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set( + self.module.keys() + ) + if unknown_module_ids: + raise ValueError( + f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that " + f"are not in this Learner!" + ) + + # Clone the parameters for the differentiable updates. + params = self._clone_named_parameters() + + # Update all differentiable learners. + others_loss_per_module = [] + others_results = {} + for other, other_training_data in zip(self.others, others_training_data): + # TODO (simon): Maybe return all losses from each other's update step. + params, other_loss_per_module, other_results = other.update( + training_data=other_training_data, + params=params, + # TODO (simon): Check, if this is still needed. + _no_metrics_reduce=_no_metrics_reduce, + **kwargs, + ) + others_loss_per_module.append(other_loss_per_module) + # TODO (simon): Find a more elegant way for naming. + others_results[to_snake_case(other.__class__.__name__)] = other_results + + # Log training results from the `DifferentiableLearner`s. + # TODO (simon): Right now metrics are not carried over b/c of + # the double tensormode problem. + self.metrics.aggregate( + stats_dicts=[others_results], key=DIFFERENTIABLE_LEARNER_RESULTS + ) + + # Make the actual in-graph/traced meta-`_update` call. This should return + # all tensor values (no numpy). + fwd_out, loss_per_module, _ = self._update( + tensor_minibatch.policy_batches, + params, + others_loss_per_module, + ) + + # TODO (sven): Maybe move this into loop above to get metrics more accuratcely + # cover the minibatch/epoch logic. + # Log all timesteps (env, agent, modules) based on given episodes/batch. + self._log_steps_trained_metrics(tensor_minibatch) + + self._set_slicing_by_batch_id(tensor_minibatch, value=False) + + if self.iterator: + # Record the number of batches pulled from the dataset. + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED), + iteration + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_TRAINED_LIFETIME), + iteration + 1, + reduce="sum", + ) + # Log all individual RLModules' loss terms and its registered optimizers' + # current learning rates. + # Note: We do this only once for the last of the minibatch updates, b/c the + # window is only 1 anyways. + for mid, loss in loss_per_module.items(): + self.metrics.log_value( + key=(mid, self.TOTAL_LOSS_KEY), + value=loss, + window=1, + ) + + # Call `after_gradient_based_update` to allow for non-gradient based + # cleanups-, logging-, and update logic to happen. + # TODO (simon): Check, if this should stay here, when running multiple + # gradient steps inside the iterator loop above (could be a complete epoch) + # the target networks might need to be updated earlier. + self.after_gradient_based_update(timesteps=timesteps or {}) + + self.metrics.deactivate_tensor_mode() + + # Reduce results across all minibatch update steps. + if not _no_metrics_reduce: + return self.metrics.reduce() + + def _clone_named_parameters(self): + """Clone named parameters for functional updates.""" + return self.module.foreach_module( + lambda _, m: {name: p.clone() for name, p in m.named_parameters()}, + return_dict=True, + ) + + @OverrideToImplementCustomLogic + @override(TorchLearner) + def _update( + self, + batch: Dict[str, Any], + params: NamedParamDict, + others_loss_per_module: List[Dict[ModuleID, TensorType]] = None, + ) -> Tuple[Any, Any, Any]: + # The first time we call _update after building the learner or + # adding/removing models, we update with the uncompiled update method. + # This makes it so that any variables that may be created during the first + # update step are already there when compiling. More specifically, + # this avoids errors that occur around using defaultdicts with + # torch.compile(). + if ( + self._torch_compile_complete_update + and not self._compiled_update_initialized + ): + self._compiled_update_initialized = True + return self._uncompiled_update(batch, params, others_loss_per_module) + else: + return self._possibly_compiled_update(batch, params, others_loss_per_module) + + def _uncompiled_update( + self, + batch: Dict, + params: NamedParamDict, + others_loss_per_module: List[Dict[ModuleID, TensorType]] = None, + **kwargs, + ): + """Performs a single functional update using a batch of data. + + This update utilizes a functional forward pass via PyTorch 2.0's `func` module. + The updated parameters provided by `TorchDifferentiableLearner` instances are + functions of the `MultiRLModule`'s parameters, allowing for differentiation. + + To ensure compatibility, the `MultiRLModule`'s `forward` method must encapsulate + all logic from `forward_train` and support additional keyword arguments (`kwargs`). + """ + # TODO (sven): Causes weird cuda error when WandB is used. + # Diagnosis thus far: + # - All peek values during metrics.reduce are non-tensors. + # - However, in impala.py::training_step(), a tensor does arrive after learner + # group.update(), so somehow, there is still a race condition + # possible (learner, which performs the reduce() and learner thread, which + # performs the logging of tensors into metrics logger). + self._compute_off_policyness(batch) + + # TODO (simon): For this to work, the `RLModule.forward` must run the + # the `forward_train`. Passing in arguments which makes the `forward` modular + # could be a workaround. + # Make a functional forward call to include higher-order gradients in the meta + # update. + fwd_out = self._make_functional_call(params, batch) + loss_per_module = self.compute_losses( + fwd_out=fwd_out, batch=batch, others_loss_per_module=others_loss_per_module + ) + gradients = self.compute_gradients(loss_per_module) + + with contextlib.ExitStack() as stack: + if self.config.num_learners > 1: + for mod in self.module.values(): + # Skip non-torch modules, b/c they may not have the `no_sync` API. + if isinstance(mod, torch.nn.Module): + stack.enter_context(mod.no_sync()) + postprocessed_gradients = self.postprocess_gradients(gradients) + self.apply_gradients(postprocessed_gradients) + + # Deactivate tensor-mode on our MetricsLogger and collect the (tensor) + # results. + return fwd_out, loss_per_module, {} + + @override(Learner) + def compute_losses( + self, + *, + fwd_out: Dict[str, Any], + batch: Dict[str, Any], + others_loss_per_module: List[Dict[ModuleID, TensorType]] = None, + **kwargs, + ) -> Dict[ModuleID, TensorType]: + """Computes the loss(es) for the module being optimized. + + This method must be overridden by MultiRLModule-specific Learners in order to + define the specific loss computation logic. If the algorithm is single-agent, + only `compute_loss_for_module()` should be overridden instead. If the algorithm + uses independent multi-agent learning (default behavior for RLlib's multi-agent + setups), also only `compute_loss_for_module()` should be overridden, but it will + be called for each individual RLModule inside the MultiRLModule. + For the functional update to work, no `forward` call should be made + within this method, especially not a non-functional one. Instead, use + the model outputs provided by `fwd_out`. + + Losses from `DifferentiableLearner` instances in `others_loss_per_module` + can be leveraged for more advanced module-wise loss calculations. + + Args: + fwd_out: Output from a call to the `forward_train()` method of the + underlying MultiRLModule (`self.module`) during training + (`self.update()`). + batch: The train batch that was used to compute `fwd_out`. + others_loss_per_module: A list of losses per module id from the contained + `DifferentiableLearner` instances in`self.others`. + + Returns: + A dictionary mapping module IDs to individual loss terms. + """ + loss_per_module = {} + for module_id in fwd_out: + module_batch = batch[module_id] + module_fwd_out = fwd_out[module_id] + + module = self.module[module_id].unwrapped() + if isinstance(module, SelfSupervisedLossAPI): + loss = module.compute_self_supervised_loss( + learner=self, + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + others_loss_per_module=others_loss_per_module, + ) + else: + loss = self.compute_loss_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + others_loss_per_module=others_loss_per_module, + ) + loss_per_module[module_id] = loss + + return loss_per_module + + @OverrideToImplementCustomLogic + @override(Learner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + others_loss_per_module: List[Dict[ModuleID, TensorType]] = None, + ) -> TensorType: + """Computes the loss for a single module. + + This method calculates the loss for an individual agent. In multi-agent + scenarios requiring more complex loss computations, consider overriding the + `compute_losses` method instead. + + Losses from `DifferentiableLearner` instances in `others_loss_per_module` + can be leveraged for more advanced module-wise loss calculations. + + Args: + module_id: The id of the module. + config: The AlgorithmConfig specific to the given `module_id`. + batch: The train batch for this particular module. + fwd_out: The output of the forward pass for this particular module. + others_loss_per_module: A list of losses per module id from the contained + `DifferentiableLearner` instances in`self.others`. + + Returns: + A single total loss tensor. If you have more than one optimizer on the + provided `module_id` and would like to compute gradients separately using + these different optimizers, simply add up the individual loss terms for + each optimizer and return the sum. Also, for recording/logging any + individual loss terms, you can use the `Learner.metrics.log_value( + key=..., value=...)` or `Learner.metrics.log_dict()` APIs. See: + :py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for more + information. + """ + # By default, `others_loss_per_module` is not used; instead, the superclass + # method is called. + return super().compute_loss_for_module( + module_id=module_id, config=config, batch=batch, fwd_out=fwd_out + ) + + def _make_functional_call( + self, params: Dict[ModuleID, NamedParamDict], batch: MultiAgentBatch + ) -> Dict[ModuleID, NamedParamDict]: + """Make a functional forward call to all modules in the `MultiRLModule`.""" + return self._module.foreach_module( + lambda mid, m: torch.func.functional_call(m, params[mid], batch[mid]), + return_dict=True, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/training_data.py b/lib/python3.12/site-packages/ray/rllib/core/learner/training_data.py new file mode 100644 index 0000000000000000000000000000000000000000..dc2bb1bd3271925d59338815b94cbb4272a18b85 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/training_data.py @@ -0,0 +1,157 @@ +import dataclasses +from collections import defaultdict +from typing import List, Optional + +import tree # pip install dm_tree + +import ray +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.minibatch_utils import ( + ShardBatchIterator, + ShardEpisodesIterator, + ShardObjectRefIterator, +) +from ray.rllib.utils.typing import EpisodeType + + +# TODO (sven): Switch to dataclass(slots=True) once on py >= 3.10. +@dataclasses.dataclass +class TrainingData: + batch: Optional[MultiAgentBatch] = None + batches: Optional[List[MultiAgentBatch]] = None + batch_refs: Optional[List[ray.ObjectRef]] = None + + episodes: Optional[List[EpisodeType]] = None + episodes_refs: Optional[List[ray.ObjectRef]] = None + + data_iterators: Optional[List[ray.data.iterator.DataIterator]] = None + + def validate(self): + # Exactly one training data type must be provided. + if ( + sum( + td is not None + for td in [ + self.batch, + self.batches, + self.batch_refs, + self.episodes, + self.episodes_refs, + self.data_iterators, + ] + ) + != 1 + ): + raise ValueError("Exactly one training data type must be provided!") + + def shard( + self, + num_shards: int, + len_lookback_buffer: Optional[int] = None, + **kwargs, + ): + # Single batch -> Split into n smaller batches. + if self.batch is not None: + return [ + (TrainingData(batch=b), {}) + for b in ShardBatchIterator(self.batch, num_shards=num_shards) + ] + + # TODO (sven): Do we need a more sohpisticated shard mechanism for this case? + elif self.batches is not None: + assert num_shards == len(self.batches) + return [(TrainingData(batch=b), {}) for b in self.batches] + + # List of batch refs. + elif self.batch_refs is not None: + return [ + (TrainingData(batch_refs=b), {}) + for b in ShardObjectRefIterator(self.batch_refs, num_shards=num_shards) + ] + + # List of episodes -> Split into n equally sized shards (based on the lengths + # of the episodes). + elif self.episodes is not None: + num_total_minibatches = 0 + if "minibatch_size" in kwargs and num_shards > 1: + num_total_minibatches = self._compute_num_total_minibatches( + self.episodes, + num_shards, + kwargs["minibatch_size"], + kwargs.get("num_epochs", 1), + ) + return [ + ( + TrainingData(episodes=e), + {"num_total_minibatches": num_total_minibatches}, + ) + for e in ShardEpisodesIterator( + self.episodes, + num_shards=num_shards, + len_lookback_buffer=len_lookback_buffer, + ) + ] + # List of episodes refs. + elif self.episodes_refs is not None: + return [ + (TrainingData(episodes_refs=e), {}) + for e in ShardObjectRefIterator(self.episodes_refs, num_shards) + ] + # List of data iterators. + else: + assert self.data_iterators and len(self.data_iterators) == num_shards + return [ + (TrainingData(data_iterators=[di]), {}) for di in self.data_iterators + ] + + def solve_refs(self): + # Batch references. + if self.batch_refs is not None: + # Solve the ray.ObjRefs. + batches = tree.flatten(ray.get(self.batch_refs)) + # If only a single batch, set `self.batch`. + if len(batches) == 1: + self.batch = batches[0] + # Otherwise, set `self.batches`. + else: + self.batches = batches + # Empty `self.batch_refs`. + self.batch_refs = None + + # Episode references. + elif self.episodes_refs is not None: + # It's possible that individual refs are invalid due to the EnvRunner + # that produced the ref has crashed or had its entire node go down. + # In this case, try each ref individually and collect only valid results. + try: + episodes = tree.flatten(ray.get(self.episodes_refs)) + except ray.exceptions.OwnerDiedError: + episode_refs = self.episodes_refs + episodes = [] + for ref in episode_refs: + try: + episodes.extend(ray.get(ref)) + except ray.exceptions.OwnerDiedError: + pass + self.episodes = episodes + self.episodes_refs = None + + @staticmethod + def _compute_num_total_minibatches( + episodes, + num_shards, + minibatch_size, + num_epochs, + ): + # Count total number of timesteps per module ID. + if isinstance(episodes[0], MultiAgentEpisode): + per_mod_ts = defaultdict(int) + for ma_episode in episodes: + for sa_episode in ma_episode.agent_episodes.values(): + per_mod_ts[sa_episode.module_id] += len(sa_episode) + max_ts = max(per_mod_ts.values()) + else: + max_ts = sum(map(len, episodes)) + + return int((num_epochs * max_ts) / (num_shards * minibatch_size)) diff --git a/lib/python3.12/site-packages/ray/rllib/core/learner/utils.py b/lib/python3.12/site-packages/ray/rllib/core/learner/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a511dd71c337c0d87f564e45b16ab73f6ee7e2fd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/learner/utils.py @@ -0,0 +1,58 @@ +import copy + +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import NetworkType +from ray.util import PublicAPI + +torch, _ = try_import_torch() + + +def make_target_network(main_net: NetworkType) -> NetworkType: + """Creates a (deep) copy of `main_net` (including synched weights) and returns it. + + Args: + main_net: The main network to return a target network for + + Returns: + The copy of `main_net` that can be used as a target net. Note that the weights + of the returned net are already synched (identical) with `main_net`. + """ + # Deepcopy the main net (this should already take care of synching all weights). + target_net = copy.deepcopy(main_net) + # Make the target net not trainable. + if isinstance(main_net, torch.nn.Module): + target_net.requires_grad_(False) + else: + raise ValueError(f"Unsupported framework for given `main_net` {main_net}!") + + return target_net + + +@PublicAPI(stability="beta") +def update_target_network( + *, + main_net: NetworkType, + target_net: NetworkType, + tau: float, +) -> None: + """Updates a target network (from a "main" network) using Polyak averaging. + + Thereby: + new_target_net_weight = ( + tau * main_net_weight + (1.0 - tau) * current_target_net_weight + ) + + Args: + main_net: The nn.Module to update from. + target_net: The target network to update. + tau: The tau value to use in the Polyak averaging formula. Use 1.0 for a + complete sync of the weights (target and main net will be the exact same + after updating). + """ + if isinstance(main_net, torch.nn.Module): + from ray.rllib.utils.torch_utils import update_target_network as _update_target + + else: + raise ValueError(f"Unsupported framework for given `main_net` {main_net}!") + + _update_target(main_net=main_net, target_net=target_net, tau=tau) diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4508636d4c2a5d10bb7dd98628f5de17c9e3147a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__init__.py @@ -0,0 +1,53 @@ +import logging +import re + +from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, +) +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.util import log_once +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger("ray.rllib") + + +@DeveloperAPI +def validate_module_id(policy_id: str, error: bool = False) -> None: + """Makes sure the given `policy_id` is valid. + + Args: + policy_id: The Policy ID to check. + IMPORTANT: Must not contain characters that + are also not allowed in Unix/Win filesystems, such as: `<>:"/\\|?*` + or a dot `.` or space ` ` at the end of the ID. + error: Whether to raise an error (ValueError) or a warning in case of an + invalid `policy_id`. + + Raises: + ValueError: If the given `policy_id` is not a valid one and `error` is True. + """ + if ( + not isinstance(policy_id, str) + or len(policy_id) == 0 + or re.search('[<>:"/\\\\|?]', policy_id) + or policy_id[-1] in (" ", ".") + ): + msg = ( + f"PolicyID `{policy_id}` not valid! IDs must be a non-empty string, " + "must not contain characters that are also disallowed file- or directory " + "names on Unix/Windows and must not end with a dot `.` or a space ` `." + ) + if error: + raise ValueError(msg) + elif log_once("invalid_policy_id"): + logger.warning(msg) + + +__all__ = [ + "MultiRLModule", + "MultiRLModuleSpec", + "RLModule", + "RLModuleSpec", + "validate_module_id", +] diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf3a6e811d3caed4cc14416518e916c71e3510b3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/default_model_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/default_model_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bea3b08d4ee605905f17f64c1b6435a28756901b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/default_model_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/multi_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/multi_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddf848c72e0ebf86ca3b875849f9f3d331fd9dfb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/multi_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dfb02007c9459478c4062485cf700edecf5758ef Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/__pycache__/rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d9d89bd07a2359c8864b4fa05a38d4b77cabc62b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__init__.py @@ -0,0 +1,17 @@ +from ray.rllib.core.rl_module.apis.inference_only_api import InferenceOnlyAPI +from ray.rllib.core.rl_module.apis.q_net_api import QNetAPI +from ray.rllib.core.rl_module.apis.self_supervised_loss_api import SelfSupervisedLossAPI +from ray.rllib.core.rl_module.apis.target_network_api import ( + TARGET_NETWORK_ACTION_DIST_INPUTS, + TargetNetworkAPI, +) +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI + +__all__ = [ + "InferenceOnlyAPI", + "QNetAPI", + "SelfSupervisedLossAPI", + "TargetNetworkAPI", + "TARGET_NETWORK_ACTION_DIST_INPUTS", + "ValueFunctionAPI", +] diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f41d5a10e61d1014aa71a8532c0d5925b742d6fc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/inference_only_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/inference_only_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db97722d8a3015ca875b0ba75d68979b4209feb6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/inference_only_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/q_net_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/q_net_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18520aef9949d4313d65b7c4642823821b099e4e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/q_net_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/self_supervised_loss_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/self_supervised_loss_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65834d60c3c5d913eebcd944e9afcae335e1d453 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/self_supervised_loss_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/target_network_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/target_network_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d54cb374fc8cd0c963824cdb782d4a6088e8f022 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/target_network_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/value_function_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/value_function_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca01205e39e558694495cc2c572e109b4e3b5ad6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/__pycache__/value_function_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/inference_only_api.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/inference_only_api.py new file mode 100644 index 0000000000000000000000000000000000000000..723c1e3480c236764828c3e5ef4cd48e6189b46a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/inference_only_api.py @@ -0,0 +1,64 @@ +import abc +from typing import List + +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class InferenceOnlyAPI(abc.ABC): + """An API to be implemented by RLModules that have an inference-only mode. + + Only the `get_non_inference_attributes` method needs to get implemented for + an RLModule to have the following functionality: + - On EnvRunners (or when self.inference_only=True), RLlib will remove + those parts of the model not required for action computation. + - An RLModule on a Learner (where `self.inference_only=False`) will + return only those weights from `get_state()` that are part of its inference-only + version, thus possibly saving network traffic/time. + """ + + @abc.abstractmethod + def get_non_inference_attributes(self) -> List[str]: + """Returns a list of attribute names (str) of components NOT used for inference. + + RLlib will use this information to remove those attributes/components from an + RLModule, whose `config.inference_only` is set to True. This so-called + "inference-only setup" is activated. Normally, all RLModules located on + EnvRunners are constructed this way (because they are only used for computing + actions). Similarly, when deployed into a production environment, users should + consider building their RLModules with this flag set to True as well. + + For example: + + .. code-block:: python + + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + + spec = RLModuleSpec(module_class=..., inference_only=True) + + If an RLModule has the following setup() implementation: + + .. testcode:: + :skipif: True + + class MyRLModule(RLModule): + + def setup(self): + self._policy_head = [some NN component] + self._value_function_head = [some NN component] + + self._encoder = [some NN component with attributes: pol and vf + (policy- and value func. encoder)] + + Then its `get_non_inference_attributes()` should return: + ["_value_function_head", "_encoder.vf"]. + + Note the "." notation to separate attributes and their sub-attributes in case + you need more fine-grained control over which exact sub-attributes to exclude in + the inference-only setup. + + Returns: + A list of names (str) of those attributes (or sub-attributes) that should be + excluded (deleted) from this RLModule in case it's setup in + `inference_only` mode. + """ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/q_net_api.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/q_net_api.py new file mode 100644 index 0000000000000000000000000000000000000000..88fb1a00c0fe7c27e4ca72b594d6f548940fe56f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/q_net_api.py @@ -0,0 +1,56 @@ +import abc +from typing import Dict + +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class QNetAPI(abc.ABC): + """An API to be implemented by RLModules used for (distributional) Q-learning. + + RLModules implementing this API must override the `compute_q_values` and the + `compute_advantage_distribution` methods. + """ + + @abc.abstractmethod + def compute_q_values( + self, + batch: Dict[str, TensorType], + ) -> Dict[str, TensorType]: + """Computes Q-values, given encoder, q-net and (optionally), advantage net. + + Note, these can be accompanied by logits and probabilities + in case of distributional Q-learning, i.e. `self.num_atoms > 1`. + + Args: + batch: The batch received in the forward pass. + + Results: + A dictionary containing the Q-value predictions ("qf_preds") + and in case of distributional Q-learning - in addition to the Q-value + predictions ("qf_preds") - the support atoms ("atoms"), the Q-logits + ("qf_logits"), and the probabilities ("qf_probs"). + """ + + def compute_advantage_distribution( + self, + batch: Dict[str, TensorType], + ) -> Dict[str, TensorType]: + """Computes the advantage distribution. + + Note this distribution is identical to the Q-distribution in case no dueling + architecture is used. + + Args: + batch: A dictionary containing a tensor with the outputs of the + forward pass of the Q-head or advantage stream head. + + Returns: + A `dict` containing the support of the discrete distribution for + either Q-values or advantages (in case of a dueling architecture), + ("atoms"), the logits per action and atom and the probabilities + of the discrete distribution (per action and atom of the support). + """ + # Return the Q-distribution by default. + return self.compute_q_values(batch) diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/self_supervised_loss_api.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/self_supervised_loss_api.py new file mode 100644 index 0000000000000000000000000000000000000000..c907a896b01a66a5c56965d3b2fb215132b87c5f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/self_supervised_loss_api.py @@ -0,0 +1,55 @@ +import abc +from typing import TYPE_CHECKING, Any, Dict + +from ray.rllib.utils.typing import ModuleID, TensorType +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.core.learner.learner import Learner + + +@PublicAPI(stability="alpha") +class SelfSupervisedLossAPI(abc.ABC): + """An API to be implemented by RLModules that bring their own self-supervised loss. + + Learners will call these model's `compute_self_supervised_loss()` method instead of + the Learner's own `compute_loss_for_module()` method. + The call signature is identical to the Learner's `compute_loss_for_module()` method + except of an additional mandatory `learner` kwarg. + """ + + @abc.abstractmethod + def compute_self_supervised_loss( + self, + *, + learner: "Learner", + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + **kwargs, + ) -> TensorType: + """Computes the loss for a single module. + + Think of this as computing loss for a single agent. For multi-agent use-cases + that require more complicated computation for loss, consider overriding the + `compute_losses` method instead. + + Args: + learner: The Learner calling this loss method on the RLModule. + module_id: The ID of the RLModule (within a MultiRLModule). + config: The AlgorithmConfig specific to the given `module_id`. + batch: The sample batch for this particular RLModule. + fwd_out: The output of the forward pass for this particular RLModule. + + Returns: + A single total loss tensor. If you have more than one optimizer on the + provided `module_id` and would like to compute gradients separately using + these different optimizers, simply add up the individual loss terms for + each optimizer and return the sum. Also, for recording/logging any + individual loss terms, you can use the `Learner.metrics.log_value( + key=..., value=...)` or `Learner.metrics.log_dict()` APIs. See: + :py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for more + information. + """ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/target_network_api.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/target_network_api.py new file mode 100644 index 0000000000000000000000000000000000000000..a368d4f7c6d3a47988a0b0c533b306d1c792fd6c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/target_network_api.py @@ -0,0 +1,56 @@ +import abc +from typing import Any, Dict, List, Tuple + +from ray.rllib.utils.typing import NetworkType +from ray.util.annotations import PublicAPI + +TARGET_NETWORK_ACTION_DIST_INPUTS = "target_network_action_dist_inputs" + + +@PublicAPI(stability="alpha") +class TargetNetworkAPI(abc.ABC): + """An API to be implemented by RLModules for handling target networks. + + RLModules implementing this API must override the `make_target_networks`, + `get_target_network_pairs`, and the `forward_target` methods. + + Note that the respective Learner that owns the implementing RLModule handles all + target syncing logic. + """ + + @abc.abstractmethod + def make_target_networks(self) -> None: + """Creates the required target nets for this RLModule. + + Use the convenience `ray.rllib.core.learner.utils.make_target_network()` utility + when implementing this method. Pass in an already existing, corresponding "main" + net (for which you need a target net). + This function already takes care of initialization (from the "main" net). + """ + + @abc.abstractmethod + def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]: + """Returns a list of 2-tuples of (main_net, target_net). + + For example, if your RLModule has a property: `self.q_net` and this network + has a corresponding target net `self.target_q_net`, return from this + (overridden) method: [(self.q_net, self.target_q_net)]. + + Note that you need to create all target nets in your overridden + `make_target_networks` method and store the target nets in any properly of your + choice. + + Returns: + A list of 2-tuples of (main_net, target_net) + """ + + @abc.abstractmethod + def forward_target(self, batch: Dict[str, Any]) -> Dict[str, Any]: + """Performs the forward pass through the target net(s). + + Args: + batch: The batch to use for the forward pass. + + Returns: + The results from the forward pass(es) through the target net(s). + """ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/value_function_api.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/value_function_api.py new file mode 100644 index 0000000000000000000000000000000000000000..46e33d2fd315fe75e555e1629d4a1e4a0e44324b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/apis/value_function_api.py @@ -0,0 +1,35 @@ +import abc +from typing import Any, Dict, Optional + +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class ValueFunctionAPI(abc.ABC): + """An API to be implemented by RLModules for handling value function-based learning. + + RLModules implementing this API must override the `compute_values` method. + """ + + @abc.abstractmethod + def compute_values( + self, + batch: Dict[str, Any], + embeddings: Optional[Any] = None, + ) -> TensorType: + """Computes the value estimates given `batch`. + + Args: + batch: The batch to compute value function estimates for. + embeddings: Optional embeddings already computed from the `batch` (by + another forward pass through the model's encoder (or other subcomponent + that computes an embedding). For example, the caller of thie method + should provide `embeddings` - if available - to avoid duplicate passes + through a shared encoder. + + Returns: + A tensor of shape (B,) or (B, T) (in case the input `batch` has a + time dimension. Note that the last value dimension should already be + squeezed out (not 1!). + """ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/default_model_config.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/default_model_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e53b852a4e7b5ec3c25a2b713e8779d8cdec834e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/default_model_config.py @@ -0,0 +1,204 @@ +from dataclasses import dataclass, field +from typing import Callable, List, Optional, Union + +from ray.rllib.utils.typing import ConvFilterSpec +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +@dataclass +class DefaultModelConfig: + """Dataclass to configure all default RLlib RLModules. + + Users should NOT use this class for configuring their own custom RLModules, but + use a custom `model_config` dict with arbitrary (str) keys passed into the + `RLModuleSpec` used to define the custom RLModule. + For example: + + .. testcode:: + + import gymnasium as gym + import numpy as np + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.examples.rl_modules.classes.tiny_atari_cnn_rlm import ( + TinyAtariCNN + ) + + my_rl_module = RLModuleSpec( + module_class=TinyAtariCNN, + observation_space=gym.spaces.Box(-1.0, 1.0, (64, 64, 4), np.float32), + action_space=gym.spaces.Discrete(7), + # DreamerV3-style stack working on a 64x64, color or 4x-grayscale-stacked, + # normalized image. + model_config={ + "conv_filters": [[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]], + }, + ).build() + + Only RLlib's default RLModules (defined by the various algorithms) should use + this dataclass. Pass an instance of it into your algorithm config like so: + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig + + config = ( + PPOConfig() + .rl_module( + model_config=DefaultModelConfig(fcnet_hiddens=[32, 32]), + ) + ) + """ + + # ==================================================== + # MLP stacks + # ==================================================== + # __sphinx_doc_default_model_config_fcnet_begin__ + #: List containing the sizes (number of nodes) of a fully connected (MLP) stack. + #: Note that in an encoder-based default architecture with a policy head (and + #: possible value head), this setting only affects the encoder component. To set the + #: policy (and value) head sizes, use `post_fcnet_hiddens`, instead. For example, + #: if you set `fcnet_hiddens=[32, 32]` and `post_fcnet_hiddens=[64]`, you would get + #: an RLModule with a [32, 32] encoder, a [64, act-dim] policy head, and a [64, 1] + #: value head (if applicable). + fcnet_hiddens: List[int] = field(default_factory=lambda: [256, 256]) + #: Activation function descriptor for the stack configured by `fcnet_hiddens`. + #: Supported values are: 'tanh', 'relu', 'swish' (or 'silu', which is the same), + #: and 'linear' (or None). + fcnet_activation: str = "tanh" + #: Initializer function or class descriptor for the weight/kernel matrices in the + #: stack configured by `fcnet_hiddens`. Supported values are the initializer names + #: (str), classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + fcnet_kernel_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `fcnet_kernel_initializer`. + fcnet_kernel_initializer_kwargs: Optional[dict] = None + #: Initializer function or class descriptor for the bias vectors in the stack + #: configured by `fcnet_hiddens`. Supported values are the initializer names (str), + #: classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + fcnet_bias_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `fcnet_bias_initializer`. + fcnet_bias_initializer_kwargs: Optional[dict] = None + # __sphinx_doc_default_model_config_fcnet_end__ + + # ==================================================== + # Conv2D stacks + # ==================================================== + # __sphinx_doc_default_model_config_conv_begin__ + #: List of lists of format [num_out_channels, kernel, stride] defining a Conv2D + #: stack if the input space is 2D. Each item in the outer list represents one Conv2D + #: layer. `kernel` and `stride` may be single ints (width and height have same + #: value) or 2-tuples (int, int) specifying width and height dimensions separately. + #: If None (default) and the input space is 2D, RLlib tries to find a default filter + #: setup given the exact input dimensions. + conv_filters: Optional[ConvFilterSpec] = None + #: Activation function descriptor for the stack configured by `conv_filters`. + #: Supported values are: 'tanh', 'relu', 'swish' (or 'silu', which is the same), and + #: 'linear' (or None). + conv_activation: str = "relu" + #: Initializer function or class descriptor for the weight/kernel matrices in the + #: stack configured by `conv_filters`. Supported values are the initializer names + #: (str), classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + conv_kernel_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `conv_kernel_initializer`. + conv_kernel_initializer_kwargs: Optional[dict] = None + #: Initializer function or class descriptor for the bias vectors in the stack + #: configured by `conv_filters`. Supported values are the initializer names (str), + #: classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + conv_bias_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `conv_bias_initializer`. + conv_bias_initializer_kwargs: Optional[dict] = None + # __sphinx_doc_default_model_config_conv_end__ + + # ==================================================== + # Head configs (e.g. policy- or value function heads) + # ==================================================== + #: List containing the sizes (number of nodes) of a fully connected (MLP) head (ex. + #: policy-, value-, or Q-head). Note that in order to configure the encoder + #: architecture, use `fcnet_hiddens`, instead. + head_fcnet_hiddens: List[int] = field(default_factory=lambda: []) + #: Activation function descriptor for the stack configured by `head_fcnet_hiddens`. + #: Supported values are: 'tanh', 'relu', 'swish' (or 'silu', which is the same), + #: and 'linear' (or None). + head_fcnet_activation: str = "relu" + #: Initializer function or class descriptor for the weight/kernel matrices in the + #: stack configured by `head_fcnet_hiddens`. Supported values are the initializer + #: names (str), classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + head_fcnet_kernel_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `head_fcnet_kernel_initializer`. + head_fcnet_kernel_initializer_kwargs: Optional[dict] = None + #: Initializer function or class descriptor for the bias vectors in the stack + #: configured by `head_fcnet_hiddens`. Supported values are the initializer names + #: (str), classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + head_fcnet_bias_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `head_fcnet_bias_initializer`. + head_fcnet_bias_initializer_kwargs: Optional[dict] = None + + # ==================================================== + # Continuous action settings + # ==================================================== + #: If True, for DiagGaussian action distributions (or any other continuous control + #: distribution), make the second half of the policy's outputs a "free" bias + #: parameter, rather than state-/NN-dependent nodes. In this case, the number of + #: nodes of the policy head have the same dimension as the action space as no slots + #: for log(stddev) are required (only for the mean values). + free_log_std: bool = False + #: Whether to clip the log(stddev) when using a DiagGaussian action distribution + #: (or any other continuous control distribution). This can stabilize training and + #: avoid very small or large log(stddev) values leading to numerical instabilities + #: turning outputs to `nan`. The default is to clamp the log(stddev) in between + #: -20 and 20. Set to float("inf") for no clamping. + log_std_clip_param: float = 20.0 + #: Whether encoder layers (defined by `fcnet_hiddens` or `conv_filters`) should be + #: shared between policy- and value function. + vf_share_layers: bool = True + + # ==================================================== + # LSTM settings + # ==================================================== + #: Whether to wrap the encoder component (defined by `fcnet_hiddens` or + #: `conv_filters`) with an LSTM. + use_lstm: bool = False + #: The maximum seq len for building the train batch for an LSTM model. + #: Defaults to 20. + max_seq_len: int = 20 + #: The size of the LSTM cell. + lstm_cell_size: int = 256 + lstm_use_prev_action: bool = False + lstm_use_prev_reward: bool = False + #: Initializer function or class descriptor for the weight/kernel matrices in the + #: LSTM layer. Supported values are the initializer names (str), classes or + #: functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + lstm_kernel_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `lstm_kernel_initializer`. + lstm_kernel_initializer_kwargs: Optional[dict] = None + #: Initializer function or class descriptor for the bias vectors in the stack + #: configured by the LSTM layer. Supported values are the initializer names (str), + #: classes or functions listed by the frameworks (`torch`). See + #: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default), + #: the default initializer defined by `torch` is used. + lstm_bias_initializer: Optional[Union[str, Callable]] = None + #: Kwargs passed into the initializer function defined through + #: `lstm_bias_initializer`. + lstm_bias_initializer_kwargs: Optional[dict] = None diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/multi_rl_module.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/multi_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..d9d04b58866ad954976e4822a179873804263e36 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/multi_rl_module.py @@ -0,0 +1,832 @@ +import copy +import dataclasses +import logging +import pprint +from typing import ( + Any, + Callable, + Collection, + Dict, + ItemsView, + KeysView, + List, + Optional, + Set, + Tuple, + Type, + Union, + ValuesView, +) + +import gymnasium as gym + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + Deprecated, + deprecation_warning, +) +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + override, +) +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.serialization import ( + deserialize_type, + gym_space_from_dict, + gym_space_to_dict, + serialize_type, +) +from ray.rllib.utils.typing import ModuleID, StateDict, T +from ray.util.annotations import PublicAPI + +logger = logging.getLogger("ray.rllib") + + +@PublicAPI(stability="alpha") +class MultiRLModule(RLModule): + """Base class for an RLModule that contains n sub-RLModules. + + This class holds a mapping from ModuleID to underlying RLModules. It provides + a convenient way of accessing each individual module, as well as accessing all of + them with only one API call. Whether a given module is trainable is + determined by the caller of this class (not the instance of this class itself). + + The extension of this class can include any arbitrary neural networks as part of + the MultiRLModule. For example, a MultiRLModule can include a shared encoder network + that is used by all the individual (single-agent) RLModules. It is up to the user + to decide how to implement this class. + + The default implementation assumes the data communicated as input and output of + the APIs in this class are `Dict[ModuleID, Dict[str, Any]]` types. The + `MultiRLModule` by default loops through each `module_id`, and runs the forward pass + of the corresponding `RLModule` object with the associated `batch` within the + input. + It also assumes that the underlying RLModules do not share any parameters or + communication with one another. The behavior of modules with such advanced + communication would be undefined by default. To share parameters or communication + between the underlying RLModules, you should implement your own + `MultiRLModule` subclass. + """ + + def __init__( + self, + config=DEPRECATED_VALUE, + *, + observation_space: Optional[gym.Space] = None, + action_space: Optional[gym.Space] = None, + inference_only: Optional[bool] = None, + # TODO (sven): Ignore learner_only setting for now on MultiRLModule. + learner_only: Optional[bool] = None, + model_config: Optional[dict] = None, + rl_module_specs: Optional[Dict[ModuleID, RLModuleSpec]] = None, + **kwargs, + ) -> None: + """Initializes a MultiRLModule instance. + + Args: + observation_space: The MultiRLModule's observation space. + action_space: The MultiRLModule's action space. + inference_only: The MultiRLModule's `inference_only` setting. If True, force + sets all inference_only flags inside `rl_module_specs` also to True. + If None, infers the value for `inference_only` by setting it to True, + iff all `inference_only` flags inside `rl_module_specs`, otherwise to + False. + model_config: The MultiRLModule's `model_config` dict. + rl_module_specs: A dict mapping ModuleIDs to `RLModuleSpec` instances used + to create the submodules. + """ + if config != DEPRECATED_VALUE and isinstance(config, MultiRLModuleConfig): + deprecation_warning( + old="MultiRLModule(config=..)", + new="MultiRLModule(*, observation_space=.., action_space=.., " + "inference_only=.., model_config=.., rl_module_specs=..)", + error=True, + ) + + # Make sure we don't alter incoming module specs in this c'tor. + rl_module_specs = copy.deepcopy(rl_module_specs or {}) + # Figure out global inference_only setting. + # If not provided (None), only if all submodules are + # inference_only, this MultiRLModule will be inference_only. + inference_only = ( + inference_only + if inference_only is not None + else all(spec.inference_only for spec in rl_module_specs.values()) + ) + # If given inference_only=True, make all submodules also inference_only (before + # creating them). + if inference_only is True: + for rl_module_spec in rl_module_specs.values(): + rl_module_spec.inference_only = True + self._check_module_specs(rl_module_specs) + self.rl_module_specs = rl_module_specs + + super().__init__( + observation_space=observation_space, + action_space=action_space, + inference_only=inference_only, + learner_only=None, + catalog_class=None, + model_config=model_config, + **kwargs, + ) + + @OverrideToImplementCustomLogic + @override(RLModule) + def setup(self): + """Sets up the underlying, individual RLModules.""" + self._rl_modules = {} + # Make sure all individual RLModules have the same framework OR framework=None. + framework = None + for module_id, rl_module_spec in self.rl_module_specs.items(): + self._rl_modules[module_id] = rl_module_spec.build() + if framework is None: + framework = self._rl_modules[module_id].framework + else: + assert self._rl_modules[module_id].framework in [None, framework] + self.framework = framework + + @override(RLModule) + def _forward( + self, + batch: Dict[ModuleID, Any], + **kwargs, + ) -> Dict[ModuleID, Dict[str, Any]]: + """Generic forward pass method, used in all phases of training and evaluation. + + If you need a more nuanced distinction between forward passes in the different + phases of training and evaluation, override the following methods instead: + For distinct action computation logic w/o exploration, override the + `self._forward_inference()` method. + For distinct action computation logic with exploration, override the + `self._forward_exploration()` method. + For distinct forward pass logic before loss computation, override the + `self._forward_train()` method. + + Args: + batch: The input batch, a dict mapping from ModuleID to individual modules' + batches. + **kwargs: Additional keyword arguments. + + Returns: + The output of the forward pass. + """ + return { + mid: self._rl_modules[mid]._forward(batch[mid], **kwargs) + for mid in batch.keys() + if mid in self + } + + @override(RLModule) + def _forward_inference( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + """Forward-pass used for action computation without exploration behavior. + + Override this method only, if you need specific behavior for non-exploratory + action computation behavior. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return { + mid: self._rl_modules[mid]._forward_inference(batch[mid], **kwargs) + for mid in batch.keys() + if mid in self + } + + @override(RLModule) + def _forward_exploration( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + """Forward-pass used for action computation with exploration behavior. + + Override this method only, if you need specific behavior for exploratory + action computation behavior. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return { + mid: self._rl_modules[mid]._forward_exploration(batch[mid], **kwargs) + for mid in batch.keys() + if mid in self + } + + @override(RLModule) + def _forward_train( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + """Forward-pass used before the loss computation (training). + + Override this method only, if you need specific behavior and outputs for your + loss computations. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return { + mid: self._rl_modules[mid]._forward_train(batch[mid], **kwargs) + for mid in batch.keys() + if mid in self + } + + @OverrideToImplementCustomLogic + @override(RLModule) + def get_initial_state(self) -> Any: + # TODO (sven): Replace by call to `self.foreach_module`, but only if this method + # supports returning dicts. + ret = {} + for module_id, module in self._rl_modules.items(): + ret[module_id] = module.get_initial_state() + return ret + + @OverrideToImplementCustomLogic + @override(RLModule) + def is_stateful(self) -> bool: + initial_state = self.get_initial_state() + assert isinstance(initial_state, dict), ( + "The initial state of an RLModule must be a dict, but is " + f"{type(initial_state)} instead." + ) + return bool(any(sa_init_state for sa_init_state in initial_state.values())) + + def add_module( + self, + module_id: ModuleID, + module: RLModule, + *, + override: bool = False, + ) -> None: + """Adds a module at run time to the multi-agent module. + + Args: + module_id: The module ID to add. If the module ID already exists and + override is False, an error is raised. If override is True, the module + is replaced. + module: The module to add. + override: Whether to override the module if it already exists. + + Raises: + ValueError: If the module ID already exists and override is False. + Warnings are raised if the module id is not valid according to the + logic of ``validate_module_id()``. + """ + from ray.rllib.core.rl_module import validate_module_id + + validate_module_id(module_id) + + if module_id in self._rl_modules and not override: + raise ValueError( + f"Module ID {module_id} already exists. If your intention is to " + "override, set override=True." + ) + # Set our own inference_only flag to False as soon as any added Module + # has `inference_only=False`. + if not module.inference_only: + self.inference_only = False + + # Check framework of incoming RLModule against `self.framework`. + if module.framework is not None: + if self.framework is None: + self.framework = module.framework + elif module.framework != self.framework: + raise ValueError( + f"Framework ({module.framework}) of incoming RLModule does NOT " + f"match framework ({self.framework}) of MultiRLModule! If the " + f"added module should not be trained, try setting its framework " + f"to None." + ) + + self._rl_modules[module_id] = module + # Update our RLModuleSpecs dict, such that - if written to disk - + # it'll allow for proper restoring this instance through `.from_checkpoint()`. + self.rl_module_specs[module_id] = RLModuleSpec.from_module(module) + + def remove_module( + self, module_id: ModuleID, *, raise_err_if_not_found: bool = True + ) -> None: + """Removes a module at runtime from the multi-agent module. + + Args: + module_id: The module ID to remove. + raise_err_if_not_found: Whether to raise an error if the module ID is not + found. + Raises: + ValueError: If the module ID does not exist and raise_err_if_not_found is + True. + """ + if raise_err_if_not_found: + self._check_module_exists(module_id) + del self._rl_modules[module_id] + del self.rl_module_specs[module_id] + + def foreach_module( + self, + func: Callable[[ModuleID, RLModule, Optional[Any]], T], + *, + return_dict: bool = False, + **kwargs, + ) -> Union[List[T], Dict[ModuleID, T]]: + """Calls the given function with each (module_id, module). + + Args: + func: The function to call with each (module_id, module) tuple. + return_dict: Whether to return a dict mapping ModuleID to the individual + module's return values of calling `func`. If False (default), return + a list. + + Returns: + The list of return values of all calls to + `func([module_id, module, **kwargs])` or a dict (if `return_dict=True`) + mapping ModuleIDs to the respective models' return values. + """ + ret_dict = { + module_id: func(module_id, module.unwrapped(), **kwargs) + for module_id, module in self._rl_modules.items() + } + if return_dict: + return ret_dict + return list(ret_dict.values()) + + def __contains__(self, item) -> bool: + """Returns whether the given `item` (ModuleID) is present in self.""" + return item in self._rl_modules + + def __getitem__(self, module_id: ModuleID) -> RLModule: + """Returns the RLModule with the given module ID. + + Args: + module_id: The module ID to get. + + Returns: + The RLModule with the given module ID. + + Raises: + KeyError: If `module_id` cannot be found in self. + """ + self._check_module_exists(module_id) + return self._rl_modules[module_id] + + def get( + self, + module_id: ModuleID, + default: Optional[RLModule] = None, + ) -> Optional[RLModule]: + """Returns the module with the given module ID or default if not found in self. + + Args: + module_id: The module ID to get. + + Returns: + The RLModule with the given module ID or `default` if `module_id` not found + in `self`. + """ + if module_id not in self._rl_modules: + return default + return self._rl_modules[module_id] + + def items(self) -> ItemsView[ModuleID, RLModule]: + """Returns an ItemsView over the module IDs in this MultiRLModule.""" + return self._rl_modules.items() + + def keys(self) -> KeysView[ModuleID]: + """Returns a KeysView over the module IDs in this MultiRLModule.""" + return self._rl_modules.keys() + + def values(self) -> ValuesView[ModuleID]: + """Returns a ValuesView over the module IDs in this MultiRLModule.""" + return self._rl_modules.values() + + def __len__(self) -> int: + """Returns the number of RLModules within this MultiRLModule.""" + return len(self._rl_modules) + + def __repr__(self) -> str: + return f"MARL({pprint.pformat(self._rl_modules)})" + + @override(RLModule) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + inference_only: bool = False, + **kwargs, + ) -> StateDict: + state = {} + + for module_id, rl_module in self.get_checkpointable_components(): + if self._check_component(module_id, components, not_components): + state[module_id] = rl_module.get_state( + components=self._get_subcomponents(module_id, components), + not_components=self._get_subcomponents(module_id, not_components), + inference_only=inference_only, + ) + return state + + @override(RLModule) + def set_state(self, state: StateDict) -> None: + """Sets the state of the multi-agent module. + + It is assumed that the state_dict is a mapping from module IDs to the + corresponding module's state. This method sets the state of each module by + calling their set_state method. If you want to set the state of some of the + RLModules within this MultiRLModule your state_dict can only include the + state of those RLModules. Override this method to customize the state_dict for + custom more advanced multi-agent use cases. + + Args: + state: The state dict to set. + """ + # Now, set the individual states + for module_id, module_state in state.items(): + if module_id in self: + self._rl_modules[module_id].set_state(module_state) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + { + "observation_space": self.observation_space, + "action_space": self.action_space, + "inference_only": self.inference_only, + "learner_only": self.learner_only, + "model_config": self.model_config, + "rl_module_specs": self.rl_module_specs, + }, # **kwargs + ) + + @override(Checkpointable) + def get_checkpointable_components(self) -> List[Tuple[str, Checkpointable]]: + return list(self._rl_modules.items()) + + @override(RLModule) + def as_multi_rl_module(self) -> "MultiRLModule": + """Returns self in order to match `RLModule.as_multi_rl_module()` behavior. + + This method is overridden to avoid double wrapping. + + Returns: + The instance itself. + """ + return self + + @classmethod + def _check_module_specs(cls, rl_module_specs: Dict[ModuleID, RLModuleSpec]): + """Checks the individual RLModuleSpecs for validity. + + Args: + rl_module_specs: Dict mapping ModuleIDs to the respective RLModuleSpec. + + Raises: + ValueError: If any RLModuleSpec is invalid. + """ + for module_id, rl_module_spec in rl_module_specs.items(): + if not isinstance(rl_module_spec, RLModuleSpec): + raise ValueError(f"Module {module_id} is not a RLModuleSpec object.") + + def _check_module_exists(self, module_id: ModuleID) -> None: + if module_id not in self._rl_modules: + raise KeyError( + f"Module with module_id {module_id} not found. " + f"Available modules: {set(self.keys())}" + ) + + @Deprecated(error=False) + def output_specs_train(self): + pass + + @Deprecated(error=False) + def output_specs_inference(self): + pass + + @Deprecated(error=False) + def output_specs_exploration(self): + pass + + @Deprecated(error=False) + def _default_input_specs(self): + pass + + +@PublicAPI(stability="alpha") +@dataclasses.dataclass +class MultiRLModuleSpec: + """A utility spec class to make it constructing MultiRLModules easier. + + Users can extend this class to modify the behavior of base class. For example to + share neural networks across the modules, the build method can be overridden to + create the shared module first and then pass it to custom module classes that would + then use it as a shared module. + """ + + #: The class of the MultiRLModule to construct. By default, + #: this is the base `MultiRLModule` class. + multi_rl_module_class: Type[MultiRLModule] = MultiRLModule + #: Optional global observation space for the MultiRLModule. + #: Useful for shared network components that live only inside the MultiRLModule + #: and don't have their own ModuleID and own RLModule within + #: `self._rl_modules`. + observation_space: Optional[gym.Space] = None + #: Optional global action space for the MultiRLModule. Useful for + #: shared network components that live only inside the MultiRLModule and don't + #: have their own ModuleID and own RLModule within `self._rl_modules`. + action_space: Optional[gym.Space] = None + #: An optional global inference_only flag. If not set (None by + #: default), considers the MultiRLModule to be inference_only=True, only if all + #: submodules also have their own inference_only flags set to True. + inference_only: Optional[bool] = None + # TODO (sven): Once we support MultiRLModules inside other MultiRLModules, we would + # need this flag in here as well, but for now, we'll leave it out for simplicity. + # learner_only: bool = False + #: An optional global model_config dict. Useful to configure shared + #: network components that only live inside the MultiRLModule and don't have + #: their own ModuleID and own RLModule within `self._rl_modules`. + model_config: Optional[dict] = None + #: The module specs for each individual module. It can be either + #: an RLModuleSpec used for all module_ids or a dictionary mapping from module + #: IDs to RLModuleSpecs for each individual module. + rl_module_specs: Union[RLModuleSpec, Dict[ModuleID, RLModuleSpec]] = None + + # TODO (sven): Deprecate these in favor of using the pure Checkpointable APIs for + # loading and saving state. + load_state_path: Optional[str] = None + modules_to_load: Optional[Set[ModuleID]] = None + + # Deprecated: Do not use anymore. + module_specs: Optional[Union[RLModuleSpec, Dict[ModuleID, RLModuleSpec]]] = None + + def __post_init__(self): + if self.module_specs is not None: + deprecation_warning( + old="MultiRLModuleSpec(module_specs=..)", + new="MultiRLModuleSpec(rl_module_specs=..)", + error=True, + ) + if self.rl_module_specs is None: + raise ValueError( + "Module_specs cannot be None. It should be either a " + "RLModuleSpec or a dictionary mapping from module IDs to " + "RLModuleSpecs for each individual module." + ) + self.module_specs = self.rl_module_specs + # Figure out global inference_only setting. + # If not provided (None), only if all submodules are + # inference_only, this MultiRLModule will be inference_only. + self.inference_only = ( + self.inference_only + if self.inference_only is not None + else all(spec.inference_only for spec in self.rl_module_specs.values()) + ) + + @OverrideToImplementCustomLogic + def build(self, module_id: Optional[ModuleID] = None) -> RLModule: + """Builds either the MultiRLModule or a (single) sub-RLModule under `module_id`. + + Args: + module_id: Optional ModuleID of a single RLModule to be built. If None + (default), builds the MultiRLModule. + + Returns: + The built RLModule if `module_id` is provided, otherwise the built + MultiRLModule. + """ + self._check_before_build() + + # ModuleID provided, return single-agent RLModule. + if module_id: + return self.rl_module_specs[module_id].build() + + # Return MultiRLModule. + try: + module = self.multi_rl_module_class( + observation_space=self.observation_space, + action_space=self.action_space, + inference_only=self.inference_only, + model_config=( + dataclasses.asdict(self.model_config) + if dataclasses.is_dataclass(self.model_config) + else self.model_config + ), + rl_module_specs=self.rl_module_specs, + ) + # Older custom model might still require the old `MultiRLModuleConfig` under + # the `config` arg. + except AttributeError as e: + if self.multi_rl_module_class is not MultiRLModule: + multi_rl_module_config = self.get_rl_module_config() + module = self.multi_rl_module_class(multi_rl_module_config) + else: + raise e + + return module + + def add_modules( + self, + module_specs: Dict[ModuleID, RLModuleSpec], + override: bool = True, + ) -> None: + """Add new module specs to the spec or updates existing ones. + + Args: + module_specs: The mapping for the module_id to the single-agent module + specs to be added to this multi-agent module spec. + override: Whether to override the existing module specs if they already + exist. If False, they are only updated. + """ + if self.rl_module_specs is None: + self.rl_module_specs = {} + for module_id, module_spec in module_specs.items(): + if override or module_id not in self.rl_module_specs: + # Disable our `inference_only` as soon as any single-agent module has + # `inference_only=False`. + if not module_spec.inference_only: + self.inference_only = False + self.rl_module_specs[module_id] = module_spec + else: + self.rl_module_specs[module_id].update(module_spec) + + def remove_modules(self, module_ids: Union[ModuleID, Collection[ModuleID]]) -> None: + """Removes the provided ModuleIDs from this MultiRLModuleSpec. + + Args: + module_ids: Collection of the ModuleIDs to remove from this spec. + """ + for module_id in force_list(module_ids): + self.rl_module_specs.pop(module_id, None) + + @classmethod + def from_module(self, module: MultiRLModule) -> "MultiRLModuleSpec": + """Creates a MultiRLModuleSpec from a MultiRLModule. + + Args: + module: The MultiRLModule to create the spec from. + + Returns: + The MultiRLModuleSpec. + """ + # we want to get the spec of the underlying unwrapped module that way we can + # easily reconstruct it. The only wrappers that we expect to support today are + # wrappers that allow us to do distributed training. Those will be added back + # by the learner if necessary. + rl_module_specs = { + module_id: RLModuleSpec.from_module(rl_module.unwrapped()) + for module_id, rl_module in module._rl_modules.items() + } + multi_rl_module_class = module.__class__ + return MultiRLModuleSpec( + multi_rl_module_class=multi_rl_module_class, + observation_space=module.observation_space, + action_space=module.action_space, + inference_only=module.inference_only, + model_config=module.model_config, + rl_module_specs=rl_module_specs, + ) + + def _check_before_build(self): + if not isinstance(self.rl_module_specs, dict): + raise ValueError( + f"When build() is called on {self.__class__}, the `rl_module_specs` " + "attribute should be a dictionary mapping ModuleIDs to " + "RLModuleSpecs for each individual RLModule." + ) + + def to_dict(self) -> Dict[str, Any]: + """Converts the MultiRLModuleSpec to a dictionary.""" + return { + "multi_rl_module_class": serialize_type(self.multi_rl_module_class), + "observation_space": gym_space_to_dict(self.observation_space), + "action_space": gym_space_to_dict(self.action_space), + "inference_only": self.inference_only, + "model_config": self.model_config, + "rl_module_specs": { + module_id: rl_module_spec.to_dict() + for module_id, rl_module_spec in self.rl_module_specs.items() + }, + } + + @classmethod + def from_dict(cls, d) -> "MultiRLModuleSpec": + """Creates a MultiRLModuleSpec from a dictionary.""" + return MultiRLModuleSpec( + multi_rl_module_class=deserialize_type(d["multi_rl_module_class"]), + observation_space=gym_space_from_dict(d.get("observation_space")), + action_space=gym_space_from_dict(d.get("action_space")), + model_config=d.get("model_config"), + inference_only=d["inference_only"], + rl_module_specs={ + module_id: RLModuleSpec.from_dict(rl_module_spec) + for module_id, rl_module_spec in ( + d.get("rl_module_specs", d.get("module_specs")).items() + ) + }, + ) + + def update( + self, + other: Union["MultiRLModuleSpec", RLModuleSpec], + override: bool = False, + ) -> None: + """Updates this spec with the other spec. + + Traverses this MultiRLModuleSpec's module_specs and updates them with + the module specs from the `other` (Multi)RLModuleSpec. + + Args: + other: The other spec to update this spec with. + override: Whether to override the existing module specs if they already + exist. If False, they are only updated. + """ + if isinstance(other, RLModuleSpec): + # Disable our `inference_only` as soon as any single-agent module has + # `inference_only=False`. + if not other.inference_only: + self.inference_only = False + for mid, spec in self.rl_module_specs.items(): + self.rl_module_specs[mid].update(other, override=False) + elif isinstance(other.module_specs, dict): + self.add_modules(other.rl_module_specs, override=override) + else: + assert isinstance(other, MultiRLModuleSpec) + if not self.rl_module_specs: + self.inference_only = other.inference_only + self.rl_module_specs = other.rl_module_specs + else: + if not other.inference_only: + self.inference_only = False + self.rl_module_specs.update(other.rl_module_specs) + + def as_multi_rl_module_spec(self) -> "MultiRLModuleSpec": + """Returns self in order to match `RLModuleSpec.as_multi_rl_module_spec()`.""" + return self + + def __contains__(self, item) -> bool: + """Returns whether the given `item` (ModuleID) is present in self.""" + return item in self.rl_module_specs + + def __getitem__(self, item) -> RLModuleSpec: + """Returns the RLModuleSpec under the ModuleID.""" + return self.rl_module_specs[item] + + @Deprecated( + new="MultiRLModule(*, module_specs={module1: [RLModuleSpec], " + "module2: [RLModuleSpec], ..}, inference_only=..)", + error=True, + ) + def get_multi_rl_module_config(self): + pass + + @Deprecated(new="MultiRLModuleSpec.as_multi_rl_module_spec()", error=True) + def as_multi_agent(self): + pass + + @Deprecated(new="MultiRLModuleSpec.get_multi_rl_module_config", error=True) + def get_marl_config(self, *args, **kwargs): + pass + + @Deprecated( + new="MultiRLModule(*, observation_space=.., action_space=.., ....)", + error=False, + ) + def get_rl_module_config(self): + return MultiRLModuleConfig( + inference_only=self.inference_only, + modules=self.rl_module_specs, + ) + + +@Deprecated( + new="MultiRLModule(*, rl_module_specs={module1: [RLModuleSpec], " + "module2: [RLModuleSpec], ..}, inference_only=..)", + error=False, +) +@dataclasses.dataclass +class MultiRLModuleConfig: + inference_only: bool = False + modules: Dict[ModuleID, RLModuleSpec] = dataclasses.field(default_factory=dict) + + def to_dict(self): + return { + "inference_only": self.inference_only, + "modules": { + module_id: module_spec.to_dict() + for module_id, module_spec in self.modules.items() + }, + } + + @classmethod + def from_dict(cls, d) -> "MultiRLModuleConfig": + return cls( + inference_only=d["inference_only"], + modules={ + module_id: RLModuleSpec.from_dict(module_spec) + for module_id, module_spec in d["modules"].items() + }, + ) + + def get_catalog(self) -> None: + return None diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/rl_module.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..7e63c9cd96621833e10ba99fdef31dffbcdd0e71 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/rl_module.py @@ -0,0 +1,835 @@ +import abc +import dataclasses +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Collection, Dict, Optional, Type, Union + +import gymnasium as gym + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + Deprecated, + deprecation_warning, +) +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.distribution.distribution import Distribution +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + override, +) +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.serialization import ( + deserialize_type, + gym_space_from_dict, + gym_space_to_dict, + serialize_type, +) +from ray.rllib.utils.typing import StateDict +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.core.models.catalog import Catalog + from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, + ) + +logger = logging.getLogger("ray.rllib") +torch, _ = try_import_torch() + + +@PublicAPI(stability="beta") +@dataclass +class RLModuleSpec: + """Utility spec class to make constructing RLModules (in single-agent case) easier. + + Args: + module_class: The RLModule class to use. + observation_space: The observation space of the RLModule. This may differ + from the observation space of the environment. For example, a discrete + observation space of an environment, would usually correspond to a + one-hot encoded observation space of the RLModule because of preprocessing. + action_space: The action space of the RLModule. + inference_only: Whether the RLModule should be configured in its inference-only + state, in which those components not needed for action computing (for + example a value function or a target network) might be missing. + Note that `inference_only=True` AND `learner_only=True` is not allowed. + learner_only: Whether this RLModule should only be built on Learner workers, but + NOT on EnvRunners. Useful for RLModules inside a MultiRLModule that are only + used for training, for example a shared value function in a multi-agent + setup or a world model in a curiosity-learning setup. + Note that `inference_only=True` AND `learner_only=True` is not allowed. + model_config: The model config dict or default RLlib dataclass to use. + catalog_class: The Catalog class to use. + load_state_path: The path to the module state to load from. NOTE: This must be + an absolute path. + """ + + module_class: Optional[Type["RLModule"]] = None + observation_space: Optional[gym.Space] = None + action_space: Optional[gym.Space] = None + inference_only: bool = False + learner_only: bool = False + model_config: Optional[Union[Dict[str, Any], DefaultModelConfig]] = None + catalog_class: Optional[Type["Catalog"]] = None + load_state_path: Optional[str] = None + + # Deprecated field. + model_config_dict: Optional[Union[dict, int]] = None + + def __post_init__(self): + if self.model_config_dict is not None: + deprecation_warning( + old="RLModuleSpec(model_config_dict=..)", + new="RLModuleSpec(model_config=..)", + error=True, + ) + + def build(self) -> "RLModule": + """Builds the RLModule from this spec.""" + if self.module_class is None: + raise ValueError("RLModule class is not set.") + if self.observation_space is None: + raise ValueError("Observation space is not set.") + + try: + module = self.module_class( + observation_space=self.observation_space, + action_space=self.action_space, + inference_only=self.inference_only, + model_config=self._get_model_config(), + catalog_class=self.catalog_class, + ) + # Older custom model might still require the old `RLModuleConfig` under + # the `config` arg. + except AttributeError: + module_config = self.get_rl_module_config() + module = self.module_class(module_config) + return module + + @classmethod + def from_module(cls, module: "RLModule") -> "RLModuleSpec": + from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule + + if isinstance(module, MultiRLModule): + raise ValueError("MultiRLModule cannot be converted to RLModuleSpec.") + + # Try instantiating a new RLModule from the spec using the new c'tor args. + try: + rl_module_spec = RLModuleSpec( + module_class=type(module), + observation_space=module.observation_space, + action_space=module.action_space, + inference_only=module.inference_only, + learner_only=module.learner_only, + model_config=module.model_config, + catalog_class=( + type(module.catalog) if module.catalog is not None else None + ), + ) + + # Old path through deprecated `RLModuleConfig` class. Used only if `module` + # still has a valid `config` attribute. + except AttributeError: + rl_module_spec = RLModuleSpec( + module_class=type(module), + observation_space=module.config.observation_space, + action_space=module.config.action_space, + inference_only=module.config.inference_only, + learner_only=module.config.learner_only, + model_config=module.config.model_config_dict, + catalog_class=module.config.catalog_class, + ) + return rl_module_spec + + def to_dict(self): + """Returns a serialized representation of the spec.""" + return { + "module_class": serialize_type(self.module_class), + "observation_space": gym_space_to_dict(self.observation_space), + "action_space": gym_space_to_dict(self.action_space), + "inference_only": self.inference_only, + "learner_only": self.learner_only, + "model_config": self._get_model_config(), + "catalog_class": serialize_type(self.catalog_class) + if self.catalog_class is not None + else None, + } + + @classmethod + def from_dict(cls, d): + """Returns a single agent RLModule spec from a serialized representation.""" + module_class = deserialize_type(d["module_class"]) + try: + spec = RLModuleSpec( + module_class=module_class, + observation_space=gym_space_from_dict(d["observation_space"]), + action_space=gym_space_from_dict(d["action_space"]), + inference_only=d["inference_only"], + learner_only=d["learner_only"], + model_config=d["model_config"], + catalog_class=deserialize_type(d["catalog_class"]) + if d["catalog_class"] is not None + else None, + ) + + # Old path through deprecated `RLModuleConfig` class. + except KeyError: + module_config = RLModuleConfig.from_dict(d["module_config"]) + spec = RLModuleSpec( + module_class=module_class, + observation_space=module_config.observation_space, + action_space=module_config.action_space, + inference_only=module_config.inference_only, + learner_only=module_config.learner_only, + model_config=module_config.model_config_dict, + catalog_class=module_config.catalog_class, + ) + return spec + + def update(self, other, override: bool = True) -> None: + """Updates this spec with the given other spec. Works like dict.update(). + + Args: + other: The other SingleAgentRLModule spec to update this one from. + override: Whether to update all properties in `self` with those of `other. + If False, only update those properties in `self` that are not None. + """ + if not isinstance(other, RLModuleSpec): + raise ValueError("Can only update with another RLModuleSpec.") + + # If the field is None in the other, keep the current field, otherwise update + # with the new value. + if override: + self.module_class = other.module_class or self.module_class + self.observation_space = other.observation_space or self.observation_space + self.action_space = other.action_space or self.action_space + self.inference_only = other.inference_only or self.inference_only + self.learner_only = other.learner_only and self.learner_only + self.model_config = other.model_config or self.model_config + self.catalog_class = other.catalog_class or self.catalog_class + self.load_state_path = other.load_state_path or self.load_state_path + # Only override, if the field is None in `self`. + # Do NOT override the boolean settings: `inference_only` and `learner_only`. + else: + self.module_class = self.module_class or other.module_class + self.observation_space = self.observation_space or other.observation_space + self.action_space = self.action_space or other.action_space + self.model_config = self.model_config or other.model_config + self.catalog_class = self.catalog_class or other.catalog_class + self.load_state_path = self.load_state_path or other.load_state_path + + def as_multi_rl_module_spec(self) -> "MultiRLModuleSpec": + """Returns a MultiRLModuleSpec (`self` under DEFAULT_MODULE_ID key).""" + from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec + + return MultiRLModuleSpec( + rl_module_specs={DEFAULT_MODULE_ID: self}, + load_state_path=self.load_state_path, + ) + + def _get_model_config(self): + return ( + dataclasses.asdict(self.model_config) + if dataclasses.is_dataclass(self.model_config) + else (self.model_config or {}) + ) + + @Deprecated( + new="RLModule(*, observation_space=.., action_space=.., ....)", + error=False, + ) + def get_rl_module_config(self): + return RLModuleConfig( + observation_space=self.observation_space, + action_space=self.action_space, + inference_only=self.inference_only, + learner_only=self.learner_only, + model_config_dict=self._get_model_config(), + catalog_class=self.catalog_class, + ) + + +@PublicAPI(stability="beta") +class RLModule(Checkpointable, abc.ABC): + """Base class for RLlib modules. + + Subclasses should call `super().__init__(observation_space=.., action_space=.., + inference_only=.., learner_only=.., model_config={..})` in their __init__ methods. + + Here is the pseudocode for how the forward methods are called: + + Example for creating a (inference-only) sampling loop: + + .. testcode:: + + from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule + ) + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig + import gymnasium as gym + import torch + + env = gym.make("CartPole-v1") + + # Create an instance of the default RLModule used by PPO. + module = DefaultPPOTorchRLModule( + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[128, 128]), + catalog_class=PPOCatalog, + ) + action_dist_class = module.get_inference_action_dist_cls() + obs, info = env.reset() + terminated = False + + while not terminated: + fwd_ins = {"obs": torch.Tensor([obs])} + fwd_outputs = module.forward_exploration(fwd_ins) + # This can be either deterministic or stochastic distribution. + action_dist = action_dist_class.from_logits( + fwd_outputs["action_dist_inputs"] + ) + action = action_dist.sample()[0].numpy() + obs, reward, terminated, truncated, info = env.step(action) + + + Example for training: + + .. testcode:: + + from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import ( + PPOTorchRLModule + ) + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig + import gymnasium as gym + import torch + + from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule + ) + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + + env = gym.make("CartPole-v1") + + # Create an instance of the default RLModule used by PPO. + module = DefaultPPOTorchRLModule( + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[128, 128]), + catalog_class=PPOCatalog, + ) + + fwd_ins = {"obs": torch.Tensor([obs])} + fwd_outputs = module.forward_train(fwd_ins) + # loss = compute_loss(fwd_outputs, fwd_ins) + # update_params(module, loss) + + Example for inference: + + .. testcode:: + + from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import ( + PPOTorchRLModule + ) + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig + import gymnasium as gym + import torch + + from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule + ) + from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog + + env = gym.make("CartPole-v1") + + # Create an instance of the default RLModule used by PPO. + module = DefaultPPOTorchRLModule( + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[128, 128]), + catalog_class=PPOCatalog, + ) + + while not terminated: + fwd_ins = {"obs": torch.Tensor([obs])} + fwd_outputs = module.forward_inference(fwd_ins) + # this can be either deterministic or stochastic distribution + action_dist = action_dist_class.from_logits( + fwd_outputs["action_dist_inputs"] + ) + action = action_dist.sample()[0].numpy() + obs, reward, terminated, truncated, info = env.step(action) + + + Args: + observation_space: The observation space of the model. Note that in multi-agent + setups, this is typically the observation space of an agent that maps to + this RLModule. + action_space: The action space of the model. Note that in multi-agent + setups, this is typically the action space of an agent that maps to + this RLModule. + inference_only: If True, this RLModule should construct itself in an inference- + only fashion. This is done automatically, if the user implements the + `InferenceOnlyAPI` with their custom RLModule subclass. False by default. + learner_only: If True, RLlib won't built this RLModule on EnvRunner actors. + False by default. + model_config: A config dict to specify features of this RLModule. + + Attributes: + action_dist_cls: An optional ray.rllib.core.distribution.distribution. + Distribution subclass to use for sampling actions, given parameters from + a batch (`Columns.ACTION_DIST_INPUTS`). + + Abstract Methods: + ``~_forward_train``: Forward pass during training. + + ``~_forward_exploration``: Forward pass during training for exploration. + + ``~_forward_inference``: Forward pass during inference. + """ + + framework: str = None + + STATE_FILE_NAME = "module_state" + + def __init__( + self, + config=DEPRECATED_VALUE, + *, + observation_space: Optional[gym.Space] = None, + action_space: Optional[gym.Space] = None, + inference_only: Optional[bool] = None, + learner_only: bool = False, + model_config: Optional[Union[dict, DefaultModelConfig]] = None, + catalog_class=None, + **kwargs, + ): + # TODO (sven): Deprecate Catalog and replace with utility functions to create + # primitive components based on obs- and action spaces. + self.catalog = None + self._catalog_ctor_error = None + + # Deprecated + self.config = config + if self.config != DEPRECATED_VALUE: + deprecation_warning( + old="RLModule(config=[RLModuleConfig])", + new="RLModule(observation_space=.., action_space=.., inference_only=..," + " learner_only=.., model_config=..)", + help="See https://github.com/ray-project/ray/blob/master/rllib/examples/rl_modules/custom_cnn_rl_module.py " # noqa + "for how to write a custom RLModule.", + error=True, + ) + else: + self.observation_space = observation_space + self.action_space = action_space + self.inference_only = inference_only + self.learner_only = learner_only + self.model_config = model_config + if catalog_class is not None: + try: + self.catalog = catalog_class( + observation_space=self.observation_space, + action_space=self.action_space, + model_config_dict=self.model_config, + ) + except Exception as e: + logger.warning( + "Didn't create a Catalog object for your RLModule! If you are " + "not using the new API stack yet, make sure to switch it off in" + " your config: `config.api_stack(enable_rl_module_and_learner=" + "False, enable_env_runner_and_connector_v2=False)`. All algos " + "use the new stack by default. Ignore this message, if your " + "RLModule does not use a Catalog to build its sub-components." + ) + self._catalog_ctor_error = e + + # TODO (sven): Deprecate this. We keep it here for now in case users + # still have custom models (or subclasses of RLlib default models) + # into which they pass in a `config` argument. + self.config = RLModuleConfig( + observation_space=self.observation_space, + action_space=self.action_space, + inference_only=self.inference_only, + learner_only=self.learner_only, + model_config_dict=self.model_config, + catalog_class=catalog_class, + ) + + self.action_dist_cls = None + if self.catalog is not None: + self.action_dist_cls = self.catalog.get_action_dist_cls( + framework=self.framework + ) + + # Make sure, `setup()` is only called once, no matter what. + if hasattr(self, "_is_setup") and self._is_setup: + raise RuntimeError( + "`RLModule.setup()` called twice within your RLModule implementation " + f"{self}! Make sure you are using the proper inheritance order " + "(TorchRLModule before [Algo]RLModule) or (TfRLModule before " + "[Algo]RLModule) and that you are NOT overriding the constructor, but " + "only the `setup()` method of your subclass." + ) + try: + self.setup() + except AttributeError as e: + if "'NoneType' object has no attribute " in e.args[0]: + raise (self._catalog_ctor_error or e) + self._is_setup = True + # Cache value for returning from `is_stateful` so we don't have to call + # the module's `get_initial_state()` method all the time (might be expensive). + self._is_stateful = None + + @OverrideToImplementCustomLogic + def setup(self): + """Sets up the components of the module. + + This is called automatically during the __init__ method of this class, + therefore, the subclass should call super.__init__() in its constructor. This + abstraction can be used to create any components (e.g. NN layers) that your + RLModule needs. + """ + pass + + @OverrideToImplementCustomLogic + def get_inference_action_dist_cls(self) -> Type[Distribution]: + """Returns the action distribution class for this RLModule used for inference. + + This class is used to create action distributions from outputs of the forward + inference method. If the case that no action distribution class is needed, + this method can return None. + + Note that RLlib's distribution classes all implement the `Distribution` + interface. This requires two special methods: `Distribution.from_logits()` and + `Distribution.to_deterministic()`. See the documentation of the + :py:class:`~ray.rllib.core.distribution.distribution.Distribution` class for + more details. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def get_exploration_action_dist_cls(self) -> Type[Distribution]: + """Returns the action distribution class for this RLModule used for exploration. + + This class is used to create action distributions from outputs of the + forward_exploration method. If the case that no action distribution class is + needed, this method can return None. + + Note that RLlib's distribution classes all implement the `Distribution` + interface. This requires two special methods: `Distribution.from_logits()` and + `Distribution.to_deterministic()`. See the documentation of the + :py:class:`~ray.rllib.core.distribution.distribution.Distribution` class for + more details. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def get_train_action_dist_cls(self) -> Type[Distribution]: + """Returns the action distribution class for this RLModule used for training. + + This class is used to get the correct action distribution class to be used by + the training components. In case that no action distribution class is needed, + this method can return None. + + Note that RLlib's distribution classes all implement the `Distribution` + interface. This requires two special methods: `Distribution.from_logits()` and + `Distribution.to_deterministic()`. See the documentation of the + :py:class:`~ray.rllib.core.distribution.distribution.Distribution` class for + more details. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def _forward(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Generic forward pass method, used in all phases of training and evaluation. + + If you need a more nuanced distinction between forward passes in the different + phases of training and evaluation, override the following methods instead: + + For distinct action computation logic w/o exploration, override the + `self._forward_inference()` method. + For distinct action computation logic with exploration, override the + `self._forward_exploration()` method. + For distinct forward pass logic before loss computation, override the + `self._forward_train()` method. + + Args: + batch: The input batch. + **kwargs: Additional keyword arguments. + + Returns: + The output of the forward pass. + """ + return {} + + def forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """DO NOT OVERRIDE! Forward-pass during evaluation, called from the sampler. + + This method should not be overridden. Override the `self._forward_inference()` + method instead. + + Args: + batch: The input batch. + **kwargs: Additional keyword arguments. + + Returns: + The output of the forward pass. This output should comply with the + ouptut_specs_inference(). + """ + return self._forward_inference(batch, **kwargs) + + @OverrideToImplementCustomLogic + def _forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Forward-pass used for action computation without exploration behavior. + + Override this method only, if you need specific behavior for non-exploratory + action computation behavior. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return self._forward(batch, **kwargs) + + def forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """DO NOT OVERRIDE! Forward-pass during exploration, called from the sampler. + + This method should not be overridden. Override the `self._forward_exploration()` + method instead. + + Args: + batch: The input batch. + **kwargs: Additional keyword arguments. + + Returns: + The output of the forward pass. + """ + return self._forward_exploration(batch, **kwargs) + + @OverrideToImplementCustomLogic + def _forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Forward-pass used for action computation with exploration behavior. + + Override this method only, if you need specific behavior for exploratory + action computation behavior. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return self._forward(batch, **kwargs) + + def forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """DO NOT OVERRIDE! Forward-pass during training called from the learner. + + This method should not be overridden. Override the `self._forward_train()` + method instead. + + Args: + batch: The input batch. + **kwargs: Additional keyword arguments. + + Returns: + The output of the forward pass. + """ + if self.inference_only: + raise RuntimeError( + "Calling `forward_train` on an inference_only module is not allowed! " + "Set the `inference_only=False` flag in the RLModuleSpec (or the " + "RLModule's constructor)." + ) + return self._forward_train(batch, **kwargs) + + @OverrideToImplementCustomLogic + def _forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Forward-pass used before the loss computation (training). + + Override this method only, if you need specific behavior and outputs for your + loss computations. If you have only one generic behavior for all + phases of training and evaluation, override `self._forward()` instead. + + By default, this calls the generic `self._forward()` method. + """ + return self._forward(batch, **kwargs) + + @OverrideToImplementCustomLogic + def get_initial_state(self) -> Any: + """Returns the initial state of the RLModule, in case this is a stateful module. + + Returns: + A tensor or any nested struct of tensors, representing an initial state for + this (stateful) RLModule. + """ + return {} + + @OverrideToImplementCustomLogic + def is_stateful(self) -> bool: + """By default, returns False if the initial state is an empty dict (or None). + + By default, RLlib assumes that the module is non-recurrent, if the initial + state is an empty dict and recurrent otherwise. + This behavior can be customized by overriding this method. + """ + if self._is_stateful is None: + initial_state = self.get_initial_state() + assert isinstance(initial_state, dict), ( + "The initial state of an RLModule must be a dict, but is " + f"{type(initial_state)} instead." + ) + self._is_stateful = bool(initial_state) + return self._is_stateful + + @OverrideToImplementCustomLogic + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + inference_only: bool = False, + **kwargs, + ) -> StateDict: + """Returns the state dict of the module. + + Args: + inference_only: Whether the returned state should be an inference-only + state (w/o those model components that are not needed for action + computations, such as a value function or a target network). + Note that setting this to `False` might raise an error if + `self.inference_only` is True. + + Returns: + This RLModule's state dict. + """ + return {} + + @OverrideToImplementCustomLogic + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + pass + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + { + "observation_space": self.observation_space, + "action_space": self.action_space, + "inference_only": self.inference_only, + "learner_only": self.learner_only, + "model_config": self.model_config, + "catalog_class": ( + type(self.catalog) if self.catalog is not None else None + ), + }, # **kwargs + ) + + def as_multi_rl_module(self) -> "MultiRLModule": + """Returns a multi-agent wrapper around this module.""" + from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule + + multi_rl_module = MultiRLModule( + rl_module_specs={DEFAULT_MODULE_ID: RLModuleSpec.from_module(self)} + ) + return multi_rl_module + + def unwrapped(self) -> "RLModule": + """Returns the underlying module if this module is a wrapper. + + An example of a wrapped is the TorchDDPRLModule class, which wraps + a TorchRLModule. + + Returns: + The underlying module. + """ + return self + + @Deprecated(error=False) + def output_specs_train(self): + pass + + @Deprecated(error=False) + def output_specs_inference(self): + pass + + @Deprecated(error=False) + def output_specs_exploration(self): + pass + + @Deprecated(error=False) + def input_specs_inference(self): + pass + + @Deprecated(error=False) + def input_specs_exploration(self): + pass + + @Deprecated(error=False) + def input_specs_train(self): + pass + + @Deprecated(error=False) + def _default_input_specs(self): + pass + + +@Deprecated( + old="RLModule(config=[RLModuleConfig object])", + new="RLModule(observation_space=.., action_space=.., inference_only=.., " + "model_config=.., catalog_class=..)", + error=False, +) +@dataclass +class RLModuleConfig: + observation_space: gym.Space = None + action_space: gym.Space = None + inference_only: bool = False + learner_only: bool = False + model_config_dict: Dict[str, Any] = field(default_factory=dict) + catalog_class: Type["Catalog"] = None + + def get_catalog(self) -> Optional["Catalog"]: + if self.catalog_class is not None: + return self.catalog_class( + observation_space=self.observation_space, + action_space=self.action_space, + model_config_dict=self.model_config_dict, + ) + return None + + def to_dict(self): + catalog_class_path = ( + serialize_type(self.catalog_class) if self.catalog_class else "" + ) + return { + "observation_space": gym_space_to_dict(self.observation_space), + "action_space": gym_space_to_dict(self.action_space), + "inference_only": self.inference_only, + "learner_only": self.learner_only, + "model_config_dict": self.model_config_dict, + "catalog_class_path": catalog_class_path, + } + + @classmethod + def from_dict(cls, d: Dict[str, Any]): + catalog_class = ( + None + if d["catalog_class_path"] == "" + else deserialize_type(d["catalog_class_path"]) + ) + return cls( + observation_space=gym_space_from_dict(d["observation_space"]), + action_space=gym_space_from_dict(d["action_space"]), + inference_only=d["inference_only"], + learner_only=d["learner_only"], + model_config_dict=d["model_config_dict"], + catalog_class=catalog_class, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..217c5a5abc8aa4ec29708e81d0b2884256d8fae7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__init__.py @@ -0,0 +1,3 @@ +from .torch_rl_module import TorchRLModule + +__all__ = ["TorchRLModule"] diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2f0e501aa7793a6d3d264889c0f94a96ab8be2c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_compile_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_compile_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..873e16823aae26d222abf1f1ea16aa2b324db66b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_compile_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bdcb53aa7fd0cfe9df4fe72b1da0e6514ce6246f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/__pycache__/torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_compile_config.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_compile_config.py new file mode 100644 index 0000000000000000000000000000000000000000..7c95c62fb52acab3c054c5cce7cc429046f1c84e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_compile_config.py @@ -0,0 +1,39 @@ +import sys +from dataclasses import dataclass, field + + +@dataclass +class TorchCompileConfig: + """Configuration options for RLlib's usage of torch.compile in RLModules. + + # On `torch.compile` in Torch RLModules + `torch.compile` invokes torch's dynamo JIT compiler that can potentially bring + speedups to RL Module's forward methods. + This is a performance optimization that should be disabled for debugging. + + General usage: + - Usually, you only want to `RLModule._forward_train` to be compiled on + instances of RLModule used for learning. (e.g. the learner) + - In some cases, it can bring speedups to also compile + `RLModule._forward_exploration` on instances used for exploration. (e.g. + RolloutWorker) + - In some cases, it can bring speedups to also compile + `RLModule._forward_inference` on instances used for inference. (e.g. + RolloutWorker) + + Note that different backends are available on different platforms. + Also note that the default backend for torch dynamo is "aot_eager" on macOS. + This is a debugging backend that is expected not to improve performance because + the inductor backend is not supported on OSX so far. + + Args: + torch_dynamo_backend: The torch.dynamo backend to use. + torch_dynamo_mode: The torch.dynamo mode to use. + kwargs: Additional keyword arguments to pass to `torch.compile()` + """ + + torch_dynamo_backend: str = ( + "aot_eager" if sys.platform == "darwin" else "cudagraphs" + ) + torch_dynamo_mode: str = None + kwargs: dict = field(default_factory=lambda: dict()) diff --git a/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..b35c9a6572fa2e4bedc8155a4e6ca4fac1a21b27 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/rl_module/torch/torch_rl_module.py @@ -0,0 +1,306 @@ +from typing import Any, Collection, Dict, Optional, Type, Union + +import gymnasium as gym +from packaging import version + +from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical, + TorchDiagGaussian, + TorchDistribution, +) +from ray.rllib.core.rl_module.apis import InferenceOnlyAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.core.rl_module.torch.torch_compile_config import TorchCompileConfig +from ray.rllib.utils.annotations import OverrideToImplementCustomLogic, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.torch_utils import ( + TORCH_COMPILE_REQUIRED_VERSION, + convert_to_torch_tensor, +) +from ray.rllib.utils.typing import StateDict + +torch, nn = try_import_torch() + + +class TorchRLModule(nn.Module, RLModule): + """A base class for RLlib PyTorch RLModules. + + Note that the `_forward` methods of this class can be 'torch.compiled' individually: + - `TorchRLModule._forward_train()` + - `TorchRLModule._forward_inference()` + - `TorchRLModule._forward_exploration()` + + As a rule of thumb, they should only contain torch-native tensor manipulations, + or otherwise they may yield wrong outputs. In particular, the creation of RLlib + distributions inside these methods should be avoided when using `torch.compile`. + When in doubt, you can use `torch.dynamo.explain()` to check whether a compiled + method has broken up into multiple sub-graphs. + + Compiling these methods can bring speedups under certain conditions. + """ + + framework: str = "torch" + + # Stick with torch default. + STATE_FILE_NAME = "module_state" + + def __init__(self, *args, **kwargs) -> None: + nn.Module.__init__(self) + RLModule.__init__(self, *args, **kwargs) + + # If an inference-only class AND self.inference_only is True, + # remove all attributes that are returned by + # `self.get_non_inference_attributes()`. + if self.inference_only and isinstance(self, InferenceOnlyAPI): + for attr in self.get_non_inference_attributes(): + parts = attr.split(".") + if not hasattr(self, parts[0]): + continue + target_name = parts[0] + target_obj = getattr(self, target_name) + # Traverse from the next part on (if nested). + for part in parts[1:]: + if not hasattr(target_obj, part): + target_obj = None + break + target_name = part + target_obj = getattr(target_obj, target_name) + # Delete, if target is valid. + if target_obj is not None: + delattr(self, target_name) + + def compile(self, compile_config: TorchCompileConfig): + """Compile the forward methods of this module. + + This is a convenience method that calls `compile_wrapper` with the given + compile_config. + + Args: + compile_config: The compile config to use. + """ + return compile_wrapper(self, compile_config) + + @OverrideToImplementCustomLogic + def _forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + # By default, calls the generic `_forward()` method, but with a no-grad context + # for performance reasons. + with torch.no_grad(): + return self._forward(batch, **kwargs) + + @OverrideToImplementCustomLogic + def _forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + # By default, calls the generic `_forward()` method, but with a no-grad context + # for performance reasons. + with torch.no_grad(): + return self._forward(batch, **kwargs) + + @OverrideToImplementCustomLogic + @override(RLModule) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + inference_only: bool = False, + **kwargs, + ) -> StateDict: + state_dict = self.state_dict() + # Filter out `inference_only` keys from the state dict if `inference_only` and + # this RLModule is NOT `inference_only` (but does implement the + # InferenceOnlyAPI). + if ( + inference_only + and not self.inference_only + and isinstance(self, InferenceOnlyAPI) + ): + attr = self.get_non_inference_attributes() + for key in list(state_dict.keys()): + if any( + key.startswith(a) and (len(key) == len(a) or key[len(a)] == ".") + for a in attr + ): + del state_dict[key] + return convert_to_numpy(state_dict) + + @OverrideToImplementCustomLogic + @override(RLModule) + def set_state(self, state: StateDict) -> None: + # If state contains more keys than `self.state_dict()`, then we simply ignore + # these keys (strict=False). This is most likely due to `state` coming from + # an `inference_only=False` RLModule, while `self` is an `inference_only=True` + # RLModule. + self.load_state_dict(convert_to_torch_tensor(state), strict=False) + + @OverrideToImplementCustomLogic + @override(RLModule) + def get_inference_action_dist_cls(self) -> Type[TorchDistribution]: + if self.action_dist_cls is not None: + return self.action_dist_cls + elif isinstance(self.action_space, gym.spaces.Discrete): + return TorchCategorical + elif isinstance(self.action_space, gym.spaces.Box): + return TorchDiagGaussian + else: + raise ValueError( + f"Default action distribution for action space " + f"{self.action_space} not supported! Either set the " + f"`self.action_dist_cls` property in your RLModule's `setup()` method " + f"to a subclass of `ray.rllib.core.distribution.torch." + f"torch_distribution.TorchDistribution` or - if you need different " + f"distributions for inference and training - override the three " + f"methods: `get_inference_action_dist_cls`," + f"`get_exploration_action_dist_cls`, and `get_train_action_dist_cls` " + f"in your RLModule." + ) + + @OverrideToImplementCustomLogic + @override(RLModule) + def get_exploration_action_dist_cls(self) -> Type[TorchDistribution]: + return self.get_inference_action_dist_cls() + + @OverrideToImplementCustomLogic + @override(RLModule) + def get_train_action_dist_cls(self) -> Type[TorchDistribution]: + return self.get_inference_action_dist_cls() + + def forward(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """DO NOT OVERRIDE! + + This is aliased to `self.forward_train` because Torch DDP requires a forward + method to be implemented for backpropagation to work. + + Instead, override: + `_forward()` to define a generic forward pass for all phases (exploration, + inference, training) + `_forward_inference()` to define the forward pass for action inference in + deployment/production (no exploration). + `_forward_exploration()` to define the forward pass for action inference during + training sample collection (w/ exploration behavior). + `_forward_train()` to define the forward pass prior to loss computation. + """ + # TODO (sven): Experimental to make ONNX exported models work. + if self.config.inference_only: + return self.forward_exploration(batch, **kwargs) + else: + return self.forward_train(batch, **kwargs) + + +class TorchDDPRLModule(RLModule, nn.parallel.DistributedDataParallel): + def __init__(self, *args, **kwargs) -> None: + nn.parallel.DistributedDataParallel.__init__(self, *args, **kwargs) + # We do not want to call RLModule.__init__ here because all we need is + # the interface of that base-class not the actual implementation. + # RLModule.__init__(self, *args, **kwargs) + self.observation_space = self.unwrapped().observation_space + self.action_space = self.unwrapped().action_space + self.inference_only = self.unwrapped().inference_only + self.learner_only = self.unwrapped().learner_only + self.model_config = self.unwrapped().model_config + self.catalog = self.unwrapped().catalog + + # Deprecated. + self.config = self.unwrapped().config + + @override(RLModule) + def get_inference_action_dist_cls(self, *args, **kwargs) -> Type[TorchDistribution]: + return self.unwrapped().get_inference_action_dist_cls(*args, **kwargs) + + @override(RLModule) + def get_exploration_action_dist_cls( + self, *args, **kwargs + ) -> Type[TorchDistribution]: + return self.unwrapped().get_exploration_action_dist_cls(*args, **kwargs) + + @override(RLModule) + def get_train_action_dist_cls(self, *args, **kwargs) -> Type[TorchDistribution]: + return self.unwrapped().get_train_action_dist_cls(*args, **kwargs) + + @override(RLModule) + def get_initial_state(self) -> Any: + return self.unwrapped().get_initial_state() + + @override(RLModule) + def is_stateful(self) -> bool: + return self.unwrapped().is_stateful() + + @override(RLModule) + def _forward(self, *args, **kwargs): + return self.unwrapped()._forward(*args, **kwargs) + + @override(RLModule) + def _forward_inference(self, *args, **kwargs) -> Dict[str, Any]: + return self.unwrapped()._forward_inference(*args, **kwargs) + + @override(RLModule) + def _forward_exploration(self, *args, **kwargs) -> Dict[str, Any]: + return self.unwrapped()._forward_exploration(*args, **kwargs) + + @override(RLModule) + def _forward_train(self, *args, **kwargs): + return self(*args, **kwargs) + + @override(RLModule) + def get_state(self, *args, **kwargs): + return self.unwrapped().get_state(*args, **kwargs) + + @override(RLModule) + def set_state(self, *args, **kwargs): + self.unwrapped().set_state(*args, **kwargs) + + @override(RLModule) + def save_to_path(self, *args, **kwargs): + self.unwrapped().save_to_path(*args, **kwargs) + + @override(RLModule) + def restore_from_path(self, *args, **kwargs): + self.unwrapped().restore_from_path(*args, **kwargs) + + @override(RLModule) + def get_metadata(self, *args, **kwargs): + self.unwrapped().get_metadata(*args, **kwargs) + + @override(RLModule) + def unwrapped(self) -> "RLModule": + return self.module + + +def compile_wrapper(rl_module: "TorchRLModule", compile_config: TorchCompileConfig): + """A wrapper that compiles the forward methods of a TorchRLModule.""" + + # TODO(Artur): Remove this once our requirements enforce torch >= 2.0.0 + # Check if torch framework supports torch.compile. + if ( + torch is not None + and version.parse(torch.__version__) < TORCH_COMPILE_REQUIRED_VERSION + ): + raise ValueError("torch.compile is only supported from torch 2.0.0") + + compiled_forward_train = torch.compile( + rl_module._forward_train, + backend=compile_config.torch_dynamo_backend, + mode=compile_config.torch_dynamo_mode, + **compile_config.kwargs, + ) + + rl_module._forward_train = compiled_forward_train + + compiled_forward_inference = torch.compile( + rl_module._forward_inference, + backend=compile_config.torch_dynamo_backend, + mode=compile_config.torch_dynamo_mode, + **compile_config.kwargs, + ) + + rl_module._forward_inference = compiled_forward_inference + + compiled_forward_exploration = torch.compile( + rl_module._forward_exploration, + backend=compile_config.torch_dynamo_backend, + mode=compile_config.torch_dynamo_mode, + **compile_config.kwargs, + ) + + rl_module._forward_exploration = compiled_forward_exploration + + return rl_module diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7478318b33357118a66eb21598d9cd4e79bd5a22 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/bc_algorithm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/bc_algorithm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c3d01d4ac8c3b74ae043f2c0c9341fb39d57616 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/bc_algorithm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/testing_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/testing_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29486c8a4f198841efbc0dfd68968d8dc5fbacd2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/__pycache__/testing_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..326fab3f1bd325ae13ae0006d4978f2df8475a81 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c0f1ddde30b7db5e53e577a52d1c9699f912d22 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc4087922d3463ded13d9578650198f0acd9f447 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/__pycache__/bc_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_learner.py b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..6e1963a93038f3edd767c3d1f121f183c022d95e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_learner.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING, Any, Dict + +import torch + +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.core.testing.testing_learner import BaseTestingLearner +from ray.rllib.utils.typing import ModuleID, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + +class BCTorchLearner(TorchLearner, BaseTestingLearner): + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + BaseTestingLearner.compute_loss_for_module( + self, + module_id=module_id, + config=config, + batch=batch, + fwd_out=fwd_out, + ) + action_dist_inputs = fwd_out[Columns.ACTION_DIST_INPUTS] + action_dist_class = self._module[module_id].get_train_action_dist_cls() + action_dist = action_dist_class.from_logits(action_dist_inputs) + loss = -torch.mean(action_dist.logp(batch[Columns.ACTIONS])) + + return loss diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_module.py b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_module.py new file mode 100644 index 0000000000000000000000000000000000000000..afaedd7ae581f7da505fb15222b6eea91e43c827 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/testing/torch/bc_module.py @@ -0,0 +1,149 @@ +from typing import Any, Dict + +from ray.rllib.core.columns import Columns +from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class DiscreteBCTorchModule(TorchRLModule): + def setup(self): + input_dim = self.observation_space.shape[0] + hidden_dim = self.model_config["fcnet_hiddens"][0] + output_dim = self.action_space.n + + self.policy = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, output_dim), + ) + + self.input_dim = input_dim + + def get_train_action_dist_cls(self): + return TorchCategorical + + def get_exploration_action_dist_cls(self): + return TorchCategorical + + def get_inference_action_dist_cls(self): + return TorchCategorical + + @override(RLModule) + def _forward_inference(self, batch: Dict[str, Any]) -> Dict[str, Any]: + with torch.no_grad(): + return self._forward_train(batch) + + @override(RLModule) + def _forward_exploration(self, batch: Dict[str, Any]) -> Dict[str, Any]: + with torch.no_grad(): + return self._forward_train(batch) + + @override(RLModule) + def _forward_train(self, batch: Dict[str, Any]) -> Dict[str, Any]: + action_logits = self.policy(batch["obs"]) + return {Columns.ACTION_DIST_INPUTS: action_logits} + + +class BCTorchRLModuleWithSharedGlobalEncoder(TorchRLModule): + """An example of an RLModule that uses an encoder shared with other things. + + For example, we could consider a multi-agent case where for inference each agent + needs to know the global state of the environment, as well as the local state of + itself. For better representation learning we would like to share the encoder + across all the modules. So this module simply accepts the encoder object as its + input argument and uses it to encode the global state. The local state is passed + through as is. The policy head is then a simple MLP that takes the concatenation of + the global and local state as input and outputs the action logits. + + """ + + def __init__( + self, + encoder: nn.Module, + local_dim: int, + hidden_dim: int, + action_dim: int, + config=None, + ) -> None: + super().__init__(config=config) + + self.encoder = encoder + self.policy_head = nn.Sequential( + nn.Linear(hidden_dim + local_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, action_dim), + ) + + def get_train_action_dist_cls(self): + return TorchCategorical + + def get_exploration_action_dist_cls(self): + return TorchCategorical + + def get_inference_action_dist_cls(self): + return TorchCategorical + + @override(RLModule) + def _default_input_specs(self): + return [("obs", "global"), ("obs", "local")] + + @override(RLModule) + def _forward_inference(self, batch): + with torch.no_grad(): + return self._common_forward(batch) + + @override(RLModule) + def _forward_exploration(self, batch): + with torch.no_grad(): + return self._common_forward(batch) + + @override(RLModule) + def _forward_train(self, batch): + return self._common_forward(batch) + + def _common_forward(self, batch): + obs = batch["obs"] + global_enc = self.encoder(obs["global"]) + policy_in = torch.cat([global_enc, obs["local"]], dim=-1) + action_logits = self.policy_head(policy_in) + + return {Columns.ACTION_DIST_INPUTS: action_logits} + + +class BCTorchMultiAgentModuleWithSharedEncoder(MultiRLModule): + def setup(self): + module_specs = self.config.modules + module_spec = next(iter(module_specs.values())) + global_dim = module_spec.observation_space["global"].shape[0] + hidden_dim = module_spec.model_config_dict["fcnet_hiddens"][0] + shared_encoder = nn.Sequential( + nn.Linear(global_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + + rl_modules = {} + for module_id, module_spec in module_specs.items(): + rl_modules[module_id] = module_spec.module_class( + config=self.config.modules[module_id].get_rl_module_config(), + encoder=shared_encoder, + local_dim=module_spec.observation_space["local"].shape[0], + hidden_dim=hidden_dim, + action_dim=module_spec.action_space.n, + ) + + self._rl_modules = rl_modules + + def serialize(self): + # TODO (Kourosh): Implement when needed. + raise NotImplementedError + + def deserialize(self, data): + # TODO (Kourosh): Implement when needed. + raise NotImplementedError diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__init__.py b/lib/python3.12/site-packages/ray/rllib/offline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c58b423c77babda6f6d016b7a75629d506b8ca8e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/__init__.py @@ -0,0 +1,29 @@ +from ray.rllib.offline.d4rl_reader import D4RLReader +from ray.rllib.offline.dataset_reader import DatasetReader, get_dataset_and_shards +from ray.rllib.offline.dataset_writer import DatasetWriter +from ray.rllib.offline.feature_importance import FeatureImportance +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.offline.json_reader import JsonReader +from ray.rllib.offline.json_writer import JsonWriter +from ray.rllib.offline.mixed_input import MixedInput +from ray.rllib.offline.output_writer import NoopOutput, OutputWriter +from ray.rllib.offline.resource import get_offline_io_resource_bundles +from ray.rllib.offline.shuffled_input import ShuffledInput + +__all__ = [ + "IOContext", + "JsonReader", + "JsonWriter", + "NoopOutput", + "OutputWriter", + "InputReader", + "MixedInput", + "ShuffledInput", + "D4RLReader", + "DatasetReader", + "DatasetWriter", + "get_dataset_and_shards", + "get_offline_io_resource_bundles", + "FeatureImportance", +] diff --git a/lib/python3.12/site-packages/ray/rllib/offline/d4rl_reader.py b/lib/python3.12/site-packages/ray/rllib/offline/d4rl_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..2800bf08c04e80ad96b011d44c415f5a1003baf3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/d4rl_reader.py @@ -0,0 +1,52 @@ +import logging +from typing import Dict + +import gymnasium as gym + +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import PublicAPI, override +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger(__name__) + + +@PublicAPI +class D4RLReader(InputReader): + """Reader object that loads the dataset from the D4RL dataset.""" + + @PublicAPI + def __init__(self, inputs: str, ioctx: IOContext = None): + """Initializes a D4RLReader instance. + + Args: + inputs: String corresponding to the D4RL environment name. + ioctx: Current IO context object. + """ + import d4rl + + self.env = gym.make(inputs) + self.dataset = _convert_to_batch(d4rl.qlearning_dataset(self.env)) + assert self.dataset.count >= 1 + self.counter = 0 + + @override(InputReader) + def next(self) -> SampleBatchType: + if self.counter >= self.dataset.count: + self.counter = 0 + + self.counter += 1 + return self.dataset.slice(start=self.counter, end=self.counter + 1) + + +def _convert_to_batch(dataset: Dict) -> SampleBatchType: + # Converts D4RL dataset to SampleBatch + d = {} + d[SampleBatch.OBS] = dataset["observations"] + d[SampleBatch.ACTIONS] = dataset["actions"] + d[SampleBatch.NEXT_OBS] = dataset["next_observations"] + d[SampleBatch.REWARDS] = dataset["rewards"] + d[SampleBatch.TERMINATEDS] = dataset["terminals"] + + return SampleBatch(d) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/dataset_reader.py b/lib/python3.12/site-packages/ray/rllib/offline/dataset_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..cf5abffc9c40048f4c3320f53743795df4854874 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/dataset_reader.py @@ -0,0 +1,290 @@ +import logging +import math +import re +import zipfile +from pathlib import Path +from typing import TYPE_CHECKING, List, Optional, Tuple + +import numpy as np + +import ray.data +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.offline.json_reader import from_json_data, postprocess_actions +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch, concat_samples +from ray.rllib.utils.annotations import PublicAPI, override +from ray.rllib.utils.typing import SampleBatchType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +DEFAULT_NUM_CPUS_PER_TASK = 0.5 + +logger = logging.getLogger(__name__) + + +def _unzip_this_path(fpath: Path, extract_path: str): + with zipfile.ZipFile(str(fpath), "r") as zip_ref: + zip_ref.extractall(extract_path) + + +def _unzip_if_needed(paths: List[str], format: str): + """If a path in paths is a zip file, unzip it and use path of the unzipped file""" + ret_paths = [] + for path in paths: + if re.search("\\.zip$", str(path)): + # TODO: We need to add unzip support for s3 + if str(path).startswith("s3://"): + raise ValueError( + "unzip_if_needed currently does not support remote paths from s3" + ) + extract_path = "./" + try: + _unzip_this_path(str(path), extract_path) + except FileNotFoundError: + # intrepreted as a relative path to rllib folder + try: + # TODO: remove this later when we replace all tests with s3 paths + _unzip_this_path(Path(__file__).parent.parent / path, extract_path) + except FileNotFoundError: + raise FileNotFoundError(f"File not found: {path}") + + unzipped_path = str( + Path(extract_path).absolute() / f"{Path(path).stem}.{format}" + ) + ret_paths.append(unzipped_path) + else: + # TODO: We can get rid of this logic when we replace all tests with s3 paths + if str(path).startswith("s3://"): + ret_paths.append(path) + else: + if not Path(path).exists(): + relative_path = str(Path(__file__).parent.parent / path) + if not Path(relative_path).exists(): + raise FileNotFoundError(f"File not found: {path}") + path = relative_path + ret_paths.append(path) + return ret_paths + + +@PublicAPI +def get_dataset_and_shards( + config: "AlgorithmConfig", num_workers: int = 0 +) -> Tuple[ray.data.Dataset, List[ray.data.Dataset]]: + """Returns a dataset and a list of shards. + + This function uses algorithm configs to create a dataset and a list of shards. + The following config keys are used to create the dataset: + input: The input type should be "dataset". + input_config: A dict containing the following key and values: + `format`: str, speciifies the format of the input data. This will be the + format that ray dataset supports. See ray.data.Dataset for + supported formats. Only "parquet" or "json" are supported for now. + `paths`: str, a single string or a list of strings. Each string is a path + to a file or a directory holding the dataset. It can be either a local path + or a remote path (e.g. to an s3 bucket). + `loader_fn`: Callable[None, ray.data.Dataset], Instead of + specifying paths and format, you can specify a function to load the dataset. + `parallelism`: int, The number of tasks to use for loading the dataset. + If not specified, it will be set to the number of workers. + `num_cpus_per_read_task`: float, The number of CPUs to use for each read + task. If not specified, it will be set to 0.5. + + Args: + config: The config dict for the algorithm. + num_workers: The number of shards to create for remote workers. + + Returns: + dataset: The dataset object. + shards: A list of dataset shards. For num_workers > 0 the first returned + shared would be a dummy None shard for local_worker. + """ + # check input and input config keys + assert config.input_ == "dataset", ( + f"Must specify config.input_ as 'dataset' if" + f" calling `get_dataset_and_shards`. Got {config.input_}" + ) + + # check input config format + input_config = config.input_config + format = input_config.get("format") + + supported_fmts = ["json", "parquet"] + if format is not None and format not in supported_fmts: + raise ValueError( + f"Unsupported format {format}. Supported formats are {supported_fmts}" + ) + + # check paths and loader_fn since only one of them is required. + paths = input_config.get("paths") + loader_fn = input_config.get("loader_fn") + if loader_fn and (format or paths): + raise ValueError( + "When using a `loader_fn`, you cannot specify a `format` or `path`." + ) + + # check if at least loader_fn or format + path is specified. + if not (format and paths) and not loader_fn: + raise ValueError( + "Must specify either a `loader_fn` or a `format` and `path` in " + "`input_config`." + ) + + # check paths to be a str or list[str] if not None + if paths is not None: + if isinstance(paths, str): + paths = [paths] + elif isinstance(paths, list): + assert isinstance(paths[0], str), "Paths must be a list of path strings." + else: + raise ValueError("Paths must be a path string or a list of path strings.") + paths = _unzip_if_needed(paths, format) + + # TODO (Kourosh): num_workers is not necessary since we can use parallelism for + # everything. Having two parameters is confusing here. Remove num_workers later. + parallelism = input_config.get("parallelism", num_workers or 1) + cpus_per_task = input_config.get( + "num_cpus_per_read_task", DEFAULT_NUM_CPUS_PER_TASK + ) + + if loader_fn: + dataset = loader_fn() + elif format == "json": + dataset = ray.data.read_json( + paths, parallelism=parallelism, ray_remote_args={"num_cpus": cpus_per_task} + ) + elif format == "parquet": + dataset = ray.data.read_parquet( + paths, parallelism=parallelism, ray_remote_args={"num_cpus": cpus_per_task} + ) + else: + raise ValueError("Un-supported Ray dataset format: ", format) + + # Local worker will be responsible for sampling. + if num_workers == 0: + # Dataset is the only shard we need. + return dataset, [dataset] + # Remote workers are responsible for sampling: + else: + # Each remote worker gets 1 shard. + remote_shards = dataset.repartition( + num_blocks=num_workers, shuffle=False + ).split(num_workers) + + # The first None shard is for the local worker, which + # shouldn't be doing rollout work anyways. + return dataset, [None] + remote_shards + + +@PublicAPI +class DatasetReader(InputReader): + """Reader object that loads data from Ray Dataset. + + Examples: + config = { + "input": "dataset", + "input_config": { + "format": "json", + # A single data file, a directory, or anything + # that ray.data.dataset recognizes. + "paths": "/tmp/sample_batches/", + # By default, parallelism=num_workers. + "parallelism": 3, + # Dataset allocates 0.5 CPU for each reader by default. + # Adjust this value based on the size of your offline dataset. + "num_cpus_per_read_task": 0.5, + } + } + """ + + @PublicAPI + def __init__(self, ds: ray.data.Dataset, ioctx: Optional[IOContext] = None): + """Initializes a DatasetReader instance. + + Args: + ds: Ray dataset to sample from. + """ + self._ioctx = ioctx or IOContext() + self._default_policy = self.policy_map = None + self.preprocessor = None + self._dataset = ds + self.count = None if not self._dataset else self._dataset.count() + # do this to disable the ray data stdout logging + ray.data.DataContext.get_current().enable_progress_bars = False + + # the number of steps to return per call to next() + self.batch_size = self._ioctx.config.get("train_batch_size", 1) + num_workers = self._ioctx.config.get("num_env_runners", 0) + seed = self._ioctx.config.get("seed", None) + if num_workers: + self.batch_size = max(math.ceil(self.batch_size / num_workers), 1) + # We allow the creation of a non-functioning None DatasetReader. + # It's useful for example for a non-rollout local worker. + if ds: + if self._ioctx.worker is not None: + self._policy_map = self._ioctx.worker.policy_map + self._default_policy = self._policy_map.get(DEFAULT_POLICY_ID) + self.preprocessor = ( + self._ioctx.worker.preprocessors.get(DEFAULT_POLICY_ID) + if not self._ioctx.config.get("_disable_preprocessors", False) + else None + ) + print( + f"DatasetReader {self._ioctx.worker_index} has {ds.count()}, samples." + ) + + def iterator(): + while True: + ds = self._dataset.random_shuffle(seed=seed) + yield from ds.iter_rows() + + self._iter = iterator() + else: + self._iter = None + + @override(InputReader) + def next(self) -> SampleBatchType: + # next() should not get called on None DatasetReader. + assert self._iter is not None + ret = [] + count = 0 + while count < self.batch_size: + d = next(self._iter) + # Columns like obs are compressed when written by DatasetWriter. + d = from_json_data(d, self._ioctx.worker) + count += d.count + d = self._preprocess_if_needed(d) + d = postprocess_actions(d, self._ioctx) + d = self._postprocess_if_needed(d) + ret.append(d) + ret = concat_samples(ret) + return ret + + def _preprocess_if_needed(self, batch: SampleBatchType) -> SampleBatchType: + # TODO: @kourosh, preprocessor is only supported for single agent case. + if self.preprocessor: + for key in (SampleBatch.CUR_OBS, SampleBatch.NEXT_OBS): + if key in batch: + batch[key] = np.stack( + [self.preprocessor.transform(s) for s in batch[key]] + ) + return batch + + def _postprocess_if_needed(self, batch: SampleBatchType) -> SampleBatchType: + if not self._ioctx.config.get("postprocess_inputs"): + return batch + + if isinstance(batch, SampleBatch): + out = [] + for sub_batch in batch.split_by_episode(): + if self._default_policy is not None: + out.append(self._default_policy.postprocess_trajectory(sub_batch)) + else: + out.append(sub_batch) + return concat_samples(out) + else: + # TODO(ekl) this is trickier since the alignments between agent + # trajectories in the episode are not available any more. + raise NotImplementedError( + "Postprocessing of multi-agent data not implemented yet." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/dataset_writer.py b/lib/python3.12/site-packages/ray/rllib/offline/dataset_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..6b3ac7a153735dface27283d4941800b87c25abe --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/dataset_writer.py @@ -0,0 +1,82 @@ +import logging +import os +import time +from typing import Dict, List + +from ray import data +from ray.rllib.offline.io_context import IOContext +from ray.rllib.offline.json_writer import _to_json_dict +from ray.rllib.offline.output_writer import OutputWriter +from ray.rllib.utils.annotations import PublicAPI, override +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger(__name__) + + +@PublicAPI +class DatasetWriter(OutputWriter): + """Writer object that saves experiences using Datasets.""" + + @PublicAPI + def __init__( + self, + ioctx: IOContext = None, + compress_columns: List[str] = frozenset(["obs", "new_obs"]), + ): + """Initializes a DatasetWriter instance. + + Examples: + config = { + "output": "dataset", + "output_config": { + "format": "json", + "path": "/tmp/test_samples/", + "max_num_samples_per_file": 100000, + } + } + + Args: + ioctx: current IO context object. + compress_columns: list of sample batch columns to compress. + """ + self.ioctx = ioctx or IOContext() + + output_config: Dict = ioctx.output_config + assert ( + "format" in output_config + ), "output_config.format must be specified when using Dataset output." + assert ( + "path" in output_config + ), "output_config.path must be specified when using Dataset output." + + self.format = output_config["format"] + self.path = os.path.abspath(os.path.expanduser(output_config["path"])) + self.max_num_samples_per_file = ( + output_config["max_num_samples_per_file"] + if "max_num_samples_per_file" in output_config + else 100000 + ) + self.compress_columns = compress_columns + + self.samples = [] + + @override(OutputWriter) + def write(self, sample_batch: SampleBatchType): + start = time.time() + + # Make sure columns like obs are compressed and writable. + d = _to_json_dict(sample_batch, self.compress_columns) + self.samples.append(d) + + # Todo: We should flush at the end of sampling even if this + # condition was not reached. + if len(self.samples) >= self.max_num_samples_per_file: + ds = data.from_items(self.samples).repartition(num_blocks=1, shuffle=False) + if self.format == "json": + ds.write_json(self.path, try_create_dir=True) + elif self.format == "parquet": + ds.write_parquet(self.path, try_create_dir=True) + else: + raise ValueError("Unknown output type: ", self.format) + self.samples = [] + logger.debug("Wrote dataset in {}s".format(time.time() - start)) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__init__.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f2561b64877680bce3ea272208b8f18cb893fd0f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__init__.py @@ -0,0 +1,15 @@ +from ray.rllib.offline.estimators.direct_method import DirectMethod +from ray.rllib.offline.estimators.doubly_robust import DoublyRobust +from ray.rllib.offline.estimators.importance_sampling import ImportanceSampling +from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator +from ray.rllib.offline.estimators.weighted_importance_sampling import ( + WeightedImportanceSampling, +) + +__all__ = [ + "OffPolicyEstimator", + "ImportanceSampling", + "WeightedImportanceSampling", + "DirectMethod", + "DoublyRobust", +] diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/doubly_robust.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/doubly_robust.py new file mode 100644 index 0000000000000000000000000000000000000000..3d17ea6c22f1ed4918a362d316417be1fdd18007 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/doubly_robust.py @@ -0,0 +1,251 @@ +import logging +import math +from typing import Any, Dict, List, Optional + +import numpy as np +import pandas as pd + +from ray.data import Dataset +from ray.rllib.offline.estimators.fqe_torch_model import FQETorchModel +from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator +from ray.rllib.offline.offline_evaluation_utils import ( + compute_is_weights, + compute_q_and_v_values, +) +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch +from ray.rllib.utils.annotations import DeveloperAPI, override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger() + + +@DeveloperAPI +class DoublyRobust(OffPolicyEstimator): + r"""The Doubly Robust estimator. + + Let s_t, a_t, and r_t be the state, action, and reward at timestep t. + + This method trains a Q-model for the evaluation policy \pi_e on behavior + data generated by \pi_b. Currently, RLlib implements this using + Fitted-Q Evaluation (FQE). You can also implement your own model + and pass it in as `q_model_config = {"type": your_model_class, **your_kwargs}`. + + For behavior policy \pi_b and evaluation policy \pi_e, define the + cumulative importance ratio at timestep t as: + p_t = \sum_{t'=0}^t (\pi_e(a_{t'} | s_{t'}) / \pi_b(a_{t'} | s_{t'})). + + Consider an episode with length T. Let V_T = 0. + For all t in {0, T - 1}, use the following recursive update: + V_t^DR = (\sum_{a \in A} \pi_e(a | s_t) Q(s_t, a)) + + p_t * (r_t + \gamma * V_{t+1}^DR - Q(s_t, a_t)) + + This estimator computes the expected return for \pi_e for an episode as: + V^{\pi_e}(s_0) = V_0^DR + and returns the mean and standard deviation over episodes. + + For more information refer to https://arxiv.org/pdf/1911.06854.pdf""" + + @override(OffPolicyEstimator) + def __init__( + self, + policy: Policy, + gamma: float, + epsilon_greedy: float = 0.0, + normalize_weights: bool = True, + q_model_config: Optional[Dict] = None, + ): + """Initializes a Doubly Robust OPE Estimator. + + Args: + policy: Policy to evaluate. + gamma: Discount factor of the environment. + epsilon_greedy: The probability by which we act acording to a fully random + policy during deployment. With 1-epsilon_greedy we act + according the target policy. + normalize_weights: If True, the inverse propensity scores are normalized to + their sum across the entire dataset. The effect of this is similar to + weighted importance sampling compared to standard importance sampling. + q_model_config: Arguments to specify the Q-model. Must specify + a `type` key pointing to the Q-model class. + This Q-model is trained in the train() method and is used + to compute the state-value and Q-value estimates + for the DoublyRobust estimator. + It must implement `train`, `estimate_q`, and `estimate_v`. + TODO (Rohan138): Unify this with RLModule API. + """ + + super().__init__(policy, gamma, epsilon_greedy) + q_model_config = q_model_config or {} + q_model_config["gamma"] = gamma + + self._model_cls = q_model_config.pop("type", FQETorchModel) + self._model_configs = q_model_config + self._normalize_weights = normalize_weights + + self.model = self._model_cls( + policy=policy, + **q_model_config, + ) + assert hasattr( + self.model, "estimate_v" + ), "self.model must implement `estimate_v`!" + assert hasattr( + self.model, "estimate_q" + ), "self.model must implement `estimate_q`!" + + @override(OffPolicyEstimator) + def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, Any]: + estimates_per_epsiode = {} + + rewards, old_prob = episode["rewards"], episode["action_prob"] + new_prob = self.compute_action_probs(episode) + + weight = new_prob / old_prob + + v_behavior = 0.0 + v_target = 0.0 + q_values = self.model.estimate_q(episode) + q_values = convert_to_numpy(q_values) + v_values = self.model.estimate_v(episode) + v_values = convert_to_numpy(v_values) + assert q_values.shape == v_values.shape == (episode.count,) + + for t in reversed(range(episode.count)): + v_behavior = rewards[t] + self.gamma * v_behavior + v_target = v_values[t] + weight[t] * ( + rewards[t] + self.gamma * v_target - q_values[t] + ) + v_target = v_target.item() + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def estimate_on_single_step_samples( + self, batch: SampleBatch + ) -> Dict[str, List[float]]: + estimates_per_epsiode = {} + + rewards, old_prob = batch["rewards"], batch["action_prob"] + new_prob = self.compute_action_probs(batch) + + q_values = self.model.estimate_q(batch) + q_values = convert_to_numpy(q_values) + v_values = self.model.estimate_v(batch) + v_values = convert_to_numpy(v_values) + + v_behavior = rewards + + weight = new_prob / old_prob + v_target = v_values + weight * (rewards - q_values) + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def train(self, batch: SampleBatchType) -> Dict[str, Any]: + """Trains self.model on the given batch. + + Args: + batch: A SampleBatch or MultiAgentbatch to train on + + Returns: + A dict with key "loss" and value as the mean training loss. + """ + batch = convert_ma_batch_to_sample_batch(batch) + losses = self.model.train(batch) + return {"loss": np.mean(losses)} + + @override(OfflineEvaluator) + def estimate_on_dataset( + self, dataset: Dataset, *, n_parallelism: int = ... + ) -> Dict[str, Any]: + """Estimates the policy value using the Doubly Robust estimator. + + The doubly robust estimator uses normalization of importance sampling weights + (aka. propensity ratios) to the average of the importance weights across the + entire dataset. This is done to reduce the variance of the estimate (similar to + weighted importance sampling). You can disable this by setting + `normalize_weights=False` in the constructor. + + Note: This estimate works for only discrete action spaces for now. + + Args: + dataset: Dataset to compute the estimate on. Each record in dataset should + include the following columns: `obs`, `actions`, `action_prob` and + `rewards`. The `obs` on each row shoud be a vector of D dimensions. + n_parallelism: Number of parallelism to use for the computation. + + Returns: + A dict with the following keys: + v_target: The estimated value of the target policy. + v_behavior: The estimated value of the behavior policy. + v_gain: The estimated gain of the target policy over the behavior + policy. + v_std: The standard deviation of the estimated value of the target. + """ + + # step 1: compute the weights and weighted rewards + batch_size = max(dataset.count() // n_parallelism, 1) + updated_ds = dataset.map_batches( + compute_is_weights, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "policy_state": self.policy.get_state(), + "estimator_class": self.__class__, + }, + ) + + # step 2: compute q_values and v_values + batch_size = max(updated_ds.count() // n_parallelism, 1) + updated_ds = updated_ds.map_batches( + compute_q_and_v_values, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "model_class": self.model.__class__, + "model_state": self.model.get_state(), + }, + ) + + # step 3: compute the v_target + def compute_v_target(batch: pd.DataFrame, normalizer: float = 1.0): + weights = batch["weights"] / normalizer + batch["v_target"] = batch["v_values"] + weights * ( + batch["rewards"] - batch["q_values"] + ) + batch["v_behavior"] = batch["rewards"] + return batch + + normalizer = updated_ds.mean("weights") if self._normalize_weights else 1.0 + updated_ds = updated_ds.map_batches( + compute_v_target, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={"normalizer": normalizer}, + ) + + v_behavior = updated_ds.mean("v_behavior") + v_target = updated_ds.mean("v_target") + v_gain_mean = v_target / v_behavior + v_gain_ste = ( + updated_ds.std("v_target") + / normalizer + / v_behavior + / math.sqrt(dataset.count()) + ) + + return { + "v_behavior": v_behavior, + "v_target": v_target, + "v_gain_mean": v_gain_mean, + "v_gain_ste": v_gain_ste, + } diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/fqe_torch_model.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/fqe_torch_model.py new file mode 100644 index 0000000000000000000000000000000000000000..b95417e3169bcaec2c753f17f49aff0cf87bfe4a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/fqe_torch_model.py @@ -0,0 +1,297 @@ +from typing import Any, Dict + +from gymnasium.spaces import Discrete + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.models.utils import get_initializer +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI, is_overridden +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType + +torch, nn = try_import_torch() + +# TODO: Create a config object for FQE and unify it with the RLModule API + + +@DeveloperAPI +class FQETorchModel: + """Pytorch implementation of the Fitted Q-Evaluation (FQE) model from + https://arxiv.org/abs/1911.06854 + """ + + def __init__( + self, + policy: Policy, + gamma: float, + model_config: ModelConfigDict = None, + n_iters: int = 1, + lr: float = 1e-3, + min_loss_threshold: float = 1e-4, + clip_grad_norm: float = 100.0, + minibatch_size: int = None, + polyak_coef: float = 1.0, + ) -> None: + """ + Args: + policy: Policy to evaluate. + gamma: Discount factor of the environment. + model_config: The ModelConfigDict for self.q_model, defaults to: + { + "fcnet_hiddens": [8, 8], + "fcnet_activation": "relu", + "vf_share_layers": True, + }, + n_iters: Number of gradient steps to run on batch, defaults to 1 + lr: Learning rate for Adam optimizer + min_loss_threshold: Early stopping if mean loss < min_loss_threshold + clip_grad_norm: Clip loss gradients to this maximum value + minibatch_size: Minibatch size for training Q-function; + if None, train on the whole batch + polyak_coef: Polyak averaging factor for target Q-function + """ + self.policy = policy + assert isinstance( + policy.action_space, Discrete + ), f"{self.__class__.__name__} only supports discrete action spaces!" + self.gamma = gamma + self.observation_space = policy.observation_space + self.action_space = policy.action_space + + if model_config is None: + model_config = { + "fcnet_hiddens": [32, 32, 32], + "fcnet_activation": "relu", + "vf_share_layers": True, + } + self.model_config = model_config + + self.device = self.policy.device + self.q_model: TorchModelV2 = ModelCatalog.get_model_v2( + self.observation_space, + self.action_space, + self.action_space.n, + model_config, + framework="torch", + name="TorchQModel", + ).to(self.device) + + self.target_q_model: TorchModelV2 = ModelCatalog.get_model_v2( + self.observation_space, + self.action_space, + self.action_space.n, + model_config, + framework="torch", + name="TargetTorchQModel", + ).to(self.device) + + self.n_iters = n_iters + self.lr = lr + self.min_loss_threshold = min_loss_threshold + self.clip_grad_norm = clip_grad_norm + self.minibatch_size = minibatch_size + self.polyak_coef = polyak_coef + self.optimizer = torch.optim.Adam(self.q_model.variables(), self.lr) + initializer = get_initializer("xavier_uniform", framework="torch") + # Hard update target + self.update_target(polyak_coef=1.0) + + def f(m): + if isinstance(m, nn.Linear): + initializer(m.weight) + + self.initializer = f + + def train(self, batch: SampleBatch) -> TensorType: + """Trains self.q_model using FQE loss on given batch. + + Args: + batch: A SampleBatch of episodes to train on + + Returns: + A list of losses for each training iteration + """ + losses = [] + minibatch_size = self.minibatch_size or batch.count + # Copy batch for shuffling + batch = batch.copy(shallow=True) + for _ in range(self.n_iters): + minibatch_losses = [] + batch.shuffle() + for idx in range(0, batch.count, minibatch_size): + minibatch = batch[idx : idx + minibatch_size] + obs = torch.tensor(minibatch[SampleBatch.OBS], device=self.device) + actions = torch.tensor( + minibatch[SampleBatch.ACTIONS], + device=self.device, + dtype=int, + ) + rewards = torch.tensor( + minibatch[SampleBatch.REWARDS], device=self.device + ) + next_obs = torch.tensor( + minibatch[SampleBatch.NEXT_OBS], device=self.device + ) + dones = torch.tensor( + minibatch[SampleBatch.TERMINATEDS], device=self.device, dtype=float + ) + + # Compute Q-values for current obs + q_values, _ = self.q_model({"obs": obs}, [], None) + q_acts = torch.gather(q_values, -1, actions.unsqueeze(-1)).squeeze(-1) + + next_action_probs = self._compute_action_probs(next_obs) + + # Compute Q-values for next obs + with torch.no_grad(): + next_q_values, _ = self.target_q_model({"obs": next_obs}, [], None) + + # Compute estimated state value next_v = E_{a ~ pi(s)} [Q(next_obs,a)] + next_v = torch.sum(next_q_values * next_action_probs, axis=-1) + targets = rewards + (1 - dones) * self.gamma * next_v + loss = (targets - q_acts) ** 2 + loss = torch.mean(loss) + self.optimizer.zero_grad() + loss.backward() + nn.utils.clip_grad.clip_grad_norm_( + self.q_model.variables(), self.clip_grad_norm + ) + self.optimizer.step() + minibatch_losses.append(loss.item()) + iter_loss = sum(minibatch_losses) / len(minibatch_losses) + losses.append(iter_loss) + if iter_loss < self.min_loss_threshold: + break + self.update_target() + return losses + + def estimate_q(self, batch: SampleBatch) -> TensorType: + obs = torch.tensor(batch[SampleBatch.OBS], device=self.device) + with torch.no_grad(): + q_values, _ = self.q_model({"obs": obs}, [], None) + actions = torch.tensor( + batch[SampleBatch.ACTIONS], device=self.device, dtype=int + ) + q_values = torch.gather(q_values, -1, actions.unsqueeze(-1)).squeeze(-1) + return q_values + + def estimate_v(self, batch: SampleBatch) -> TensorType: + obs = torch.tensor(batch[SampleBatch.OBS], device=self.device) + with torch.no_grad(): + q_values, _ = self.q_model({"obs": obs}, [], None) + # Compute pi(a | s) for each action a in policy.action_space + action_probs = self._compute_action_probs(obs) + v_values = torch.sum(q_values * action_probs, axis=-1) + return v_values + + def update_target(self, polyak_coef=None): + # Update_target will be called periodically to copy Q network to + # target Q network, using (soft) polyak_coef-synching. + polyak_coef = polyak_coef or self.polyak_coef + model_state_dict = self.q_model.state_dict() + # Support partial (soft) synching. + # If polyak_coef == 1.0: Full sync from Q-model to target Q-model. + target_state_dict = self.target_q_model.state_dict() + model_state_dict = { + k: polyak_coef * model_state_dict[k] + (1 - polyak_coef) * v + for k, v in target_state_dict.items() + } + + self.target_q_model.load_state_dict(model_state_dict) + + def _compute_action_probs(self, obs: TensorType) -> TensorType: + """Compute action distribution over the action space. + + Args: + obs: A tensor of observations of shape (batch_size * obs_dim) + + Returns: + action_probs: A tensor of action probabilities + of shape (batch_size * action_dim) + """ + input_dict = {SampleBatch.OBS: obs} + seq_lens = torch.ones(len(obs), device=self.device, dtype=int) + state_batches = [] + if is_overridden(self.policy.action_distribution_fn): + try: + # TorchPolicyV2 function signature + dist_inputs, dist_class, _ = self.policy.action_distribution_fn( + self.policy.model, + obs_batch=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=False, + is_training=False, + ) + except TypeError: + # TorchPolicyV1 function signature for compatibility with DQN + # TODO: Remove this once DQNTorchPolicy is migrated to PolicyV2 + dist_inputs, dist_class, _ = self.policy.action_distribution_fn( + self.policy, + self.policy.model, + input_dict=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=False, + is_training=False, + ) + else: + dist_class = self.policy.dist_class + dist_inputs, _ = self.policy.model(input_dict, state_batches, seq_lens) + action_dist = dist_class(dist_inputs, self.policy.model) + assert isinstance( + action_dist.dist, torch.distributions.categorical.Categorical + ), "FQE only supports Categorical or MultiCategorical distributions!" + action_probs = action_dist.dist.probs + return action_probs + + def get_state(self) -> Dict[str, Any]: + """Returns the current state of the FQE Model.""" + return { + "policy_state": self.policy.get_state(), + "model_config": self.model_config, + "n_iters": self.n_iters, + "lr": self.lr, + "min_loss_threshold": self.min_loss_threshold, + "clip_grad_norm": self.clip_grad_norm, + "minibatch_size": self.minibatch_size, + "polyak_coef": self.polyak_coef, + "gamma": self.gamma, + "q_model_state": self.q_model.state_dict(), + "target_q_model_state": self.target_q_model.state_dict(), + } + + def set_state(self, state: Dict[str, Any]) -> None: + """Sets the current state of the FQE Model. + Args: + state: A state dict returned by `get_state()`. + """ + self.n_iters = state["n_iters"] + self.lr = state["lr"] + self.min_loss_threshold = state["min_loss_threshold"] + self.clip_grad_norm = state["clip_grad_norm"] + self.minibatch_size = state["minibatch_size"] + self.polyak_coef = state["polyak_coef"] + self.gamma = state["gamma"] + self.policy.set_state(state["policy_state"]) + self.q_model.load_state_dict(state["q_model_state"]) + self.target_q_model.load_state_dict(state["target_q_model_state"]) + + @classmethod + def from_state(cls, state: Dict[str, Any]) -> "FQETorchModel": + """Creates a FQE Model from a state dict. + + Args: + state: A state dict returned by `get_state`. + + Returns: + An instance of the FQETorchModel. + """ + policy = Policy.from_state(state["policy_state"]) + model = cls( + policy=policy, gamma=state["gamma"], model_config=state["model_config"] + ) + model.set_state(state) + return model diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/importance_sampling.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/importance_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..0d62a902b9f49c20d15f4f7a03cc8b6784dd6441 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/importance_sampling.py @@ -0,0 +1,125 @@ +import math +from typing import Any, Dict, List + +from ray.data import Dataset +from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator +from ray.rllib.offline.offline_evaluation_utils import ( + compute_is_weights, + remove_time_dim, +) +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI, override + + +@DeveloperAPI +class ImportanceSampling(OffPolicyEstimator): + r"""The step-wise IS estimator. + + Let s_t, a_t, and r_t be the state, action, and reward at timestep t. + + For behavior policy \pi_b and evaluation policy \pi_e, define the + cumulative importance ratio at timestep t as: + p_t = \sum_{t'=0}^t (\pi_e(a_{t'} | s_{t'}) / \pi_b(a_{t'} | s_{t'})). + + This estimator computes the expected return for \pi_e for an episode as: + V^{\pi_e}(s_0) = \sum_t \gamma ^ {t} * p_t * r_t + and returns the mean and standard deviation over episodes. + + For more information refer to https://arxiv.org/pdf/1911.06854.pdf""" + + @override(OffPolicyEstimator) + def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, float]: + estimates_per_epsiode = {} + + rewards, old_prob = episode["rewards"], episode["action_prob"] + new_prob = self.compute_action_probs(episode) + + # calculate importance ratios + p = [] + for t in range(episode.count): + if t == 0: + pt_prev = 1.0 + else: + pt_prev = p[t - 1] + p.append(pt_prev * new_prob[t] / old_prob[t]) + + # calculate stepwise IS estimate + v_behavior = 0.0 + v_target = 0.0 + for t in range(episode.count): + v_behavior += rewards[t] * self.gamma**t + v_target += p[t] * rewards[t] * self.gamma**t + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def estimate_on_single_step_samples( + self, batch: SampleBatch + ) -> Dict[str, List[float]]: + estimates_per_epsiode = {} + + rewards, old_prob = batch["rewards"], batch["action_prob"] + new_prob = self.compute_action_probs(batch) + + weights = new_prob / old_prob + v_behavior = rewards + v_target = weights * rewards + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OfflineEvaluator) + def estimate_on_dataset( + self, dataset: Dataset, *, n_parallelism: int = ... + ) -> Dict[str, Any]: + """Computes the Importance sampling estimate on the given dataset. + + Note: This estimate works for both continuous and discrete action spaces. + + Args: + dataset: Dataset to compute the estimate on. Each record in dataset should + include the following columns: `obs`, `actions`, `action_prob` and + `rewards`. The `obs` on each row shoud be a vector of D dimensions. + n_parallelism: The number of parallel workers to use. + + Returns: + A dictionary containing the following keys: + v_target: The estimated value of the target policy. + v_behavior: The estimated value of the behavior policy. + v_gain_mean: The mean of the gain of the target policy over the + behavior policy. + v_gain_ste: The standard error of the gain of the target policy over + the behavior policy. + """ + batch_size = max(dataset.count() // n_parallelism, 1) + dataset = dataset.map_batches( + remove_time_dim, batch_size=batch_size, batch_format="pandas" + ) + updated_ds = dataset.map_batches( + compute_is_weights, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "policy_state": self.policy.get_state(), + "estimator_class": self.__class__, + }, + ) + v_target = updated_ds.mean("weighted_rewards") + v_behavior = updated_ds.mean("rewards") + v_gain_mean = v_target / v_behavior + v_gain_ste = ( + updated_ds.std("weighted_rewards") / v_behavior / math.sqrt(dataset.count()) + ) + + return { + "v_target": v_target, + "v_behavior": v_behavior, + "v_gain_mean": v_gain_mean, + "v_gain_ste": v_gain_ste, + } diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/weighted_importance_sampling.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/weighted_importance_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..67d14682a9969a9ad61c93505be09d50ba7954eb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/weighted_importance_sampling.py @@ -0,0 +1,185 @@ +import math +from typing import Any, Dict, List + +import numpy as np + +from ray.data import Dataset +from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator +from ray.rllib.offline.offline_evaluation_utils import ( + compute_is_weights, + remove_time_dim, +) +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI, override + + +@DeveloperAPI +class WeightedImportanceSampling(OffPolicyEstimator): + r"""The step-wise WIS estimator. + + Let s_t, a_t, and r_t be the state, action, and reward at timestep t. + + For behavior policy \pi_b and evaluation policy \pi_e, define the + cumulative importance ratio at timestep t as: + p_t = \sum_{t'=0}^t (\pi_e(a_{t'} | s_{t'}) / \pi_b(a_{t'} | s_{t'})). + + Define the average importance ratio over episodes i in the dataset D as: + w_t = \sum_{i \in D} p^(i)_t / |D| + + This estimator computes the expected return for \pi_e for an episode as: + V^{\pi_e}(s_0) = \E[\sum_t \gamma ^ {t} * (p_t / w_t) * r_t] + and returns the mean and standard deviation over episodes. + + For more information refer to https://arxiv.org/pdf/1911.06854.pdf""" + + @override(OffPolicyEstimator) + def __init__(self, policy: Policy, gamma: float, epsilon_greedy: float = 0.0): + super().__init__(policy, gamma, epsilon_greedy) + # map from time to cummulative propensity values + self.cummulative_ips_values = [] + # map from time to number of episodes that reached this time + self.episode_timestep_count = [] + # map from eps id to mapping from time to propensity values + self.p = {} + + @override(OffPolicyEstimator) + def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, Any]: + estimates_per_epsiode = {} + rewards = episode["rewards"] + + eps_id = episode[SampleBatch.EPS_ID][0] + if eps_id not in self.p: + raise ValueError( + f"Cannot find target weight for episode {eps_id}. " + f"Did it go though the peek_on_single_episode() function?" + ) + + # calculate stepwise weighted IS estimate + v_behavior = 0.0 + v_target = 0.0 + episode_p = self.p[eps_id] + for t in range(episode.count): + v_behavior += rewards[t] * self.gamma**t + w_t = self.cummulative_ips_values[t] / self.episode_timestep_count[t] + v_target += episode_p[t] / w_t * rewards[t] * self.gamma**t + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def estimate_on_single_step_samples( + self, batch: SampleBatch + ) -> Dict[str, List[float]]: + estimates_per_epsiode = {} + rewards, old_prob = batch["rewards"], batch["action_prob"] + new_prob = self.compute_action_probs(batch) + + weights = new_prob / old_prob + v_behavior = rewards + v_target = weights * rewards / np.mean(weights) + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + estimates_per_epsiode["weights"] = weights + estimates_per_epsiode["new_prob"] = new_prob + estimates_per_epsiode["old_prob"] = old_prob + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def on_before_split_batch_by_episode( + self, sample_batch: SampleBatch + ) -> SampleBatch: + self.cummulative_ips_values = [] + self.episode_timestep_count = [] + self.p = {} + + return sample_batch + + @override(OffPolicyEstimator) + def peek_on_single_episode(self, episode: SampleBatch) -> None: + old_prob = episode["action_prob"] + new_prob = self.compute_action_probs(episode) + + # calculate importance ratios + episode_p = [] + for t in range(episode.count): + if t == 0: + pt_prev = 1.0 + else: + pt_prev = episode_p[t - 1] + episode_p.append(pt_prev * new_prob[t] / old_prob[t]) + + for t, p_t in enumerate(episode_p): + if t >= len(self.cummulative_ips_values): + self.cummulative_ips_values.append(p_t) + self.episode_timestep_count.append(1.0) + else: + self.cummulative_ips_values[t] += p_t + self.episode_timestep_count[t] += 1.0 + + eps_id = episode[SampleBatch.EPS_ID][0] + if eps_id in self.p: + raise ValueError( + f"eps_id {eps_id} was already passed to the peek function. " + f"Make sure dataset contains only unique episodes with unique ids." + ) + self.p[eps_id] = episode_p + + @override(OfflineEvaluator) + def estimate_on_dataset( + self, dataset: Dataset, *, n_parallelism: int = ... + ) -> Dict[str, Any]: + """Computes the weighted importance sampling estimate on a dataset. + + Note: This estimate works for both continuous and discrete action spaces. + + Args: + dataset: Dataset to compute the estimate on. Each record in dataset should + include the following columns: `obs`, `actions`, `action_prob` and + `rewards`. The `obs` on each row shoud be a vector of D dimensions. + n_parallelism: Number of parallel workers to use for the computation. + + Returns: + Dictionary with the following keys: + v_target: The weighted importance sampling estimate. + v_behavior: The behavior policy estimate. + v_gain_mean: The mean of the gain of the target policy over the + behavior policy. + v_gain_ste: The standard error of the gain of the target policy over + the behavior policy. + """ + # compute the weights and weighted rewards + batch_size = max(dataset.count() // n_parallelism, 1) + dataset = dataset.map_batches( + remove_time_dim, batch_size=batch_size, batch_format="pandas" + ) + updated_ds = dataset.map_batches( + compute_is_weights, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "policy_state": self.policy.get_state(), + "estimator_class": self.__class__, + }, + ) + v_target = updated_ds.mean("weighted_rewards") / updated_ds.mean("weights") + v_behavior = updated_ds.mean("rewards") + v_gain_mean = v_target / v_behavior + v_gain_ste = ( + updated_ds.std("weighted_rewards") + / updated_ds.mean("weights") + / v_behavior + / math.sqrt(dataset.count()) + ) + + return { + "v_target": v_target, + "v_behavior": v_behavior, + "v_gain_mean": v_gain_mean, + "v_gain_ste": v_gain_ste, + } diff --git a/lib/python3.12/site-packages/ray/rllib/offline/feature_importance.py b/lib/python3.12/site-packages/ray/rllib/offline/feature_importance.py new file mode 100644 index 0000000000000000000000000000000000000000..bc520f8e5cbe9578fdefa8fb6dcecf70914acccb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/feature_importance.py @@ -0,0 +1,283 @@ +import copy +from typing import Any, Callable, Dict + +import numpy as np +import pandas as pd + +import ray +from ray.data import Dataset +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch +from ray.rllib.utils.annotations import DeveloperAPI, ExperimentalAPI, override +from ray.rllib.utils.typing import SampleBatchType + + +@DeveloperAPI +def _perturb_fn(batch: np.ndarray, index: int): + # shuffle the indexth column features + random_inds = np.random.permutation(batch.shape[0]) + batch[:, index] = batch[random_inds, index] + + +@ExperimentalAPI +def _perturb_df(batch: pd.DataFrame, index: int): + obs_batch = np.vstack(batch["obs"].values) + _perturb_fn(obs_batch, index) + batch["perturbed_obs"] = list(obs_batch) + return batch + + +def _compute_actions( + batch: pd.DataFrame, + policy_state: Dict[str, Any], + input_key: str = "", + output_key: str = "", +): + """A custom local function to do batch prediction of a policy. + + Given the policy state the action predictions are computed as a function of + `input_key` and stored in the `output_key` column. + + Args: + batch: A sub-batch from the dataset. + policy_state: The state of the policy to use for the prediction. + input_key: The key to use for the input to the policy. If not given, the + default is SampleBatch.OBS. + output_key: The key to use for the output of the policy. If not given, the + default is "predicted_actions". + + Returns: + The modified batch with the predicted actions added as a column. + """ + if not input_key: + input_key = SampleBatch.OBS + + policy = Policy.from_state(policy_state) + sample_batch = SampleBatch( + { + SampleBatch.OBS: np.vstack(batch[input_key].values), + } + ) + actions, _, _ = policy.compute_actions_from_input_dict(sample_batch, explore=False) + + if not output_key: + output_key = "predicted_actions" + batch[output_key] = actions + + return batch + + +@ray.remote +def get_feature_importance_on_index( + dataset: ray.data.Dataset, + *, + index: int, + perturb_fn: Callable[[pd.DataFrame, int], None], + batch_size: int, + policy_state: Dict[str, Any], +): + """A remote function to compute the feature importance of a given index. + + Args: + dataset: The dataset to use for the computation. The dataset should have `obs` + and `actions` columns. Each record should be flat d-dimensional array. + index: The index of the feature to compute the importance for. + perturb_fn: The function to use for perturbing the dataset at the given index. + batch_size: The batch size to use for the computation. + policy_state: The state of the policy to use for the computation. + + Returns: + The modified dataset that contains a `delta` column which is the absolute + difference between the expected output and the output due to the perturbation. + """ + perturbed_ds = dataset.map_batches( + perturb_fn, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={"index": index}, + ) + perturbed_actions = perturbed_ds.map_batches( + _compute_actions, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "output_key": "perturbed_actions", + "input_key": "perturbed_obs", + "policy_state": policy_state, + }, + ) + + def delta_fn(batch): + # take the abs difference between columns 'ref_actions` and `perturbed_actions` + # and store it in `diff` + batch["delta"] = np.abs(batch["ref_actions"] - batch["perturbed_actions"]) + return batch + + delta = perturbed_actions.map_batches( + delta_fn, batch_size=batch_size, batch_format="pandas" + ) + + return delta + + +@DeveloperAPI +class FeatureImportance(OfflineEvaluator): + @override(OfflineEvaluator) + def __init__( + self, + policy: Policy, + repeat: int = 1, + limit_fraction: float = 1.0, + perturb_fn: Callable[[pd.DataFrame, int], pd.DataFrame] = _perturb_df, + ): + """Feature importance in a model inspection technique that can be used for any + fitted predictor when the data is tablular. + + This implementation is also known as permutation importance that is defined to + be the variation of the model's prediction when a single feature value is + randomly shuffled. In RLlib it is implemented as a custom OffPolicyEstimator + which is used to evaluate RLlib policies without performing environment + interactions. + + Example usage: In the example below the feature importance module is used to + evaluate the policy and the each feature's importance is computed after each + training iteration. The permutation are repeated `self.repeat` times and the + results are averages across repeats. + + ```python + config = ( + AlgorithmConfig() + .offline_data( + off_policy_estimation_methods= + { + "feature_importance": { + "type": FeatureImportance, + "repeat": 10, + "limit_fraction": 0.1, + } + } + ) + ) + + algorithm = DQN(config=config) + results = algorithm.train() + ``` + + Args: + policy: the policy to use for feature importance. + repeat: number of times to repeat the perturbation. + perturb_fn: function to perturb the features. By default reshuffle the + features within the batch. + limit_fraction: fraction of the dataset to use for feature importance + This is only used in estimate_on_dataset when the dataset is too large + to compute feature importance on. + """ + super().__init__(policy) + self.repeat = repeat + self.perturb_fn = perturb_fn + self.limit_fraction = limit_fraction + + def estimate(self, batch: SampleBatchType) -> Dict[str, Any]: + """Estimate the feature importance of the policy. + + Given a batch of tabular observations, the importance of each feature is + computed by perturbing each feature and computing the difference between the + perturbed policy and the reference policy. The importance is computed for each + feature and each perturbation is repeated `self.repeat` times. + + Args: + batch: the batch of data to use for feature importance. + + Returns: + A dict mapping each feature index string to its importance. + """ + batch = convert_ma_batch_to_sample_batch(batch) + obs_batch = batch["obs"] + n_features = obs_batch.shape[-1] + importance = np.zeros((self.repeat, n_features)) + + ref_actions, _, _ = self.policy.compute_actions(obs_batch, explore=False) + for r in range(self.repeat): + for i in range(n_features): + copy_obs_batch = copy.deepcopy(obs_batch) + _perturb_fn(copy_obs_batch, index=i) + perturbed_actions, _, _ = self.policy.compute_actions( + copy_obs_batch, explore=False + ) + importance[r, i] = np.mean(np.abs(perturbed_actions - ref_actions)) + + # take an average across repeats + importance = importance.mean(0) + metrics = {f"feature_{i}": importance[i] for i in range(len(importance))} + + return metrics + + @override(OfflineEvaluator) + def estimate_on_dataset( + self, dataset: Dataset, *, n_parallelism: int = ... + ) -> Dict[str, Any]: + """Estimate the feature importance of the policy given a dataset. + + For each feature in the dataset, the importance is computed by applying + perturbations to each feature and computing the difference between the + perturbed prediction and the reference prediction. The importance + computation for each feature and each perturbation is repeated `self.repeat` + times. If dataset is large the user can initialize the estimator with a + `limit_fraction` to limit the dataset to a fraction of the original dataset. + + The dataset should include a column named `obs` where each row is a vector of D + dimensions. The importance is computed for each dimension of the vector. + + Note (Implementation detail): The computation across features are distributed + with ray workers since each feature is independent of each other. + + Args: + dataset: the dataset to use for feature importance. + n_parallelism: number of parallel workers to use for feature importance. + + Returns: + A dict mapping each feature index string to its importance. + """ + + policy_state = self.policy.get_state() + # step 1: limit the dataset to a few first rows + ds = dataset.limit(int(self.limit_fraction * dataset.count())) + + # step 2: compute the reference actions + bsize = max(1, ds.count() // n_parallelism) + actions_ds = ds.map_batches( + _compute_actions, + batch_size=bsize, + fn_kwargs={ + "output_key": "ref_actions", + "policy_state": policy_state, + }, + ) + + # step 3: compute the feature importance + n_features = ds.take(1)[0][SampleBatch.OBS].shape[-1] + importance = np.zeros((self.repeat, n_features)) + for r in range(self.repeat): + # shuffle the entire dataset + shuffled_ds = actions_ds.random_shuffle() + bsize_per_task = max(1, (shuffled_ds.count() * n_features) // n_parallelism) + + # for each index perturb the dataset and compute the feat importance score + remote_fns = [ + get_feature_importance_on_index.remote( + dataset=shuffled_ds, + index=i, + perturb_fn=self.perturb_fn, + bsize=bsize_per_task, + policy_state=policy_state, + ) + for i in range(n_features) + ] + ds_w_fi_scores = ray.get(remote_fns) + importance[r] = np.array([d.mean("delta") for d in ds_w_fi_scores]) + + importance = importance.mean(0) + metrics = {f"feature_{i}": importance[i] for i in range(len(importance))} + + return metrics diff --git a/lib/python3.12/site-packages/ray/rllib/offline/input_reader.py b/lib/python3.12/site-packages/ray/rllib/offline/input_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..18f40176072ee6afa081bc5b43e9b89e9138a61f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/input_reader.py @@ -0,0 +1,133 @@ +import logging +import threading +from abc import ABCMeta, abstractmethod +from typing import Dict, List + +import numpy as np + +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import SampleBatchType, TensorType + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +@PublicAPI +class InputReader(metaclass=ABCMeta): + """API for collecting and returning experiences during policy evaluation.""" + + @abstractmethod + @PublicAPI + def next(self) -> SampleBatchType: + """Returns the next batch of read experiences. + + Returns: + The experience read (SampleBatch or MultiAgentBatch). + """ + raise NotImplementedError + + @PublicAPI + def tf_input_ops(self, queue_size: int = 1) -> Dict[str, TensorType]: + """Returns TensorFlow queue ops for reading inputs from this reader. + + The main use of these ops is for integration into custom model losses. + For example, you can use tf_input_ops() to read from files of external + experiences to add an imitation learning loss to your model. + + This method creates a queue runner thread that will call next() on this + reader repeatedly to feed the TensorFlow queue. + + Args: + queue_size: Max elements to allow in the TF queue. + + .. testcode:: + :skipif: True + + from ray.rllib.models.modelv2 import ModelV2 + from ray.rllib.offline.json_reader import JsonReader + imitation_loss = ... + class MyModel(ModelV2): + def custom_loss(self, policy_loss, loss_inputs): + reader = JsonReader(...) + input_ops = reader.tf_input_ops() + logits, _ = self._build_layers_v2( + {"obs": input_ops["obs"]}, + self.num_outputs, self.options) + il_loss = imitation_loss(logits, input_ops["action"]) + return policy_loss + il_loss + + You can find a runnable version of this in examples/custom_loss.py. + + Returns: + Dict of Tensors, one for each column of the read SampleBatch. + """ + + if hasattr(self, "_queue_runner"): + raise ValueError( + "A queue runner already exists for this input reader. " + "You can only call tf_input_ops() once per reader." + ) + + logger.info("Reading initial batch of data from input reader.") + batch = self.next() + if isinstance(batch, MultiAgentBatch): + raise NotImplementedError( + "tf_input_ops() is not implemented for multi agent batches" + ) + + # Note on casting to `np.array(batch[k])`: In order to get all keys that + # are numbers, we need to convert to numpy everything that is not a numpy array. + # This is because SampleBatches used to only hold numpy arrays, but since our + # RNN efforts under RLModules, we also allow lists. + keys = [ + k + for k in sorted(batch.keys()) + if np.issubdtype(np.array(batch[k]).dtype, np.number) + ] + dtypes = [batch[k].dtype for k in keys] + shapes = {k: (-1,) + s[1:] for (k, s) in [(k, batch[k].shape) for k in keys]} + queue = tf1.FIFOQueue(capacity=queue_size, dtypes=dtypes, names=keys) + tensors = queue.dequeue() + + logger.info("Creating TF queue runner for {}".format(self)) + self._queue_runner = _QueueRunner(self, queue, keys, dtypes) + self._queue_runner.enqueue(batch) + self._queue_runner.start() + + out = {k: tf.reshape(t, shapes[k]) for k, t in tensors.items()} + return out + + +class _QueueRunner(threading.Thread): + """Thread that feeds a TF queue from a InputReader.""" + + def __init__( + self, + input_reader: InputReader, + queue: "tf1.FIFOQueue", + keys: List[str], + dtypes: "tf.dtypes.DType", + ): + threading.Thread.__init__(self) + self.sess = tf1.get_default_session() + self.daemon = True + self.input_reader = input_reader + self.keys = keys + self.queue = queue + self.placeholders = [tf1.placeholder(dtype) for dtype in dtypes] + self.enqueue_op = queue.enqueue(dict(zip(keys, self.placeholders))) + + def enqueue(self, batch: SampleBatchType): + data = {self.placeholders[i]: batch[key] for i, key in enumerate(self.keys)} + self.sess.run(self.enqueue_op, feed_dict=data) + + def run(self): + while True: + try: + batch = self.input_reader.next() + self.enqueue(batch) + except Exception: + logger.exception("Error reading from input") diff --git a/lib/python3.12/site-packages/ray/rllib/offline/io_context.py b/lib/python3.12/site-packages/ray/rllib/offline/io_context.py new file mode 100644 index 0000000000000000000000000000000000000000..72576d3730f84e7569e4970c45860e98d9b3a39b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/io_context.py @@ -0,0 +1,72 @@ +import os +from typing import TYPE_CHECKING, Optional + +from ray.rllib.utils.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.evaluation.sampler import SamplerInput + + +@PublicAPI +class IOContext: + """Class containing attributes to pass to input/output class constructors. + + RLlib auto-sets these attributes when constructing input/output classes, + such as InputReaders and OutputWriters. + """ + + @PublicAPI + def __init__( + self, + log_dir: Optional[str] = None, + config: Optional["AlgorithmConfig"] = None, + worker_index: int = 0, + worker: Optional["RolloutWorker"] = None, + ): + """Initializes a IOContext object. + + Args: + log_dir: The logging directory to read from/write to. + config: The (main) AlgorithmConfig object. + worker_index: When there are multiple workers created, this + uniquely identifies the current worker. 0 for the local + worker, >0 for any of the remote workers. + worker: The RolloutWorker object reference. + """ + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + self.log_dir = log_dir or os.getcwd() + # In case no config is provided, use the default one, but set + # `actions_in_input_normalized=True` if we don't have a worker. + # Not having a worker and/or a config should only be the case in some test + # cases, though. + self.config = config or AlgorithmConfig().offline_data( + actions_in_input_normalized=worker is None + ).training(train_batch_size=1) + self.worker_index = worker_index + self.worker = worker + + @PublicAPI + def default_sampler_input(self) -> Optional["SamplerInput"]: + """Returns the RolloutWorker's SamplerInput object, if any. + + Returns None if the RolloutWorker has no SamplerInput. Note that local + workers in case there are also one or more remote workers by default + do not create a SamplerInput object. + + Returns: + The RolloutWorkers' SamplerInput object or None if none exists. + """ + return self.worker.sampler + + @property + @PublicAPI + def input_config(self): + return self.config.get("input_config", {}) + + @property + @PublicAPI + def output_config(self): + return self.config.get("output_config", {}) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/is_estimator.py b/lib/python3.12/site-packages/ray/rllib/offline/is_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..871120ad676bcae81c47626672ba991c16a16fbd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/is_estimator.py @@ -0,0 +1,10 @@ +from ray._common.deprecation import Deprecated +from ray.rllib.offline.estimators.importance_sampling import ImportanceSampling + + +@Deprecated( + new="ray.rllib.offline.estimators.importance_sampling::ImportanceSampling", + error=True, +) +class ImportanceSamplingEstimator(ImportanceSampling): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/offline/json_reader.py b/lib/python3.12/site-packages/ray/rllib/offline/json_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..076791716d8232aab44a23e3d04b05b9524bfbdb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/json_reader.py @@ -0,0 +1,438 @@ +import glob +import json +import logging +import math +import os +import random +import re +import zipfile +from pathlib import Path +from typing import TYPE_CHECKING, List, Optional, Union +from urllib.parse import urlparse + +import numpy as np +import tree # pip install dm_tree + +try: + from smart_open import smart_open +except ImportError: + smart_open = None + +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import ( + DEFAULT_POLICY_ID, + MultiAgentBatch, + SampleBatch, + concat_samples, + convert_ma_batch_to_sample_batch, +) +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI, override +from ray.rllib.utils.compression import unpack_if_needed +from ray.rllib.utils.spaces.space_utils import clip_action, normalize_action +from ray.rllib.utils.typing import Any, FileType, SampleBatchType + +if TYPE_CHECKING: + from ray.rllib.evaluation import RolloutWorker + +logger = logging.getLogger(__name__) + +WINDOWS_DRIVES = [chr(i) for i in range(ord("c"), ord("z") + 1)] + + +def _adjust_obs_actions_for_policy(json_data: dict, policy: Policy) -> dict: + """Handle nested action/observation spaces for policies. + + Translates nested lists/dicts from the json into proper + np.ndarrays, according to the (nested) observation- and action- + spaces of the given policy. + + Providing nested lists w/o this preprocessing step would + confuse a SampleBatch constructor. + """ + for k, v in json_data.items(): + data_col = ( + policy.view_requirements[k].data_col + if k in policy.view_requirements + else "" + ) + # No action flattening -> Process nested (leaf) action(s). + if policy.config.get("_disable_action_flattening") and ( + k == SampleBatch.ACTIONS + or data_col == SampleBatch.ACTIONS + or k == SampleBatch.PREV_ACTIONS + or data_col == SampleBatch.PREV_ACTIONS + ): + json_data[k] = tree.map_structure_up_to( + policy.action_space_struct, + lambda comp: np.array(comp), + json_data[k], + check_types=False, + ) + # No preprocessing -> Process nested (leaf) observation(s). + elif policy.config.get("_disable_preprocessor_api") and ( + k == SampleBatch.OBS + or data_col == SampleBatch.OBS + or k == SampleBatch.NEXT_OBS + or data_col == SampleBatch.NEXT_OBS + ): + json_data[k] = tree.map_structure_up_to( + policy.observation_space_struct, + lambda comp: np.array(comp), + json_data[k], + check_types=False, + ) + return json_data + + +@DeveloperAPI +def _adjust_dones(json_data: dict) -> dict: + """Make sure DONES in json data is properly translated into TERMINATEDS.""" + new_json_data = {} + for k, v in json_data.items(): + # Translate DONES into TERMINATEDS. + if k == SampleBatch.DONES: + new_json_data[SampleBatch.TERMINATEDS] = v + # Leave everything else as-is. + else: + new_json_data[k] = v + + return new_json_data + + +@DeveloperAPI +def postprocess_actions(batch: SampleBatchType, ioctx: IOContext) -> SampleBatchType: + # Clip actions (from any values into env's bounds), if necessary. + cfg = ioctx.config + # TODO(jungong): We should not clip_action in input reader. + # Use connector to handle this. + if cfg.get("clip_actions"): + if ioctx.worker is None: + raise ValueError( + "clip_actions is True but cannot clip actions since no workers exist" + ) + + if isinstance(batch, SampleBatch): + policy = ioctx.worker.policy_map.get(DEFAULT_POLICY_ID) + if policy is None: + assert len(ioctx.worker.policy_map) == 1 + policy = next(iter(ioctx.worker.policy_map.values())) + batch[SampleBatch.ACTIONS] = clip_action( + batch[SampleBatch.ACTIONS], policy.action_space_struct + ) + else: + for pid, b in batch.policy_batches.items(): + b[SampleBatch.ACTIONS] = clip_action( + b[SampleBatch.ACTIONS], + ioctx.worker.policy_map[pid].action_space_struct, + ) + # Re-normalize actions (from env's bounds to zero-centered), if + # necessary. + if ( + cfg.get("actions_in_input_normalized") is False + and cfg.get("normalize_actions") is True + ): + if ioctx.worker is None: + raise ValueError( + "actions_in_input_normalized is False but" + "cannot normalize actions since no workers exist" + ) + + # If we have a complex action space and actions were flattened + # and we have to normalize -> Error. + error_msg = ( + "Normalization of offline actions that are flattened is not " + "supported! Make sure that you record actions into offline " + "file with the `_disable_action_flattening=True` flag OR " + "as already normalized (between -1.0 and 1.0) values. " + "Also, when reading already normalized action values from " + "offline files, make sure to set " + "`actions_in_input_normalized=True` so that RLlib will not " + "perform normalization on top." + ) + + if isinstance(batch, SampleBatch): + policy = ioctx.worker.policy_map.get(DEFAULT_POLICY_ID) + if policy is None: + assert len(ioctx.worker.policy_map) == 1 + policy = next(iter(ioctx.worker.policy_map.values())) + if isinstance( + policy.action_space_struct, (tuple, dict) + ) and not policy.config.get("_disable_action_flattening"): + raise ValueError(error_msg) + batch[SampleBatch.ACTIONS] = normalize_action( + batch[SampleBatch.ACTIONS], policy.action_space_struct + ) + else: + for pid, b in batch.policy_batches.items(): + policy = ioctx.worker.policy_map[pid] + if isinstance( + policy.action_space_struct, (tuple, dict) + ) and not policy.config.get("_disable_action_flattening"): + raise ValueError(error_msg) + b[SampleBatch.ACTIONS] = normalize_action( + b[SampleBatch.ACTIONS], + ioctx.worker.policy_map[pid].action_space_struct, + ) + + return batch + + +@DeveloperAPI +def from_json_data(json_data: Any, worker: Optional["RolloutWorker"]): + # Try to infer the SampleBatchType (SampleBatch or MultiAgentBatch). + if "type" in json_data: + data_type = json_data.pop("type") + else: + raise ValueError("JSON record missing 'type' field") + + if data_type == "SampleBatch": + if worker is not None and len(worker.policy_map) != 1: + raise ValueError( + "Found single-agent SampleBatch in input file, but our " + "PolicyMap contains more than 1 policy!" + ) + for k, v in json_data.items(): + json_data[k] = unpack_if_needed(v) + if worker is not None: + policy = next(iter(worker.policy_map.values())) + json_data = _adjust_obs_actions_for_policy(json_data, policy) + json_data = _adjust_dones(json_data) + return SampleBatch(json_data) + elif data_type == "MultiAgentBatch": + policy_batches = {} + for policy_id, policy_batch in json_data["policy_batches"].items(): + inner = {} + for k, v in policy_batch.items(): + # Translate DONES into TERMINATEDS. + if k == SampleBatch.DONES: + k = SampleBatch.TERMINATEDS + inner[k] = unpack_if_needed(v) + if worker is not None: + policy = worker.policy_map[policy_id] + inner = _adjust_obs_actions_for_policy(inner, policy) + inner = _adjust_dones(inner) + policy_batches[policy_id] = SampleBatch(inner) + return MultiAgentBatch(policy_batches, json_data["count"]) + else: + raise ValueError( + "Type field must be one of ['SampleBatch', 'MultiAgentBatch']", data_type + ) + + +# TODO(jungong) : use DatasetReader to back JsonReader, so we reduce +# codebase complexity without losing existing functionality. +@PublicAPI +class JsonReader(InputReader): + """Reader object that loads experiences from JSON file chunks. + + The input files will be read from in random order. + """ + + @PublicAPI + def __init__( + self, inputs: Union[str, List[str]], ioctx: Optional[IOContext] = None + ): + """Initializes a JsonReader instance. + + Args: + inputs: Either a glob expression for files, e.g. `/tmp/**/*.json`, + or a list of single file paths or URIs, e.g., + ["s3://bucket/file.json", "s3://bucket/file2.json"]. + ioctx: Current IO context object or None. + """ + logger.info( + "You are using JSONReader. It is recommended to use " + + "DatasetReader instead for better sharding support." + ) + + self.ioctx = ioctx or IOContext() + self.default_policy = self.policy_map = None + self.batch_size = 1 + if self.ioctx: + self.batch_size = self.ioctx.config.get("train_batch_size", 1) + num_workers = self.ioctx.config.get("num_env_runners", 0) + if num_workers: + self.batch_size = max(math.ceil(self.batch_size / num_workers), 1) + + if self.ioctx.worker is not None: + self.policy_map = self.ioctx.worker.policy_map + self.default_policy = self.policy_map.get(DEFAULT_POLICY_ID) + if self.default_policy is None: + assert len(self.policy_map) == 1 + self.default_policy = next(iter(self.policy_map.values())) + + if isinstance(inputs, str): + inputs = os.path.abspath(os.path.expanduser(inputs)) + if os.path.isdir(inputs): + inputs = [os.path.join(inputs, "*.json"), os.path.join(inputs, "*.zip")] + logger.warning(f"Treating input directory as glob patterns: {inputs}") + else: + inputs = [inputs] + + if any(urlparse(i).scheme not in [""] + WINDOWS_DRIVES for i in inputs): + raise ValueError( + "Don't know how to glob over `{}`, ".format(inputs) + + "please specify a list of files to read instead." + ) + else: + self.files = [] + for i in inputs: + self.files.extend(glob.glob(i)) + elif isinstance(inputs, (list, tuple)): + self.files = list(inputs) + else: + raise ValueError( + "type of inputs must be list or str, not {}".format(inputs) + ) + if self.files: + logger.info("Found {} input files.".format(len(self.files))) + else: + raise ValueError("No files found matching {}".format(inputs)) + self.cur_file = None + + @override(InputReader) + def next(self) -> SampleBatchType: + ret = [] + count = 0 + while count < self.batch_size: + batch = self._try_parse(self._next_line()) + tries = 0 + while not batch and tries < 100: + tries += 1 + logger.debug("Skipping empty line in {}".format(self.cur_file)) + batch = self._try_parse(self._next_line()) + if not batch: + raise ValueError( + "Failed to read valid experience batch from file: {}".format( + self.cur_file + ) + ) + batch = self._postprocess_if_needed(batch) + count += batch.count + ret.append(batch) + ret = concat_samples(ret) + return ret + + def read_all_files(self) -> SampleBatchType: + """Reads through all files and yields one SampleBatchType per line. + + When reaching the end of the last file, will start from the beginning + again. + + Yields: + One SampleBatch or MultiAgentBatch per line in all input files. + """ + for path in self.files: + file = self._try_open_file(path) + while True: + line = file.readline() + if not line: + break + batch = self._try_parse(line) + if batch is None: + break + yield batch + + def _postprocess_if_needed(self, batch: SampleBatchType) -> SampleBatchType: + if not self.ioctx.config.get("postprocess_inputs"): + return batch + + batch = convert_ma_batch_to_sample_batch(batch) + + if isinstance(batch, SampleBatch): + out = [] + for sub_batch in batch.split_by_episode(): + out.append(self.default_policy.postprocess_trajectory(sub_batch)) + return concat_samples(out) + else: + # TODO(ekl) this is trickier since the alignments between agent + # trajectories in the episode are not available any more. + raise NotImplementedError( + "Postprocessing of multi-agent data not implemented yet." + ) + + def _try_open_file(self, path): + if urlparse(path).scheme not in [""] + WINDOWS_DRIVES: + if smart_open is None: + raise ValueError( + "You must install the `smart_open` module to read " + "from URIs like {}".format(path) + ) + ctx = smart_open + else: + # Allow shortcut for home directory ("~/" -> env[HOME]). + if path.startswith("~/"): + path = os.path.join(os.environ.get("HOME", ""), path[2:]) + + # If path doesn't exist, try to interpret is as relative to the + # rllib directory (located ../../ from this very module). + path_orig = path + if not os.path.exists(path): + path = os.path.join(Path(__file__).parent.parent, path) + if not os.path.exists(path): + raise FileNotFoundError(f"Offline file {path_orig} not found!") + + # Unzip files, if necessary and re-point to extracted json file. + if re.search("\\.zip$", path): + with zipfile.ZipFile(path, "r") as zip_ref: + zip_ref.extractall(Path(path).parent) + path = re.sub("\\.zip$", ".json", path) + assert os.path.exists(path) + ctx = open + file = ctx(path, "r") + return file + + def _try_parse(self, line: str) -> Optional[SampleBatchType]: + line = line.strip() + if not line: + return None + try: + batch = self._from_json(line) + except Exception: + logger.exception( + "Ignoring corrupt json record in {}: {}".format(self.cur_file, line) + ) + return None + + batch = postprocess_actions(batch, self.ioctx) + + return batch + + def _next_line(self) -> str: + if not self.cur_file: + self.cur_file = self._next_file() + line = self.cur_file.readline() + tries = 0 + while not line and tries < 100: + tries += 1 + if hasattr(self.cur_file, "close"): # legacy smart_open impls + self.cur_file.close() + self.cur_file = self._next_file() + line = self.cur_file.readline() + if not line: + logger.debug("Ignoring empty file {}".format(self.cur_file)) + if not line: + raise ValueError( + "Failed to read next line from files: {}".format(self.files) + ) + return line + + def _next_file(self) -> FileType: + # If this is the first time, we open a file, make sure all workers + # start with a different one if possible. + if self.cur_file is None and self.ioctx.worker is not None: + idx = self.ioctx.worker.worker_index + total = self.ioctx.worker.num_workers or 1 + path = self.files[round((len(self.files) - 1) * (idx / total))] + # After the first file, pick all others randomly. + else: + path = random.choice(self.files) + return self._try_open_file(path) + + def _from_json(self, data: str) -> SampleBatchType: + if isinstance(data, bytes): # smart_open S3 doesn't respect "r" + data = data.decode("utf-8") + json_data = json.loads(data) + return from_json_data(json_data, self.ioctx.worker) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/json_writer.py b/lib/python3.12/site-packages/ray/rllib/offline/json_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..da7b49e5b17b30d573684a681c52930435075ba9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/json_writer.py @@ -0,0 +1,144 @@ +import json +import logging +import os +import time +from datetime import datetime +from urllib.parse import urlparse + +import numpy as np + +try: + from smart_open import smart_open +except ImportError: + smart_open = None + +from typing import Any, Dict, List + +from ray.air._internal.json import SafeFallbackEncoder +from ray.rllib.offline.io_context import IOContext +from ray.rllib.offline.output_writer import OutputWriter +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import PublicAPI, override +from ray.rllib.utils.compression import compression_supported, pack +from ray.rllib.utils.typing import FileType, SampleBatchType + +logger = logging.getLogger(__name__) + +WINDOWS_DRIVES = [chr(i) for i in range(ord("c"), ord("z") + 1)] + + +# TODO(jungong): use DatasetWriter to back JsonWriter, so we reduce codebase complexity +# without losing existing functionality. +@PublicAPI +class JsonWriter(OutputWriter): + """Writer object that saves experiences in JSON file chunks.""" + + @PublicAPI + def __init__( + self, + path: str, + ioctx: IOContext = None, + max_file_size: int = 64 * 1024 * 1024, + compress_columns: List[str] = frozenset(["obs", "new_obs"]), + ): + """Initializes a JsonWriter instance. + + Args: + path: a path/URI of the output directory to save files in. + ioctx: current IO context object. + max_file_size: max size of single files before rolling over. + compress_columns: list of sample batch columns to compress. + """ + logger.info( + "You are using JSONWriter. It is recommended to use " + + "DatasetWriter instead." + ) + + self.ioctx = ioctx or IOContext() + self.max_file_size = max_file_size + self.compress_columns = compress_columns + if urlparse(path).scheme not in [""] + WINDOWS_DRIVES: + self.path_is_uri = True + else: + path = os.path.abspath(os.path.expanduser(path)) + # Try to create local dirs if they don't exist + os.makedirs(path, exist_ok=True) + assert os.path.exists(path), "Failed to create {}".format(path) + self.path_is_uri = False + self.path = path + self.file_index = 0 + self.bytes_written = 0 + self.cur_file = None + + @override(OutputWriter) + def write(self, sample_batch: SampleBatchType): + start = time.time() + data = _to_json(sample_batch, self.compress_columns) + f = self._get_file() + f.write(data) + f.write("\n") + if hasattr(f, "flush"): # legacy smart_open impls + f.flush() + self.bytes_written += len(data) + logger.debug( + "Wrote {} bytes to {} in {}s".format(len(data), f, time.time() - start) + ) + + def _get_file(self) -> FileType: + if not self.cur_file or self.bytes_written >= self.max_file_size: + if self.cur_file: + self.cur_file.close() + timestr = datetime.today().strftime("%Y-%m-%d_%H-%M-%S") + path = os.path.join( + self.path, + "output-{}_worker-{}_{}.json".format( + timestr, self.ioctx.worker_index, self.file_index + ), + ) + if self.path_is_uri: + if smart_open is None: + raise ValueError( + "You must install the `smart_open` module to write " + "to URIs like {}".format(path) + ) + self.cur_file = smart_open(path, "w") + else: + self.cur_file = open(path, "w") + self.file_index += 1 + self.bytes_written = 0 + logger.info("Writing to new output file {}".format(self.cur_file)) + return self.cur_file + + +def _to_jsonable(v, compress: bool) -> Any: + if compress and compression_supported(): + return str(pack(v)) + elif isinstance(v, np.ndarray): + return v.tolist() + + return v + + +def _to_json_dict(batch: SampleBatchType, compress_columns: List[str]) -> Dict: + out = {} + if isinstance(batch, MultiAgentBatch): + out["type"] = "MultiAgentBatch" + out["count"] = batch.count + policy_batches = {} + for policy_id, sub_batch in batch.policy_batches.items(): + policy_batches[policy_id] = {} + for k, v in sub_batch.items(): + policy_batches[policy_id][k] = _to_jsonable( + v, compress=k in compress_columns + ) + out["policy_batches"] = policy_batches + else: + out["type"] = "SampleBatch" + for k, v in batch.items(): + out[k] = _to_jsonable(v, compress=k in compress_columns) + return out + + +def _to_json(batch: SampleBatchType, compress_columns: List[str]) -> str: + out = _to_json_dict(batch, compress_columns) + return json.dumps(out, cls=SafeFallbackEncoder) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/mixed_input.py b/lib/python3.12/site-packages/ray/rllib/offline/mixed_input.py new file mode 100644 index 0000000000000000000000000000000000000000..171a7a2c85c412d61b8658b69fba321bc3cd8f93 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/mixed_input.py @@ -0,0 +1,59 @@ +from types import FunctionType +from typing import Dict + +import numpy as np + +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.offline.json_reader import JsonReader +from ray.rllib.utils.annotations import DeveloperAPI, override +from ray.rllib.utils.typing import SampleBatchType +from ray.tune.registry import registry_contains_input, registry_get_input + + +@DeveloperAPI +class MixedInput(InputReader): + """Mixes input from a number of other input sources. + + .. testcode:: + :skipif: True + + from ray.rllib.offline.io_context import IOContext + from ray.rllib.offline.mixed_input import MixedInput + ioctx = IOContext(...) + MixedInput({ + "sampler": 0.4, + "/tmp/experiences/*.json": 0.4, + "s3://bucket/expert.json": 0.2, + }, ioctx) + """ + + @DeveloperAPI + def __init__(self, dist: Dict[JsonReader, float], ioctx: IOContext): + """Initialize a MixedInput. + + Args: + dist: dict mapping JSONReader paths or "sampler" to + probabilities. The probabilities must sum to 1.0. + ioctx: current IO context object. + """ + if sum(dist.values()) != 1.0: + raise ValueError("Values must sum to 1.0: {}".format(dist)) + self.choices = [] + self.p = [] + for k, v in dist.items(): + if k == "sampler": + self.choices.append(ioctx.default_sampler_input()) + elif isinstance(k, FunctionType): + self.choices.append(k(ioctx)) + elif isinstance(k, str) and registry_contains_input(k): + input_creator = registry_get_input(k) + self.choices.append(input_creator(ioctx)) + else: + self.choices.append(JsonReader(k, ioctx)) + self.p.append(v) + + @override(InputReader) + def next(self) -> SampleBatchType: + source = np.random.choice(self.choices, p=self.p) + return source.next() diff --git a/lib/python3.12/site-packages/ray/rllib/offline/off_policy_estimator.py b/lib/python3.12/site-packages/ray/rllib/offline/off_policy_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..71d3a80f01487b37850cf617f397d1201001a22d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/off_policy_estimator.py @@ -0,0 +1,10 @@ +from ray._common.deprecation import deprecation_warning +from ray.rllib.offline.estimators.off_policy_estimator import ( # noqa: F401 + OffPolicyEstimator, +) + +deprecation_warning( + old="ray.rllib.offline.off_policy_estimator", + new="ray.rllib.offline.estimators.off_policy_estimator", + error=True, +) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_data.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_data.py new file mode 100644 index 0000000000000000000000000000000000000000..64c55a6c33d947389060a3db99e48fb1c36b8089 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_data.py @@ -0,0 +1,289 @@ +import logging +import time +import types +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict + +import numpy as np +import pyarrow.fs + +import ray +from ray.rllib.core import COMPONENT_RL_MODULE +from ray.rllib.env import INPUT_ENV_SPACES +from ray.rllib.offline.offline_prelearner import OfflinePreLearner +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils import force_list, unflatten_dict +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, +) +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="alpha") +class OfflineData: + @OverrideToImplementCustomLogic_CallToSuperRecommended + def __init__(self, config: "AlgorithmConfig"): + + # TODO (simon): Define self.spaces here. + self.config = config + self.is_multi_agent = self.config.is_multi_agent + self.path = ( + self.config.input_ + if isinstance(config.input_, list) + else Path(config.input_) + ) + # Use `read_parquet` as default data read method. + self.data_read_method = self.config.input_read_method + # Override default arguments for the data read method. + self.data_read_method_kwargs = self.config.input_read_method_kwargs + # In case `EpisodeType` or `BatchType` batches are read the size + # could differ from the final `train_batch_size_per_learner`. + self.data_read_batch_size = self.config.input_read_batch_size + + # If data should be materialized. + self.materialize_data = config.materialize_data + # If mapped data should be materialized. + self.materialize_mapped_data = config.materialize_mapped_data + # Flag to identify, if data has already been mapped with the + # `OfflinePreLearner`. + self.data_is_mapped = False + + # Set the filesystem. + self.filesystem = self.config.input_filesystem + self.filesystem_kwargs = self.config.input_filesystem_kwargs + self.filesystem_object = None + + # If a specific filesystem is given, set it up. Note, this could + # be `gcsfs` for GCS, `pyarrow` for S3 or `adlfs` for Azure Blob Storage. + # this filesystem is specifically needed, if a session has to be created + # with the cloud provider. + if self.filesystem == "gcs": + import gcsfs + + self.filesystem_object = gcsfs.GCSFileSystem(**self.filesystem_kwargs) + elif self.filesystem == "s3": + self.filesystem_object = pyarrow.fs.S3FileSystem(**self.filesystem_kwargs) + elif self.filesystem == "abs": + import adlfs + + self.filesystem_object = adlfs.AzureBlobFileSystem(**self.filesystem_kwargs) + elif isinstance(self.filesystem, pyarrow.fs.FileSystem): + self.filesystem_object = self.filesystem + elif self.filesystem is not None: + raise ValueError( + f"Unknown `config.input_filesystem` {self.filesystem}! Filesystems " + "can be None for local, any instance of `pyarrow.fs.FileSystem`, " + "'gcs' for GCS, 's3' for S3, or 'abs' for adlfs.AzureBlobFileSystem." + ) + # Add the filesystem object to the write method kwargs. + if self.filesystem_object: + self.data_read_method_kwargs.update( + { + "filesystem": self.filesystem_object, + } + ) + + try: + # Load the dataset. + start_time = time.perf_counter() + self.data = getattr(ray.data, self.data_read_method)( + self.path, **self.data_read_method_kwargs + ) + if self.materialize_data: + self.data = self.data.materialize() + stop_time = time.perf_counter() + logger.debug( + "===> [OfflineData] - Time for loading dataset: " + f"{stop_time - start_time}s." + ) + logger.info("Reading data from {}".format(self.path)) + except Exception as e: + logger.error(e) + # Avoids reinstantiating the batch iterator each time we sample. + self.batch_iterators = None + self.map_batches_kwargs = ( + self.default_map_batches_kwargs | self.config.map_batches_kwargs + ) + self.iter_batches_kwargs = ( + self.default_iter_batches_kwargs | self.config.iter_batches_kwargs + ) + self.returned_streaming_split = False + # Defines the prelearner class. Note, this could be user-defined. + self.prelearner_class = self.config.prelearner_class or OfflinePreLearner + # For remote learner setups. + self.locality_hints = None + self.learner_handles = None + self.module_spec = None + + @OverrideToImplementCustomLogic + def sample( + self, + num_samples: int, + return_iterator: bool = False, + num_shards: int = 1, + module_state: Dict[str, Any] = None, + ): + # Materialize the mapped data, if necessary. This runs for all the + # data the `OfflinePreLearner` logic and maps them to `MultiAgentBatch`es. + # TODO (simon, sven): This would never update the module nor the + # the connectors. If this is needed we have to check, if we give + # (a) only an iterator and let the learner and OfflinePreLearner + # communicate through the object storage. This only works when + # not materializing. + # (b) Rematerialize the data every couple of iterations. This is + # is costly. + if not self.data_is_mapped: + + if not module_state: + # Get the RLModule state from learners. + if num_shards >= 1: + # Call here the learner to get an up-to-date module state. + # TODO (simon): This is a workaround as along as learners cannot + # receive any calls from another actor. + module_state = ray.get( + self.learner_handles[0].get_state.remote( + component=COMPONENT_RL_MODULE, + ) + )[COMPONENT_RL_MODULE] + # Provide the `Learner`(s) GPU devices, if needed. + # if not self.map_batches_uses_gpus(self.config) and self.config._validate_config: + # devices = ray.get(self.learner_handles[0].get_device.remote()) + # devices = [devices] if not isinstance(devices, list) else devices + # device_strings = [ + # f"{device.type}:{str(device.index)}" + # if device.type == "cuda" + # else device.type + # for device in devices + # ] + # # Otherwise, set the GPU strings to `None`. + # # TODO (simon): Check inside 'OfflinePreLearner'. + # else: + # device_strings = None + else: + # Get the module state from the `Learner`(S). + module_state = self.learner_handles[0].get_state( + component=COMPONENT_RL_MODULE, + )[COMPONENT_RL_MODULE] + # Provide the `Learner`(s) GPU devices, if needed. + # if not self.map_batches_uses_gpus(self.config) and self.config._validate_config: + # device = self.learner_handles[0].get_device() + # device_strings = [ + # f"{device.type}:{str(device.index)}" + # if device.type == "cuda" + # else device.type + # ] + # else: + # device_strings = None + # Constructor `kwargs` for the `OfflinePreLearner`. + fn_constructor_kwargs = { + "config": self.config, + "spaces": self.spaces[INPUT_ENV_SPACES], + "module_spec": self.module_spec, + "module_state": module_state, + # "device_strings": self.get_devices(), + } + + # Map the data to run the `OfflinePreLearner`s in the data pipeline + # for training. + self.data = self.data.map_batches( + self.prelearner_class, + fn_constructor_kwargs=fn_constructor_kwargs, + batch_size=self.data_read_batch_size or num_samples, + **self.map_batches_kwargs, + ) + # Set the flag to `True`. + self.data_is_mapped = True + # If the user wants to materialize the data in memory. + if self.materialize_mapped_data: + self.data = self.data.materialize() + # Build an iterator, if necessary. Note, in case that an iterator should be + # returned now and we have already generated from the iterator, i.e. + # `isinstance(self.batch_iterators, types.GeneratorType) == True`, we need + # to create here a new iterator. + # TODO (simon): Check, if this iterator could potentially exhaust. + if not self.batch_iterators or ( + return_iterator and isinstance(self.batch_iterators, types.GeneratorType) + ): + # If we have more than one learner create an iterator for each of them + # by splitting the data stream. + if num_shards > 1: + logger.debug("===> [OfflineData]: Return streaming_split ... ") + # In case of multiple shards, we return multiple + # `StreamingSplitIterator` instances. + self.batch_iterators = self.data.streaming_split( + n=num_shards, + # Note, `equal` must be `True`, i.e. the batch size must + # be the same for all batches b/c otherwise remote learners + # could block each others. + equal=True, + locality_hints=self.locality_hints, + ) + # Otherwise we create a simple iterator and - if necessary - initialize + # it here. + else: + # Should an iterator be returned? + if return_iterator: + self.batch_iterators = self.data.iterator() + # Otherwise, the user wants batches returned. + else: + # Define a collate (last-mile) transformation that maps batches + # to RLlib's `MultiAgentBatch`. + def _collate_fn(_batch: Dict[str, np.ndarray]) -> MultiAgentBatch: + _batch = unflatten_dict(_batch) + return MultiAgentBatch( + { + module_id: SampleBatch(module_data) + for module_id, module_data in _batch.items() + }, + env_steps=sum( + len(next(iter(module_data.values()))) + for module_data in _batch.values() + ), + ) + + # If no iterator should be returned, or if we want to return a single + # batch iterator, we instantiate the batch iterator once, here. + self.batch_iterators = self.data.iter_batches( + batch_size=num_samples, + _collate_fn=_collate_fn, + **self.iter_batches_kwargs, + ) + self.batch_iterators = iter(self.batch_iterators) + + # Do we want to return an iterator or a single batch? + if return_iterator: + return force_list(self.batch_iterators) + else: + # Return a single batch from the iterator. + try: + return next(self.batch_iterators) + except StopIteration: + # If the batch iterator is exhausted, reinitiate a new one. + logger.debug( + "===> [OfflineData]: Batch iterator exhausted. Reinitiating ..." + ) + self.batch_iterators = None + return self.sample( + num_samples=num_samples, + return_iterator=return_iterator, + num_shards=num_shards, + ) + + @property + def default_map_batches_kwargs(self): + return { + "concurrency": max(2, self.config.num_learners), + "zero_copy_batch": True, + } + + @property + def default_iter_batches_kwargs(self): + return { + "prefetch_batches": 2, + } diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_env_runner.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_env_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..64eae96cb6c244e1cb7c0c3b456d45d086d68272 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_env_runner.py @@ -0,0 +1,305 @@ +import logging +from pathlib import Path +from typing import List + +import ray +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.core.columns import Columns +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.env.single_agent_env_runner import SingleAgentEnvRunner +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.compression import pack_if_needed +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI +from ray.util.debug import log_once + +logger = logging.Logger(__file__) + +# TODO (simon): This class can be agnostic to the episode type as it +# calls only get_state. + + +@PublicAPI(stability="alpha") +class OfflineSingleAgentEnvRunner(SingleAgentEnvRunner): + """The environment runner to record the single agent case.""" + + @override(SingleAgentEnvRunner) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def __init__(self, *, config: AlgorithmConfig, **kwargs): + # Initialize the parent. + super().__init__(config=config, **kwargs) + + # override SingleAgentEnvRunner + self.episodes_to_numpy = False + + # Get the data context for this `EnvRunner`. + data_context = ray.data.DataContext.get_current() + # Limit the resources for Ray Data to the CPUs given to this `EnvRunner`. + data_context.execution_options.resource_limits = ( + data_context.execution_options.resource_limits.copy( + cpu=config.num_cpus_per_env_runner + ) + ) + + # Set the output write method. + self.output_write_method = self.config.output_write_method + self.output_write_method_kwargs = self.config.output_write_method_kwargs + + # Set the filesystem. + self.filesystem = self.config.output_filesystem + self.filesystem_kwargs = self.config.output_filesystem_kwargs + self.filesystem_object = None + + # Set the output base path. + self.output_path = self.config.output + # Set the subdir (environment specific). + self.subdir_path = self.config.env.lower() + # Set the worker-specific path name. Note, this is + # specifically to enable multi-threaded writing into + # the same directory. + self.worker_path = "run-" + f"{self.worker_index}".zfill(6) + + # If a specific filesystem is given, set it up. Note, this could + # be `gcsfs` for GCS, `pyarrow` for S3 or `adlfs` for Azure Blob Storage. + # this filesystem is specifically needed, if a session has to be created + # with the cloud provider. + + if self.filesystem == "gcs": + import gcsfs + + self.filesystem_object = gcsfs.GCSFileSystem(**self.filesystem_kwargs) + elif self.filesystem == "s3": + from pyarrow import fs + + self.filesystem_object = fs.S3FileSystem(**self.filesystem_kwargs) + elif self.filesystem == "abs": + import adlfs + + self.filesystem_object = adlfs.AzureBlobFileSystem(**self.filesystem_kwargs) + elif self.filesystem is not None: + raise ValueError( + f"Unknown filesystem: {self.filesystem}. Filesystems can be " + "'gcs' for GCS, 's3' for S3, or 'abs'" + ) + # Add the filesystem object to the write method kwargs. + self.output_write_method_kwargs.update( + { + "filesystem": self.filesystem_object, + } + ) + + # If we should store `SingleAgentEpisodes` or column data. + self.output_write_episodes = self.config.output_write_episodes + # Which columns should be compressed in the output data. + self.output_compress_columns = self.config.output_compress_columns + + # Buffer these many rows before writing to file. + self.output_max_rows_per_file = self.config.output_max_rows_per_file + # If the user defines a maximum number of rows per file, set the + # event to `False` and check during sampling. + if self.output_max_rows_per_file: + self.write_data_this_iter = False + # Otherwise the event is always `True` and we write always sampled + # data immediately to disk. + else: + self.write_data_this_iter = True + + # If the remaining data should be stored. Note, this is only + # relevant in case `output_max_rows_per_file` is defined. + self.write_remaining_data = self.config.output_write_remaining_data + + # Counts how often `sample` is called to define the output path for + # each file. + self._sample_counter = 0 + + # Define the buffer for experiences stored until written to disk. + self._samples = [] + + @override(SingleAgentEnvRunner) + @OverrideToImplementCustomLogic + def sample( + self, + *, + num_timesteps: int = None, + num_episodes: int = None, + explore: bool = None, + random_actions: bool = False, + force_reset: bool = False, + ) -> List[SingleAgentEpisode]: + """Samples from environments and writes data to disk.""" + + # Call the super sample method. + samples = super().sample( + num_timesteps=num_timesteps, + num_episodes=num_episodes, + explore=explore, + random_actions=random_actions, + force_reset=force_reset, + ) + + self._sample_counter += 1 + + # Add data to the buffers. + if self.output_write_episodes: + + import msgpack + import msgpack_numpy as mnp + + if log_once("msgpack"): + logger.info( + "Packing episodes with `msgpack` and encode array with " + "`msgpack_numpy` for serialization. This is needed for " + "recording episodes." + ) + # Note, we serialize episodes with `msgpack` and `msgpack_numpy` to + # ensure version compatibility. + assert all(eps.is_numpy is False for eps in samples) + self._samples.extend( + [msgpack.packb(eps.get_state(), default=mnp.encode) for eps in samples] + ) + else: + self._map_episodes_to_data(samples) + + # If the user defined the maximum number of rows to write. + if self.output_max_rows_per_file: + # Check, if this number is reached. + if len(self._samples) >= self.output_max_rows_per_file: + # Start the recording of data. + self.write_data_this_iter = True + + if self.write_data_this_iter: + # If the user wants a maximum number of experiences per file, + # cut the samples to write to disk from the buffer. + if self.output_max_rows_per_file: + # Reset the event. + self.write_data_this_iter = False + # Ensure that all data ready to be written is released from + # the buffer. Note, this is important in case we have many + # episodes sampled and a relatively small `output_max_rows_per_file`. + while len(self._samples) >= self.output_max_rows_per_file: + # Extract the number of samples to be written to disk this + # iteration. + samples_to_write = self._samples[: self.output_max_rows_per_file] + # Reset the buffer to the remaining data. This only makes sense, if + # `rollout_fragment_length` is smaller `output_max_rows_per_file` or + # a 2 x `output_max_rows_per_file`. + self._samples = self._samples[self.output_max_rows_per_file :] + samples_ds = ray.data.from_items(samples_to_write) + # Otherwise, write the complete data. + else: + samples_ds = ray.data.from_items(self._samples) + try: + # Setup the path for writing data. Each run will be written to + # its own file. A run is a writing event. The path will look + # like. 'base_path/env-name/00000-00000'. + path = ( + Path(self.output_path) + .joinpath(self.subdir_path) + .joinpath(self.worker_path + f"-{self._sample_counter}".zfill(6)) + ) + getattr(samples_ds, self.output_write_method)( + path.as_posix(), **self.output_write_method_kwargs + ) + logger.info(f"Wrote samples to storage at {path}.") + except Exception as e: + logger.error(e) + + self.metrics.log_value( + key="recording_buffer_size", + value=len(self._samples), + ) + + # Finally return the samples as usual. + return samples + + @override(EnvRunner) + @OverrideToImplementCustomLogic + def stop(self) -> None: + """Writes the reamining samples to disk + + Note, if the user defined `max_rows_per_file` the + number of rows for the remaining samples could be + less than the defined maximum row number by the user. + """ + # If there are samples left over we have to write htem to disk. them + # to a dataset. + if self._samples and self.write_remaining_data: + # Convert them to a `ray.data.Dataset`. + samples_ds = ray.data.from_items(self._samples) + # Increase the sample counter for the folder/file name. + self._sample_counter += 1 + # Try to write the dataset to disk/cloud storage. + try: + # Setup the path for writing data. Each run will be written to + # its own file. A run is a writing event. The path will look + # like. 'base_path/env-name/00000-00000'. + path = ( + Path(self.output_path) + .joinpath(self.subdir_path) + .joinpath(self.worker_path + f"-{self._sample_counter}".zfill(6)) + ) + getattr(samples_ds, self.output_write_method)( + path.as_posix(), **self.output_write_method_kwargs + ) + logger.info( + f"Wrote final samples to storage at {path}. Note " + "Note, final samples could be smaller in size than " + f"`max_rows_per_file`, if defined." + ) + except Exception as e: + logger.error(e) + + logger.debug(f"Experience buffer length: {len(self._samples)}") + + @OverrideToImplementCustomLogic + def _map_episodes_to_data(self, samples: List[EpisodeType]) -> None: + """Converts list of episodes to list of single dict experiences. + + Note, this method also appends all sampled experiences to the + buffer. + + Args: + samples: List of episodes to be converted. + """ + # Loop through all sampled episodes. + for sample in samples: + # Loop through all items of the episode. + for i in range(len(sample)): + sample_data = { + Columns.EPS_ID: sample.id_, + Columns.AGENT_ID: sample.agent_id, + Columns.MODULE_ID: sample.module_id, + # Compress observations, if requested. + Columns.OBS: pack_if_needed(sample.get_observations(i)) + if Columns.OBS in self.output_compress_columns + else sample.get_observations(i), + # Compress actions, if requested. + Columns.ACTIONS: pack_if_needed(sample.get_actions(i)) + if Columns.ACTIONS in self.output_compress_columns + else sample.get_actions(i), + Columns.REWARDS: sample.get_rewards(i), + # Compress next observations, if requested. + Columns.NEXT_OBS: pack_if_needed(sample.get_observations(i + 1)) + if Columns.OBS in self.output_compress_columns + else sample.get_observations(i + 1), + Columns.TERMINATEDS: False + if i < len(sample) - 1 + else sample.is_terminated, + Columns.TRUNCATEDS: False + if i < len(sample) - 1 + else sample.is_truncated, + **{ + # Compress any extra model output, if requested. + k: pack_if_needed(sample.get_extra_model_outputs(k, i)) + if k in self.output_compress_columns + else sample.get_extra_model_outputs(k, i) + for k in sample.extra_model_outputs.keys() + }, + } + # Finally append to the data buffer. + self._samples.append(sample_data) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..c49de5a061e972794e9914ca7a501bdc1e4f406a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner.py @@ -0,0 +1,466 @@ +import types +from typing import TYPE_CHECKING, Any, Collection, Dict, Iterable, Optional, Union + +import ray +from ray.data.iterator import DataIterator +from ray.rllib.core import ( + ALL_MODULES, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.framework import get_device, try_import_torch +from ray.rllib.utils.metrics import ( + DATASET_NUM_ITERS_EVALUATED, + DATASET_NUM_ITERS_EVALUATED_LIFETIME, + MODULE_SAMPLE_BATCH_SIZE_MEAN, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_MODULE_STEPS_SAMPLED, + NUM_MODULE_STEPS_SAMPLED_LIFETIME, + OFFLINE_SAMPLING_TIMER, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.minibatch_utils import MiniBatchRayDataIterator +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.runners.runner import Runner +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import DeviceType, ModuleID, StateDict, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +torch, _ = try_import_torch() + +TOTAL_EVAL_LOSS_KEY = "total_eval_loss" + + +class OfflineEvaluationRunner(Runner, Checkpointable): + def __init__( + self, + config: "AlgorithmConfig", + module_spec: Optional[MultiRLModuleSpec] = None, + **kwargs, + ): + + # This needs to be defined before we call the `Runner.__init__` + # b/c the latter calls the `make_module` and then needs the spec. + # TODO (simon): Check, if we make this a generic attribute. + self.__module_spec: MultiRLModuleSpec = module_spec + self.__dataset_iterator = None + self.__batch_iterator = None + + Runner.__init__(self, config=config, **kwargs) + Checkpointable.__init__(self) + + # This has to be defined after we have a `self.config`. + self._loss_for_module_fn = types.MethodType(self.get_loss_for_module_fn(), self) + + @override(Runner) + def run( + self, + explore: bool = False, + train: bool = True, + **kwargs, + ) -> None: + + if self.__dataset_iterator is None: + raise ValueError( + f"{self} doesn't have a data iterator. Can't call `run` on " + "`OfflineEvaluationRunner`." + ) + + if not self._batch_iterator: + self.__batch_iterator = self._create_batch_iterator( + **self.config.iter_batches_kwargs + ) + + # Log current weight seq no. + self.metrics.log_value( + key=WEIGHTS_SEQ_NO, + value=self._weights_seq_no, + window=1, + ) + + with self.metrics.log_time(OFFLINE_SAMPLING_TIMER): + if explore is None: + explore = self.config.explore + + # Evaluate on offline data. + return self._evaluate( + explore=explore, + train=train, + ) + + def _create_batch_iterator(self, **kwargs) -> Iterable: + + # Return a minibatch iterator. + return MiniBatchRayDataIterator( + iterator=self._dataset_iterator, + device=self._device, + minibatch_size=self.config.offline_eval_batch_size_per_runner, + num_iters=self.config.dataset_num_iters_per_eval_runner, + **kwargs, + ) + + def _evaluate( + self, + explore: bool, + train: bool, + ) -> None: + + self.metrics.activate_tensor_mode() + + for iteration, tensor_minibatch in enumerate(self._batch_iterator): + # Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs + # found in this batch. If not, throw an error. + unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set( + self.module.keys() + ) + if unknown_module_ids: + raise ValueError( + f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that " + f"are not in this Learner!" + ) + + if explore: + fwd_out = self.module.forward_exploration( + tensor_minibatch.policy_batches + ) + elif train: + fwd_out = self.module.forward_train(tensor_minibatch.policy_batches) + else: + fwd_out = self.module.forward_inference(tensor_minibatch.policy_batches) + + eval_loss_per_module = self.compute_eval_losses( + fwd_out=fwd_out, batch=tensor_minibatch.policy_batches + ) + + self._log_steps_evaluated_metrics(tensor_minibatch) + + # Record the number of batches pulled from the dataset. + self.metrics.log_value( + # TODO (simon): Create extra eval metrics. + (ALL_MODULES, DATASET_NUM_ITERS_EVALUATED), + iteration + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_EVALUATED_LIFETIME), + iteration + 1, + reduce="sum", + ) + # Log all individual RLModules' loss terms + # Note: We do this only once for the last of the minibatch updates, b/c the + # window is only 1 anyways. + for mid, loss in convert_to_numpy(eval_loss_per_module).items(): + self.metrics.log_value( + key=(mid, TOTAL_EVAL_LOSS_KEY), + value=loss, + window=1, + ) + + self.metrics.deactivate_tensor_mode() + + return self.metrics.reduce() + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + {"config": self.config}, # **kwargs + ) + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + state = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + ), + } + + if self._check_component(COMPONENT_RL_MODULE, components, not_components): + state[COMPONENT_RL_MODULE] = self.module.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=self._get_subcomponents( + COMPONENT_RL_MODULE, not_components + ), + **kwargs, + ) + state[WEIGHTS_SEQ_NO] = self._weights_seq_no + + return state + + def _convert_to_tensor(self, struct) -> TensorType: + """Converts structs to a framework-specific tensor.""" + return convert_to_torch_tensor(struct) + + @override(Runner) + def stop(self) -> None: + """Releases all resources used by this EnvRunner. + + For example, when using a gym.Env in this EnvRunner, you should make sure + that its `close()` method is called. + """ + pass + + @override(Runner) + def __del__(self) -> None: + """If this Actor is deleted, clears all resources used by it.""" + pass + + @override(Runner) + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Ensures that the instances has a `MultiRLModule` and an + environment defined. + + Raises: + AssertionError: If the EnvRunner Actor has NOT been properly initialized. + """ + # Make sure, we have built our RLModule properly and assigned a dataset iterator. + assert self._dataset_iterator and hasattr(self, "module") + + @override(Runner) + def get_metrics(self): + return self.metrics.reduce() + + def _convert_batch_type( + self, + batch: MultiAgentBatch, + to_device: bool = True, + pin_memory: bool = False, + use_stream: bool = False, + ) -> MultiAgentBatch: + batch = convert_to_torch_tensor( + batch.policy_batches, + device=self._device if to_device else None, + pin_memory=pin_memory, + use_stream=use_stream, + ) + # TODO (sven): This computation of `env_steps` is not accurate! + length = max(len(b) for b in batch.values()) + batch = MultiAgentBatch(batch, env_steps=length) + return batch + + def compute_eval_losses( + self, *, fwd_out: Dict[str, Any], batch: Dict[str, Any] + ) -> Dict[str, Any]: + + loss_per_module = {} + for module_id in fwd_out: + module_batch = batch[module_id] + module_fwd_out = fwd_out[module_id] + + module = self.module[module_id].unwrapped() + if isinstance(module, SelfSupervisedLossAPI): + loss = module.compute_self_supervised_loss( + learner=self, + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + ) + else: + loss = self.compute_eval_loss_for_module( + module_id=module_id, + config=self.config.get_config_for_module(module_id), + batch=module_batch, + fwd_out=module_fwd_out, + ) + loss_per_module[module_id] = loss + + return loss_per_module + + def compute_eval_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + + return self._loss_for_module_fn( + module_id=module_id, + config=config, + batch=batch, + fwd_out=fwd_out, + ) + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + + # Update the RLModule state. + if COMPONENT_RL_MODULE in state: + # A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the + # update. + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Only update the weigths, if this is the first synchronization or + # if the weights of this `EnvRunner` lacks behind the actual ones. + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + rl_module_state = state[COMPONENT_RL_MODULE] + if isinstance(rl_module_state, ray.ObjectRef): + rl_module_state = ray.get(rl_module_state) + self.module.set_state(rl_module_state) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update our lifetime counters. + # TODO (simon): Create extra metrics. + if NUM_ENV_STEPS_SAMPLED_LIFETIME in state: + self.metrics.set_value( + key=NUM_ENV_STEPS_SAMPLED_LIFETIME, + value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME], + reduce="sum", + with_throughput=True, + ) + + def _log_steps_evaluated_metrics(self, batch: MultiAgentBatch) -> None: + for mid, module_batch in batch.policy_batches.items(): + # Log weights seq no for this batch. + self.metrics.log_value( + (mid, WEIGHTS_SEQ_NO), + self._weights_seq_no, + window=1, + ) + + module_batch_size = len(module_batch) + # Log average batch size (for each module). + self.metrics.log_value( + key=(mid, MODULE_SAMPLE_BATCH_SIZE_MEAN), + value=module_batch_size, + ) + # Log module steps (for each module). + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_SAMPLED), + value=module_batch_size, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(mid, NUM_MODULE_STEPS_SAMPLED_LIFETIME), + value=module_batch_size, + reduce="sum", + ) + # Log module steps (sum of all modules). + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_SAMPLED), + value=module_batch_size, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_SAMPLED_LIFETIME), + value=module_batch_size, + reduce="sum", + ) + # Log env steps (all modules). + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_SAMPLED), + batch.env_steps(), + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, NUM_ENV_STEPS_SAMPLED_LIFETIME), + batch.env_steps(), + reduce="sum", + with_throughput=True, + ) + + @override(Runner) + def set_device(self): + try: + self.__device = get_device( + self.config, + ( + 0 + if not self.worker_index + else self.config.num_gpus_per_offline_eval_runner + ), + ) + except NotImplementedError: + self.__device = None + + @override(Runner) + def make_module(self): + try: + from ray.rllib.env import INPUT_ENV_SPACES + + if not self._module_spec: + self.__module_spec = self.config.get_multi_rl_module_spec( + # Note, usually we have no environemnt in case of offline evaluation. + env=self.config.env, + spaces={ + INPUT_ENV_SPACES: ( + self.config.observation_space, + self.config.action_space, + ) + }, + inference_only=self.config.offline_eval_rl_module_inference_only, + ) + # Build the module from its spec. + self.module = self._module_spec.build() + # TODO (simon): Implement GPU inference. + # Move the RLModule to our device. + # TODO (sven): In order to make this framework-agnostic, we should maybe + # make the MultiRLModule.build() method accept a device OR create an + # additional `(Multi)RLModule.to()` override. + + self.module.foreach_module( + lambda mid, mod: ( + mod.to(self._device) if isinstance(mod, torch.nn.Module) else mod + ) + ) + + # If `AlgorithmConfig.get_multi_rl_module_spec()` is not implemented, this env runner + # will not have an RLModule, but might still be usable with random actions. + except NotImplementedError: + self.module = None + + def get_loss_for_module_fn(self): + # Either the user has provided a loss-for-module function, or we take + # the loss function from the default `Learner` class. + return ( + self.config.offline_loss_for_module_fn + or self.config.get_default_learner_class().__dict__[ + "compute_loss_for_module" + ] + ) + + @property + def _dataset_iterator(self) -> DataIterator: + """Returns the dataset iterator.""" + return self.__dataset_iterator + + def set_dataset_iterator(self, iterator): + """Sets the dataset iterator.""" + self.__dataset_iterator = iterator + + @property + def _batch_iterator(self) -> MiniBatchRayDataIterator: + return self.__batch_iterator + + @property + def _device(self) -> Union[DeviceType, None]: + return self.__device + + @property + def _module_spec(self) -> MultiRLModuleSpec: + """Returns the `MultiRLModuleSpec` of this `Runner`.""" + return self.__module_spec diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner_group.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner_group.py new file mode 100644 index 0000000000000000000000000000000000000000..470d762211c15fca7adf860b98240c6fcbaaa470 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_runner_group.py @@ -0,0 +1,191 @@ +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import ray +from ray.data.iterator import DataIterator +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.env import INPUT_ENV_SPACES +from ray.rllib.offline.offline_data import OfflineData +from ray.rllib.offline.offline_evaluation_runner import OfflineEvaluationRunner +from ray.rllib.offline.offline_policy_evaluation_runner import ( + OfflinePolicyEvaluationRunner, + OfflinePolicyPreEvaluator, +) +from ray.rllib.offline.offline_prelearner import OfflinePreLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.runners.runner_group import RunnerGroup + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + +class OfflineEvaluationRunnerGroup(RunnerGroup): + def __init__( + self, + config: "AlgorithmConfig", + local_runner: Optional[bool] = False, + logdir: Optional[str] = None, + tune_trial_id: Optional[str] = None, + pg_offset: int = 0, + _setup: bool = True, + spaces: Optional[Dict[str, Any]] = None, + module_state: Dict[str, Any] = None, + module_spec: Optional[MultiRLModuleSpec] = None, + **kwargs: Dict[str, Any], + ) -> None: + + # TODO (simon): Check, if this should happen later when the dataset + # is created. Maybe just overriding _setup. + # First initialize the super class. + super().__init__( + config=config, + local_runner=local_runner, + logdir=logdir, + tune_trial_id=tune_trial_id, + pg_offset=pg_offset, + _setup=_setup, + module_state=module_state, + module_spec=module_spec, + spaces=spaces, + ) + + @override(RunnerGroup) + def _setup( + self, + *, + config: Optional["AlgorithmConfig"] = None, + num_runners: int = 0, + local_runner: Optional[bool] = False, + module_state: Dict[str, Any] = None, + module_spec: Optional[MultiRLModuleSpec] = None, + spaces: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], + ) -> None: + + # Define the offline evaluation runner class. + self._runner_cls = config.offline_eval_runner_class or ( + OfflineEvaluationRunner + if config.offline_evaluation_type == "eval_loss" + else OfflinePolicyEvaluationRunner + ) + # Define + self._pre_learner_or_evaluator_cls = self.config.prelearner_class or ( + OfflinePreLearner + if config.offline_evaluation_type == "eval_loss" + else OfflinePolicyPreEvaluator + ) + self.config._is_frozen = False + self.config.prelearner_class = self._pre_learner_or_evaluator_cls + self.config._is_frozen = True + + # We can either run on a local runner or on remote runners only b/c + # streaming split needs remote runners. + if num_runners > 0 and local_runner: + raise ValueError( + f"Cannot run `OfflineEvaluationRunnerGroup with {num_runners=} " + "and a local runner. Either use no remote runners or only " + "remote runners." + ) + # Create all workers. + super()._setup( + config=config, + num_runners=num_runners, + local_runner=local_runner, + # Do not validate until the `DataIterators` are distributed. + validate=False, + module_spec=module_spec, + module_state=module_state, + spaces=spaces, + ) + + # Setup the evaluation offline dataset and return an iterator. + self._offline_data: OfflineData = OfflineData(config=config) + # We need the spaces to be defined for the `OfflinePreLearner`. + spaces = spaces or { + INPUT_ENV_SPACES: (config.observation_space, config.action_space) + } + self._offline_data.spaces = spaces + # The `OfflinePreLearner` also needs the module spec. + module_spec: MultiRLModuleSpec = module_spec or self.config.get_multi_rl_module_spec( + # TODO (simon): this needs merely the spaces defined via the connectors. + spaces={DEFAULT_MODULE_ID: spaces[INPUT_ENV_SPACES]}, + inference_only=self.config.offline_eval_rl_module_inference_only, + ) + self._offline_data.module_spec = module_spec + # If we have remote runners set the locality hints for the streaming split + # dataset iterators. + if self.num_remote_runners > 0: + runner_node_ids = self.foreach_runner( + lambda _: ray.get_runtime_context().get_node_id() + ) + if self.local_runner is not None: + runner_node_ids.insert(0, ray.get_runtime_context().get_node_id()) + self._offline_data.locality_hints = runner_node_ids + # Return a data iterator for each `Runner`. + self._offline_data_iterators: List[DataIterator] = self.offline_data.sample( + num_samples=self.config.offline_eval_batch_size_per_runner, + return_iterator=True, + num_shards=num_runners, + module_state=module_state, + ) + # Provide each `Runner` with a `DataIterator`. + self.foreach_runner( + func="set_dataset_iterator", + local_runner=local_runner, + kwargs=[ + {"iterator": iterator} for iterator in self._offline_data_iterators + ], + ) + # Now validate healthiness. + self.validate() + + @property + def runner_health_probe_timeout_s(self): + """Number of seconds to wait for health probe calls to `Runner`s.""" + return self.config.offline_eval_runner_health_probe_timeout_s + + @property + def runner_cls(self) -> Callable: + """Class for each runner.""" + return self._runner_cls + + @property + def num_runners(self) -> int: + """Number of runners to schedule and manage.""" + return self.config.num_offline_eval_runners + + @property + def offline_data(self) -> OfflineData: + return self._offline_data + + @property + def _remote_args(self): + """Remote arguments for each runner.""" + return { + "num_cpus": self._remote_config.num_cpus_per_offline_eval_runner, + "num_gpus": self._remote_config.num_gpus_per_offline_eval_runner, + "resources": self._remote_config.custom_resources_per_offline_eval_runner, + "max_restarts": ( + self.config.max_num_offline_eval_runner_restarts + if self.config.restart_failed_offline_eval_runners + else 0 + ), + } + + @property + def _ignore_ray_errors_on_runners(self): + """If errors in runners should be ignored.""" + return ( + self.config.ignore_offline_eval_runner_failures + or self.config.restart_failed_offline_eval_runners + ) + + @property + def _max_requests_in_flight_per_runner(self): + """Maximum requests in flight per runner.""" + return self.config.max_requests_in_flight_per_offline_eval_runner + + @property + def _validate_runners_after_construction(self): + """If runners should validated after constructed.""" + return self.config.validate_offline_eval_runners_after_construction diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_utils.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7eb8c3f728fbc46ee57155cd75784848a4edc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluation_utils.py @@ -0,0 +1,132 @@ +from typing import TYPE_CHECKING, Any, Dict, Type + +import numpy as np +import pandas as pd + +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.numpy import convert_to_numpy + +if TYPE_CHECKING: + from ray.rllib.offline.estimators.fqe_torch_model import FQETorchModel + from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator + + +@DeveloperAPI +def compute_q_and_v_values( + batch: pd.DataFrame, + model_class: Type["FQETorchModel"], + model_state: Dict[str, Any], + compute_q_values: bool = True, +) -> pd.DataFrame: + """Computes the Q and V values for the given batch of samples. + + This function is to be used with map_batches() to perform a batch prediction on a + dataset of records with `obs` and `actions` columns. + + Args: + batch: A sub-batch from the dataset. + model_class: The model class to use for the prediction. This class should be a + sub-class of FQEModel that implements the estimate_q() and estimate_v() + methods. + model_state: The state of the model to use for the prediction. + compute_q_values: Whether to compute the Q values or not. If False, only the V + is computed and returned. + + Returns: + The modified batch with the Q and V values added as columns. + """ + model = model_class.from_state(model_state) + + sample_batch = SampleBatch( + { + SampleBatch.OBS: np.vstack(batch[SampleBatch.OBS]), + SampleBatch.ACTIONS: np.vstack(batch[SampleBatch.ACTIONS]).squeeze(-1), + } + ) + + v_values = model.estimate_v(sample_batch) + v_values = convert_to_numpy(v_values) + batch["v_values"] = v_values + + if compute_q_values: + q_values = model.estimate_q(sample_batch) + q_values = convert_to_numpy(q_values) + batch["q_values"] = q_values + + return batch + + +@DeveloperAPI +def compute_is_weights( + batch: pd.DataFrame, + policy_state: Dict[str, Any], + estimator_class: Type["OffPolicyEstimator"], +) -> pd.DataFrame: + """Computes the importance sampling weights for the given batch of samples. + + For a lot of off-policy estimators, the importance sampling weights are computed as + the propensity score ratio between the new and old policies + (i.e. new_pi(act|obs) / old_pi(act|obs)). This function is to be used with + map_batches() to perform a batch prediction on a dataset of records with `obs`, + `actions`, `action_prob` and `rewards` columns. + + Args: + batch: A sub-batch from the dataset. + policy_state: The state of the policy to use for the prediction. + estimator_class: The estimator class to use for the prediction. This class + + Returns: + The modified batch with the importance sampling weights, weighted rewards, new + and old propensities added as columns. + """ + policy = Policy.from_state(policy_state) + estimator = estimator_class(policy=policy, gamma=0, epsilon_greedy=0) + sample_batch = SampleBatch( + { + SampleBatch.OBS: np.vstack(batch["obs"].values), + SampleBatch.ACTIONS: np.vstack(batch["actions"].values).squeeze(-1), + SampleBatch.ACTION_PROB: np.vstack(batch["action_prob"].values).squeeze(-1), + SampleBatch.REWARDS: np.vstack(batch["rewards"].values).squeeze(-1), + } + ) + new_prob = estimator.compute_action_probs(sample_batch) + old_prob = sample_batch[SampleBatch.ACTION_PROB] + rewards = sample_batch[SampleBatch.REWARDS] + weights = new_prob / old_prob + weighted_rewards = weights * rewards + + batch["weights"] = weights + batch["weighted_rewards"] = weighted_rewards + batch["new_prob"] = new_prob + batch["old_prob"] = old_prob + + return batch + + +@DeveloperAPI +def remove_time_dim(batch: pd.DataFrame) -> pd.DataFrame: + """Removes the time dimension from the given sub-batch of the dataset. + + If each row in a dataset has a time dimension ([T, D]), and T=1, this function will + remove the T dimension to convert each row to of shape [D]. If T > 1, the row is + left unchanged. This function is to be used with map_batches(). + + Args: + batch: The batch to remove the time dimension from. + Returns: + The modified batch with the time dimension removed (when applicable) + """ + BATCHED_KEYS = { + SampleBatch.OBS, + SampleBatch.ACTIONS, + SampleBatch.ACTION_PROB, + SampleBatch.REWARDS, + SampleBatch.NEXT_OBS, + SampleBatch.DONES, + } + for k in batch.columns: + if k in BATCHED_KEYS: + batch[k] = batch[k].apply(lambda x: x[0] if len(x) == 1 else x) + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluator.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..277c514e4582147210eb3f00d9d134d2593bb0a3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_evaluator.py @@ -0,0 +1,76 @@ +import abc +import logging +import os +from typing import Any, Dict + +from ray.data import Dataset +from ray.rllib.policy import Policy +from ray.rllib.utils.annotations import DeveloperAPI, ExperimentalAPI +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class OfflineEvaluator(abc.ABC): + """Interface for an offline evaluator of a policy""" + + @DeveloperAPI + def __init__(self, policy: Policy, **kwargs): + """Initializes an OffPolicyEstimator instance. + + Args: + policy: Policy to evaluate. + kwargs: forward compatibility placeholder. + """ + self.policy = policy + + @abc.abstractmethod + @DeveloperAPI + def estimate(self, batch: SampleBatchType, **kwargs) -> Dict[str, Any]: + """Returns the evaluation results for the given batch of episodes. + + Args: + batch: The batch to evaluate. + kwargs: forward compatibility placeholder. + + Returns: + The evaluation done on the given batch. The returned + dict can be any arbitrary mapping of strings to metrics. + """ + raise NotImplementedError + + @DeveloperAPI + def train(self, batch: SampleBatchType, **kwargs) -> Dict[str, Any]: + """Sometimes you need to train a model inside an evaluator. This method + abstracts the training process. + + Args: + batch: SampleBatch to train on + kwargs: forward compatibility placeholder. + + Returns: + Any optional metrics to return from the evaluator + """ + return {} + + @ExperimentalAPI + def estimate_on_dataset( + self, + dataset: Dataset, + *, + n_parallelism: int = os.cpu_count(), + ) -> Dict[str, Any]: + + """Calculates the estimate of the metrics based on the given offline dataset. + + Typically, the dataset is passed through only once via n_parallel tasks in + mini-batches to improve the run-time of metric estimation. + + Args: + dataset: The ray dataset object to do offline evaluation on. + n_parallelism: The number of parallelism to use for the computation. + + Returns: + Dict[str, Any]: A dictionary of the estimated values. + """ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_policy_evaluation_runner.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_policy_evaluation_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..6dec0e761ee673b0ba4c258236f1049b6076f1ed --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_policy_evaluation_runner.py @@ -0,0 +1,643 @@ +import math +from enum import Enum +from typing import ( + TYPE_CHECKING, + Collection, + Dict, + Iterable, + List, + Optional, + Union, +) + +import gymnasium as gym +import numpy + +import ray +from ray.data.iterator import DataIterator +from ray.rllib.connectors.env_to_module import EnvToModulePipeline +from ray.rllib.core import ( + ALL_MODULES, + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_RL_MODULE, + DEFAULT_AGENT_ID, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.offline.offline_prelearner import SCHEMA, OfflinePreLearner +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.framework import get_device, try_import_torch +from ray.rllib.utils.metrics import ( + DATASET_NUM_ITERS_EVALUATED, + DATASET_NUM_ITERS_EVALUATED_LIFETIME, + EPISODE_LEN_MAX, + EPISODE_LEN_MEAN, + EPISODE_LEN_MIN, + EPISODE_RETURN_MAX, + EPISODE_RETURN_MEAN, + EPISODE_RETURN_MIN, + MODULE_SAMPLE_BATCH_SIZE_MEAN, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_MODULE_STEPS_SAMPLED, + NUM_MODULE_STEPS_SAMPLED_LIFETIME, + OFFLINE_SAMPLING_TIMER, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.minibatch_utils import MiniBatchRayDataIterator +from ray.rllib.utils.runners.runner import Runner +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import ( + DeviceType, + EpisodeID, + StateDict, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +torch, _ = try_import_torch() + +TOTAL_EVAL_LOSS_KEY = "total_eval_loss" + + +# TODO (simon): Implement more ... +class OfflinePolicyEvaluationTypes(str, Enum): + """Defines the offline policy evaluation types. + + IS: Importance Sampling. + PDIS: Per-Decision Importance Sampling. In contrast to IS this method + weighs each reward and not the return as a whole. As a result it + usually exhibits lower variance. + """ + + IS = "is" + PDIS = "pdis" + + +class OfflinePolicyPreEvaluator(OfflinePreLearner): + def __call__(self, batch: Dict[str, numpy.ndarray]) -> Dict[str, numpy.ndarray]: + # If we directly read in episodes we just convert to list. + if self.input_read_episodes: + # Import `msgpack` for decoding. + import msgpack + import msgpack_numpy as mnp + + # Read the episodes and decode them. + episodes: List[SingleAgentEpisode] = [ + SingleAgentEpisode.from_state( + msgpack.unpackb(state, object_hook=mnp.decode) + ) + for state in batch["item"] + ] + # Ensure that all episodes are done and no duplicates are in the batch. + episodes = self._validate_episodes(episodes) + # Add the episodes to the buffer. + self.episode_buffer.add(episodes) + # TODO (simon): Refactor into a single code block for both cases. + episodes = self.episode_buffer.sample( + num_items=self.config.train_batch_size_per_learner, + batch_length_T=( + self.config.model_config.get("max_seq_len", 0) + if self._module.is_stateful() + else None + ), + n_step=self.config.get("n_step", 1) or 1, + # TODO (simon): This can be removed as soon as DreamerV3 has been + # cleaned up, i.e. can use episode samples for training. + sample_episodes=True, + to_numpy=True, + ) + # Else, if we have old stack `SampleBatch`es. + elif self.input_read_sample_batches: + episodes: List[ + SingleAgentEpisode + ] = OfflinePreLearner._map_sample_batch_to_episode( + self._is_multi_agent, + batch, + to_numpy=True, + schema=SCHEMA | self.config.input_read_schema, + input_compress_columns=self.config.input_compress_columns, + )[ + "episodes" + ] + # Ensure that all episodes are done and no duplicates are in the batch. + episodes = self._validate_episodes(episodes) + # Add the episodes to the buffer. + self.episode_buffer.add(episodes) + # Sample steps from the buffer. + episodes = self.episode_buffer.sample( + num_items=self.config.train_batch_size_per_learner, + batch_length_T=( + self.config.model_config.get("max_seq_len", 0) + if self._module.is_stateful() + else None + ), + n_step=self.config.get("n_step", 1) or 1, + # TODO (simon): This can be removed as soon as DreamerV3 has been + # cleaned up, i.e. can use episode samples for training. + sample_episodes=True, + to_numpy=True, + ) + # Otherwise we map the batch to episodes. + else: + episodes: List[SingleAgentEpisode] = self._map_to_episodes( + self._is_multi_agent, + batch, + schema=SCHEMA | self.config.input_read_schema, + to_numpy=False, + input_compress_columns=self.config.input_compress_columns, + observation_space=self.observation_space, + action_space=self.action_space, + )["episodes"] + + episode_dicts = [] + for episode in episodes: + # Note, we expect users to provide terminated episodes in `SingleAgentEpisode` + # or `SampleBatch` format. Otherwise computation of episode returns will be + # biased. + episode_dict = {} + episode_dict[Columns.OBS] = episode.get_observations(slice(0, len(episode))) + episode_dict[Columns.ACTIONS] = episode.get_actions() + episode_dict[Columns.REWARDS] = episode.get_rewards() + episode_dict[Columns.ACTION_LOGP] = episode.get_extra_model_outputs( + key=Columns.ACTION_LOGP + ) + episode_dicts.append(episode_dict) + + return {"episodes": episode_dicts} + + +class OfflinePolicyEvaluationRunner(Runner, Checkpointable): + def __init__( + self, + config: "AlgorithmConfig", + module_spec: Optional[MultiRLModuleSpec] = None, + **kwargs, + ): + + # This needs to be defined before we call the `Runner.__init__` + # b/c the latter calls the `make_module` and then needs the spec. + # TODO (simon): Check, if we make this a generic attribute. + self.__module_spec: MultiRLModuleSpec = module_spec + self.__dataset_iterator = None + self.__batch_iterator = None + + Runner.__init__(self, config=config, **kwargs) + Checkpointable.__init__(self) + + # This has to be defined after we have a `self.config`. + self.__spaces = kwargs.get("spaces") + self.__env_to_module = self.config.build_env_to_module_connector( + spaces=self._spaces, device=self._device + ) + self.__offline_evaluation_type = OfflinePolicyEvaluationTypes( + self.config["offline_evaluation_type"] + ) + + def run( + self, + explore: bool = False, + train: bool = True, + **kwargs, + ) -> None: + + if self.__dataset_iterator is None: + raise ValueError( + f"{self} doesn't have a data iterator. Can't call `run` on " + "`OfflinePolicyEvaluationRunner`." + ) + + if not self._batch_iterator: + self.__batch_iterator = self._create_batch_iterator( + **self.config.iter_batches_kwargs + ) + + # Log current weight seq no. + self.metrics.log_value( + key=WEIGHTS_SEQ_NO, + value=self._weights_seq_no, + window=1, + ) + + with self.metrics.log_time(OFFLINE_SAMPLING_TIMER): + if explore is None: + explore = self.config.explore + + # Evaluate on offline data. + return self._evaluate( + explore=explore, + train=train, + ) + + def _create_batch_iterator(self, **kwargs) -> Iterable: + + # Import the torch utils here b/c Ray Air imports `torch`` directly. + from ray.air._internal.torch_utils import ( + convert_ndarray_batch_to_torch_tensor_batch, + ) + + # Define the collate function that converts the flattened dictionary + # to a `MultiAgentBatch` with Tensors. + def _collate_fn( + _batch: Dict[str, numpy.ndarray], + ) -> Dict[EpisodeID, Dict[str, numpy.ndarray]]: + + return _batch["episodes"] + + # Define the finalize function that makes the host-to-device transfer. + def _finalize_fn( + _batch: Dict[EpisodeID, Dict[str, numpy.ndarray]], + ) -> Dict[EpisodeID, Dict[str, TensorType]]: + + return [ + convert_ndarray_batch_to_torch_tensor_batch( + episode, device=self._device, dtypes=torch.float32 + ) + for episode in _batch + ] + + # Return a minibatch iterator. + return MiniBatchRayDataIterator( + iterator=self._dataset_iterator, + collate_fn=_collate_fn, + finalize_fn=_finalize_fn, + minibatch_size=self.config.offline_eval_batch_size_per_runner, + num_iters=self.config.dataset_num_iters_per_eval_runner, + **kwargs, + ) + + def _evaluate( + self, + explore: bool, + train: bool, + ) -> None: + + self.metrics.activate_tensor_mode() + + num_env_steps = 0 + for iteration, tensor_minibatch in enumerate(self._batch_iterator): + for episode in tensor_minibatch: + action_dist_cls = self.module[ + DEFAULT_MODULE_ID + ].get_inference_action_dist_cls() + # TODO (simon): It needs here the `EnvToModule` pipeline. + action_logits = self.module[DEFAULT_MODULE_ID].forward_inference( + episode + )[Columns.ACTION_DIST_INPUTS] + # TODO (simon): It might need here the ModuleToEnv pipeline until the + # `GetActions` piece. + action_dist = action_dist_cls.from_logits(action_logits) + actions = action_dist.sample() + action_logp = action_dist.logp(actions) + # If we have action log-probs use them. + if Columns.ACTION_LOGP in episode: + behavior_action_logp = episode[Columns.ACTION_LOGP] + # Otherwise approximate them via the current action distribution. + else: + behavior_action_logp = action_dist.logp(episode[Columns.ACTIONS]) + + # Compute the weights. + if self.__offline_evaluation_type == OfflinePolicyEvaluationTypes.IS: + weight = torch.prod( + torch.exp(action_logp) / torch.exp(behavior_action_logp) + ) + # Note, we use the (un)-discounted return to compare with the `EnvRunner` + # returns. + episode_return = episode[Columns.REWARDS].sum() + offline_return = (weight * episode_return).item() + elif ( + self.__offline_evaluation_type == OfflinePolicyEvaluationTypes.PDIS + ): + weights = torch.exp(action_logp) / torch.exp(behavior_action_logp) + offline_return = torch.dot(weights, episode[Columns.REWARDS]).item() + + episode_len = episode[Columns.REWARDS].shape[0] + 1 + num_env_steps += episode_len + + self._log_episode_metrics(episode_len, offline_return) + + self._log_batch_metrics(len(tensor_minibatch), num_env_steps) + + # Record the number of batches pulled from the dataset. + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_EVALUATED), + iteration + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (ALL_MODULES, DATASET_NUM_ITERS_EVALUATED_LIFETIME), + iteration + 1, + reduce="sum", + ) + + self.metrics.deactivate_tensor_mode() + + return self.metrics.reduce() + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + {"config": self.config}, # **kwargs + ) + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + state = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + ), + } + + if self._check_component(COMPONENT_RL_MODULE, components, not_components): + state[COMPONENT_RL_MODULE] = self.module.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=self._get_subcomponents( + COMPONENT_RL_MODULE, not_components + ), + **kwargs, + ) + state[WEIGHTS_SEQ_NO] = self._weights_seq_no + if self._check_component( + COMPONENT_ENV_TO_MODULE_CONNECTOR, components, not_components + ): + state[COMPONENT_ENV_TO_MODULE_CONNECTOR] = self._env_to_module.get_state() + + return state + + def _convert_to_tensor(self, struct) -> TensorType: + """Converts structs to a framework-specific tensor.""" + return convert_to_torch_tensor(struct) + + def stop(self) -> None: + """Releases all resources used by this EnvRunner. + + For example, when using a gym.Env in this EnvRunner, you should make sure + that its `close()` method is called. + """ + pass + + def __del__(self) -> None: + """If this Actor is deleted, clears all resources used by it.""" + pass + + @override(Runner) + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Ensures that the instances has a `MultiRLModule` and an + environment defined. + + Raises: + AssertionError: If the EnvRunner Actor has NOT been properly initialized. + """ + # Make sure, we have built our RLModule properly and assigned a dataset iterator. + assert self._dataset_iterator and hasattr(self, "module") + + @override(Runner) + def get_metrics(self): + return self.metrics.reduce() + + def _convert_batch_type( + self, + batch: MultiAgentBatch, + to_device: bool = True, + pin_memory: bool = False, + use_stream: bool = False, + ) -> MultiAgentBatch: + batch = convert_to_torch_tensor( + batch.policy_batches, + device=self._device if to_device else None, + pin_memory=pin_memory, + use_stream=use_stream, + ) + # TODO (sven): This computation of `env_steps` is not accurate! + length = max(len(b) for b in batch.values()) + batch = MultiAgentBatch(batch, env_steps=length) + return batch + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + if COMPONENT_ENV_TO_MODULE_CONNECTOR in state: + self._env_to_module.set_state(state[COMPONENT_ENV_TO_MODULE_CONNECTOR]) + + # Update the RLModule state. + if COMPONENT_RL_MODULE in state: + # A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the + # update. + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Only update the weigths, if this is the first synchronization or + # if the weights of this `EnvRunner` lacks behind the actual ones. + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + rl_module_state = state[COMPONENT_RL_MODULE] + if isinstance(rl_module_state, ray.ObjectRef): + rl_module_state = ray.get(rl_module_state) + self.module.set_state(rl_module_state) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update our lifetime counters. + # TODO (simon): Create extra metrics. + if NUM_ENV_STEPS_SAMPLED_LIFETIME in state: + self.metrics.set_value( + key=NUM_ENV_STEPS_SAMPLED_LIFETIME, + value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME], + reduce="sum", + with_throughput=True, + ) + + def _log_episode_metrics(self, episode_len: int, episode_return: float) -> None: + """Logs episode metrics for each episode.""" + + # Log general episode metrics. + # Use the configured window, but factor in the parallelism of the + # `OfflinePolicyEvaluationRunners`. As a result, we only log the last + # `window / num_env_runners` steps here, b/c everything gets + # parallel-merged in the Algorithm process. + win = max( + 1, + int( + math.ceil( + self.config.metrics_num_episodes_for_smoothing + / (self.config.num_offline_eval_runners or 1) + ) + ), + ) + self.metrics.log_value(EPISODE_LEN_MEAN, episode_len, window=win) + self.metrics.log_value(EPISODE_RETURN_MEAN, episode_return, window=win) + # Per-agent returns. + self.metrics.log_value( + ("agent_episode_return_mean", DEFAULT_AGENT_ID), episode_return, window=win + ) + # Per-RLModule returns. + self.metrics.log_value( + ("module_episode_return_mean", DEFAULT_MODULE_ID), + episode_return, + window=win, + ) + + # For some metrics, log min/max as well. + self.metrics.log_value(EPISODE_LEN_MIN, episode_len, reduce="min", window=win) + self.metrics.log_value( + EPISODE_RETURN_MIN, episode_return, reduce="min", window=win + ) + self.metrics.log_value(EPISODE_LEN_MAX, episode_len, reduce="max", window=win) + self.metrics.log_value( + EPISODE_RETURN_MAX, episode_return, reduce="max", window=win + ) + + def _log_batch_metrics(self, batch_size: int, num_env_steps: int): + """Logs batch metrics for each mini batch.""" + + # Log weights seq no for this batch. + self.metrics.log_value( + (DEFAULT_MODULE_ID, WEIGHTS_SEQ_NO), + self._weights_seq_no, + window=1, + ) + + # Log average batch size (for each module). + self.metrics.log_value( + key=(DEFAULT_MODULE_ID, MODULE_SAMPLE_BATCH_SIZE_MEAN), + value=batch_size, + ) + # Log module steps (for each module). + self.metrics.log_value( + key=(DEFAULT_MODULE_ID, NUM_MODULE_STEPS_SAMPLED), + value=num_env_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(DEFAULT_MODULE_ID, NUM_MODULE_STEPS_SAMPLED_LIFETIME), + value=num_env_steps, + reduce="sum", + ) + # Log module steps (sum of all modules). + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_SAMPLED), + value=num_env_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(ALL_MODULES, NUM_MODULE_STEPS_SAMPLED_LIFETIME), + value=num_env_steps, + reduce="sum", + ) + # Log env steps (all modules). + self.metrics.log_value( + key=(ALL_MODULES, NUM_ENV_STEPS_SAMPLED), + value=num_env_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + key=(ALL_MODULES, NUM_ENV_STEPS_SAMPLED_LIFETIME), + value=num_env_steps, + reduce="sum", + with_throughput=True, + ) + + @override(Runner) + def set_device(self): + try: + self.__device = get_device( + self.config, + ( + 0 + if not self.worker_index + else self.config.num_gpus_per_offline_eval_runner + ), + ) + except NotImplementedError: + self.__device = None + + @override(Runner) + def make_module(self): + try: + from ray.rllib.env import INPUT_ENV_SPACES + + if not self._module_spec: + self.__module_spec = self.config.get_multi_rl_module_spec( + # Note, usually we have no environemnt in case of offline evaluation. + env=self.config.env, + spaces={ + INPUT_ENV_SPACES: ( + self.config.observation_space, + self.config.action_space, + ) + }, + inference_only=self.config.offline_eval_rl_module_inference_only, + ) + # Build the module from its spec. + self.module = self._module_spec.build() + # TODO (simon): Implement GPU inference. + # Move the RLModule to our device. + # TODO (sven): In order to make this framework-agnostic, we should maybe + # make the MultiRLModule.build() method accept a device OR create an + # additional `(Multi)RLModule.to()` override. + + self.module.foreach_module( + lambda mid, mod: ( + mod.to(self._device) if isinstance(mod, torch.nn.Module) else mod + ) + ) + + # If `AlgorithmConfig.get_multi_rl_module_spec()` is not implemented, this env runner + # will not have an RLModule, but might still be usable with random actions. + except NotImplementedError: + self.module = None + + @property + def _dataset_iterator(self) -> DataIterator: + """Returns the dataset iterator.""" + return self.__dataset_iterator + + def set_dataset_iterator(self, iterator): + """Sets the dataset iterator.""" + self.__dataset_iterator = iterator + + @property + def _batch_iterator(self) -> MiniBatchRayDataIterator: + return self.__batch_iterator + + @property + def _device(self) -> Union[DeviceType, None]: + return self.__device + + @property + def _module_spec(self) -> MultiRLModuleSpec: + """Returns the `MultiRLModuleSpec` of this `Runner`.""" + return self.__module_spec + + @property + def _spaces(self) -> Dict[str, gym.spaces.Space]: + """Returns the spaces of thsi `Runner`.""" + return self.__spaces + + @property + def _env_to_module(self) -> EnvToModulePipeline: + """Returns the env-to-module pipeline of this `Runner`.""" + return self.__env_to_module + + @property + def _offline_evaluation_type(self) -> Enum: + """Returns the offline evaluation type of this `Runner`.""" + return self.__offline_evaluation_type diff --git a/lib/python3.12/site-packages/ray/rllib/offline/offline_prelearner.py b/lib/python3.12/site-packages/ray/rllib/offline/offline_prelearner.py new file mode 100644 index 0000000000000000000000000000000000000000..18d66974dce8926aecbbb798d52fb496d88fb8dd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/offline_prelearner.py @@ -0,0 +1,622 @@ +import copy +import logging +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union + +import gymnasium as gym +import numpy as np +import tree + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule, MultiRLModuleSpec +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils import flatten_dict +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, +) +from ray.rllib.utils.compression import unpack_if_needed +from ray.rllib.utils.replay_buffers.replay_buffer import ReplayBuffer +from ray.rllib.utils.typing import EpisodeType, ModuleID +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +#: This is the default schema used if no `input_read_schema` is set in +#: the config. If a user passes in a schema into `input_read_schema` +#: this user-defined schema has to comply with the keys of `SCHEMA`, +#: while values correspond to the columns in the user's dataset. Note +#: that only the user-defined values will be overridden while all +#: other values from SCHEMA remain as defined here. +SCHEMA = { + Columns.EPS_ID: Columns.EPS_ID, + Columns.AGENT_ID: Columns.AGENT_ID, + Columns.MODULE_ID: Columns.MODULE_ID, + Columns.OBS: Columns.OBS, + Columns.ACTIONS: Columns.ACTIONS, + Columns.REWARDS: Columns.REWARDS, + Columns.INFOS: Columns.INFOS, + Columns.NEXT_OBS: Columns.NEXT_OBS, + Columns.TERMINATEDS: Columns.TERMINATEDS, + Columns.TRUNCATEDS: Columns.TRUNCATEDS, + Columns.T: Columns.T, + # TODO (simon): Add remove as soon as we are new stack only. + "agent_index": "agent_index", + "dones": "dones", + "unroll_id": "unroll_id", +} + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="alpha") +class OfflinePreLearner: + """Class that coordinates data transformation from dataset to learner. + + This class is an essential part of the new `Offline RL API` of `RLlib`. + It is a callable class that is run in `ray.data.Dataset.map_batches` + when iterating over batches for training. It's basic function is to + convert data in batch from rows to episodes (`SingleAGentEpisode`s + for now) and to then run the learner connector pipeline to convert + further to trainable batches. These batches are used directly in the + `Learner`'s `update` method. + + The main reason to run these transformations inside of `map_batches` + is for better performance. Batches can be pre-fetched in `ray.data` + and therefore batch trransformation can be run highly parallelized to + the `Learner''s `update`. + + This class can be overridden to implement custom logic for transforming + batches and make them 'Learner'-ready. When deriving from this class + the `__call__` method and `_map_to_episodes` can be overridden to induce + custom logic for the complete transformation pipeline (`__call__`) or + for converting to episodes only ('_map_to_episodes`). + + Custom `OfflinePreLearner` classes can be passed into + `AlgorithmConfig.offline`'s `prelearner_class`. The `OfflineData` class + will then use the custom class in its data pipeline. + """ + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def __init__( + self, + *, + config: "AlgorithmConfig", + spaces: Optional[Tuple[gym.Space, gym.Space]] = None, + module_spec: Optional[MultiRLModuleSpec] = None, + module_state: Optional[Dict[ModuleID, Any]] = None, + **kwargs: Dict[str, Any], + ): + + self.config: AlgorithmConfig = config + self.input_read_episodes: bool = self.config.input_read_episodes + self.input_read_sample_batches: bool = self.config.input_read_sample_batches + # Build the module from spec. + self._module: MultiRLModule = module_spec.build() + self._module.set_state(module_state) + # Map the module to the device, if necessary. + # TODO (simon): Check here if we already have a list. + # self._set_device(device_strings) + + # Store the observation and action space if defined, otherwise we + # set them to `None`. Note, if `None` the `convert_from_jsonable` + # will not convert the input space samples. + self.observation_space, self.action_space = spaces or (None, None) + + # Build the learner connector pipeline. + self._learner_connector = self.config.build_learner_connector( + input_observation_space=self.observation_space, + input_action_space=self.action_space, + ) + # Cache the policies to be trained to update weights only for these. + self._policies_to_train = self.config.policies_to_train + self._is_multi_agent: bool = config.is_multi_agent + # Set the counter to zero. + self.iter_since_last_module_update: int = 0 + # self._future = None + + # Set up an episode buffer, if the module is stateful or we sample from + # `SampleBatch` types. + if ( + self.input_read_sample_batches + or self._module.is_stateful() + or self.input_read_episodes + ): + # Either the user defined a buffer class or we fall back to the default. + prelearner_buffer_class = ( + self.config.prelearner_buffer_class + or self.default_prelearner_buffer_class + ) + prelearner_buffer_kwargs = ( + self.default_prelearner_buffer_kwargs + | self.config.prelearner_buffer_kwargs + ) + # Initialize the buffer. + self.episode_buffer = prelearner_buffer_class( + **prelearner_buffer_kwargs, + ) + + @OverrideToImplementCustomLogic + def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: + """Prepares plain data batches for training with `Learner`'s. + + Args: + batch: A dictionary of numpy arrays containing either column data + with `self.config.input_read_schema`, `EpisodeType` data, or + `BatchType` data. + + Returns: + A `MultiAgentBatch` that can be passed to `Learner.update` methods. + """ + # If we directly read in episodes we just convert to list. + if self.input_read_episodes: + # Import `msgpack` for decoding. + import msgpack + import msgpack_numpy as mnp + + # Read the episodes and decode them. + episodes: List[SingleAgentEpisode] = [ + SingleAgentEpisode.from_state( + msgpack.unpackb(state, object_hook=mnp.decode) + ) + for state in batch["item"] + ] + # Ensure that all episodes are done and no duplicates are in the batch. + episodes = self._validate_episodes(episodes) + # Add the episodes to the buffer. + self.episode_buffer.add(episodes) + # TODO (simon): Refactor into a single code block for both cases. + episodes = self.episode_buffer.sample( + num_items=self.config.train_batch_size_per_learner, + batch_length_T=self.config.model_config.get("max_seq_len", 0) + if self._module.is_stateful() + else None, + n_step=self.config.get("n_step", 1) or 1, + # TODO (simon): This can be removed as soon as DreamerV3 has been + # cleaned up, i.e. can use episode samples for training. + sample_episodes=True, + to_numpy=True, + ) + # Else, if we have old stack `SampleBatch`es. + elif self.input_read_sample_batches: + episodes: List[ + SingleAgentEpisode + ] = OfflinePreLearner._map_sample_batch_to_episode( + self._is_multi_agent, + batch, + to_numpy=True, + schema=SCHEMA | self.config.input_read_schema, + input_compress_columns=self.config.input_compress_columns, + )[ + "episodes" + ] + # Ensure that all episodes are done and no duplicates are in the batch. + episodes = self._validate_episodes(episodes) + # Add the episodes to the buffer. + self.episode_buffer.add(episodes) + # Sample steps from the buffer. + episodes = self.episode_buffer.sample( + num_items=self.config.train_batch_size_per_learner, + batch_length_T=self.config.model_config.get("max_seq_len", 0) + if self._module.is_stateful() + else None, + n_step=self.config.get("n_step", 1) or 1, + # TODO (simon): This can be removed as soon as DreamerV3 has been + # cleaned up, i.e. can use episode samples for training. + sample_episodes=True, + to_numpy=True, + ) + # Otherwise we map the batch to episodes. + else: + episodes: List[SingleAgentEpisode] = self._map_to_episodes( + self._is_multi_agent, + batch, + schema=SCHEMA | self.config.input_read_schema, + to_numpy=False, + input_compress_columns=self.config.input_compress_columns, + observation_space=self.observation_space, + action_space=self.action_space, + )["episodes"] + + # TODO (simon): Make synching work. Right now this becomes blocking or never + # receives weights. Learners appear to be non accessi ble via other actors. + # Increase the counter for updating the module. + # self.iter_since_last_module_update += 1 + + # if self._future: + # refs, _ = ray.wait([self._future], timeout=0) + # print(f"refs: {refs}") + # if refs: + # module_state = ray.get(self._future) + # + # self._module.set_state(module_state) + # self._future = None + + # # Synch the learner module, if necessary. Note, in case of a local learner + # # we have a reference to the module and therefore an up-to-date module. + # if self.learner_is_remote and self.iter_since_last_module_update + # > self.config.prelearner_module_synch_period: + # # Reset the iteration counter. + # self.iter_since_last_module_update = 0 + # # Request the module weights from the remote learner. + # self._future = + # self._learner.get_module_state.remote(inference_only=False) + # # module_state = + # ray.get(self._learner.get_module_state.remote(inference_only=False)) + # # self._module.set_state(module_state) + + # Run the `Learner`'s connector pipeline. + batch = self._learner_connector( + rl_module=self._module, + batch={}, + episodes=episodes, + shared_data={}, + # TODO (sven): Add MetricsLogger to non-Learner components that have a + # LearnerConnector pipeline. + metrics=None, + ) + # Remove all data from modules that should not be trained. We do + # not want to pass around more data than necessary. + for module_id in batch: + if not self._should_module_be_updated(module_id, batch): + del batch[module_id] + + # Flatten the dictionary to increase serialization performance. + return flatten_dict(batch) + + @property + def default_prelearner_buffer_class(self) -> ReplayBuffer: + """Sets the default replay buffer.""" + from ray.rllib.utils.replay_buffers.episode_replay_buffer import ( + EpisodeReplayBuffer, + ) + + # Return the buffer. + return EpisodeReplayBuffer + + @property + def default_prelearner_buffer_kwargs(self) -> Dict[str, Any]: + """Sets the default arguments for the replay buffer. + + Note, the `capacity` might vary with the size of the episodes or + sample batches in the offline dataset. + """ + return { + "capacity": self.config.train_batch_size_per_learner * 10, + "batch_size_B": self.config.train_batch_size_per_learner, + } + + def _validate_episodes( + self, episodes: List[SingleAgentEpisode] + ) -> Set[SingleAgentEpisode]: + """Validate episodes sampled from the dataset. + + Note, our episode buffers cannot handle either duplicates nor + non-ordered fragmentations, i.e. fragments from episodes that do + not arrive in timestep order. + + Args: + episodes: A list of `SingleAgentEpisode` instances sampled + from a dataset. + + Returns: + A set of `SingleAgentEpisode` instances. + + Raises: + ValueError: If not all episodes are `done`. + """ + # Ensure that episodes are all done. + if not all(eps.is_done for eps in episodes): + raise ValueError( + "When sampling from episodes (`input_read_episodes=True`) all " + "recorded episodes must be done (i.e. either `terminated=True`) " + "or `truncated=True`)." + ) + # Ensure that episodes do not contain duplicates. Note, this can happen + # if the dataset is small and pulled batches contain multiple episodes. + unique_episode_ids = set() + cleaned_episodes = set() + for eps in episodes: + if ( + eps.id_ not in unique_episode_ids + and eps.id_ not in self.episode_buffer.episode_id_to_index + ): + unique_episode_ids.add(eps.id_) + cleaned_episodes.add(eps) + return cleaned_episodes + + def _should_module_be_updated(self, module_id, multi_agent_batch=None) -> bool: + """Checks which modules in a MultiRLModule should be updated.""" + if not self._policies_to_train: + # In case of no update information, the module is updated. + return True + elif not callable(self._policies_to_train): + return module_id in set(self._policies_to_train) + else: + return self._policies_to_train(module_id, multi_agent_batch) + + @OverrideToImplementCustomLogic + @staticmethod + def _map_to_episodes( + is_multi_agent: bool, + batch: Dict[str, Union[list, np.ndarray]], + schema: Dict[str, str] = SCHEMA, + to_numpy: bool = False, + input_compress_columns: Optional[List[str]] = None, + ignore_final_observation: Optional[bool] = False, + observation_space: gym.Space = None, + action_space: gym.Space = None, + **kwargs: Dict[str, Any], + ) -> Dict[str, List[EpisodeType]]: + """Maps a batch of data to episodes.""" + + # Set to empty list, if `None`. + input_compress_columns = input_compress_columns or [] + + episodes = [] + for i, obs in enumerate(batch[schema[Columns.OBS]]): + + # If multi-agent we need to extract the agent ID. + # TODO (simon): Check, what happens with the module ID. + if is_multi_agent: + agent_id = ( + batch[schema[Columns.AGENT_ID]][i] + if Columns.AGENT_ID in batch + # The old stack uses "agent_index" instead of "agent_id". + # TODO (simon): Remove this as soon as we are new stack only. + else ( + batch[schema["agent_index"]][i] + if schema["agent_index"] in batch + else None + ) + ) + else: + agent_id = None + + if is_multi_agent: + # TODO (simon): Add support for multi-agent episodes. + pass + else: + # Unpack observations, if needed. + unpacked_obs = ( + unpack_if_needed(obs) + if Columns.OBS in input_compress_columns + else obs + ) + # Set the next observation. + if ignore_final_observation: + unpacked_next_obs = tree.map_structure( + lambda x: 0 * x, copy.deepcopy(unpacked_obs) + ) + else: + unpacked_next_obs = ( + unpack_if_needed(batch[schema[Columns.NEXT_OBS]][i]) + if Columns.OBS in input_compress_columns + else batch[schema[Columns.NEXT_OBS]][i] + ) + # Build a single-agent episode with a single row of the batch. + episode = SingleAgentEpisode( + id_=str(batch[schema[Columns.EPS_ID]][i]) + if schema[Columns.EPS_ID] in batch + else uuid.uuid4().hex, + agent_id=agent_id, + # Observations might be (a) serialized and/or (b) converted + # to a JSONable (when a composite space was used). We unserialize + # and then reconvert from JSONable to space sample. + observations=[unpacked_obs, unpacked_next_obs], + infos=[ + {}, + batch[schema[Columns.INFOS]][i] + if schema[Columns.INFOS] in batch + else {}, + ], + # Actions might be (a) serialized and/or (b) converted to a JSONable + # (when a composite space was used). We unserializer and then + # reconvert from JSONable to space sample. + actions=[ + unpack_if_needed(batch[schema[Columns.ACTIONS]][i]) + if Columns.ACTIONS in input_compress_columns + else batch[schema[Columns.ACTIONS]][i] + ], + rewards=[batch[schema[Columns.REWARDS]][i]], + terminated=batch[ + schema[Columns.TERMINATEDS] + if schema[Columns.TERMINATEDS] in batch + else "dones" + ][i], + truncated=batch[schema[Columns.TRUNCATEDS]][i] + if schema[Columns.TRUNCATEDS] in batch + else False, + # TODO (simon): Results in zero-length episodes in connector. + # t_started=batch[Columns.T if Columns.T in batch else + # "unroll_id"][i][0], + # TODO (simon): Single-dimensional columns are not supported. + # Extra model outputs might be serialized. We unserialize them here + # if needed. + # TODO (simon): Check, if we need here also reconversion from + # JSONable in case of composite spaces. + extra_model_outputs={ + k: [ + unpack_if_needed(v[i]) + if k in input_compress_columns + else v[i] + ] + for k, v in batch.items() + if ( + k not in schema + and k not in schema.values() + and k not in ["dones", "agent_index", "type"] + ) + }, + len_lookback_buffer=0, + ) + + if to_numpy: + episode.to_numpy() + episodes.append(episode) + # Note, `map_batches` expects a `Dict` as return value. + return {"episodes": episodes} + + @OverrideToImplementCustomLogic + @staticmethod + def _map_sample_batch_to_episode( + is_multi_agent: bool, + batch: Dict[str, Union[list, np.ndarray]], + schema: Dict[str, str] = SCHEMA, + to_numpy: bool = False, + input_compress_columns: Optional[List[str]] = None, + ) -> Dict[str, List[EpisodeType]]: + """Maps an old stack `SampleBatch` to new stack episodes.""" + # Set `input_compress_columns` to an empty `list` if `None`. + input_compress_columns = input_compress_columns or [] + + # TODO (simon): CHeck, if needed. It could possibly happen that a batch contains + # data from different episodes. Merging and resplitting the batch would then + # be the solution. + # Check, if batch comes actually from multiple episodes. + # episode_begin_indices = np.where(np.diff(np.hstack(batch["eps_id"])) != 0) + 1 + + # Define a container to collect episodes. + episodes = [] + # Loop over `SampleBatch`es in the `ray.data` batch (a dict). + for i, obs in enumerate(batch[schema[Columns.OBS]]): + + # If multi-agent we need to extract the agent ID. + # TODO (simon): Check, what happens with the module ID. + if is_multi_agent: + agent_id = ( + # The old stack uses "agent_index" instead of "agent_id". + batch[schema["agent_index"]][i][0] + if schema["agent_index"] in batch + else None + ) + else: + agent_id = None + + if is_multi_agent: + # TODO (simon): Add support for multi-agent episodes. + pass + else: + # Unpack observations, if needed. Note, observations could + # be either compressed by their entirety (the complete batch + # column) or individually (each column entry). + if isinstance(obs, str): + # Decompress the observations if we have a string, i.e. + # observations are compressed in their entirety. + obs = unpack_if_needed(obs) + # Convert to a list of arrays. This is needed as input by + # the `SingleAgentEpisode`. + obs = [obs[i, ...] for i in range(obs.shape[0])] + # Otherwise observations are only compressed inside of the + # batch column (if at all). + elif isinstance(obs, np.ndarray): + # Unpack observations, if they are compressed otherwise we + # simply convert to a list, which is needed by the + # `SingleAgentEpisode`. + obs = ( + unpack_if_needed(obs.tolist()) + if schema[Columns.OBS] in input_compress_columns + else obs.tolist() + ) + else: + raise TypeError( + f"Unknown observation type: {type(obs)}. When mapping " + "from old recorded `SampleBatches` batched " + "observations should be either of type `np.array` " + "or - if the column is compressed - of `str` type." + ) + + if schema[Columns.NEXT_OBS] in batch: + # Append the last `new_obs` to get the correct length of + # observations. + obs.append( + unpack_if_needed(batch[schema[Columns.NEXT_OBS]][i][-1]) + if schema[Columns.OBS] in input_compress_columns + else batch[schema[Columns.NEXT_OBS]][i][-1] + ) + else: + # Otherwise we duplicate the last observation. + obs.append(obs[-1]) + + # Check, if we have `done`, `truncated`, or `terminated`s in + # the batch. + if ( + schema[Columns.TRUNCATEDS] in batch + and schema[Columns.TERMINATEDS] in batch + ): + truncated = batch[schema[Columns.TRUNCATEDS]][i][-1] + terminated = batch[schema[Columns.TERMINATEDS]][i][-1] + elif ( + schema[Columns.TRUNCATEDS] in batch + and schema[Columns.TERMINATEDS] not in batch + ): + truncated = batch[schema[Columns.TRUNCATEDS]][i][-1] + terminated = False + elif ( + schema[Columns.TRUNCATEDS] not in batch + and schema[Columns.TERMINATEDS] in batch + ): + terminated = batch[schema[Columns.TERMINATEDS]][i][-1] + truncated = False + elif "done" in batch: + terminated = batch["done"][i][-1] + truncated = False + # Otherwise, if no `terminated`, nor `truncated` nor `done` + # is given, we consider the episode as terminated. + else: + terminated = True + truncated = False + + # Create a `SingleAgentEpisode`. + episode = SingleAgentEpisode( + # If the recorded episode has an ID we use this ID, + # otherwise we generate a new one. + id_=str(batch[schema[Columns.EPS_ID]][i][0]) + if schema[Columns.EPS_ID] in batch + else uuid.uuid4().hex, + agent_id=agent_id, + observations=obs, + infos=( + batch[schema[Columns.INFOS]][i] + if schema[Columns.INFOS] in batch + else [{}] * len(obs) + ), + # Actions might be (a) serialized. We unserialize them here. + actions=( + unpack_if_needed(batch[schema[Columns.ACTIONS]][i]) + if Columns.ACTIONS in input_compress_columns + else batch[schema[Columns.ACTIONS]][i] + ), + rewards=batch[schema[Columns.REWARDS]][i], + terminated=terminated, + truncated=truncated, + # TODO (simon): Results in zero-length episodes in connector. + # t_started=batch[Columns.T if Columns.T in batch else + # "unroll_id"][i][0], + # TODO (simon): Single-dimensional columns are not supported. + # Extra model outputs might be serialized. We unserialize them here + # if needed. + # TODO (simon): Check, if we need here also reconversion from + # JSONable in case of composite spaces. + extra_model_outputs={ + k: unpack_if_needed(v[i]) + if k in input_compress_columns + else v[i] + for k, v in batch.items() + if ( + k not in schema + and k not in schema.values() + and k not in ["dones", "agent_index", "type"] + ) + }, + len_lookback_buffer=0, + ) + # Numpy'ized, if necessary. + # TODO (simon, sven): Check, if we should convert all data to lists + # before. Right now only obs are lists. + if to_numpy: + episode.to_numpy() + episodes.append(episode) + # Note, `map_batches` expects a `Dict` as return value. + return {"episodes": episodes} diff --git a/lib/python3.12/site-packages/ray/rllib/offline/output_writer.py b/lib/python3.12/site-packages/ray/rllib/offline/output_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..b278ca90ae91fc1ecb8e5e9c402cf3ec40bfd29f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/output_writer.py @@ -0,0 +1,26 @@ +from ray.rllib.utils.annotations import PublicAPI, override +from ray.rllib.utils.typing import SampleBatchType + + +@PublicAPI +class OutputWriter: + """Writer API for saving experiences from policy evaluation.""" + + @PublicAPI + def write(self, sample_batch: SampleBatchType): + """Saves a batch of experiences. + + Args: + sample_batch: SampleBatch or MultiAgentBatch to save. + """ + raise NotImplementedError + + +@PublicAPI +class NoopOutput(OutputWriter): + """Output writer that discards its outputs.""" + + @override(OutputWriter) + def write(self, sample_batch: SampleBatchType): + # Do nothing. + pass diff --git a/lib/python3.12/site-packages/ray/rllib/offline/resource.py b/lib/python3.12/site-packages/ray/rllib/offline/resource.py new file mode 100644 index 0000000000000000000000000000000000000000..ff01b2fb7b89450e6658bab96673e4bcceba9776 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/resource.py @@ -0,0 +1,31 @@ +from typing import TYPE_CHECKING, Dict, List + +from ray.rllib.utils.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +DEFAULT_NUM_CPUS_PER_TASK = 0.5 + + +@PublicAPI +def get_offline_io_resource_bundles( + config: "AlgorithmConfig", +) -> List[Dict[str, float]]: + # DatasetReader is the only offline I/O component today that + # requires compute resources. + if config.input_ == "dataset": + input_config = config.input_config + # TODO (Kourosh): parallelism is use for reading the dataset, which defaults to + # num_workers. This logic here relies on the information that dataset reader + # will have the same logic. So to remove the information leakage, inside + # Algorithm config, we should set parallelism to num_workers if not specified + # and only deal with parallelism here or in dataset_reader.py. same thing is + # true with cpus_per_task. + parallelism = input_config.get("parallelism", config.get("num_env_runners", 1)) + cpus_per_task = input_config.get( + "num_cpus_per_read_task", DEFAULT_NUM_CPUS_PER_TASK + ) + return [{"CPU": cpus_per_task} for _ in range(parallelism)] + else: + return [] diff --git a/lib/python3.12/site-packages/ray/rllib/offline/shuffled_input.py b/lib/python3.12/site-packages/ray/rllib/offline/shuffled_input.py new file mode 100644 index 0000000000000000000000000000000000000000..a6633fa76e5c972582fb5e98d83c50f9651f3107 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/shuffled_input.py @@ -0,0 +1,42 @@ +import logging +import random + +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.utils.annotations import DeveloperAPI, override +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class ShuffledInput(InputReader): + """Randomizes data over a sliding window buffer of N batches. + + This increases the randomization of the data, which is useful if the + batches were not in random order to start with. + """ + + @DeveloperAPI + def __init__(self, child: InputReader, n: int = 0): + """Initializes a ShuffledInput instance. + + Args: + child: child input reader to shuffle. + n: If positive, shuffle input over this many batches. + """ + self.n = n + self.child = child + self.buffer = [] + + @override(InputReader) + def next(self) -> SampleBatchType: + if self.n <= 1: + return self.child.next() + if len(self.buffer) < self.n: + logger.info("Filling shuffle buffer to {} batches".format(self.n)) + while len(self.buffer) < self.n: + self.buffer.append(self.child.next()) + logger.info("Shuffle buffer filled") + i = random.randint(0, len(self.buffer) - 1) + self.buffer[i] = self.child.next() + return random.choice(self.buffer) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/wis_estimator.py b/lib/python3.12/site-packages/ray/rllib/offline/wis_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..d207d7e904281cce1a101cb035765217f8752fd7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/wis_estimator.py @@ -0,0 +1,13 @@ +from ray._common.deprecation import Deprecated +from ray.rllib.offline.estimators.weighted_importance_sampling import ( + WeightedImportanceSampling, +) + + +@Deprecated( + new="ray.rllib.offline.estimators.weighted_importance_sampling::" + "WeightedImportanceSampling", + error=True, +) +class WeightedImportanceSamplingEstimator(WeightedImportanceSampling): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__init__.py b/lib/python3.12/site-packages/ray/rllib/policy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8e164d1dcf856f446682fcd744f4d31de2ee9edf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/__init__.py @@ -0,0 +1,13 @@ +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.policy_template import build_policy_class +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.policy.tf_policy_template import build_tf_policy +from ray.rllib.policy.torch_policy import TorchPolicy + +__all__ = [ + "Policy", + "TFPolicy", + "TorchPolicy", + "build_policy_class", + "build_tf_policy", +] diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f61955037e994339fcdfbea54dca06898d07d04 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ce6394d084692b6ec1a6c0297fa3165da0f52e2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1c2bdb0c7e0ec54f6c687029d9700914970131e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/dynamic_tf_policy_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cebbf6b18e2ff15868f421386b22f75124040c25 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5d648a32648c0ddad2fa23469152c4f12289440 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/eager_tf_policy_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7bba4ed25e72204d47298f786c55826d5f7cfcb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_map.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_map.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a774064a0960f88dded39cddc6ec7193bd77353 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_map.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_template.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_template.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..420ac3c3eded17828a3441d4065cdfcef89a24f2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/policy_template.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/rnn_sequencing.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/rnn_sequencing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7af1f1d196e5fe441ae22cbdb7a6535f87e1c77 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/rnn_sequencing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/sample_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/sample_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d89f4517751974c2990d758d33e917930e626a81 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/sample_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_mixins.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_mixins.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b32fb49c5f96c9b5ee35bf205e4eeec05ef7447 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_mixins.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b89c4ce0fc0841daface3258eaa9ac0d8609c37e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy_template.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy_template.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2638a3e7f03e52eee442be46c6813ab3d6cb81c9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/tf_policy_template.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_mixins.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_mixins.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8524165947cde0b2956abcf42d5817b2dc1f7d91 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_mixins.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71a3e5e6187959ecc47bc46cffd8230bbd70d120 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..945a4b1b83c283e7b31b31f7a73ac6e984f6b10f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/torch_policy_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/view_requirement.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/view_requirement.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a17f5e10c5a97bbb9fb56e4f1d206ed78fe39ac Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/policy/__pycache__/view_requirement.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..172860c17931d35bad03ae45cd44e20e45ba6f2c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy.py @@ -0,0 +1,1359 @@ +import logging +import re +from collections import OrderedDict, namedtuple +from typing import Callable, Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym +import tree # pip install dm_tree + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space +from ray.rllib.utils.tf_utils import get_placeholder +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + +# Variable scope in which created variables will be placed under. +TOWER_SCOPE_NAME = "tower" + + +@OldAPIStack +class DynamicTFPolicy(TFPolicy): + """A TFPolicy that auto-defines placeholders dynamically at runtime. + + Do not sub-class this class directly (neither should you sub-class + TFPolicy), but rather use rllib.policy.tf_policy_template.build_tf_policy + to generate your custom tf (graph-mode or eager) Policy classes. + """ + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + loss_fn: Callable[ + [Policy, ModelV2, Type[TFActionDistribution], SampleBatch], TensorType + ], + *, + stats_fn: Optional[ + Callable[[Policy, SampleBatch], Dict[str, TensorType]] + ] = None, + grad_stats_fn: Optional[ + Callable[[Policy, SampleBatch, ModelGradients], Dict[str, TensorType]] + ] = None, + before_loss_init: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], None + ] + ] = None, + make_model: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], + ModelV2, + ] + ] = None, + action_sampler_fn: Optional[ + Callable[ + [TensorType, List[TensorType]], + Union[ + Tuple[TensorType, TensorType], + Tuple[TensorType, TensorType, TensorType, List[TensorType]], + ], + ] + ] = None, + action_distribution_fn: Optional[ + Callable[ + [Policy, ModelV2, TensorType, TensorType, TensorType], + Tuple[TensorType, type, List[TensorType]], + ] + ] = None, + existing_inputs: Optional[Dict[str, "tf1.placeholder"]] = None, + existing_model: Optional[ModelV2] = None, + get_batch_divisibility_req: Optional[Callable[[Policy], int]] = None, + obs_include_prev_action_reward=DEPRECATED_VALUE, + ): + """Initializes a DynamicTFPolicy instance. + + Initialization of this class occurs in two phases and defines the + static graph. + + Phase 1: The model is created and model variables are initialized. + + Phase 2: A fake batch of data is created, sent to the trajectory + postprocessor, and then used to create placeholders for the loss + function. The loss and stats functions are initialized with these + placeholders. + + Args: + observation_space: Observation space of the policy. + action_space: Action space of the policy. + config: Policy-specific configuration data. + loss_fn: Function that returns a loss tensor for the policy graph. + stats_fn: Optional callable that - given the policy and batch + input tensors - returns a dict mapping str to TF ops. + These ops are fetched from the graph after loss calculations + and the resulting values can be found in the results dict + returned by e.g. `Algorithm.train()` or in tensorboard (if TB + logging is enabled). + grad_stats_fn: Optional callable that - given the policy, batch + input tensors, and calculated loss gradient tensors - returns + a dict mapping str to TF ops. These ops are fetched from the + graph after loss and gradient calculations and the resulting + values can be found in the results dict returned by e.g. + `Algorithm.train()` or in tensorboard (if TB logging is + enabled). + before_loss_init: Optional function to run prior to + loss init that takes the same arguments as __init__. + make_model: Optional function that returns a ModelV2 object + given policy, obs_space, action_space, and policy config. + All policy variables should be created in this function. If not + specified, a default model will be created. + action_sampler_fn: A callable returning either a sampled action and + its log-likelihood or a sampled action, its log-likelihood, + action distribution inputs and updated state given Policy, + ModelV2, observation inputs, explore, and is_training. + Provide `action_sampler_fn` if you would like to have full + control over the action computation step, including the + model forward pass, possible sampling from a distribution, + and exploration logic. + Note: If `action_sampler_fn` is given, `action_distribution_fn` + must be None. If both `action_sampler_fn` and + `action_distribution_fn` are None, RLlib will simply pass + inputs through `self.model` to get distribution inputs, create + the distribution object, sample from it, and apply some + exploration logic to the results. + The callable takes as inputs: Policy, ModelV2, obs_batch, + state_batches (optional), seq_lens (optional), + prev_actions_batch (optional), prev_rewards_batch (optional), + explore, and is_training. + action_distribution_fn: A callable returning distribution inputs + (parameters), a dist-class to generate an action distribution + object from, and internal-state outputs (or an empty list if + not applicable). + Provide `action_distribution_fn` if you would like to only + customize the model forward pass call. The resulting + distribution parameters are then used by RLlib to create a + distribution object, sample from it, and execute any + exploration logic. + Note: If `action_distribution_fn` is given, `action_sampler_fn` + must be None. If both `action_sampler_fn` and + `action_distribution_fn` are None, RLlib will simply pass + inputs through `self.model` to get distribution inputs, create + the distribution object, sample from it, and apply some + exploration logic to the results. + The callable takes as inputs: Policy, ModelV2, input_dict, + explore, timestep, is_training. + existing_inputs: When copying a policy, this specifies an existing + dict of placeholders to use instead of defining new ones. + existing_model: When copying a policy, this specifies an existing + model to clone and share weights with. + get_batch_divisibility_req: Optional callable that returns the + divisibility requirement for sample batches. If None, will + assume a value of 1. + """ + if obs_include_prev_action_reward != DEPRECATED_VALUE: + deprecation_warning(old="obs_include_prev_action_reward", error=True) + self.observation_space = obs_space + self.action_space = action_space + self.config = config + self.framework = "tf" + self._loss_fn = loss_fn + self._stats_fn = stats_fn + self._grad_stats_fn = grad_stats_fn + self._seq_lens = None + self._is_tower = existing_inputs is not None + + dist_class = None + if action_sampler_fn or action_distribution_fn: + if not make_model: + raise ValueError( + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given" + ) + else: + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"] + ) + + # Setup self.model. + if existing_model: + if isinstance(existing_model, list): + self.model = existing_model[0] + # TODO: (sven) hack, but works for `target_[q_]?model`. + for i in range(1, len(existing_model)): + setattr(self, existing_model[i][0], existing_model[i][1]) + elif make_model: + self.model = make_model(self, obs_space, action_space, config) + else: + self.model = ModelCatalog.get_model_v2( + obs_space=obs_space, + action_space=action_space, + num_outputs=logit_dim, + model_config=self.config["model"], + framework="tf", + ) + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + + # Input placeholders already given -> Use these. + if existing_inputs: + self._state_inputs = [ + v for k, v in existing_inputs.items() if k.startswith("state_in_") + ] + # Placeholder for RNN time-chunk valid lengths. + if self._state_inputs: + self._seq_lens = existing_inputs[SampleBatch.SEQ_LENS] + # Create new input placeholders. + else: + self._state_inputs = [ + get_placeholder( + space=vr.space, + time_axis=not isinstance(vr.shift, int), + name=k, + ) + for k, vr in self.model.view_requirements.items() + if k.startswith("state_in_") + ] + # Placeholder for RNN time-chunk valid lengths. + if self._state_inputs: + self._seq_lens = tf1.placeholder( + dtype=tf.int32, shape=[None], name="seq_lens" + ) + + # Use default settings. + # Add NEXT_OBS, STATE_IN_0.., and others. + self.view_requirements = self._get_default_view_requirements() + # Combine view_requirements for Model and Policy. + self.view_requirements.update(self.model.view_requirements) + # Disable env-info placeholder. + if SampleBatch.INFOS in self.view_requirements: + self.view_requirements[SampleBatch.INFOS].used_for_training = False + + # Setup standard placeholders. + if self._is_tower: + timestep = existing_inputs["timestep"] + explore = False + self._input_dict, self._dummy_batch = self._get_input_dict_and_dummy_batch( + self.view_requirements, existing_inputs + ) + else: + if not self.config.get("_disable_action_flattening"): + action_ph = ModelCatalog.get_action_placeholder(action_space) + prev_action_ph = {} + if SampleBatch.PREV_ACTIONS not in self.view_requirements: + prev_action_ph = { + SampleBatch.PREV_ACTIONS: ModelCatalog.get_action_placeholder( + action_space, "prev_action" + ) + } + ( + self._input_dict, + self._dummy_batch, + ) = self._get_input_dict_and_dummy_batch( + self.view_requirements, + dict({SampleBatch.ACTIONS: action_ph}, **prev_action_ph), + ) + else: + ( + self._input_dict, + self._dummy_batch, + ) = self._get_input_dict_and_dummy_batch(self.view_requirements, {}) + # Placeholder for (sampling steps) timestep (int). + timestep = tf1.placeholder_with_default( + tf.zeros((), dtype=tf.int64), (), name="timestep" + ) + # Placeholder for `is_exploring` flag. + explore = tf1.placeholder_with_default(True, (), name="is_exploring") + + # Placeholder for `is_training` flag. + self._input_dict.set_training(self._get_is_training_placeholder()) + + # Multi-GPU towers do not need any action computing/exploration + # graphs. + sampled_action = None + sampled_action_logp = None + dist_inputs = None + extra_action_fetches = {} + self._state_out = None + if not self._is_tower: + # Create the Exploration object to use for this Policy. + self.exploration = self._create_exploration() + + # Fully customized action generation (e.g., custom policy). + if action_sampler_fn: + action_sampler_outputs = action_sampler_fn( + self, + self.model, + obs_batch=self._input_dict[SampleBatch.CUR_OBS], + state_batches=self._state_inputs, + seq_lens=self._seq_lens, + prev_action_batch=self._input_dict.get(SampleBatch.PREV_ACTIONS), + prev_reward_batch=self._input_dict.get(SampleBatch.PREV_REWARDS), + explore=explore, + is_training=self._input_dict.is_training, + ) + if len(action_sampler_outputs) == 4: + ( + sampled_action, + sampled_action_logp, + dist_inputs, + self._state_out, + ) = action_sampler_outputs + else: + dist_inputs = None + self._state_out = [] + sampled_action, sampled_action_logp = action_sampler_outputs + # Distribution generation is customized, e.g., DQN, DDPG. + else: + if action_distribution_fn: + # Try new action_distribution_fn signature, supporting + # state_batches and seq_lens. + in_dict = self._input_dict + try: + ( + dist_inputs, + dist_class, + self._state_out, + ) = action_distribution_fn( + self, + self.model, + input_dict=in_dict, + state_batches=self._state_inputs, + seq_lens=self._seq_lens, + explore=explore, + timestep=timestep, + is_training=in_dict.is_training, + ) + # Trying the old way (to stay backward compatible). + # TODO: Remove in future. + except TypeError as e: + if ( + "positional argument" in e.args[0] + or "unexpected keyword argument" in e.args[0] + ): + ( + dist_inputs, + dist_class, + self._state_out, + ) = action_distribution_fn( + self, + self.model, + obs_batch=in_dict[SampleBatch.CUR_OBS], + state_batches=self._state_inputs, + seq_lens=self._seq_lens, + prev_action_batch=in_dict.get(SampleBatch.PREV_ACTIONS), + prev_reward_batch=in_dict.get(SampleBatch.PREV_REWARDS), + explore=explore, + is_training=in_dict.is_training, + ) + else: + raise e + + # Default distribution generation behavior: + # Pass through model. E.g., PG, PPO. + else: + if isinstance(self.model, tf.keras.Model): + dist_inputs, self._state_out, extra_action_fetches = self.model( + self._input_dict + ) + else: + dist_inputs, self._state_out = self.model(self._input_dict) + + action_dist = dist_class(dist_inputs, self.model) + + # Using exploration to get final action (e.g. via sampling). + ( + sampled_action, + sampled_action_logp, + ) = self.exploration.get_exploration_action( + action_distribution=action_dist, timestep=timestep, explore=explore + ) + + if dist_inputs is not None: + extra_action_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + + if sampled_action_logp is not None: + extra_action_fetches[SampleBatch.ACTION_LOGP] = sampled_action_logp + extra_action_fetches[SampleBatch.ACTION_PROB] = tf.exp( + tf.cast(sampled_action_logp, tf.float32) + ) + + # Phase 1 init. + sess = tf1.get_default_session() or tf1.Session( + config=tf1.ConfigProto(**self.config["tf_session_args"]) + ) + + batch_divisibility_req = ( + get_batch_divisibility_req(self) + if callable(get_batch_divisibility_req) + else (get_batch_divisibility_req or 1) + ) + + prev_action_input = ( + self._input_dict[SampleBatch.PREV_ACTIONS] + if SampleBatch.PREV_ACTIONS in self._input_dict.accessed_keys + else None + ) + prev_reward_input = ( + self._input_dict[SampleBatch.PREV_REWARDS] + if SampleBatch.PREV_REWARDS in self._input_dict.accessed_keys + else None + ) + + super().__init__( + observation_space=obs_space, + action_space=action_space, + config=config, + sess=sess, + obs_input=self._input_dict[SampleBatch.OBS], + action_input=self._input_dict[SampleBatch.ACTIONS], + sampled_action=sampled_action, + sampled_action_logp=sampled_action_logp, + dist_inputs=dist_inputs, + dist_class=dist_class, + loss=None, # dynamically initialized on run + loss_inputs=[], + model=self.model, + state_inputs=self._state_inputs, + state_outputs=self._state_out, + prev_action_input=prev_action_input, + prev_reward_input=prev_reward_input, + seq_lens=self._seq_lens, + max_seq_len=config["model"]["max_seq_len"], + batch_divisibility_req=batch_divisibility_req, + explore=explore, + timestep=timestep, + ) + + # Phase 2 init. + if before_loss_init is not None: + before_loss_init(self, obs_space, action_space, config) + if hasattr(self, "_extra_action_fetches"): + self._extra_action_fetches.update(extra_action_fetches) + else: + self._extra_action_fetches = extra_action_fetches + + # Loss initialization and model/postprocessing test calls. + if not self._is_tower: + self._initialize_loss_from_dummy_batch(auto_remove_unneeded_view_reqs=True) + + # Create MultiGPUTowerStacks, if we have at least one actual + # GPU or >1 CPUs (fake GPUs). + if len(self.devices) > 1 or any("gpu" in d for d in self.devices): + # Per-GPU graph copies created here must share vars with the + # policy. Therefore, `reuse` is set to tf1.AUTO_REUSE because + # Adam nodes are created after all of the device copies are + # created. + with tf1.variable_scope("", reuse=tf1.AUTO_REUSE): + self.multi_gpu_tower_stacks = [ + TFMultiGPUTowerStack(policy=self) + for i in range(self.config.get("num_multi_gpu_tower_stacks", 1)) + ] + + # Initialize again after loss and tower init. + self.get_session().run(tf1.global_variables_initializer()) + + @override(TFPolicy) + def copy(self, existing_inputs: List[Tuple[str, "tf1.placeholder"]]) -> TFPolicy: + """Creates a copy of self using existing input placeholders.""" + + flat_loss_inputs = tree.flatten(self._loss_input_dict) + flat_loss_inputs_no_rnn = tree.flatten(self._loss_input_dict_no_rnn) + + # Note that there might be RNN state inputs at the end of the list + if len(flat_loss_inputs) != len(existing_inputs): + raise ValueError( + "Tensor list mismatch", + self._loss_input_dict, + self._state_inputs, + existing_inputs, + ) + for i, v in enumerate(flat_loss_inputs_no_rnn): + if v.shape.as_list() != existing_inputs[i].shape.as_list(): + raise ValueError( + "Tensor shape mismatch", i, v.shape, existing_inputs[i].shape + ) + # By convention, the loss inputs are followed by state inputs and then + # the seq len tensor. + rnn_inputs = [] + for i in range(len(self._state_inputs)): + rnn_inputs.append( + ( + "state_in_{}".format(i), + existing_inputs[len(flat_loss_inputs_no_rnn) + i], + ) + ) + if rnn_inputs: + rnn_inputs.append((SampleBatch.SEQ_LENS, existing_inputs[-1])) + existing_inputs_unflattened = tree.unflatten_as( + self._loss_input_dict_no_rnn, + existing_inputs[: len(flat_loss_inputs_no_rnn)], + ) + input_dict = OrderedDict( + [("is_exploring", self._is_exploring), ("timestep", self._timestep)] + + [ + (k, existing_inputs_unflattened[k]) + for i, k in enumerate(self._loss_input_dict_no_rnn.keys()) + ] + + rnn_inputs + ) + + instance = self.__class__( + self.observation_space, + self.action_space, + self.config, + existing_inputs=input_dict, + existing_model=[ + self.model, + # Deprecated: Target models should all reside under + # `policy.target_model` now. + ("target_q_model", getattr(self, "target_q_model", None)), + ("target_model", getattr(self, "target_model", None)), + ], + ) + + instance._loss_input_dict = input_dict + losses = instance._do_loss_init(SampleBatch(input_dict)) + loss_inputs = [ + (k, existing_inputs_unflattened[k]) + for i, k in enumerate(self._loss_input_dict_no_rnn.keys()) + ] + + TFPolicy._initialize_loss(instance, losses, loss_inputs) + if instance._grad_stats_fn: + instance._stats_fetches.update( + instance._grad_stats_fn(instance, input_dict, instance._grads) + ) + return instance + + @override(Policy) + def get_initial_state(self) -> List[TensorType]: + if self.model: + return self.model.get_initial_state() + else: + return [] + + @override(Policy) + def load_batch_into_buffer( + self, + batch: SampleBatch, + buffer_index: int = 0, + ) -> int: + # Set the is_training flag of the batch. + batch.set_training(True) + + # Shortcut for 1 CPU only: Store batch in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + self._loaded_single_cpu_batch = batch + return len(batch) + + input_dict = self._get_loss_inputs_dict(batch, shuffle=False) + data_keys = tree.flatten(self._loss_input_dict_no_rnn) + if self._state_inputs: + state_keys = self._state_inputs + [self._seq_lens] + else: + state_keys = [] + inputs = [input_dict[k] for k in data_keys] + state_inputs = [input_dict[k] for k in state_keys] + + return self.multi_gpu_tower_stacks[buffer_index].load_data( + sess=self.get_session(), + inputs=inputs, + state_inputs=state_inputs, + num_grad_updates=batch.num_grad_updates, + ) + + @override(Policy) + def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int: + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + return ( + len(self._loaded_single_cpu_batch) + if self._loaded_single_cpu_batch is not None + else 0 + ) + + return self.multi_gpu_tower_stacks[buffer_index].num_tuples_loaded + + @override(Policy) + def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0): + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + if self._loaded_single_cpu_batch is None: + raise ValueError( + "Must call Policy.load_batch_into_buffer() before " + "Policy.learn_on_loaded_batch()!" + ) + # Get the correct slice of the already loaded batch to use, + # based on offset and batch size. + batch_size = self.config.get("minibatch_size") + if batch_size is None: + batch_size = self.config.get( + "sgd_minibatch_size", self.config["train_batch_size"] + ) + if batch_size >= len(self._loaded_single_cpu_batch): + sliced_batch = self._loaded_single_cpu_batch + else: + sliced_batch = self._loaded_single_cpu_batch.slice( + start=offset, end=offset + batch_size + ) + return self.learn_on_batch(sliced_batch) + + tower_stack = self.multi_gpu_tower_stacks[buffer_index] + results = tower_stack.optimize(self.get_session(), offset) + self.num_grad_updates += 1 + + results.update( + { + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates - 1 - (tower_stack.num_grad_updates or 0) + ), + } + ) + + return results + + def _get_input_dict_and_dummy_batch(self, view_requirements, existing_inputs): + """Creates input_dict and dummy_batch for loss initialization. + + Used for managing the Policy's input placeholders and for loss + initialization. + Input_dict: Str -> tf.placeholders, dummy_batch: str -> np.arrays. + + Args: + view_requirements: The view requirements dict. + existing_inputs (Dict[str, tf.placeholder]): A dict of already + existing placeholders. + + Returns: + Tuple[Dict[str, tf.placeholder], Dict[str, np.ndarray]]: The + input_dict/dummy_batch tuple. + """ + input_dict = {} + for view_col, view_req in view_requirements.items(): + # Point state_in to the already existing self._state_inputs. + mo = re.match(r"state_in_(\d+)", view_col) + if mo is not None: + input_dict[view_col] = self._state_inputs[int(mo.group(1))] + # State-outs (no placeholders needed). + elif view_col.startswith("state_out_"): + continue + # Skip action dist inputs placeholder (do later). + elif view_col == SampleBatch.ACTION_DIST_INPUTS: + continue + # This is a tower: Input placeholders already exist. + elif view_col in existing_inputs: + input_dict[view_col] = existing_inputs[view_col] + # All others. + else: + time_axis = not isinstance(view_req.shift, int) + if view_req.used_for_training: + # Create a +time-axis placeholder if the shift is not an + # int (range or list of ints). + # Do not flatten actions if action flattening disabled. + if self.config.get("_disable_action_flattening") and view_col in [ + SampleBatch.ACTIONS, + SampleBatch.PREV_ACTIONS, + ]: + flatten = False + # Do not flatten observations if no preprocessor API used. + elif ( + view_col in [SampleBatch.OBS, SampleBatch.NEXT_OBS] + and self.config["_disable_preprocessor_api"] + ): + flatten = False + # Flatten everything else. + else: + flatten = True + input_dict[view_col] = get_placeholder( + space=view_req.space, + name=view_col, + time_axis=time_axis, + flatten=flatten, + ) + dummy_batch = self._get_dummy_batch_from_view_requirements(batch_size=32) + + return SampleBatch(input_dict, seq_lens=self._seq_lens), dummy_batch + + @override(Policy) + def _initialize_loss_from_dummy_batch( + self, auto_remove_unneeded_view_reqs: bool = True, stats_fn=None + ) -> None: + # Create the optimizer/exploration optimizer here. Some initialization + # steps (e.g. exploration postprocessing) may need this. + if not self._optimizers: + self._optimizers = force_list(self.optimizer()) + # Backward compatibility. + self._optimizer = self._optimizers[0] + + # Test calls depend on variable init, so initialize model first. + self.get_session().run(tf1.global_variables_initializer()) + + # Fields that have not been accessed are not needed for action + # computations -> Tag them as `used_for_compute_actions=False`. + for key, view_req in self.view_requirements.items(): + if ( + not key.startswith("state_in_") + and key not in self._input_dict.accessed_keys + ): + view_req.used_for_compute_actions = False + for key, value in self._extra_action_fetches.items(): + self._dummy_batch[key] = get_dummy_batch_for_space( + gym.spaces.Box( + -1.0, 1.0, shape=value.shape.as_list()[1:], dtype=value.dtype.name + ), + batch_size=len(self._dummy_batch), + ) + self._input_dict[key] = get_placeholder(value=value, name=key) + if key not in self.view_requirements: + logger.info("Adding extra-action-fetch `{}` to view-reqs.".format(key)) + self.view_requirements[key] = ViewRequirement( + space=gym.spaces.Box( + -1.0, + 1.0, + shape=value.shape.as_list()[1:], + dtype=value.dtype.name, + ), + used_for_compute_actions=False, + ) + dummy_batch = self._dummy_batch + + logger.info("Testing `postprocess_trajectory` w/ dummy batch.") + self.exploration.postprocess_trajectory(self, dummy_batch, self.get_session()) + _ = self.postprocess_trajectory(dummy_batch) + # Add new columns automatically to (loss) input_dict. + for key in dummy_batch.added_keys: + if key not in self._input_dict: + self._input_dict[key] = get_placeholder( + value=dummy_batch[key], name=key + ) + if key not in self.view_requirements: + self.view_requirements[key] = ViewRequirement( + space=gym.spaces.Box( + -1.0, + 1.0, + shape=dummy_batch[key].shape[1:], + dtype=dummy_batch[key].dtype, + ), + used_for_compute_actions=False, + ) + + train_batch = SampleBatch( + dict(self._input_dict, **self._loss_input_dict), + _is_training=True, + ) + + if self._state_inputs: + train_batch[SampleBatch.SEQ_LENS] = self._seq_lens + self._loss_input_dict.update( + {SampleBatch.SEQ_LENS: train_batch[SampleBatch.SEQ_LENS]} + ) + + self._loss_input_dict.update(dict(train_batch)) + + if log_once("loss_init"): + logger.debug( + "Initializing loss function with dummy input:\n\n{}\n".format( + summarize(train_batch) + ) + ) + + losses = self._do_loss_init(train_batch) + + all_accessed_keys = ( + train_batch.accessed_keys + | dummy_batch.accessed_keys + | dummy_batch.added_keys + | set(self.model.view_requirements.keys()) + ) + + TFPolicy._initialize_loss( + self, + losses, + [(k, v) for k, v in train_batch.items() if k in all_accessed_keys] + + ( + [(SampleBatch.SEQ_LENS, train_batch[SampleBatch.SEQ_LENS])] + if SampleBatch.SEQ_LENS in train_batch + else [] + ), + ) + + if "is_training" in self._loss_input_dict: + del self._loss_input_dict["is_training"] + + # Call the grads stats fn. + # TODO: (sven) rename to simply stats_fn to match eager and torch. + if self._grad_stats_fn: + self._stats_fetches.update( + self._grad_stats_fn(self, train_batch, self._grads) + ) + + # Add new columns automatically to view-reqs. + if auto_remove_unneeded_view_reqs: + # Add those needed for postprocessing and training. + all_accessed_keys = train_batch.accessed_keys | dummy_batch.accessed_keys + # Tag those only needed for post-processing (with some exceptions). + for key in dummy_batch.accessed_keys: + if ( + key not in train_batch.accessed_keys + and key not in self.model.view_requirements + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + SampleBatch.OBS_EMBEDS, + ] + ): + if key in self.view_requirements: + self.view_requirements[key].used_for_training = False + if key in self._loss_input_dict: + del self._loss_input_dict[key] + # Remove those not needed at all (leave those that are needed + # by Sampler to properly execute sample collection). + # Also always leave TERMINATEDS, TRUNCATEDS, REWARDS, and INFOS, + # no matter what. + for key in list(self.view_requirements.keys()): + if ( + key not in all_accessed_keys + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + ] + and key not in self.model.view_requirements + ): + # If user deleted this key manually in postprocessing + # fn, warn about it and do not remove from + # view-requirements. + if key in dummy_batch.deleted_keys: + logger.warning( + "SampleBatch key '{}' was deleted manually in " + "postprocessing function! RLlib will " + "automatically remove non-used items from the " + "data stream. Remove the `del` from your " + "postprocessing function.".format(key) + ) + # If we are not writing output to disk, safe to erase + # this key to save space in the sample batch. + elif self.config["output"] is None: + del self.view_requirements[key] + + if key in self._loss_input_dict: + del self._loss_input_dict[key] + # Add those data_cols (again) that are missing and have + # dependencies by view_cols. + for key in list(self.view_requirements.keys()): + vr = self.view_requirements[key] + if ( + vr.data_col is not None + and vr.data_col not in self.view_requirements + ): + used_for_training = vr.data_col in train_batch.accessed_keys + self.view_requirements[vr.data_col] = ViewRequirement( + space=vr.space, used_for_training=used_for_training + ) + + self._loss_input_dict_no_rnn = { + k: v + for k, v in self._loss_input_dict.items() + if (v not in self._state_inputs and v != self._seq_lens) + } + + def _do_loss_init(self, train_batch: SampleBatch): + losses = self._loss_fn(self, self.model, self.dist_class, train_batch) + losses = force_list(losses) + if self._stats_fn: + self._stats_fetches.update(self._stats_fn(self, train_batch)) + # Override the update ops to be those of the model. + self._update_ops = [] + if not isinstance(self.model, tf.keras.Model): + self._update_ops = self.model.update_ops() + return losses + + +@OldAPIStack +class TFMultiGPUTowerStack: + """Optimizer that runs in parallel across multiple local devices. + + TFMultiGPUTowerStack automatically splits up and loads training data + onto specified local devices (e.g. GPUs) with `load_data()`. During a call + to `optimize()`, the devices compute gradients over slices of the data in + parallel. The gradients are then averaged and applied to the shared + weights. + + The data loaded is pinned in device memory until the next call to + `load_data`, so you can make multiple passes (possibly in randomized order) + over the same data once loaded. + + This is similar to tf1.train.SyncReplicasOptimizer, but works within a + single TensorFlow graph, i.e. implements in-graph replicated training: + + https://www.tensorflow.org/api_docs/python/tf/train/SyncReplicasOptimizer + """ + + def __init__( + self, + # Deprecated. + optimizer=None, + devices=None, + input_placeholders=None, + rnn_inputs=None, + max_per_device_batch_size=None, + build_graph=None, + grad_norm_clipping=None, + # Use only `policy` argument from here on. + policy: TFPolicy = None, + ): + """Initializes a TFMultiGPUTowerStack instance. + + Args: + policy: The TFPolicy object that this tower stack + belongs to. + """ + # Obsoleted usage, use only `policy` arg from here on. + if policy is None: + deprecation_warning( + old="TFMultiGPUTowerStack(...)", + new="TFMultiGPUTowerStack(policy=[Policy])", + error=True, + ) + self.policy = None + self.optimizers = optimizer + self.devices = devices + self.max_per_device_batch_size = max_per_device_batch_size + self.policy_copy = build_graph + else: + self.policy: TFPolicy = policy + self.optimizers: List[LocalOptimizer] = self.policy._optimizers + self.devices = self.policy.devices + self.max_per_device_batch_size = ( + max_per_device_batch_size + or policy.config.get( + "minibatch_size", policy.config.get("train_batch_size", 999999) + ) + ) // len(self.devices) + input_placeholders = tree.flatten(self.policy._loss_input_dict_no_rnn) + rnn_inputs = [] + if self.policy._state_inputs: + rnn_inputs = self.policy._state_inputs + [self.policy._seq_lens] + grad_norm_clipping = self.policy.config.get("grad_clip") + self.policy_copy = self.policy.copy + + assert len(self.devices) > 1 or "gpu" in self.devices[0] + self.loss_inputs = input_placeholders + rnn_inputs + + shared_ops = tf1.get_collection( + tf1.GraphKeys.UPDATE_OPS, scope=tf1.get_variable_scope().name + ) + + # Then setup the per-device loss graphs that use the shared weights + self._batch_index = tf1.placeholder(tf.int32, name="batch_index") + + # Dynamic batch size, which may be shrunk if there isn't enough data + self._per_device_batch_size = tf1.placeholder( + tf.int32, name="per_device_batch_size" + ) + self._loaded_per_device_batch_size = max_per_device_batch_size + + # When loading RNN input, we dynamically determine the max seq len + self._max_seq_len = tf1.placeholder(tf.int32, name="max_seq_len") + self._loaded_max_seq_len = 1 + + device_placeholders = [[] for _ in range(len(self.devices))] + + for t in tree.flatten(self.loss_inputs): + # Split on the CPU in case the data doesn't fit in GPU memory. + with tf.device("/cpu:0"): + splits = tf.split(t, len(self.devices)) + for i, d in enumerate(self.devices): + device_placeholders[i].append(splits[i]) + + self._towers = [] + for tower_i, (device, placeholders) in enumerate( + zip(self.devices, device_placeholders) + ): + self._towers.append( + self._setup_device( + tower_i, device, placeholders, len(tree.flatten(input_placeholders)) + ) + ) + + if self.policy.config["_tf_policy_handles_more_than_one_loss"]: + avgs = [] + for i, optim in enumerate(self.optimizers): + avg = _average_gradients([t.grads[i] for t in self._towers]) + if grad_norm_clipping: + clipped = [] + for grad, _ in avg: + clipped.append(grad) + clipped, _ = tf.clip_by_global_norm(clipped, grad_norm_clipping) + for i, (grad, var) in enumerate(avg): + avg[i] = (clipped[i], var) + avgs.append(avg) + + # Gather update ops for any batch norm layers. + # TODO(ekl) here we + # will use all the ops found which won't work for DQN / DDPG, but + # those aren't supported with multi-gpu right now anyways. + self._update_ops = tf1.get_collection( + tf1.GraphKeys.UPDATE_OPS, scope=tf1.get_variable_scope().name + ) + for op in shared_ops: + self._update_ops.remove(op) # only care about tower update ops + if self._update_ops: + logger.debug( + "Update ops to run on apply gradient: {}".format(self._update_ops) + ) + + with tf1.control_dependencies(self._update_ops): + self._train_op = tf.group( + [o.apply_gradients(a) for o, a in zip(self.optimizers, avgs)] + ) + else: + avg = _average_gradients([t.grads for t in self._towers]) + if grad_norm_clipping: + clipped = [] + for grad, _ in avg: + clipped.append(grad) + clipped, _ = tf.clip_by_global_norm(clipped, grad_norm_clipping) + for i, (grad, var) in enumerate(avg): + avg[i] = (clipped[i], var) + + # Gather update ops for any batch norm layers. + # TODO(ekl) here we + # will use all the ops found which won't work for DQN / DDPG, but + # those aren't supported with multi-gpu right now anyways. + self._update_ops = tf1.get_collection( + tf1.GraphKeys.UPDATE_OPS, scope=tf1.get_variable_scope().name + ) + for op in shared_ops: + self._update_ops.remove(op) # only care about tower update ops + if self._update_ops: + logger.debug( + "Update ops to run on apply gradient: {}".format(self._update_ops) + ) + + with tf1.control_dependencies(self._update_ops): + self._train_op = self.optimizers[0].apply_gradients(avg) + + # The lifetime number of gradient updates that the policy having sent + # some data (SampleBatchType) into this tower stack's GPU buffer(s) has already + # undergone. + self.num_grad_updates = 0 + + def load_data(self, sess, inputs, state_inputs, num_grad_updates=None): + """Bulk loads the specified inputs into device memory. + + The shape of the inputs must conform to the shapes of the input + placeholders this optimizer was constructed with. + + The data is split equally across all the devices. If the data is not + evenly divisible by the batch size, excess data will be discarded. + + Args: + sess: TensorFlow session. + inputs: List of arrays matching the input placeholders, of shape + [BATCH_SIZE, ...]. + state_inputs: List of RNN input arrays. These arrays have size + [BATCH_SIZE / MAX_SEQ_LEN, ...]. + num_grad_updates: The lifetime number of gradient updates that the + policy having collected the data has already undergone. + + Returns: + The number of tuples loaded per device. + """ + self.num_grad_updates = num_grad_updates + + if log_once("load_data"): + logger.info( + "Training on concatenated sample batches:\n\n{}\n".format( + summarize( + { + "placeholders": self.loss_inputs, + "inputs": inputs, + "state_inputs": state_inputs, + } + ) + ) + ) + + feed_dict = {} + assert len(self.loss_inputs) == len(inputs + state_inputs), ( + self.loss_inputs, + inputs, + state_inputs, + ) + + # Let's suppose we have the following input data, and 2 devices: + # 1 2 3 4 5 6 7 <- state inputs shape + # A A A B B B C C C D D D E E E F F F G G G <- inputs shape + # The data is truncated and split across devices as follows: + # |---| seq len = 3 + # |---------------------------------| seq batch size = 6 seqs + # |----------------| per device batch size = 9 tuples + + if len(state_inputs) > 0: + smallest_array = state_inputs[0] + seq_len = len(inputs[0]) // len(state_inputs[0]) + self._loaded_max_seq_len = seq_len + else: + smallest_array = inputs[0] + self._loaded_max_seq_len = 1 + + sequences_per_minibatch = ( + self.max_per_device_batch_size + // self._loaded_max_seq_len + * len(self.devices) + ) + if sequences_per_minibatch < 1: + logger.warning( + ( + "Target minibatch size is {}, however the rollout sequence " + "length is {}, hence the minibatch size will be raised to " + "{}." + ).format( + self.max_per_device_batch_size, + self._loaded_max_seq_len, + self._loaded_max_seq_len * len(self.devices), + ) + ) + sequences_per_minibatch = 1 + + if len(smallest_array) < sequences_per_minibatch: + # Dynamically shrink the batch size if insufficient data + sequences_per_minibatch = _make_divisible_by( + len(smallest_array), len(self.devices) + ) + + if log_once("data_slicing"): + logger.info( + ( + "Divided {} rollout sequences, each of length {}, among " + "{} devices." + ).format( + len(smallest_array), self._loaded_max_seq_len, len(self.devices) + ) + ) + + if sequences_per_minibatch < len(self.devices): + raise ValueError( + "Must load at least 1 tuple sequence per device. Try " + "increasing `minibatch_size` or reducing `max_seq_len` " + "to ensure that at least one sequence fits per device." + ) + self._loaded_per_device_batch_size = ( + sequences_per_minibatch // len(self.devices) * self._loaded_max_seq_len + ) + + if len(state_inputs) > 0: + # First truncate the RNN state arrays to the sequences_per_minib. + state_inputs = [ + _make_divisible_by(arr, sequences_per_minibatch) for arr in state_inputs + ] + # Then truncate the data inputs to match + inputs = [arr[: len(state_inputs[0]) * seq_len] for arr in inputs] + assert len(state_inputs[0]) * seq_len == len(inputs[0]), ( + len(state_inputs[0]), + sequences_per_minibatch, + seq_len, + len(inputs[0]), + ) + for ph, arr in zip(self.loss_inputs, inputs + state_inputs): + feed_dict[ph] = arr + truncated_len = len(inputs[0]) + else: + truncated_len = 0 + for ph, arr in zip(self.loss_inputs, inputs): + truncated_arr = _make_divisible_by(arr, sequences_per_minibatch) + feed_dict[ph] = truncated_arr + if truncated_len == 0: + truncated_len = len(truncated_arr) + + sess.run([t.init_op for t in self._towers], feed_dict=feed_dict) + + self.num_tuples_loaded = truncated_len + samples_per_device = truncated_len // len(self.devices) + assert samples_per_device > 0, "No data loaded?" + assert samples_per_device % self._loaded_per_device_batch_size == 0 + # Return loaded samples per-device. + return samples_per_device + + def optimize(self, sess, batch_index): + """Run a single step of SGD. + + Runs a SGD step over a slice of the preloaded batch with size given by + self._loaded_per_device_batch_size and offset given by the batch_index + argument. + + Updates shared model weights based on the averaged per-device + gradients. + + Args: + sess: TensorFlow session. + batch_index: Offset into the preloaded data. This value must be + between `0` and `tuples_per_device`. The amount of data to + process is at most `max_per_device_batch_size`. + + Returns: + The outputs of extra_ops evaluated over the batch. + """ + feed_dict = { + self._batch_index: batch_index, + self._per_device_batch_size: self._loaded_per_device_batch_size, + self._max_seq_len: self._loaded_max_seq_len, + } + for tower in self._towers: + feed_dict.update(tower.loss_graph.extra_compute_grad_feed_dict()) + + fetches = {"train": self._train_op} + for tower_num, tower in enumerate(self._towers): + tower_fetch = tower.loss_graph._get_grad_and_stats_fetches() + fetches["tower_{}".format(tower_num)] = tower_fetch + + return sess.run(fetches, feed_dict=feed_dict) + + def get_device_losses(self): + return [t.loss_graph for t in self._towers] + + def _setup_device(self, tower_i, device, device_input_placeholders, num_data_in): + assert num_data_in <= len(device_input_placeholders) + with tf.device(device): + with tf1.name_scope(TOWER_SCOPE_NAME + f"_{tower_i}"): + device_input_batches = [] + device_input_slices = [] + for i, ph in enumerate(device_input_placeholders): + current_batch = tf1.Variable( + ph, trainable=False, validate_shape=False, collections=[] + ) + device_input_batches.append(current_batch) + if i < num_data_in: + scale = self._max_seq_len + granularity = self._max_seq_len + else: + scale = self._max_seq_len + granularity = 1 + current_slice = tf.slice( + current_batch, + ( + [self._batch_index // scale * granularity] + + [0] * len(ph.shape[1:]) + ), + ( + [self._per_device_batch_size // scale * granularity] + + [-1] * len(ph.shape[1:]) + ), + ) + current_slice.set_shape(ph.shape) + device_input_slices.append(current_slice) + graph_obj = self.policy_copy(device_input_slices) + device_grads = graph_obj.gradients(self.optimizers, graph_obj._losses) + return _Tower( + tf.group(*[batch.initializer for batch in device_input_batches]), + device_grads, + graph_obj, + ) + + +# Each tower is a copy of the loss graph pinned to a specific device. +_Tower = namedtuple("Tower", ["init_op", "grads", "loss_graph"]) + + +def _make_divisible_by(a, n): + if type(a) is int: + return a - a % n + return a[0 : a.shape[0] - a.shape[0] % n] + + +def _average_gradients(tower_grads): + """Averages gradients across towers. + + Calculate the average gradient for each shared variable across all towers. + Note that this function provides a synchronization point across all towers. + + Args: + tower_grads: List of lists of (gradient, variable) tuples. The outer + list is over individual gradients. The inner list is over the + gradient calculation for each tower. + + Returns: + List of pairs of (gradient, variable) where the gradient has been + averaged across all towers. + + TODO(ekl): We could use NCCL if this becomes a bottleneck. + """ + + average_grads = [] + for grad_and_vars in zip(*tower_grads): + # Note that each grad_and_vars looks like the following: + # ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN)) + grads = [] + for g, _ in grad_and_vars: + if g is not None: + # Add 0 dimension to the gradients to represent the tower. + expanded_g = tf.expand_dims(g, 0) + + # Append on a 'tower' dimension which we will average over + # below. + grads.append(expanded_g) + + if not grads: + continue + + # Average over the 'tower' dimension. + grad = tf.concat(axis=0, values=grads) + grad = tf.reduce_mean(grad, 0) + + # Keep in mind that the Variables are redundant because they are shared + # across towers. So .. we will just return the first tower's pointer to + # the Variable. + v = grad_and_vars[0][1] + grad_and_var = (grad, v) + average_grads.append(grad_and_var) + + return average_grads diff --git a/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy_v2.py b/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..51fe10ce8ecbf045a3dcd1ab9bb5aaf2a7d3c5bb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/dynamic_tf_policy_v2.py @@ -0,0 +1,1048 @@ +import logging +import re +from collections import OrderedDict +from typing import Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym +import tree # pip install dm_tree + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy.dynamic_tf_policy import TFMultiGPUTowerStack +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + is_overridden, + override, +) +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space +from ray.rllib.utils.tf_utils import get_placeholder +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class DynamicTFPolicyV2(TFPolicy): + """A TFPolicy that auto-defines placeholders dynamically at runtime. + + This class is intended to be used and extended by sub-classing. + """ + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + *, + existing_inputs: Optional[Dict[str, "tf1.placeholder"]] = None, + existing_model: Optional[ModelV2] = None, + ): + self.observation_space = obs_space + self.action_space = action_space + self.config = config + self.framework = "tf" + self._seq_lens = None + self._is_tower = existing_inputs is not None + + self.validate_spaces(obs_space, action_space, config) + + self.dist_class = self._init_dist_class() + # Setup self.model. + if existing_model and isinstance(existing_model, list): + self.model = existing_model[0] + # TODO: (sven) hack, but works for `target_[q_]?model`. + for i in range(1, len(existing_model)): + setattr(self, existing_model[i][0], existing_model[i][1]) + else: + self.model = self.make_model() + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + + self._init_state_inputs(existing_inputs) + self._init_view_requirements() + timestep, explore = self._init_input_dict_and_dummy_batch(existing_inputs) + ( + sampled_action, + sampled_action_logp, + dist_inputs, + self._policy_extra_action_fetches, + ) = self._init_action_fetches(timestep, explore) + + # Phase 1 init. + sess = tf1.get_default_session() or tf1.Session( + config=tf1.ConfigProto(**self.config["tf_session_args"]) + ) + + batch_divisibility_req = self.get_batch_divisibility_req() + + prev_action_input = ( + self._input_dict[SampleBatch.PREV_ACTIONS] + if SampleBatch.PREV_ACTIONS in self._input_dict.accessed_keys + else None + ) + prev_reward_input = ( + self._input_dict[SampleBatch.PREV_REWARDS] + if SampleBatch.PREV_REWARDS in self._input_dict.accessed_keys + else None + ) + + super().__init__( + observation_space=obs_space, + action_space=action_space, + config=config, + sess=sess, + obs_input=self._input_dict[SampleBatch.OBS], + action_input=self._input_dict[SampleBatch.ACTIONS], + sampled_action=sampled_action, + sampled_action_logp=sampled_action_logp, + dist_inputs=dist_inputs, + dist_class=self.dist_class, + loss=None, # dynamically initialized on run + loss_inputs=[], + model=self.model, + state_inputs=self._state_inputs, + state_outputs=self._state_out, + prev_action_input=prev_action_input, + prev_reward_input=prev_reward_input, + seq_lens=self._seq_lens, + max_seq_len=config["model"].get("max_seq_len", 20), + batch_divisibility_req=batch_divisibility_req, + explore=explore, + timestep=timestep, + ) + + @staticmethod + def enable_eager_execution_if_necessary(): + # This is static graph TF policy. + # Simply do nothing. + pass + + @OverrideToImplementCustomLogic + def validate_spaces( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + ): + return {} + + @OverrideToImplementCustomLogic + @override(Policy) + def loss( + self, + model: Union[ModelV2, "tf.keras.Model"], + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + """Constructs loss computation graph for this TF1 policy. + + Args: + model: The Model to calculate the loss for. + dist_class: The action distr. class. + train_batch: The training data. + + Returns: + A single loss tensor or a list of loss tensors. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function. Returns a dict of statistics. + + Args: + train_batch: The SampleBatch (already) used for training. + + Returns: + The stats dict. + """ + return {} + + @OverrideToImplementCustomLogic + def grad_stats_fn( + self, train_batch: SampleBatch, grads: ModelGradients + ) -> Dict[str, TensorType]: + """Gradient stats function. Returns a dict of statistics. + + Args: + train_batch: The SampleBatch (already) used for training. + + Returns: + The stats dict. + """ + return {} + + @OverrideToImplementCustomLogic + def make_model(self) -> ModelV2: + """Build underlying model for this Policy. + + Returns: + The Model for the Policy to use. + """ + # Default ModelV2 model. + _, logit_dim = ModelCatalog.get_action_dist( + self.action_space, self.config["model"] + ) + return ModelCatalog.get_model_v2( + obs_space=self.observation_space, + action_space=self.action_space, + num_outputs=logit_dim, + model_config=self.config["model"], + framework="tf", + ) + + @OverrideToImplementCustomLogic + def compute_gradients_fn( + self, optimizer: LocalOptimizer, loss: TensorType + ) -> ModelGradients: + """Gradients computing function (from loss tensor, using local optimizer). + + Args: + policy: The Policy object that generated the loss tensor and + that holds the given local optimizer. + optimizer: The tf (local) optimizer object to + calculate the gradients with. + loss: The loss tensor for which gradients should be + calculated. + + Returns: + ModelGradients: List of the possibly clipped gradients- and variable + tuples. + """ + return None + + @OverrideToImplementCustomLogic + def apply_gradients_fn( + self, + optimizer: "tf.keras.optimizers.Optimizer", + grads: ModelGradients, + ) -> "tf.Operation": + """Gradients computing function (from loss tensor, using local optimizer). + + Args: + optimizer: The tf (local) optimizer object to + calculate the gradients with. + grads: The gradient tensor to be applied. + + Returns: + "tf.Operation": TF operation that applies supplied gradients. + """ + return None + + @OverrideToImplementCustomLogic + def action_sampler_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, TensorType, TensorType, List[TensorType]]: + """Custom function for sampling new actions given policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Sampled action + Log-likelihood + Action distribution inputs + Updated state + """ + return None, None, None, None + + @OverrideToImplementCustomLogic + def action_distribution_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, type, List[TensorType]]: + """Action distribution function for this Policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Distribution input. + ActionDistribution class. + State outs. + """ + return None, None, None + + @OverrideToImplementCustomLogic + def get_batch_divisibility_req(self) -> int: + """Get batch divisibility request. + + Returns: + Size N. A sample batch must be of size K*N. + """ + # By default, any sized batch is ok, so simply return 1. + return 1 + + @override(TFPolicy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_action_out_fn(self) -> Dict[str, TensorType]: + """Extra values to fetch and return from compute_actions(). + + Returns: + Dict[str, TensorType]: An extra fetch-dict to be passed to and + returned from the compute_actions() call. + """ + extra_action_fetches = super().extra_action_out_fn() + extra_action_fetches.update(self._policy_extra_action_fetches) + return extra_action_fetches + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_learn_fetches_fn(self) -> Dict[str, TensorType]: + """Extra stats to be reported after gradient computation. + + Returns: + Dict[str, TensorType]: An extra fetch-dict. + """ + return {} + + @override(TFPolicy) + def extra_compute_grad_fetches(self): + return dict({LEARNER_STATS_KEY: {}}, **self.extra_learn_fetches_fn()) + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[SampleBatch] = None, + episode=None, + ): + """Post process trajectory in the format of a SampleBatch. + + Args: + sample_batch: sample_batch: batch of experiences for the policy, + which will contain at most one episode trajectory. + other_agent_batches: In a multi-agent env, this contains a + mapping of agent ids to (policy, agent_batch) tuples + containing the policy and experiences of the other agents. + episode: An optional multi-agent episode object to provide + access to all of the internal episode state, which may + be useful for model-based or multi-agent algorithms. + + Returns: + The postprocessed sample batch. + """ + return Policy.postprocess_trajectory(self, sample_batch) + + @override(TFPolicy) + @OverrideToImplementCustomLogic + def optimizer( + self, + ) -> Union["tf.keras.optimizers.Optimizer", List["tf.keras.optimizers.Optimizer"]]: + """TF optimizer to use for policy optimization. + + Returns: + A local optimizer or a list of local optimizers to use for this + Policy's Model. + """ + return super().optimizer() + + def _init_dist_class(self): + if is_overridden(self.action_sampler_fn) or is_overridden( + self.action_distribution_fn + ): + if not is_overridden(self.make_model): + raise ValueError( + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given" + ) + return None + else: + dist_class, _ = ModelCatalog.get_action_dist( + self.action_space, self.config["model"] + ) + return dist_class + + def _init_view_requirements(self): + # If ViewRequirements are explicitly specified. + if getattr(self, "view_requirements", None): + return + + # Use default settings. + # Add NEXT_OBS, STATE_IN_0.., and others. + self.view_requirements = self._get_default_view_requirements() + # Combine view_requirements for Model and Policy. + # TODO(jungong) : models will not carry view_requirements once they + # are migrated to be organic Keras models. + self.view_requirements.update(self.model.view_requirements) + # Disable env-info placeholder. + if SampleBatch.INFOS in self.view_requirements: + self.view_requirements[SampleBatch.INFOS].used_for_training = False + + def _init_state_inputs(self, existing_inputs: Dict[str, "tf1.placeholder"]): + """Initialize input placeholders. + + Args: + existing_inputs: existing placeholders. + """ + if existing_inputs: + self._state_inputs = [ + v for k, v in existing_inputs.items() if k.startswith("state_in_") + ] + # Placeholder for RNN time-chunk valid lengths. + if self._state_inputs: + self._seq_lens = existing_inputs[SampleBatch.SEQ_LENS] + # Create new input placeholders. + else: + self._state_inputs = [ + get_placeholder( + space=vr.space, + time_axis=not isinstance(vr.shift, int), + name=k, + ) + for k, vr in self.model.view_requirements.items() + if k.startswith("state_in_") + ] + # Placeholder for RNN time-chunk valid lengths. + if self._state_inputs: + self._seq_lens = tf1.placeholder( + dtype=tf.int32, shape=[None], name="seq_lens" + ) + + def _init_input_dict_and_dummy_batch( + self, existing_inputs: Dict[str, "tf1.placeholder"] + ) -> Tuple[Union[int, TensorType], Union[bool, TensorType]]: + """Initialized input_dict and dummy_batch data. + + Args: + existing_inputs: When copying a policy, this specifies an existing + dict of placeholders to use instead of defining new ones. + + Returns: + timestep: training timestep. + explore: whether this policy should explore. + """ + # Setup standard placeholders. + if self._is_tower: + assert existing_inputs is not None + timestep = existing_inputs["timestep"] + explore = False + ( + self._input_dict, + self._dummy_batch, + ) = self._create_input_dict_and_dummy_batch( + self.view_requirements, existing_inputs + ) + else: + # Placeholder for (sampling steps) timestep (int). + timestep = tf1.placeholder_with_default( + tf.zeros((), dtype=tf.int64), (), name="timestep" + ) + # Placeholder for `is_exploring` flag. + explore = tf1.placeholder_with_default(True, (), name="is_exploring") + ( + self._input_dict, + self._dummy_batch, + ) = self._create_input_dict_and_dummy_batch(self.view_requirements, {}) + + # Placeholder for `is_training` flag. + self._input_dict.set_training(self._get_is_training_placeholder()) + + return timestep, explore + + def _create_input_dict_and_dummy_batch(self, view_requirements, existing_inputs): + """Creates input_dict and dummy_batch for loss initialization. + + Used for managing the Policy's input placeholders and for loss + initialization. + Input_dict: Str -> tf.placeholders, dummy_batch: str -> np.arrays. + + Args: + view_requirements: The view requirements dict. + existing_inputs (Dict[str, tf.placeholder]): A dict of already + existing placeholders. + + Returns: + Tuple[Dict[str, tf.placeholder], Dict[str, np.ndarray]]: The + input_dict/dummy_batch tuple. + """ + input_dict = {} + for view_col, view_req in view_requirements.items(): + # Point state_in to the already existing self._state_inputs. + mo = re.match(r"state_in_(\d+)", view_col) + if mo is not None: + input_dict[view_col] = self._state_inputs[int(mo.group(1))] + # State-outs (no placeholders needed). + elif view_col.startswith("state_out_"): + continue + # Skip action dist inputs placeholder (do later). + elif view_col == SampleBatch.ACTION_DIST_INPUTS: + continue + # This is a tower: Input placeholders already exist. + elif view_col in existing_inputs: + input_dict[view_col] = existing_inputs[view_col] + # All others. + else: + time_axis = not isinstance(view_req.shift, int) + if view_req.used_for_training: + # Create a +time-axis placeholder if the shift is not an + # int (range or list of ints). + # Do not flatten actions if action flattening disabled. + if self.config.get("_disable_action_flattening") and view_col in [ + SampleBatch.ACTIONS, + SampleBatch.PREV_ACTIONS, + ]: + flatten = False + # Do not flatten observations if no preprocessor API used. + elif ( + view_col in [SampleBatch.OBS, SampleBatch.NEXT_OBS] + and self.config["_disable_preprocessor_api"] + ): + flatten = False + # Flatten everything else. + else: + flatten = True + input_dict[view_col] = get_placeholder( + space=view_req.space, + name=view_col, + time_axis=time_axis, + flatten=flatten, + ) + dummy_batch = self._get_dummy_batch_from_view_requirements(batch_size=32) + + return SampleBatch(input_dict, seq_lens=self._seq_lens), dummy_batch + + def _init_action_fetches( + self, timestep: Union[int, TensorType], explore: Union[bool, TensorType] + ) -> Tuple[TensorType, TensorType, TensorType, type, Dict[str, TensorType]]: + """Create action related fields for base Policy and loss initialization.""" + # Multi-GPU towers do not need any action computing/exploration + # graphs. + sampled_action = None + sampled_action_logp = None + dist_inputs = None + extra_action_fetches = {} + self._state_out = None + if not self._is_tower: + # Create the Exploration object to use for this Policy. + self.exploration = self._create_exploration() + + # Fully customized action generation (e.g., custom policy). + if is_overridden(self.action_sampler_fn): + ( + sampled_action, + sampled_action_logp, + dist_inputs, + self._state_out, + ) = self.action_sampler_fn( + self.model, + obs_batch=self._input_dict[SampleBatch.OBS], + state_batches=self._state_inputs, + seq_lens=self._seq_lens, + prev_action_batch=self._input_dict.get(SampleBatch.PREV_ACTIONS), + prev_reward_batch=self._input_dict.get(SampleBatch.PREV_REWARDS), + explore=explore, + is_training=self._input_dict.is_training, + ) + # Distribution generation is customized, e.g., DQN, DDPG. + else: + if is_overridden(self.action_distribution_fn): + # Try new action_distribution_fn signature, supporting + # state_batches and seq_lens. + in_dict = self._input_dict + ( + dist_inputs, + self.dist_class, + self._state_out, + ) = self.action_distribution_fn( + self.model, + obs_batch=in_dict[SampleBatch.OBS], + state_batches=self._state_inputs, + seq_lens=self._seq_lens, + explore=explore, + timestep=timestep, + is_training=in_dict.is_training, + ) + # Default distribution generation behavior: + # Pass through model. E.g., PG, PPO. + else: + if isinstance(self.model, tf.keras.Model): + dist_inputs, self._state_out, extra_action_fetches = self.model( + self._input_dict + ) + else: + dist_inputs, self._state_out = self.model(self._input_dict) + + action_dist = self.dist_class(dist_inputs, self.model) + + # Using exploration to get final action (e.g. via sampling). + ( + sampled_action, + sampled_action_logp, + ) = self.exploration.get_exploration_action( + action_distribution=action_dist, timestep=timestep, explore=explore + ) + + if dist_inputs is not None: + extra_action_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + + if sampled_action_logp is not None: + extra_action_fetches[SampleBatch.ACTION_LOGP] = sampled_action_logp + extra_action_fetches[SampleBatch.ACTION_PROB] = tf.exp( + tf.cast(sampled_action_logp, tf.float32) + ) + + return ( + sampled_action, + sampled_action_logp, + dist_inputs, + extra_action_fetches, + ) + + def _init_optimizers(self): + # Create the optimizer/exploration optimizer here. Some initialization + # steps (e.g. exploration postprocessing) may need this. + optimizers = force_list(self.optimizer()) + if self.exploration: + optimizers = self.exploration.get_exploration_optimizer(optimizers) + + # No optimizers produced -> Return. + if not optimizers: + return + + # The list of local (tf) optimizers (one per loss term). + self._optimizers = optimizers + # Backward compatibility. + self._optimizer = optimizers[0] + + def maybe_initialize_optimizer_and_loss(self): + # We don't need to initialize loss calculation for MultiGPUTowerStack. + if self._is_tower: + self.get_session().run(tf1.global_variables_initializer()) + return + + # Loss initialization and model/postprocessing test calls. + self._init_optimizers() + self._initialize_loss_from_dummy_batch(auto_remove_unneeded_view_reqs=True) + + # Create MultiGPUTowerStacks, if we have at least one actual + # GPU or >1 CPUs (fake GPUs). + if len(self.devices) > 1 or any("gpu" in d for d in self.devices): + # Per-GPU graph copies created here must share vars with the + # policy. Therefore, `reuse` is set to tf1.AUTO_REUSE because + # Adam nodes are created after all of the device copies are + # created. + with tf1.variable_scope("", reuse=tf1.AUTO_REUSE): + self.multi_gpu_tower_stacks = [ + TFMultiGPUTowerStack(policy=self) + for _ in range(self.config.get("num_multi_gpu_tower_stacks", 1)) + ] + + # Initialize again after loss and tower init. + self.get_session().run(tf1.global_variables_initializer()) + + @override(Policy) + def _initialize_loss_from_dummy_batch( + self, auto_remove_unneeded_view_reqs: bool = True + ) -> None: + # Test calls depend on variable init, so initialize model first. + self.get_session().run(tf1.global_variables_initializer()) + + # Fields that have not been accessed are not needed for action + # computations -> Tag them as `used_for_compute_actions=False`. + for key, view_req in self.view_requirements.items(): + if ( + not key.startswith("state_in_") + and key not in self._input_dict.accessed_keys + ): + view_req.used_for_compute_actions = False + for key, value in self.extra_action_out_fn().items(): + self._dummy_batch[key] = get_dummy_batch_for_space( + gym.spaces.Box( + -1.0, 1.0, shape=value.shape.as_list()[1:], dtype=value.dtype.name + ), + batch_size=len(self._dummy_batch), + ) + self._input_dict[key] = get_placeholder(value=value, name=key) + if key not in self.view_requirements: + logger.info("Adding extra-action-fetch `{}` to view-reqs.".format(key)) + self.view_requirements[key] = ViewRequirement( + space=gym.spaces.Box( + -1.0, + 1.0, + shape=value.shape.as_list()[1:], + dtype=value.dtype.name, + ), + used_for_compute_actions=False, + ) + dummy_batch = self._dummy_batch + + logger.info("Testing `postprocess_trajectory` w/ dummy batch.") + self.exploration.postprocess_trajectory(self, dummy_batch, self.get_session()) + _ = self.postprocess_trajectory(dummy_batch) + # Add new columns automatically to (loss) input_dict. + for key in dummy_batch.added_keys: + if key not in self._input_dict: + self._input_dict[key] = get_placeholder( + value=dummy_batch[key], name=key + ) + if key not in self.view_requirements: + self.view_requirements[key] = ViewRequirement( + space=gym.spaces.Box( + -1.0, + 1.0, + shape=dummy_batch[key].shape[1:], + dtype=dummy_batch[key].dtype, + ), + used_for_compute_actions=False, + ) + + train_batch = SampleBatch( + dict(self._input_dict, **self._loss_input_dict), + _is_training=True, + ) + + if self._state_inputs: + train_batch[SampleBatch.SEQ_LENS] = self._seq_lens + self._loss_input_dict.update( + {SampleBatch.SEQ_LENS: train_batch[SampleBatch.SEQ_LENS]} + ) + + self._loss_input_dict.update(dict(train_batch)) + + if log_once("loss_init"): + logger.debug( + "Initializing loss function with dummy input:\n\n{}\n".format( + summarize(train_batch) + ) + ) + + losses = self._do_loss_init(train_batch) + + all_accessed_keys = ( + train_batch.accessed_keys + | dummy_batch.accessed_keys + | dummy_batch.added_keys + | set(self.model.view_requirements.keys()) + ) + + TFPolicy._initialize_loss( + self, + losses, + [(k, v) for k, v in train_batch.items() if k in all_accessed_keys] + + ( + [(SampleBatch.SEQ_LENS, train_batch[SampleBatch.SEQ_LENS])] + if SampleBatch.SEQ_LENS in train_batch + else [] + ), + ) + + if "is_training" in self._loss_input_dict: + del self._loss_input_dict["is_training"] + + # Call the grads stats fn. + # TODO: (sven) rename to simply stats_fn to match eager and torch. + self._stats_fetches.update(self.grad_stats_fn(train_batch, self._grads)) + + # Add new columns automatically to view-reqs. + if auto_remove_unneeded_view_reqs: + # Add those needed for postprocessing and training. + all_accessed_keys = train_batch.accessed_keys | dummy_batch.accessed_keys + # Tag those only needed for post-processing (with some exceptions). + for key in dummy_batch.accessed_keys: + if ( + key not in train_batch.accessed_keys + and key not in self.model.view_requirements + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + SampleBatch.OBS_EMBEDS, + ] + ): + if key in self.view_requirements: + self.view_requirements[key].used_for_training = False + if key in self._loss_input_dict: + del self._loss_input_dict[key] + # Remove those not needed at all (leave those that are needed + # by Sampler to properly execute sample collection). + # Also always leave TERMINATEDS, TRUNCATEDS, REWARDS, and INFOS, + # no matter what. + for key in list(self.view_requirements.keys()): + if ( + key not in all_accessed_keys + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + ] + and key not in self.model.view_requirements + ): + # If user deleted this key manually in postprocessing + # fn, warn about it and do not remove from + # view-requirements. + if key in dummy_batch.deleted_keys: + logger.warning( + "SampleBatch key '{}' was deleted manually in " + "postprocessing function! RLlib will " + "automatically remove non-used items from the " + "data stream. Remove the `del` from your " + "postprocessing function.".format(key) + ) + # If we are not writing output to disk, safe to erase + # this key to save space in the sample batch. + elif self.config["output"] is None: + del self.view_requirements[key] + + if key in self._loss_input_dict: + del self._loss_input_dict[key] + # Add those data_cols (again) that are missing and have + # dependencies by view_cols. + for key in list(self.view_requirements.keys()): + vr = self.view_requirements[key] + if ( + vr.data_col is not None + and vr.data_col not in self.view_requirements + ): + used_for_training = vr.data_col in train_batch.accessed_keys + self.view_requirements[vr.data_col] = ViewRequirement( + space=vr.space, used_for_training=used_for_training + ) + + self._loss_input_dict_no_rnn = { + k: v + for k, v in self._loss_input_dict.items() + if (v not in self._state_inputs and v != self._seq_lens) + } + + def _do_loss_init(self, train_batch: SampleBatch): + losses = self.loss(self.model, self.dist_class, train_batch) + losses = force_list(losses) + self._stats_fetches.update(self.stats_fn(train_batch)) + # Override the update ops to be those of the model. + self._update_ops = [] + if not isinstance(self.model, tf.keras.Model): + self._update_ops = self.model.update_ops() + return losses + + @override(TFPolicy) + def copy(self, existing_inputs: List[Tuple[str, "tf1.placeholder"]]) -> TFPolicy: + """Creates a copy of self using existing input placeholders.""" + + flat_loss_inputs = tree.flatten(self._loss_input_dict) + flat_loss_inputs_no_rnn = tree.flatten(self._loss_input_dict_no_rnn) + + # Note that there might be RNN state inputs at the end of the list + if len(flat_loss_inputs) != len(existing_inputs): + raise ValueError( + "Tensor list mismatch", + self._loss_input_dict, + self._state_inputs, + existing_inputs, + ) + for i, v in enumerate(flat_loss_inputs_no_rnn): + if v.shape.as_list() != existing_inputs[i].shape.as_list(): + raise ValueError( + "Tensor shape mismatch", i, v.shape, existing_inputs[i].shape + ) + # By convention, the loss inputs are followed by state inputs and then + # the seq len tensor. + rnn_inputs = [] + for i in range(len(self._state_inputs)): + rnn_inputs.append( + ( + "state_in_{}".format(i), + existing_inputs[len(flat_loss_inputs_no_rnn) + i], + ) + ) + if rnn_inputs: + rnn_inputs.append((SampleBatch.SEQ_LENS, existing_inputs[-1])) + existing_inputs_unflattened = tree.unflatten_as( + self._loss_input_dict_no_rnn, + existing_inputs[: len(flat_loss_inputs_no_rnn)], + ) + input_dict = OrderedDict( + [("is_exploring", self._is_exploring), ("timestep", self._timestep)] + + [ + (k, existing_inputs_unflattened[k]) + for i, k in enumerate(self._loss_input_dict_no_rnn.keys()) + ] + + rnn_inputs + ) + + instance = self.__class__( + self.observation_space, + self.action_space, + self.config, + existing_inputs=input_dict, + existing_model=[ + self.model, + # Deprecated: Target models should all reside under + # `policy.target_model` now. + ("target_q_model", getattr(self, "target_q_model", None)), + ("target_model", getattr(self, "target_model", None)), + ], + ) + + instance._loss_input_dict = input_dict + losses = instance._do_loss_init(SampleBatch(input_dict)) + loss_inputs = [ + (k, existing_inputs_unflattened[k]) + for i, k in enumerate(self._loss_input_dict_no_rnn.keys()) + ] + + TFPolicy._initialize_loss(instance, losses, loss_inputs) + instance._stats_fetches.update( + instance.grad_stats_fn(input_dict, instance._grads) + ) + return instance + + @override(Policy) + def get_initial_state(self) -> List[TensorType]: + if self.model: + return self.model.get_initial_state() + else: + return [] + + @override(Policy) + def load_batch_into_buffer( + self, + batch: SampleBatch, + buffer_index: int = 0, + ) -> int: + # Set the is_training flag of the batch. + batch.set_training(True) + + # Shortcut for 1 CPU only: Store batch in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + self._loaded_single_cpu_batch = batch + return len(batch) + + input_dict = self._get_loss_inputs_dict(batch, shuffle=False) + data_keys = tree.flatten(self._loss_input_dict_no_rnn) + if self._state_inputs: + state_keys = self._state_inputs + [self._seq_lens] + else: + state_keys = [] + inputs = [input_dict[k] for k in data_keys] + state_inputs = [input_dict[k] for k in state_keys] + + return self.multi_gpu_tower_stacks[buffer_index].load_data( + sess=self.get_session(), + inputs=inputs, + state_inputs=state_inputs, + num_grad_updates=batch.num_grad_updates, + ) + + @override(Policy) + def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int: + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + return ( + len(self._loaded_single_cpu_batch) + if self._loaded_single_cpu_batch is not None + else 0 + ) + + return self.multi_gpu_tower_stacks[buffer_index].num_tuples_loaded + + @override(Policy) + def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0): + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_single_cpu_batch`. + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + if self._loaded_single_cpu_batch is None: + raise ValueError( + "Must call Policy.load_batch_into_buffer() before " + "Policy.learn_on_loaded_batch()!" + ) + # Get the correct slice of the already loaded batch to use, + # based on offset and batch size. + batch_size = self.config.get("minibatch_size") + if batch_size is None: + batch_size = self.config.get( + "sgd_minibatch_size", self.config["train_batch_size"] + ) + + if batch_size >= len(self._loaded_single_cpu_batch): + sliced_batch = self._loaded_single_cpu_batch + else: + sliced_batch = self._loaded_single_cpu_batch.slice( + start=offset, end=offset + batch_size + ) + return self.learn_on_batch(sliced_batch) + + tower_stack = self.multi_gpu_tower_stacks[buffer_index] + results = tower_stack.optimize(self.get_session(), offset) + self.num_grad_updates += 1 + + results.update( + { + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates - 1 - (tower_stack.num_grad_updates or 0) + ), + } + ) + + return results + + @override(TFPolicy) + def gradients(self, optimizer, loss): + optimizers = force_list(optimizer) + losses = force_list(loss) + + if is_overridden(self.compute_gradients_fn): + # New API: Allow more than one optimizer -> Return a list of + # lists of gradients. + if self.config["_tf_policy_handles_more_than_one_loss"]: + return self.compute_gradients_fn(optimizers, losses) + # Old API: Return a single List of gradients. + else: + return self.compute_gradients_fn(optimizers[0], losses[0]) + else: + return super().gradients(optimizers, losses) diff --git a/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..0091eaf7aa182c7bd6cdd8859edb653dbe70262b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy.py @@ -0,0 +1,1051 @@ +"""Eager mode TF policy built using build_tf_policy(). + +It supports both traced and non-traced eager execution modes.""" + +import functools +import logging +import os +import threading +from typing import Dict, List, Optional, Tuple, Union + +import tree # pip install dm_tree + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.repeated_values import RepeatedValues +from ray.rllib.policy.policy import Policy, PolicyState +from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import add_mixins, force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.error import ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_AGENT_STEPS_TRAINED, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.spaces.space_utils import normalize_action +from ray.rllib.utils.tf_utils import get_gpu_devices +from ray.rllib.utils.threading import with_lock +from ray.rllib.utils.typing import ( + LocalOptimizer, + ModelGradients, + TensorStructType, + TensorType, +) +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +def _convert_to_tf(x, dtype=None): + if isinstance(x, SampleBatch): + dict_ = {k: v for k, v in x.items() if k != SampleBatch.INFOS} + return tree.map_structure(_convert_to_tf, dict_) + elif isinstance(x, Policy): + return x + # Special handling of "Repeated" values. + elif isinstance(x, RepeatedValues): + return RepeatedValues( + tree.map_structure(_convert_to_tf, x.values), x.lengths, x.max_len + ) + + if x is not None: + d = dtype + return tree.map_structure( + lambda f: _convert_to_tf(f, d) + if isinstance(f, RepeatedValues) + else tf.convert_to_tensor(f, d) + if f is not None and not tf.is_tensor(f) + else f, + x, + ) + + return x + + +def _convert_to_numpy(x): + def _map(x): + if isinstance(x, tf.Tensor): + return x.numpy() + return x + + try: + return tf.nest.map_structure(_map, x) + except AttributeError: + raise TypeError( + ("Object of type {} has no method to convert to numpy.").format(type(x)) + ) + + +def _convert_eager_inputs(func): + @functools.wraps(func) + def _func(*args, **kwargs): + if tf.executing_eagerly(): + eager_args = [_convert_to_tf(x) for x in args] + # TODO: (sven) find a way to remove key-specific hacks. + eager_kwargs = { + k: _convert_to_tf(v, dtype=tf.int64 if k == "timestep" else None) + for k, v in kwargs.items() + if k not in {"info_batch", "episodes"} + } + return func(*eager_args, **eager_kwargs) + else: + return func(*args, **kwargs) + + return _func + + +def _convert_eager_outputs(func): + @functools.wraps(func) + def _func(*args, **kwargs): + out = func(*args, **kwargs) + if tf.executing_eagerly(): + out = tf.nest.map_structure(_convert_to_numpy, out) + return out + + return _func + + +def _disallow_var_creation(next_creator, **kw): + v = next_creator(**kw) + raise ValueError( + "Detected a variable being created during an eager " + "forward pass. Variables should only be created during " + "model initialization: {}".format(v.name) + ) + + +def _check_too_many_retraces(obj): + """Asserts that a given number of re-traces is not breached.""" + + def _func(self_, *args, **kwargs): + if ( + self_.config.get("eager_max_retraces") is not None + and self_._re_trace_counter > self_.config["eager_max_retraces"] + ): + raise RuntimeError( + "Too many tf-eager re-traces detected! This could lead to" + " significant slow-downs (even slower than running in " + "tf-eager mode w/ `eager_tracing=False`). To switch off " + "these re-trace counting checks, set `eager_max_retraces`" + " in your config to None." + ) + return obj(self_, *args, **kwargs) + + return _func + + +@OldAPIStack +class EagerTFPolicy(Policy): + """Dummy class to recognize any eagerized TFPolicy by its inheritance.""" + + pass + + +def _traced_eager_policy(eager_policy_cls): + """Wrapper class that enables tracing for all eager policy methods. + + This is enabled by the `--trace`/`eager_tracing=True` config when + framework=tf2. + """ + + class TracedEagerPolicy(eager_policy_cls): + def __init__(self, *args, **kwargs): + self._traced_learn_on_batch_helper = False + self._traced_compute_actions_helper = False + self._traced_compute_gradients_helper = False + self._traced_apply_gradients_helper = False + super(TracedEagerPolicy, self).__init__(*args, **kwargs) + + @_check_too_many_retraces + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Dict[str, TensorType], + explore: bool = None, + timestep: Optional[int] = None, + episodes=None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + """Traced version of Policy.compute_actions_from_input_dict.""" + + # Create a traced version of `self._compute_actions_helper`. + if self._traced_compute_actions_helper is False and not self._no_tracing: + self._compute_actions_helper = _convert_eager_inputs( + tf.function( + super(TracedEagerPolicy, self)._compute_actions_helper, + autograph=False, + reduce_retracing=True, + ) + ) + self._traced_compute_actions_helper = True + + # Now that the helper method is traced, call super's + # `compute_actions_from_input_dict()` (which will call the traced helper). + return super(TracedEagerPolicy, self).compute_actions_from_input_dict( + input_dict=input_dict, + explore=explore, + timestep=timestep, + episodes=episodes, + **kwargs, + ) + + @_check_too_many_retraces + @override(eager_policy_cls) + def learn_on_batch(self, samples): + """Traced version of Policy.learn_on_batch.""" + + # Create a traced version of `self._learn_on_batch_helper`. + if self._traced_learn_on_batch_helper is False and not self._no_tracing: + self._learn_on_batch_helper = _convert_eager_inputs( + tf.function( + super(TracedEagerPolicy, self)._learn_on_batch_helper, + autograph=False, + reduce_retracing=True, + ) + ) + self._traced_learn_on_batch_helper = True + + # Now that the helper method is traced, call super's + # apply_gradients (which will call the traced helper). + return super(TracedEagerPolicy, self).learn_on_batch(samples) + + @_check_too_many_retraces + @override(eager_policy_cls) + def compute_gradients(self, samples: SampleBatch) -> ModelGradients: + """Traced version of Policy.compute_gradients.""" + + # Create a traced version of `self._compute_gradients_helper`. + if self._traced_compute_gradients_helper is False and not self._no_tracing: + self._compute_gradients_helper = _convert_eager_inputs( + tf.function( + super(TracedEagerPolicy, self)._compute_gradients_helper, + autograph=False, + reduce_retracing=True, + ) + ) + self._traced_compute_gradients_helper = True + + # Now that the helper method is traced, call super's + # `compute_gradients()` (which will call the traced helper). + return super(TracedEagerPolicy, self).compute_gradients(samples) + + @_check_too_many_retraces + @override(Policy) + def apply_gradients(self, grads: ModelGradients) -> None: + """Traced version of Policy.apply_gradients.""" + + # Create a traced version of `self._apply_gradients_helper`. + if self._traced_apply_gradients_helper is False and not self._no_tracing: + self._apply_gradients_helper = _convert_eager_inputs( + tf.function( + super(TracedEagerPolicy, self)._apply_gradients_helper, + autograph=False, + reduce_retracing=True, + ) + ) + self._traced_apply_gradients_helper = True + + # Now that the helper method is traced, call super's + # `apply_gradients()` (which will call the traced helper). + return super(TracedEagerPolicy, self).apply_gradients(grads) + + @classmethod + def with_tracing(cls): + # Already traced -> Return same class. + return cls + + TracedEagerPolicy.__name__ = eager_policy_cls.__name__ + "_traced" + TracedEagerPolicy.__qualname__ = eager_policy_cls.__qualname__ + "_traced" + return TracedEagerPolicy + + +class _OptimizerWrapper: + def __init__(self, tape): + self.tape = tape + + def compute_gradients(self, loss, var_list): + return list(zip(self.tape.gradient(loss, var_list), var_list)) + + +@OldAPIStack +def _build_eager_tf_policy( + name, + loss_fn, + get_default_config=None, + postprocess_fn=None, + stats_fn=None, + optimizer_fn=None, + compute_gradients_fn=None, + apply_gradients_fn=None, + grad_stats_fn=None, + extra_learn_fetches_fn=None, + extra_action_out_fn=None, + validate_spaces=None, + before_init=None, + before_loss_init=None, + after_init=None, + make_model=None, + action_sampler_fn=None, + action_distribution_fn=None, + mixins=None, + get_batch_divisibility_req=None, + # Deprecated args. + obs_include_prev_action_reward=DEPRECATED_VALUE, + extra_action_fetches_fn=None, + gradients_fn=None, +): + """Build an eager TF policy. + + An eager policy runs all operations in eager mode, which makes debugging + much simpler, but has lower performance. + + You shouldn't need to call this directly. Rather, prefer to build a TF + graph policy and use set `.framework("tf2", eager_tracing=False) in your + AlgorithmConfig to have it automatically be converted to an eager policy. + + This has the same signature as build_tf_policy().""" + + base = add_mixins(EagerTFPolicy, mixins) + + if obs_include_prev_action_reward != DEPRECATED_VALUE: + deprecation_warning(old="obs_include_prev_action_reward", error=True) + + if extra_action_fetches_fn is not None: + deprecation_warning( + old="extra_action_fetches_fn", new="extra_action_out_fn", error=True + ) + + if gradients_fn is not None: + deprecation_warning(old="gradients_fn", new="compute_gradients_fn", error=True) + + class eager_policy_cls(base): + def __init__(self, observation_space, action_space, config): + # If this class runs as a @ray.remote actor, eager mode may not + # have been activated yet. + if not tf1.executing_eagerly(): + tf1.enable_eager_execution() + self.framework = config.get("framework", "tf2") + EagerTFPolicy.__init__(self, observation_space, action_space, config) + + # Global timestep should be a tensor. + self.global_timestep = tf.Variable(0, trainable=False, dtype=tf.int64) + self.explore = tf.Variable( + self.config["explore"], trainable=False, dtype=tf.bool + ) + + # Log device and worker index. + num_gpus = self._get_num_gpus_for_policy() + if num_gpus > 0: + gpu_ids = get_gpu_devices() + logger.info(f"Found {len(gpu_ids)} visible cuda devices.") + + self._is_training = False + + # Only for `config.eager_tracing=True`: A counter to keep track of + # how many times an eager-traced method (e.g. + # `self._compute_actions_helper`) has been re-traced by tensorflow. + # We will raise an error if more than n re-tracings have been + # detected, since this would considerably slow down execution. + # The variable below should only get incremented during the + # tf.function trace operations, never when calling the already + # traced function after that. + self._re_trace_counter = 0 + + self._loss_initialized = False + # To ensure backward compatibility: + # Old way: If `loss` provided here, use as-is (as a function). + if loss_fn is not None: + self._loss = loss_fn + # New way: Convert the overridden `self.loss` into a plain + # function, so it can be called the same way as `loss` would + # be, ensuring backward compatibility. + elif self.loss.__func__.__qualname__ != "Policy.loss": + self._loss = self.loss.__func__ + # `loss` not provided nor overridden from Policy -> Set to None. + else: + self._loss = None + + self.batch_divisibility_req = ( + get_batch_divisibility_req(self) + if callable(get_batch_divisibility_req) + else (get_batch_divisibility_req or 1) + ) + self._max_seq_len = config["model"]["max_seq_len"] + + if validate_spaces: + validate_spaces(self, observation_space, action_space, config) + + if before_init: + before_init(self, observation_space, action_space, config) + + self.config = config + self.dist_class = None + if action_sampler_fn or action_distribution_fn: + if not make_model: + raise ValueError( + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given" + ) + else: + self.dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"] + ) + + if make_model: + self.model = make_model(self, observation_space, action_space, config) + else: + self.model = ModelCatalog.get_model_v2( + observation_space, + action_space, + logit_dim, + config["model"], + framework=self.framework, + ) + # Lock used for locking some methods on the object-level. + # This prevents possible race conditions when calling the model + # first, then its value function (e.g. in a loss function), in + # between of which another model call is made (e.g. to compute an + # action). + self._lock = threading.RLock() + + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + # Combine view_requirements for Model and Policy. + self.view_requirements.update(self.model.view_requirements) + + self.exploration = self._create_exploration() + self._state_inputs = self.model.get_initial_state() + self._is_recurrent = len(self._state_inputs) > 0 + + if before_loss_init: + before_loss_init(self, observation_space, action_space, config) + + if optimizer_fn: + optimizers = optimizer_fn(self, config) + else: + optimizers = tf.keras.optimizers.Adam(config["lr"]) + optimizers = force_list(optimizers) + if self.exploration: + optimizers = self.exploration.get_exploration_optimizer(optimizers) + + # The list of local (tf) optimizers (one per loss term). + self._optimizers: List[LocalOptimizer] = optimizers + # Backward compatibility: A user's policy may only support a single + # loss term and optimizer (no lists). + self._optimizer: LocalOptimizer = optimizers[0] if optimizers else None + + self._initialize_loss_from_dummy_batch( + auto_remove_unneeded_view_reqs=True, + stats_fn=stats_fn, + ) + self._loss_initialized = True + + if after_init: + after_init(self, observation_space, action_space, config) + + # Got to reset global_timestep again after fake run-throughs. + self.global_timestep.assign(0) + + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Dict[str, TensorType], + explore: bool = None, + timestep: Optional[int] = None, + episodes=None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + if not self.config.get("eager_tracing") and not tf1.executing_eagerly(): + tf1.enable_eager_execution() + + self._is_training = False + + explore = explore if explore is not None else self.explore + timestep = timestep if timestep is not None else self.global_timestep + if isinstance(timestep, tf.Tensor): + timestep = int(timestep.numpy()) + + # Pass lazy (eager) tensor dict to Model as `input_dict`. + input_dict = self._lazy_tensor_dict(input_dict) + input_dict.set_training(False) + + # Pack internal state inputs into (separate) list. + state_batches = [ + input_dict[k] for k in input_dict.keys() if "state_in" in k[:8] + ] + self._state_in = state_batches + self._is_recurrent = state_batches != [] + + # Call the exploration before_compute_actions hook. + self.exploration.before_compute_actions( + timestep=timestep, explore=explore, tf_sess=self.get_session() + ) + + ret = self._compute_actions_helper( + input_dict, + state_batches, + # TODO: Passing episodes into a traced method does not work. + None if self.config["eager_tracing"] else episodes, + explore, + timestep, + ) + # Update our global timestep by the batch size. + self.global_timestep.assign_add(tree.flatten(ret[0])[0].shape.as_list()[0]) + return convert_to_numpy(ret) + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorStructType], TensorStructType] = None, + prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None, + info_batch: Optional[Dict[str, list]] = None, + episodes: Optional[List] = None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + # Create input dict to simply pass the entire call to + # self.compute_actions_from_input_dict(). + input_dict = SampleBatch( + { + SampleBatch.CUR_OBS: obs_batch, + }, + _is_training=tf.constant(False), + ) + if state_batches is not None: + for i, s in enumerate(state_batches): + input_dict[f"state_in_{i}"] = s + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + if info_batch is not None: + input_dict[SampleBatch.INFOS] = info_batch + + return self.compute_actions_from_input_dict( + input_dict=input_dict, + explore=explore, + timestep=timestep, + episodes=episodes, + **kwargs, + ) + + @with_lock + @override(Policy) + def compute_log_likelihoods( + self, + actions, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + actions_normalized=True, + **kwargs, + ): + if action_sampler_fn and action_distribution_fn is None: + raise ValueError( + "Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!" + ) + + seq_lens = tf.ones(len(obs_batch), dtype=tf.int32) + input_batch = SampleBatch( + {SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_batch)}, + _is_training=False, + ) + if prev_action_batch is not None: + input_batch[SampleBatch.PREV_ACTIONS] = tf.convert_to_tensor( + prev_action_batch + ) + if prev_reward_batch is not None: + input_batch[SampleBatch.PREV_REWARDS] = tf.convert_to_tensor( + prev_reward_batch + ) + + if self.exploration: + # Exploration hook before each forward pass. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if action_distribution_fn: + dist_inputs, dist_class, _ = action_distribution_fn( + self, self.model, input_batch, explore=False, is_training=False + ) + # Default log-likelihood calculation. + else: + dist_inputs, _ = self.model(input_batch, state_batches, seq_lens) + dist_class = self.dist_class + + action_dist = dist_class(dist_inputs, self.model) + + # Normalize actions if necessary. + if not actions_normalized and self.config["normalize_actions"]: + actions = normalize_action(actions, self.action_space_struct) + + log_likelihoods = action_dist.logp(actions) + + return log_likelihoods + + @override(Policy) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + assert tf.executing_eagerly() + # Call super's postprocess_trajectory first. + sample_batch = EagerTFPolicy.postprocess_trajectory(self, sample_batch) + if postprocess_fn: + return postprocess_fn(self, sample_batch, other_agent_batches, episode) + return sample_batch + + @with_lock + @override(Policy) + def learn_on_batch(self, postprocessed_batch): + # Callback handling. + learn_stats = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=postprocessed_batch, result=learn_stats + ) + + pad_batch_to_sequences_of_same_size( + postprocessed_batch, + max_seq_len=self._max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + self._is_training = True + postprocessed_batch = self._lazy_tensor_dict(postprocessed_batch) + postprocessed_batch.set_training(True) + stats = self._learn_on_batch_helper(postprocessed_batch) + self.num_grad_updates += 1 + + stats.update( + { + "custom_metrics": learn_stats, + NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count, + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update + # above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates + - 1 + - (postprocessed_batch.num_grad_updates or 0) + ), + } + ) + return convert_to_numpy(stats) + + @override(Policy) + def compute_gradients( + self, postprocessed_batch: SampleBatch + ) -> Tuple[ModelGradients, Dict[str, TensorType]]: + pad_batch_to_sequences_of_same_size( + postprocessed_batch, + shuffle=False, + max_seq_len=self._max_seq_len, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + self._is_training = True + self._lazy_tensor_dict(postprocessed_batch) + postprocessed_batch.set_training(True) + grads_and_vars, grads, stats = self._compute_gradients_helper( + postprocessed_batch + ) + return convert_to_numpy((grads, stats)) + + @override(Policy) + def apply_gradients(self, gradients: ModelGradients) -> None: + self._apply_gradients_helper( + list( + zip( + [ + (tf.convert_to_tensor(g) if g is not None else None) + for g in gradients + ], + self.model.trainable_variables(), + ) + ) + ) + + @override(Policy) + def get_weights(self, as_dict=False): + variables = self.variables() + if as_dict: + return {v.name: v.numpy() for v in variables} + return [v.numpy() for v in variables] + + @override(Policy) + def set_weights(self, weights): + variables = self.variables() + assert len(weights) == len(variables), (len(weights), len(variables)) + for v, w in zip(variables, weights): + v.assign(w) + + @override(Policy) + def get_exploration_state(self): + return convert_to_numpy(self.exploration.get_state()) + + @override(Policy) + def is_recurrent(self): + return self._is_recurrent + + @override(Policy) + def num_state_tensors(self): + return len(self._state_inputs) + + @override(Policy) + def get_initial_state(self): + if hasattr(self, "model"): + return self.model.get_initial_state() + return [] + + @override(Policy) + def get_state(self) -> PolicyState: + # Legacy Policy state (w/o keras model and w/o PolicySpec). + state = super().get_state() + + state["global_timestep"] = state["global_timestep"].numpy() + if self._optimizer and len(self._optimizer.variables()) > 0: + state["_optimizer_variables"] = self._optimizer.variables() + # Add exploration state. + if self.exploration: + # This is not compatible with RLModules, which have a method + # `forward_exploration` to specify custom exploration behavior. + state["_exploration_state"] = self.exploration.get_state() + return state + + @override(Policy) + def set_state(self, state: PolicyState) -> None: + # Set optimizer vars first. + optimizer_vars = state.get("_optimizer_variables", None) + if optimizer_vars and self._optimizer.variables(): + if not type(self).__name__.endswith("_traced") and log_once( + "set_state_optimizer_vars_tf_eager_policy_v2" + ): + logger.warning( + "Cannot restore an optimizer's state for tf eager! Keras " + "is not able to save the v1.x optimizers (from " + "tf.compat.v1.train) since they aren't compatible with " + "checkpoints." + ) + for opt_var, value in zip(self._optimizer.variables(), optimizer_vars): + opt_var.assign(value) + # Set exploration's state. + if hasattr(self, "exploration") and "_exploration_state" in state: + self.exploration.set_state(state=state["_exploration_state"]) + + # Restore glbal timestep (tf vars). + self.global_timestep.assign(state["global_timestep"]) + + # Then the Policy's (NN) weights and connectors. + super().set_state(state) + + @override(Policy) + def export_model(self, export_dir, onnx: Optional[int] = None) -> None: + """Exports the Policy's Model to local directory for serving. + + Note: Since the TfModelV2 class that EagerTfPolicy uses is-NOT-a + tf.keras.Model, we need to assume that there is a `base_model` property + within this TfModelV2 class that is-a tf.keras.Model. This base model + will be used here for the export. + TODO (kourosh): This restriction will be resolved once we move Policy and + ModelV2 to the new Learner/RLModule APIs. + + Args: + export_dir: Local writable directory. + onnx: If given, will export model in ONNX format. The + value of this parameter set the ONNX OpSet version to use. + """ + if ( + hasattr(self, "model") + and hasattr(self.model, "base_model") + and isinstance(self.model.base_model, tf.keras.Model) + ): + # Store model in ONNX format. + if onnx: + try: + import tf2onnx + except ImportError as e: + raise RuntimeError( + "Converting a TensorFlow model to ONNX requires " + "`tf2onnx` to be installed. Install with " + "`pip install tf2onnx`." + ) from e + + model_proto, external_tensor_storage = tf2onnx.convert.from_keras( + self.model.base_model, + output_path=os.path.join(export_dir, "model.onnx"), + ) + # Save the tf.keras.Model (architecture and weights, so it can be + # retrieved w/o access to the original (custom) Model or Policy code). + else: + try: + self.model.base_model.save(export_dir, save_format="tf") + except Exception: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + else: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + + def variables(self): + """Return the list of all savable variables for this policy.""" + if isinstance(self.model, tf.keras.Model): + return self.model.variables + else: + return self.model.variables() + + def loss_initialized(self): + return self._loss_initialized + + @with_lock + def _compute_actions_helper( + self, input_dict, state_batches, episodes, explore, timestep + ): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + # Calculate RNN sequence lengths. + batch_size = tree.flatten(input_dict[SampleBatch.OBS])[0].shape[0] + seq_lens = tf.ones(batch_size, dtype=tf.int32) if state_batches else None + + # Add default and custom fetches. + extra_fetches = {} + + # Use Exploration object. + with tf.variable_creator_scope(_disallow_var_creation): + if action_sampler_fn: + action_sampler_outputs = action_sampler_fn( + self, + self.model, + input_dict[SampleBatch.CUR_OBS], + explore=explore, + timestep=timestep, + episodes=episodes, + ) + if len(action_sampler_outputs) == 4: + actions, logp, dist_inputs, state_out = action_sampler_outputs + else: + dist_inputs = None + state_out = [] + actions, logp = action_sampler_outputs + else: + if action_distribution_fn: + # Try new action_distribution_fn signature, supporting + # state_batches and seq_lens. + try: + ( + dist_inputs, + self.dist_class, + state_out, + ) = action_distribution_fn( + self, + self.model, + input_dict=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=explore, + timestep=timestep, + is_training=False, + ) + # Trying the old way (to stay backward compatible). + # TODO: Remove in future. + except TypeError as e: + if ( + "positional argument" in e.args[0] + or "unexpected keyword argument" in e.args[0] + ): + ( + dist_inputs, + self.dist_class, + state_out, + ) = action_distribution_fn( + self, + self.model, + input_dict[SampleBatch.OBS], + explore=explore, + timestep=timestep, + is_training=False, + ) + else: + raise e + elif isinstance(self.model, tf.keras.Model): + input_dict = SampleBatch(input_dict, seq_lens=seq_lens) + if state_batches and "state_in_0" not in input_dict: + for i, s in enumerate(state_batches): + input_dict[f"state_in_{i}"] = s + self._lazy_tensor_dict(input_dict) + dist_inputs, state_out, extra_fetches = self.model(input_dict) + else: + dist_inputs, state_out = self.model( + input_dict, state_batches, seq_lens + ) + + action_dist = self.dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = self.exploration.get_exploration_action( + action_distribution=action_dist, + timestep=timestep, + explore=explore, + ) + + # Action-logp and action-prob. + if logp is not None: + extra_fetches[SampleBatch.ACTION_PROB] = tf.exp(logp) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + # Custom extra fetches. + if extra_action_out_fn: + extra_fetches.update(extra_action_out_fn(self)) + + return actions, state_out, extra_fetches + + # TODO: Figure out, why _ray_trace_ctx=None helps to prevent a crash in + # AlphaStar w/ framework=tf2; eager_tracing=True on the policy learner actors. + # It seems there may be a clash between the traced-by-tf function and the + # traced-by-ray functions (for making the policy class a ray actor). + def _learn_on_batch_helper(self, samples, _ray_trace_ctx=None): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + with tf.variable_creator_scope(_disallow_var_creation): + grads_and_vars, _, stats = self._compute_gradients_helper(samples) + self._apply_gradients_helper(grads_and_vars) + return stats + + def _get_is_training_placeholder(self): + return tf.convert_to_tensor(self._is_training) + + @with_lock + def _compute_gradients_helper(self, samples): + """Computes and returns grads as eager tensors.""" + + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + # Gather all variables for which to calculate losses. + if isinstance(self.model, tf.keras.Model): + variables = self.model.trainable_variables + else: + variables = self.model.trainable_variables() + + # Calculate the loss(es) inside a tf GradientTape. + with tf.GradientTape(persistent=compute_gradients_fn is not None) as tape: + losses = self._loss(self, self.model, self.dist_class, samples) + losses = force_list(losses) + + # User provided a compute_gradients_fn. + if compute_gradients_fn: + # Wrap our tape inside a wrapper, such that the resulting + # object looks like a "classic" tf.optimizer. This way, custom + # compute_gradients_fn will work on both tf static graph + # and tf-eager. + optimizer = _OptimizerWrapper(tape) + # More than one loss terms/optimizers. + if self.config["_tf_policy_handles_more_than_one_loss"]: + grads_and_vars = compute_gradients_fn( + self, [optimizer] * len(losses), losses + ) + # Only one loss and one optimizer. + else: + grads_and_vars = [compute_gradients_fn(self, optimizer, losses[0])] + # Default: Compute gradients using the above tape. + else: + grads_and_vars = [ + list(zip(tape.gradient(loss, variables), variables)) + for loss in losses + ] + + if log_once("grad_vars"): + for g_and_v in grads_and_vars: + for g, v in g_and_v: + if g is not None: + logger.info(f"Optimizing variable {v.name}") + + # `grads_and_vars` is returned a list (len=num optimizers/losses) + # of lists of (grad, var) tuples. + if self.config["_tf_policy_handles_more_than_one_loss"]: + grads = [[g for g, _ in g_and_v] for g_and_v in grads_and_vars] + # `grads_and_vars` is returned as a list of (grad, var) tuples. + else: + grads_and_vars = grads_and_vars[0] + grads = [g for g, _ in grads_and_vars] + + stats = self._stats(self, samples, grads) + return grads_and_vars, grads, stats + + def _apply_gradients_helper(self, grads_and_vars): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + if apply_gradients_fn: + if self.config["_tf_policy_handles_more_than_one_loss"]: + apply_gradients_fn(self, self._optimizers, grads_and_vars) + else: + apply_gradients_fn(self, self._optimizer, grads_and_vars) + else: + if self.config["_tf_policy_handles_more_than_one_loss"]: + for i, o in enumerate(self._optimizers): + o.apply_gradients( + [(g, v) for g, v in grads_and_vars[i] if g is not None] + ) + else: + self._optimizer.apply_gradients( + [(g, v) for g, v in grads_and_vars if g is not None] + ) + + def _stats(self, outputs, samples, grads): + fetches = {} + if stats_fn: + fetches[LEARNER_STATS_KEY] = dict(stats_fn(outputs, samples)) + else: + fetches[LEARNER_STATS_KEY] = {} + + if extra_learn_fetches_fn: + fetches.update(dict(extra_learn_fetches_fn(self))) + if grad_stats_fn: + fetches.update(dict(grad_stats_fn(self, samples, grads))) + return fetches + + def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch): + # TODO: (sven): Keep for a while to ensure backward compatibility. + if not isinstance(postprocessed_batch, SampleBatch): + postprocessed_batch = SampleBatch(postprocessed_batch) + postprocessed_batch.set_get_interceptor(_convert_to_tf) + return postprocessed_batch + + @classmethod + def with_tracing(cls): + return _traced_eager_policy(cls) + + eager_policy_cls.__name__ = name + "_eager" + eager_policy_cls.__qualname__ = name + "_eager" + return eager_policy_cls diff --git a/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy_v2.py b/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..19a637ed55b72720be42b15414eacfec1a534606 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/eager_tf_policy_v2.py @@ -0,0 +1,968 @@ +"""Eager mode TF policy built using build_tf_policy(). + +It supports both traced and non-traced eager execution modes. +""" + +import logging +import os +import threading +from typing import Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym +import tree # pip install dm_tree + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy.eager_tf_policy import ( + _convert_to_tf, + _disallow_var_creation, + _OptimizerWrapper, + _traced_eager_policy, +) +from ray.rllib.policy.policy import Policy, PolicyState +from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + is_overridden, + override, +) +from ray.rllib.utils.error import ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_AGENT_STEPS_TRAINED, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.spaces.space_utils import normalize_action +from ray.rllib.utils.tf_utils import get_gpu_devices +from ray.rllib.utils.threading import with_lock +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +@OldAPIStack +class EagerTFPolicyV2(Policy): + """A TF-eager / TF2 based tensorflow policy. + + This class is intended to be used and extended by sub-classing. + """ + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + **kwargs, + ): + self.framework = config.get("framework", "tf2") + + # Log device. + logger.info( + "Creating TF-eager policy running on {}.".format( + "GPU" if get_gpu_devices() else "CPU" + ) + ) + + Policy.__init__(self, observation_space, action_space, config) + + self._is_training = False + # Global timestep should be a tensor. + self.global_timestep = tf.Variable(0, trainable=False, dtype=tf.int64) + self.explore = tf.Variable( + self.config["explore"], trainable=False, dtype=tf.bool + ) + + # Log device and worker index. + num_gpus = self._get_num_gpus_for_policy() + if num_gpus > 0: + gpu_ids = get_gpu_devices() + logger.info(f"Found {len(gpu_ids)} visible cuda devices.") + + self._is_training = False + + self._loss_initialized = False + # Backward compatibility workaround so Policy will call self.loss() directly. + # TODO(jungong): clean up after all policies are migrated to new sub-class + # implementation. + self._loss = None + + self.batch_divisibility_req = self.get_batch_divisibility_req() + self._max_seq_len = self.config["model"]["max_seq_len"] + + self.validate_spaces(observation_space, action_space, self.config) + + # If using default make_model(), dist_class will get updated when + # the model is created next. + self.dist_class = self._init_dist_class() + self.model = self.make_model() + + self._init_view_requirements() + + self.exploration = self._create_exploration() + self._state_inputs = self.model.get_initial_state() + self._is_recurrent = len(self._state_inputs) > 0 + + # Got to reset global_timestep again after fake run-throughs. + self.global_timestep.assign(0) + + # Lock used for locking some methods on the object-level. + # This prevents possible race conditions when calling the model + # first, then its value function (e.g. in a loss function), in + # between of which another model call is made (e.g. to compute an + # action). + self._lock = threading.RLock() + + # Only for `config.eager_tracing=True`: A counter to keep track of + # how many times an eager-traced method (e.g. + # `self._compute_actions_helper`) has been re-traced by tensorflow. + # We will raise an error if more than n re-tracings have been + # detected, since this would considerably slow down execution. + # The variable below should only get incremented during the + # tf.function trace operations, never when calling the already + # traced function after that. + self._re_trace_counter = 0 + + @staticmethod + def enable_eager_execution_if_necessary(): + # If this class runs as a @ray.remote actor, eager mode may not + # have been activated yet. + if tf1 and not tf1.executing_eagerly(): + tf1.enable_eager_execution() + + @OverrideToImplementCustomLogic + def validate_spaces( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + ): + return {} + + @OverrideToImplementCustomLogic + @override(Policy) + def loss( + self, + model: Union[ModelV2, "tf.keras.Model"], + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + """Compute loss for this policy using model, dist_class and a train_batch. + + Args: + model: The Model to calculate the loss for. + dist_class: The action distr. class. + train_batch: The training data. + + Returns: + A single loss tensor or a list of loss tensors. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function. Returns a dict of statistics. + + Args: + train_batch: The SampleBatch (already) used for training. + + Returns: + The stats dict. + """ + return {} + + @OverrideToImplementCustomLogic + def grad_stats_fn( + self, train_batch: SampleBatch, grads: ModelGradients + ) -> Dict[str, TensorType]: + """Gradient stats function. Returns a dict of statistics. + + Args: + train_batch: The SampleBatch (already) used for training. + + Returns: + The stats dict. + """ + return {} + + @OverrideToImplementCustomLogic + def make_model(self) -> ModelV2: + """Build underlying model for this Policy. + + Returns: + The Model for the Policy to use. + """ + # Default ModelV2 model. + _, logit_dim = ModelCatalog.get_action_dist( + self.action_space, self.config["model"] + ) + return ModelCatalog.get_model_v2( + self.observation_space, + self.action_space, + logit_dim, + self.config["model"], + framework=self.framework, + ) + + @OverrideToImplementCustomLogic + def compute_gradients_fn( + self, policy: Policy, optimizer: LocalOptimizer, loss: TensorType + ) -> ModelGradients: + """Gradients computing function (from loss tensor, using local optimizer). + + Args: + policy: The Policy object that generated the loss tensor and + that holds the given local optimizer. + optimizer: The tf (local) optimizer object to + calculate the gradients with. + loss: The loss tensor for which gradients should be + calculated. + + Returns: + ModelGradients: List of the possibly clipped gradients- and variable + tuples. + """ + return None + + @OverrideToImplementCustomLogic + def apply_gradients_fn( + self, + optimizer: "tf.keras.optimizers.Optimizer", + grads: ModelGradients, + ) -> "tf.Operation": + """Gradients computing function (from loss tensor, using local optimizer). + + Args: + optimizer: The tf (local) optimizer object to + calculate the gradients with. + grads: The gradient tensor to be applied. + + Returns: + "tf.Operation": TF operation that applies supplied gradients. + """ + return None + + @OverrideToImplementCustomLogic + def action_sampler_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, TensorType, TensorType, List[TensorType]]: + """Custom function for sampling new actions given policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Sampled action + Log-likelihood + Action distribution inputs + Updated state + """ + return None, None, None, None + + @OverrideToImplementCustomLogic + def action_distribution_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, type, List[TensorType]]: + """Action distribution function for this Policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Distribution input. + ActionDistribution class. + State outs. + """ + return None, None, None + + @OverrideToImplementCustomLogic + def get_batch_divisibility_req(self) -> int: + """Get batch divisibility request. + + Returns: + Size N. A sample batch must be of size K*N. + """ + # By default, any sized batch is ok, so simply return 1. + return 1 + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_action_out_fn(self) -> Dict[str, TensorType]: + """Extra values to fetch and return from compute_actions(). + + Returns: + Dict[str, TensorType]: An extra fetch-dict to be passed to and + returned from the compute_actions() call. + """ + return {} + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_learn_fetches_fn(self) -> Dict[str, TensorType]: + """Extra stats to be reported after gradient computation. + + Returns: + Dict[str, TensorType]: An extra fetch-dict. + """ + return {} + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[SampleBatch] = None, + episode=None, + ): + """Post process trajectory in the format of a SampleBatch. + + Args: + sample_batch: sample_batch: batch of experiences for the policy, + which will contain at most one episode trajectory. + other_agent_batches: In a multi-agent env, this contains a + mapping of agent ids to (policy, agent_batch) tuples + containing the policy and experiences of the other agents. + episode: An optional multi-agent episode object to provide + access to all of the internal episode state, which may + be useful for model-based or multi-agent algorithms. + + Returns: + The postprocessed sample batch. + """ + assert tf.executing_eagerly() + return Policy.postprocess_trajectory(self, sample_batch) + + @OverrideToImplementCustomLogic + def optimizer( + self, + ) -> Union["tf.keras.optimizers.Optimizer", List["tf.keras.optimizers.Optimizer"]]: + """TF optimizer to use for policy optimization. + + Returns: + A local optimizer or a list of local optimizers to use for this + Policy's Model. + """ + return tf.keras.optimizers.Adam(self.config["lr"]) + + def _init_dist_class(self): + if is_overridden(self.action_sampler_fn) or is_overridden( + self.action_distribution_fn + ): + if not is_overridden(self.make_model): + raise ValueError( + "`make_model` is required if `action_sampler_fn` OR " + "`action_distribution_fn` is given" + ) + return None + else: + dist_class, _ = ModelCatalog.get_action_dist( + self.action_space, self.config["model"] + ) + return dist_class + + def _init_view_requirements(self): + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + # Combine view_requirements for Model and Policy. + self.view_requirements.update(self.model.view_requirements) + + # Disable env-info placeholder. + if SampleBatch.INFOS in self.view_requirements: + self.view_requirements[SampleBatch.INFOS].used_for_training = False + + def maybe_initialize_optimizer_and_loss(self): + optimizers = force_list(self.optimizer()) + if self.exploration: + # Policies with RLModules don't have an exploration object. + optimizers = self.exploration.get_exploration_optimizer(optimizers) + + # The list of local (tf) optimizers (one per loss term). + self._optimizers: List[LocalOptimizer] = optimizers + # Backward compatibility: A user's policy may only support a single + # loss term and optimizer (no lists). + self._optimizer: LocalOptimizer = optimizers[0] if optimizers else None + + self._initialize_loss_from_dummy_batch( + auto_remove_unneeded_view_reqs=True, + ) + self._loss_initialized = True + + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Dict[str, TensorType], + explore: bool = None, + timestep: Optional[int] = None, + episodes=None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + self._is_training = False + + explore = explore if explore is not None else self.explore + timestep = timestep if timestep is not None else self.global_timestep + if isinstance(timestep, tf.Tensor): + timestep = int(timestep.numpy()) + + # Pass lazy (eager) tensor dict to Model as `input_dict`. + input_dict = self._lazy_tensor_dict(input_dict) + input_dict.set_training(False) + + # Pack internal state inputs into (separate) list. + state_batches = [ + input_dict[k] for k in input_dict.keys() if "state_in" in k[:8] + ] + self._state_in = state_batches + self._is_recurrent = len(tree.flatten(self._state_in)) > 0 + + # Call the exploration before_compute_actions hook. + if self.exploration: + # Policies with RLModules don't have an exploration object. + self.exploration.before_compute_actions( + timestep=timestep, explore=explore, tf_sess=self.get_session() + ) + + ret = self._compute_actions_helper( + input_dict, + state_batches, + # TODO: Passing episodes into a traced method does not work. + None if self.config["eager_tracing"] else episodes, + explore, + timestep, + ) + # Update our global timestep by the batch size. + self.global_timestep.assign_add(tree.flatten(ret[0])[0].shape.as_list()[0]) + return convert_to_numpy(ret) + + # TODO(jungong) : deprecate this API and make compute_actions_from_input_dict the + # only canonical entry point for inference. + @override(Policy) + def compute_actions( + self, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + info_batch=None, + episodes=None, + explore=None, + timestep=None, + **kwargs, + ): + # Create input dict to simply pass the entire call to + # self.compute_actions_from_input_dict(). + input_dict = SampleBatch( + { + SampleBatch.CUR_OBS: obs_batch, + }, + _is_training=tf.constant(False), + ) + if state_batches is not None: + for s in enumerate(state_batches): + input_dict["state_in_{i}"] = s + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + if info_batch is not None: + input_dict[SampleBatch.INFOS] = info_batch + + return self.compute_actions_from_input_dict( + input_dict=input_dict, + explore=explore, + timestep=timestep, + episodes=episodes, + **kwargs, + ) + + @with_lock + @override(Policy) + def compute_log_likelihoods( + self, + actions: Union[List[TensorType], TensorType], + obs_batch: Union[List[TensorType], TensorType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Optional[Union[List[TensorType], TensorType]] = None, + prev_reward_batch: Optional[Union[List[TensorType], TensorType]] = None, + actions_normalized: bool = True, + in_training: bool = True, + ) -> TensorType: + if is_overridden(self.action_sampler_fn) and not is_overridden( + self.action_distribution_fn + ): + raise ValueError( + "Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!" + ) + + seq_lens = tf.ones(len(obs_batch), dtype=tf.int32) + input_batch = SampleBatch( + { + SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_batch), + SampleBatch.ACTIONS: actions, + }, + _is_training=False, + ) + if prev_action_batch is not None: + input_batch[SampleBatch.PREV_ACTIONS] = tf.convert_to_tensor( + prev_action_batch + ) + if prev_reward_batch is not None: + input_batch[SampleBatch.PREV_REWARDS] = tf.convert_to_tensor( + prev_reward_batch + ) + + # Exploration hook before each forward pass. + if self.exploration: + # Policies with RLModules don't have an exploration object. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if is_overridden(self.action_distribution_fn): + dist_inputs, self.dist_class, _ = self.action_distribution_fn( + self, self.model, input_batch, explore=False, is_training=False + ) + action_dist = self.dist_class(dist_inputs, self.model) + # Default log-likelihood calculation. + else: + dist_inputs, _ = self.model(input_batch, state_batches, seq_lens) + action_dist = self.dist_class(dist_inputs, self.model) + + # Normalize actions if necessary. + if not actions_normalized and self.config["normalize_actions"]: + actions = normalize_action(actions, self.action_space_struct) + + log_likelihoods = action_dist.logp(actions) + + return log_likelihoods + + @with_lock + @override(Policy) + def learn_on_batch(self, postprocessed_batch): + # Callback handling. + learn_stats = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=postprocessed_batch, result=learn_stats + ) + + pad_batch_to_sequences_of_same_size( + postprocessed_batch, + max_seq_len=self._max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + self._is_training = True + postprocessed_batch = self._lazy_tensor_dict(postprocessed_batch) + postprocessed_batch.set_training(True) + stats = self._learn_on_batch_helper(postprocessed_batch) + self.num_grad_updates += 1 + + stats.update( + { + "custom_metrics": learn_stats, + NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count, + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates + - 1 + - (postprocessed_batch.num_grad_updates or 0) + ), + } + ) + + return convert_to_numpy(stats) + + @override(Policy) + def compute_gradients( + self, postprocessed_batch: SampleBatch + ) -> Tuple[ModelGradients, Dict[str, TensorType]]: + + pad_batch_to_sequences_of_same_size( + postprocessed_batch, + shuffle=False, + max_seq_len=self._max_seq_len, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + self._is_training = True + self._lazy_tensor_dict(postprocessed_batch) + postprocessed_batch.set_training(True) + grads_and_vars, grads, stats = self._compute_gradients_helper( + postprocessed_batch + ) + return convert_to_numpy((grads, stats)) + + @override(Policy) + def apply_gradients(self, gradients: ModelGradients) -> None: + self._apply_gradients_helper( + list( + zip( + [ + (tf.convert_to_tensor(g) if g is not None else None) + for g in gradients + ], + self.model.trainable_variables(), + ) + ) + ) + + @override(Policy) + def get_weights(self, as_dict=False): + variables = self.variables() + if as_dict: + return {v.name: v.numpy() for v in variables} + return [v.numpy() for v in variables] + + @override(Policy) + def set_weights(self, weights): + variables = self.variables() + assert len(weights) == len(variables), (len(weights), len(variables)) + for v, w in zip(variables, weights): + v.assign(w) + + @override(Policy) + def get_exploration_state(self): + return convert_to_numpy(self.exploration.get_state()) + + @override(Policy) + def is_recurrent(self): + return self._is_recurrent + + @override(Policy) + def num_state_tensors(self): + return len(self._state_inputs) + + @override(Policy) + def get_initial_state(self): + if hasattr(self, "model"): + return self.model.get_initial_state() + return [] + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def get_state(self) -> PolicyState: + # Legacy Policy state (w/o keras model and w/o PolicySpec). + state = super().get_state() + + state["global_timestep"] = state["global_timestep"].numpy() + # In the new Learner API stack, the optimizers live in the learner. + state["_optimizer_variables"] = [] + if self._optimizer and len(self._optimizer.variables()) > 0: + state["_optimizer_variables"] = self._optimizer.variables() + + # Add exploration state. + if self.exploration: + # This is not compatible with RLModules, which have a method + # `forward_exploration` to specify custom exploration behavior. + state["_exploration_state"] = self.exploration.get_state() + + return state + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def set_state(self, state: PolicyState) -> None: + # Set optimizer vars. + optimizer_vars = state.get("_optimizer_variables", None) + if optimizer_vars and self._optimizer.variables(): + if not type(self).__name__.endswith("_traced") and log_once( + "set_state_optimizer_vars_tf_eager_policy_v2" + ): + logger.warning( + "Cannot restore an optimizer's state for tf eager! Keras " + "is not able to save the v1.x optimizers (from " + "tf.compat.v1.train) since they aren't compatible with " + "checkpoints." + ) + for opt_var, value in zip(self._optimizer.variables(), optimizer_vars): + opt_var.assign(value) + # Set exploration's state. + if hasattr(self, "exploration") and "_exploration_state" in state: + self.exploration.set_state(state=state["_exploration_state"]) + + # Restore glbal timestep (tf vars). + self.global_timestep.assign(state["global_timestep"]) + + # Then the Policy's (NN) weights and connectors. + super().set_state(state) + + @override(Policy) + def export_model(self, export_dir, onnx: Optional[int] = None) -> None: + if onnx: + try: + import tf2onnx + except ImportError as e: + raise RuntimeError( + "Converting a TensorFlow model to ONNX requires " + "`tf2onnx` to be installed. Install with " + "`pip install tf2onnx`." + ) from e + + model_proto, external_tensor_storage = tf2onnx.convert.from_keras( + self.model.base_model, + output_path=os.path.join(export_dir, "model.onnx"), + ) + # Save the tf.keras.Model (architecture and weights, so it can be retrieved + # w/o access to the original (custom) Model or Policy code). + elif ( + hasattr(self, "model") + and hasattr(self.model, "base_model") + and isinstance(self.model.base_model, tf.keras.Model) + ): + try: + self.model.base_model.save(export_dir, save_format="tf") + except Exception: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + else: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + + def variables(self): + """Return the list of all savable variables for this policy.""" + if isinstance(self.model, tf.keras.Model): + return self.model.variables + else: + return self.model.variables() + + def loss_initialized(self): + return self._loss_initialized + + @with_lock + def _compute_actions_helper( + self, + input_dict, + state_batches, + episodes, + explore, + timestep, + _ray_trace_ctx=None, + ): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + # Calculate RNN sequence lengths. + if SampleBatch.SEQ_LENS in input_dict: + seq_lens = input_dict[SampleBatch.SEQ_LENS] + else: + batch_size = tree.flatten(input_dict[SampleBatch.OBS])[0].shape[0] + seq_lens = tf.ones(batch_size, dtype=tf.int32) if state_batches else None + + # Add default and custom fetches. + extra_fetches = {} + + with tf.variable_creator_scope(_disallow_var_creation): + + if is_overridden(self.action_sampler_fn): + actions, logp, dist_inputs, state_out = self.action_sampler_fn( + self.model, + obs_batch=input_dict[SampleBatch.OBS], + state_batches=state_batches, + seq_lens=seq_lens, + explore=explore, + timestep=timestep, + episodes=episodes, + ) + else: + # Try `action_distribution_fn`. + if is_overridden(self.action_distribution_fn): + ( + dist_inputs, + self.dist_class, + state_out, + ) = self.action_distribution_fn( + self.model, + obs_batch=input_dict[SampleBatch.OBS], + state_batches=state_batches, + seq_lens=seq_lens, + explore=explore, + timestep=timestep, + is_training=False, + ) + elif isinstance(self.model, tf.keras.Model): + if state_batches and "state_in_0" not in input_dict: + for i, s in enumerate(state_batches): + input_dict[f"state_in_{i}"] = s + self._lazy_tensor_dict(input_dict) + dist_inputs, state_out, extra_fetches = self.model(input_dict) + else: + dist_inputs, state_out = self.model( + input_dict, state_batches, seq_lens + ) + + action_dist = self.dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = self.exploration.get_exploration_action( + action_distribution=action_dist, + timestep=timestep, + explore=explore, + ) + + # Action-logp and action-prob. + if logp is not None: + extra_fetches[SampleBatch.ACTION_PROB] = tf.exp(logp) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + # Custom extra fetches. + extra_fetches.update(self.extra_action_out_fn()) + + return actions, state_out, extra_fetches + + # TODO: Figure out, why _ray_trace_ctx=None helps to prevent a crash in + # AlphaStar w/ framework=tf2; eager_tracing=True on the policy learner actors. + # It seems there may be a clash between the traced-by-tf function and the + # traced-by-ray functions (for making the policy class a ray actor). + def _learn_on_batch_helper(self, samples, _ray_trace_ctx=None): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + with tf.variable_creator_scope(_disallow_var_creation): + grads_and_vars, _, stats = self._compute_gradients_helper(samples) + self._apply_gradients_helper(grads_and_vars) + return stats + + def _get_is_training_placeholder(self): + return tf.convert_to_tensor(self._is_training) + + @with_lock + def _compute_gradients_helper(self, samples): + """Computes and returns grads as eager tensors.""" + + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + # Gather all variables for which to calculate losses. + if isinstance(self.model, tf.keras.Model): + variables = self.model.trainable_variables + else: + variables = self.model.trainable_variables() + + # Calculate the loss(es) inside a tf GradientTape. + with tf.GradientTape( + persistent=is_overridden(self.compute_gradients_fn) + ) as tape: + losses = self.loss(self.model, self.dist_class, samples) + losses = force_list(losses) + + # User provided a custom compute_gradients_fn. + if is_overridden(self.compute_gradients_fn): + # Wrap our tape inside a wrapper, such that the resulting + # object looks like a "classic" tf.optimizer. This way, custom + # compute_gradients_fn will work on both tf static graph + # and tf-eager. + optimizer = _OptimizerWrapper(tape) + # More than one loss terms/optimizers. + if self.config["_tf_policy_handles_more_than_one_loss"]: + grads_and_vars = self.compute_gradients_fn( + [optimizer] * len(losses), losses + ) + # Only one loss and one optimizer. + else: + grads_and_vars = [self.compute_gradients_fn(optimizer, losses[0])] + # Default: Compute gradients using the above tape. + else: + grads_and_vars = [ + list(zip(tape.gradient(loss, variables), variables)) for loss in losses + ] + + if log_once("grad_vars"): + for g_and_v in grads_and_vars: + for g, v in g_and_v: + if g is not None: + logger.info(f"Optimizing variable {v.name}") + + # `grads_and_vars` is returned a list (len=num optimizers/losses) + # of lists of (grad, var) tuples. + if self.config["_tf_policy_handles_more_than_one_loss"]: + grads = [[g for g, _ in g_and_v] for g_and_v in grads_and_vars] + # `grads_and_vars` is returned as a list of (grad, var) tuples. + else: + grads_and_vars = grads_and_vars[0] + grads = [g for g, _ in grads_and_vars] + + stats = self._stats(samples, grads) + return grads_and_vars, grads, stats + + def _apply_gradients_helper(self, grads_and_vars): + # Increase the tracing counter to make sure we don't re-trace too + # often. If eager_tracing=True, this counter should only get + # incremented during the @tf.function trace operations, never when + # calling the already traced function after that. + self._re_trace_counter += 1 + + if is_overridden(self.apply_gradients_fn): + if self.config["_tf_policy_handles_more_than_one_loss"]: + self.apply_gradients_fn(self._optimizers, grads_and_vars) + else: + self.apply_gradients_fn(self._optimizer, grads_and_vars) + else: + if self.config["_tf_policy_handles_more_than_one_loss"]: + for i, o in enumerate(self._optimizers): + o.apply_gradients( + [(g, v) for g, v in grads_and_vars[i] if g is not None] + ) + else: + self._optimizer.apply_gradients( + [(g, v) for g, v in grads_and_vars if g is not None] + ) + + def _stats(self, samples, grads): + fetches = {} + if is_overridden(self.stats_fn): + fetches[LEARNER_STATS_KEY] = dict(self.stats_fn(samples)) + else: + fetches[LEARNER_STATS_KEY] = {} + + fetches.update(dict(self.extra_learn_fetches_fn())) + fetches.update(dict(self.grad_stats_fn(samples, grads))) + return fetches + + def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch): + # TODO: (sven): Keep for a while to ensure backward compatibility. + if not isinstance(postprocessed_batch, SampleBatch): + postprocessed_batch = SampleBatch(postprocessed_batch) + postprocessed_batch.set_get_interceptor(_convert_to_tf) + return postprocessed_batch + + @classmethod + def with_tracing(cls): + return _traced_eager_policy(cls) diff --git a/lib/python3.12/site-packages/ray/rllib/policy/policy.py b/lib/python3.12/site-packages/ray/rllib/policy/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..b32fcb8098a75a10aa54e9b99c1812b337c13b7f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/policy.py @@ -0,0 +1,1696 @@ +import json +import logging +import os +import platform +from abc import ABCMeta, abstractmethod +from typing import ( + Any, + Callable, + Collection, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box +from packaging import version + +import ray +import ray.cloudpickle as pickle +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.actor import ActorHandle +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + is_overridden, +) +from ray.rllib.utils.checkpoints import ( + CHECKPOINT_VERSION, + get_checkpoint_info, + try_import_msgpack, +) +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.serialization import ( + deserialize_type, + space_from_dict, + space_to_dict, +) +from ray.rllib.utils.spaces.space_utils import ( + get_base_struct_from_space, + get_dummy_batch_for_space, + unbatch, +) +from ray.rllib.utils.tensor_dtype import get_np_dtype +from ray.rllib.utils.tf_utils import get_tf_eager_cls_if_necessary +from ray.rllib.utils.typing import ( + AgentID, + AlgorithmConfigDict, + ModelGradients, + ModelWeights, + PolicyID, + PolicyState, + T, + TensorStructType, + TensorType, +) +from ray.tune import Checkpoint + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class PolicySpec: + """A policy spec used in the "config.multiagent.policies" specification dict. + + As values (keys are the policy IDs (str)). E.g.: + config: + multiagent: + policies: { + "pol1": PolicySpec(None, Box, Discrete(2), {"lr": 0.0001}), + "pol2": PolicySpec(config={"lr": 0.001}), + } + """ + + def __init__( + self, policy_class=None, observation_space=None, action_space=None, config=None + ): + # If None, use the Algorithm's default policy class stored under + # `Algorithm._policy_class`. + self.policy_class = policy_class + # If None, use the env's observation space. If None and there is no Env + # (e.g. offline RL), an error is thrown. + self.observation_space = observation_space + # If None, use the env's action space. If None and there is no Env + # (e.g. offline RL), an error is thrown. + self.action_space = action_space + # Overrides defined keys in the main Algorithm config. + # If None, use {}. + self.config = config + + def __eq__(self, other: "PolicySpec"): + return ( + self.policy_class == other.policy_class + and self.observation_space == other.observation_space + and self.action_space == other.action_space + and self.config == other.config + ) + + def get_state(self) -> Dict[str, Any]: + """Returns the state of a `PolicyDict` as a dict.""" + return ( + self.policy_class, + self.observation_space, + self.action_space, + self.config, + ) + + @classmethod + def from_state(cls, state: Dict[str, Any]) -> "PolicySpec": + """Builds a `PolicySpec` from a state.""" + policy_spec = PolicySpec() + policy_spec.__dict__.update(state) + + return policy_spec + + def serialize(self) -> Dict: + from ray.rllib.algorithms.registry import get_policy_class_name + + # Try to figure out a durable name for this policy. + cls = get_policy_class_name(self.policy_class) + if cls is None: + logger.warning( + f"Can not figure out a durable policy name for {self.policy_class}. " + f"You are probably trying to checkpoint a custom policy. " + f"Raw policy class may cause problems when the checkpoint needs to " + "be loaded in the future. To fix this, make sure you add your " + "custom policy in rllib.algorithms.registry.POLICIES." + ) + cls = self.policy_class + + return { + "policy_class": cls, + "observation_space": space_to_dict(self.observation_space), + "action_space": space_to_dict(self.action_space), + # TODO(jungong) : try making the config dict durable by maybe + # getting rid of all the fields that are not JSON serializable. + "config": self.config, + } + + @classmethod + def deserialize(cls, spec: Dict) -> "PolicySpec": + if isinstance(spec["policy_class"], str): + # Try to recover the actual policy class from durable name. + from ray.rllib.algorithms.registry import get_policy_class + + policy_class = get_policy_class(spec["policy_class"]) + elif isinstance(spec["policy_class"], type): + # Policy spec is already a class type. Simply use it. + policy_class = spec["policy_class"] + else: + raise AttributeError(f"Unknown policy class spec {spec['policy_class']}") + + return cls( + policy_class=policy_class, + observation_space=space_from_dict(spec["observation_space"]), + action_space=space_from_dict(spec["action_space"]), + config=spec["config"], + ) + + +@OldAPIStack +class Policy(metaclass=ABCMeta): + """RLlib's base class for all Policy implementations. + + Policy is the abstract superclass for all DL-framework specific sub-classes + (e.g. TFPolicy or TorchPolicy). It exposes APIs to + + 1. Compute actions from observation (and possibly other) inputs. + + 2. Manage the Policy's NN model(s), like exporting and loading their weights. + + 3. Postprocess a given trajectory from the environment or other input via the + `postprocess_trajectory` method. + + 4. Compute losses from a train batch. + + 5. Perform updates from a train batch on the NN-models (this normally includes loss + calculations) either: + + a. in one monolithic step (`learn_on_batch`) + + b. via batch pre-loading, then n steps of actual loss computations and updates + (`load_batch_into_buffer` + `learn_on_loaded_batch`). + """ + + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + config: AlgorithmConfigDict, + ): + """Initializes a Policy instance. + + Args: + observation_space: Observation space of the policy. + action_space: Action space of the policy. + config: A complete Algorithm/Policy config dict. For the default + config keys and values, see rllib/algorithm/algorithm.py. + """ + self.observation_space: gym.Space = observation_space + self.action_space: gym.Space = action_space + # the policy id in the global context. + self.__policy_id = config.get("__policy_id") + # The base struct of the observation/action spaces. + # E.g. action-space = gym.spaces.Dict({"a": Discrete(2)}) -> + # action_space_struct = {"a": Discrete(2)} + self.observation_space_struct = get_base_struct_from_space(observation_space) + self.action_space_struct = get_base_struct_from_space(action_space) + + self.config: AlgorithmConfigDict = config + self.framework = self.config.get("framework") + + # Create the callbacks object to use for handling custom callbacks. + from ray.rllib.callbacks.callbacks import RLlibCallback + + callbacks = self.config.get("callbacks") + if isinstance(callbacks, RLlibCallback): + self.callbacks = callbacks() + elif isinstance(callbacks, (str, type)): + try: + self.callbacks: "RLlibCallback" = deserialize_type( + self.config.get("callbacks") + )() + except Exception: + pass # TEST + else: + self.callbacks: "RLlibCallback" = RLlibCallback() + + # The global timestep, broadcast down from time to time from the + # local worker to all remote workers. + self.global_timestep: int = 0 + # The number of gradient updates this policy has undergone. + self.num_grad_updates: int = 0 + + # The action distribution class to use for action sampling, if any. + # Child classes may set this. + self.dist_class: Optional[Type] = None + + # Initialize view requirements. + self.init_view_requirements() + + # Whether the Model's initial state (method) has been added + # automatically based on the given view requirements of the model. + self._model_init_state_automatically_added = False + + # Connectors. + self.agent_connectors = None + self.action_connectors = None + + @staticmethod + def from_checkpoint( + checkpoint: Union[str, Checkpoint], + policy_ids: Optional[Collection[PolicyID]] = None, + ) -> Union["Policy", Dict[PolicyID, "Policy"]]: + """Creates new Policy instance(s) from a given Policy or Algorithm checkpoint. + + Note: This method must remain backward compatible from 2.1.0 on, wrt. + checkpoints created with Ray 2.0.0 or later. + + Args: + checkpoint: The path (str) to a Policy or Algorithm checkpoint directory + or an AIR Checkpoint (Policy or Algorithm) instance to restore + from. + If checkpoint is a Policy checkpoint, `policy_ids` must be None + and only the Policy in that checkpoint is restored and returned. + If checkpoint is an Algorithm checkpoint and `policy_ids` is None, + will return a list of all Policy objects found in + the checkpoint, otherwise a list of those policies in `policy_ids`. + policy_ids: List of policy IDs to extract from a given Algorithm checkpoint. + If None and an Algorithm checkpoint is provided, will restore all + policies found in that checkpoint. If a Policy checkpoint is given, + this arg must be None. + + Returns: + An instantiated Policy, if `checkpoint` is a Policy checkpoint. A dict + mapping PolicyID to Policies, if `checkpoint` is an Algorithm checkpoint. + In the latter case, returns all policies within the Algorithm if + `policy_ids` is None, else a dict of only those Policies that are in + `policy_ids`. + """ + checkpoint_info = get_checkpoint_info(checkpoint) + + # Algorithm checkpoint: Extract one or more policies from it and return them + # in a dict (mapping PolicyID to Policy instances). + if checkpoint_info["type"] == "Algorithm": + from ray.rllib.algorithms.algorithm import Algorithm + + policies = {} + + # Old Algorithm checkpoints: State must be completely retrieved from: + # algo state file -> worker -> "state". + if checkpoint_info["checkpoint_version"] < version.Version("1.0"): + with open(checkpoint_info["state_file"], "rb") as f: + state = pickle.load(f) + # In older checkpoint versions, the policy states are stored under + # "state" within the worker state (which is pickled in itself). + worker_state = pickle.loads(state["worker"]) + policy_states = worker_state["state"] + for pid, policy_state in policy_states.items(): + # Get spec and config, merge config with + serialized_policy_spec = worker_state["policy_specs"][pid] + policy_config = Algorithm.merge_algorithm_configs( + worker_state["policy_config"], serialized_policy_spec["config"] + ) + serialized_policy_spec.update({"config": policy_config}) + policy_state.update({"policy_spec": serialized_policy_spec}) + policies[pid] = Policy.from_state(policy_state) + # Newer versions: Get policy states from "policies/" sub-dirs. + elif checkpoint_info["policy_ids"] is not None: + for policy_id in checkpoint_info["policy_ids"]: + if policy_ids is None or policy_id in policy_ids: + policy_checkpoint_info = get_checkpoint_info( + os.path.join( + checkpoint_info["checkpoint_dir"], + "policies", + policy_id, + ) + ) + assert policy_checkpoint_info["type"] == "Policy" + with open(policy_checkpoint_info["state_file"], "rb") as f: + policy_state = pickle.load(f) + policies[policy_id] = Policy.from_state(policy_state) + return policies + + # Policy checkpoint: Return a single Policy instance. + else: + msgpack = None + if checkpoint_info.get("format") == "msgpack": + msgpack = try_import_msgpack(error=True) + + with open(checkpoint_info["state_file"], "rb") as f: + if msgpack is not None: + state = msgpack.load(f) + else: + state = pickle.load(f) + return Policy.from_state(state) + + @staticmethod + def from_state(state: PolicyState) -> "Policy": + """Recovers a Policy from a state object. + + The `state` of an instantiated Policy can be retrieved by calling its + `get_state` method. This only works for the V2 Policy classes (EagerTFPolicyV2, + SynamicTFPolicyV2, and TorchPolicyV2). It contains all information necessary + to create the Policy. No access to the original code (e.g. configs, knowledge of + the policy's class, etc..) is needed. + + Args: + state: The state to recover a new Policy instance from. + + Returns: + A new Policy instance. + """ + serialized_pol_spec: Optional[dict] = state.get("policy_spec") + if serialized_pol_spec is None: + raise ValueError( + "No `policy_spec` key was found in given `state`! " + "Cannot create new Policy." + ) + pol_spec = PolicySpec.deserialize(serialized_pol_spec) + actual_class = get_tf_eager_cls_if_necessary( + pol_spec.policy_class, + pol_spec.config, + ) + + if pol_spec.config["framework"] == "tf": + from ray.rllib.policy.tf_policy import TFPolicy + + return TFPolicy._tf1_from_state_helper(state) + + # Create the new policy. + new_policy = actual_class( + # Note(jungong) : we are intentionally not using keyward arguments here + # because some policies name the observation space parameter obs_space, + # and some others name it observation_space. + pol_spec.observation_space, + pol_spec.action_space, + pol_spec.config, + ) + + # Set the new policy's state (weights, optimizer vars, exploration state, + # etc..). + new_policy.set_state(state) + # Return the new policy. + return new_policy + + def init_view_requirements(self): + """Maximal view requirements dict for `learn_on_batch()` and + `compute_actions` calls. + Specific policies can override this function to provide custom + list of view requirements. + """ + # Maximal view requirements dict for `learn_on_batch()` and + # `compute_actions` calls. + # View requirements will be automatically filtered out later based + # on the postprocessing and loss functions to ensure optimal data + # collection and transfer performance. + view_reqs = self._get_default_view_requirements() + if not hasattr(self, "view_requirements"): + self.view_requirements = view_reqs + else: + for k, v in view_reqs.items(): + if k not in self.view_requirements: + self.view_requirements[k] = v + + def get_connector_metrics(self) -> Dict: + """Get metrics on timing from connectors.""" + return { + "agent_connectors": { + name + "_ms": 1000 * timer.mean + for name, timer in self.agent_connectors.timers.items() + }, + "action_connectors": { + name + "_ms": 1000 * timer.mean + for name, timer in self.agent_connectors.timers.items() + }, + } + + def reset_connectors(self, env_id) -> None: + """Reset action- and agent-connectors for this policy.""" + self.agent_connectors.reset(env_id=env_id) + self.action_connectors.reset(env_id=env_id) + + def compute_single_action( + self, + obs: Optional[TensorStructType] = None, + state: Optional[List[TensorType]] = None, + *, + prev_action: Optional[TensorStructType] = None, + prev_reward: Optional[TensorStructType] = None, + info: dict = None, + input_dict: Optional[SampleBatch] = None, + episode=None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + # Kwars placeholder for future compatibility. + **kwargs, + ) -> Tuple[TensorStructType, List[TensorType], Dict[str, TensorType]]: + """Computes and returns a single (B=1) action value. + + Takes an input dict (usually a SampleBatch) as its main data input. + This allows for using this method in case a more complex input pattern + (view requirements) is needed, for example when the Model requires the + last n observations, the last m actions/rewards, or a combination + of any of these. + Alternatively, in case no complex inputs are required, takes a single + `obs` values (and possibly single state values, prev-action/reward + values, etc..). + + Args: + obs: Single observation. + state: List of RNN state inputs, if any. + prev_action: Previous action value, if any. + prev_reward: Previous reward, if any. + info: Info object, if any. + input_dict: A SampleBatch or input dict containing the + single (unbatched) Tensors to compute actions. If given, it'll + be used instead of `obs`, `state`, `prev_action|reward`, and + `info`. + episode: This provides access to all of the internal episode state, + which may be useful for model-based or multi-agent algorithms. + explore: Whether to pick an exploitation or + exploration action + (default: None -> use self.config["explore"]). + timestep: The current (sampling) time step. + + Keyword Args: + kwargs: Forward compatibility placeholder. + + Returns: + Tuple consisting of the action, the list of RNN state outputs (if + any), and a dictionary of extra features (if any). + """ + # Build the input-dict used for the call to + # `self.compute_actions_from_input_dict()`. + if input_dict is None: + input_dict = {SampleBatch.OBS: obs} + if state is not None: + for i, s in enumerate(state): + input_dict[f"state_in_{i}"] = s + if prev_action is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action + if prev_reward is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward + if info is not None: + input_dict[SampleBatch.INFOS] = info + + # Batch all data in input dict. + input_dict = tree.map_structure_with_path( + lambda p, s: ( + s + if p == "seq_lens" + else s.unsqueeze(0) + if torch and isinstance(s, torch.Tensor) + else np.expand_dims(s, 0) + ), + input_dict, + ) + + episodes = None + if episode is not None: + episodes = [episode] + + out = self.compute_actions_from_input_dict( + input_dict=SampleBatch(input_dict), + episodes=episodes, + explore=explore, + timestep=timestep, + ) + + # Some policies don't return a tuple, but always just a single action. + # E.g. ES and ARS. + if not isinstance(out, tuple): + single_action = out + state_out = [] + info = {} + # Normal case: Policy should return (action, state, info) tuple. + else: + batched_action, state_out, info = out + single_action = unbatch(batched_action) + assert len(single_action) == 1 + single_action = single_action[0] + + # Return action, internal state(s), infos. + return ( + single_action, + tree.map_structure(lambda x: x[0], state_out), + tree.map_structure(lambda x: x[0], info), + ) + + def compute_actions_from_input_dict( + self, + input_dict: Union[SampleBatch, Dict[str, TensorStructType]], + explore: Optional[bool] = None, + timestep: Optional[int] = None, + episodes=None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + """Computes actions from collected samples (across multiple-agents). + + Takes an input dict (usually a SampleBatch) as its main data input. + This allows for using this method in case a more complex input pattern + (view requirements) is needed, for example when the Model requires the + last n observations, the last m actions/rewards, or a combination + of any of these. + + Args: + input_dict: A SampleBatch or input dict containing the Tensors + to compute actions. `input_dict` already abides to the + Policy's as well as the Model's view requirements and can + thus be passed to the Model as-is. + explore: Whether to pick an exploitation or exploration + action (default: None -> use self.config["explore"]). + timestep: The current (sampling) time step. + episodes: This provides access to all of the internal episodes' + state, which may be useful for model-based or multi-agent + algorithms. + + Keyword Args: + kwargs: Forward compatibility placeholder. + + Returns: + actions: Batch of output actions, with shape like + [BATCH_SIZE, ACTION_SHAPE]. + state_outs: List of RNN state output + batches, if any, each with shape [BATCH_SIZE, STATE_SIZE]. + info: Dictionary of extra feature batches, if any, with shape like + {"f1": [BATCH_SIZE, ...], "f2": [BATCH_SIZE, ...]}. + """ + # Default implementation just passes obs, prev-a/r, and states on to + # `self.compute_actions()`. + state_batches = [s for k, s in input_dict.items() if k.startswith("state_in")] + return self.compute_actions( + input_dict[SampleBatch.OBS], + state_batches, + prev_action_batch=input_dict.get(SampleBatch.PREV_ACTIONS), + prev_reward_batch=input_dict.get(SampleBatch.PREV_REWARDS), + info_batch=input_dict.get(SampleBatch.INFOS), + explore=explore, + timestep=timestep, + episodes=episodes, + **kwargs, + ) + + @abstractmethod + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorStructType], TensorStructType] = None, + prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None, + info_batch: Optional[Dict[str, list]] = None, + episodes: Optional[List] = None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + """Computes actions for the current policy. + + Args: + obs_batch: Batch of observations. + state_batches: List of RNN state input batches, if any. + prev_action_batch: Batch of previous action values. + prev_reward_batch: Batch of previous rewards. + info_batch: Batch of info objects. + episodes: List of Episode objects, one for each obs in + obs_batch. This provides access to all of the internal + episode state, which may be useful for model-based or + multi-agent algorithms. + explore: Whether to pick an exploitation or exploration action. + Set to None (default) for using the value of + `self.config["explore"]`. + timestep: The current (sampling) time step. + + Keyword Args: + kwargs: Forward compatibility placeholder + + Returns: + actions: Batch of output actions, with shape like + [BATCH_SIZE, ACTION_SHAPE]. + state_outs (List[TensorType]): List of RNN state output + batches, if any, each with shape [BATCH_SIZE, STATE_SIZE]. + info (List[dict]): Dictionary of extra feature batches, if any, + with shape like + {"f1": [BATCH_SIZE, ...], "f2": [BATCH_SIZE, ...]}. + """ + raise NotImplementedError + + def compute_log_likelihoods( + self, + actions: Union[List[TensorType], TensorType], + obs_batch: Union[List[TensorType], TensorType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Optional[Union[List[TensorType], TensorType]] = None, + prev_reward_batch: Optional[Union[List[TensorType], TensorType]] = None, + actions_normalized: bool = True, + in_training: bool = True, + ) -> TensorType: + """Computes the log-prob/likelihood for a given action and observation. + + The log-likelihood is calculated using this Policy's action + distribution class (self.dist_class). + + Args: + actions: Batch of actions, for which to retrieve the + log-probs/likelihoods (given all other inputs: obs, + states, ..). + obs_batch: Batch of observations. + state_batches: List of RNN state input batches, if any. + prev_action_batch: Batch of previous action values. + prev_reward_batch: Batch of previous rewards. + actions_normalized: Is the given `actions` already normalized + (between -1.0 and 1.0) or not? If not and + `normalize_actions=True`, we need to normalize the given + actions first, before calculating log likelihoods. + in_training: Whether to use the forward_train() or forward_exploration() of + the underlying RLModule. + Returns: + Batch of log probs/likelihoods, with shape: [BATCH_SIZE]. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[ + Dict[AgentID, Tuple["Policy", SampleBatch]] + ] = None, + episode=None, + ) -> SampleBatch: + """Implements algorithm-specific trajectory postprocessing. + + This will be called on each trajectory fragment computed during policy + evaluation. Each fragment is guaranteed to be only from one episode. + The given fragment may or may not contain the end of this episode, + depending on the `batch_mode=truncate_episodes|complete_episodes`, + `rollout_fragment_length`, and other settings. + + Args: + sample_batch: batch of experiences for the policy, + which will contain at most one episode trajectory. + other_agent_batches: In a multi-agent env, this contains a + mapping of agent ids to (policy, agent_batch) tuples + containing the policy and experiences of the other agents. + episode: An optional multi-agent episode object to provide + access to all of the internal episode state, which may + be useful for model-based or multi-agent algorithms. + + Returns: + The postprocessed sample batch. + """ + # The default implementation just returns the same, unaltered batch. + return sample_batch + + @OverrideToImplementCustomLogic + def loss( + self, model: ModelV2, dist_class: ActionDistribution, train_batch: SampleBatch + ) -> Union[TensorType, List[TensorType]]: + """Loss function for this Policy. + + Override this method in order to implement custom loss computations. + + Args: + model: The model to calculate the loss(es). + dist_class: The action distribution class to sample actions + from the model's outputs. + train_batch: The input batch on which to calculate the loss. + + Returns: + Either a single loss tensor or a list of loss tensors. + """ + raise NotImplementedError + + def learn_on_batch(self, samples: SampleBatch) -> Dict[str, TensorType]: + """Perform one learning update, given `samples`. + + Either this method or the combination of `compute_gradients` and + `apply_gradients` must be implemented by subclasses. + + Args: + samples: The SampleBatch object to learn from. + + Returns: + Dictionary of extra metadata from `compute_gradients()`. + + .. testcode:: + :skipif: True + + policy, sample_batch = ... + policy.learn_on_batch(sample_batch) + """ + # The default implementation is simply a fused `compute_gradients` plus + # `apply_gradients` call. + grads, grad_info = self.compute_gradients(samples) + self.apply_gradients(grads) + return grad_info + + def learn_on_batch_from_replay_buffer( + self, replay_actor: ActorHandle, policy_id: PolicyID + ) -> Dict[str, TensorType]: + """Samples a batch from given replay actor and performs an update. + + Args: + replay_actor: The replay buffer actor to sample from. + policy_id: The ID of this policy. + + Returns: + Dictionary of extra metadata from `compute_gradients()`. + """ + # Sample a batch from the given replay actor. + # Note that for better performance (less data sent through the + # network), this policy should be co-located on the same node + # as `replay_actor`. Such a co-location step is usually done during + # the Algorithm's `setup()` phase. + batch = ray.get(replay_actor.replay.remote(policy_id=policy_id)) + if batch is None: + return {} + + # Send to own learn_on_batch method for updating. + # TODO: hack w/ `hasattr` + if hasattr(self, "devices") and len(self.devices) > 1: + self.load_batch_into_buffer(batch, buffer_index=0) + return self.learn_on_loaded_batch(offset=0, buffer_index=0) + else: + return self.learn_on_batch(batch) + + def load_batch_into_buffer(self, batch: SampleBatch, buffer_index: int = 0) -> int: + """Bulk-loads the given SampleBatch into the devices' memories. + + The data is split equally across all the Policy's devices. + If the data is not evenly divisible by the batch size, excess data + should be discarded. + + Args: + batch: The SampleBatch to load. + buffer_index: The index of the buffer (a MultiGPUTowerStack) to use + on the devices. The number of buffers on each device depends + on the value of the `num_multi_gpu_tower_stacks` config key. + + Returns: + The number of tuples loaded per device. + """ + raise NotImplementedError + + def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int: + """Returns the number of currently loaded samples in the given buffer. + + Args: + buffer_index: The index of the buffer (a MultiGPUTowerStack) + to use on the devices. The number of buffers on each device + depends on the value of the `num_multi_gpu_tower_stacks` config + key. + + Returns: + The number of tuples loaded per device. + """ + raise NotImplementedError + + def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0): + """Runs a single step of SGD on an already loaded data in a buffer. + + Runs an SGD step over a slice of the pre-loaded batch, offset by + the `offset` argument (useful for performing n minibatch SGD + updates repeatedly on the same, already pre-loaded data). + + Updates the model weights based on the averaged per-device gradients. + + Args: + offset: Offset into the preloaded data. Used for pre-loading + a train-batch once to a device, then iterating over + (subsampling through) this batch n times doing minibatch SGD. + buffer_index: The index of the buffer (a MultiGPUTowerStack) + to take the already pre-loaded data from. The number of buffers + on each device depends on the value of the + `num_multi_gpu_tower_stacks` config key. + + Returns: + The outputs of extra_ops evaluated over the batch. + """ + raise NotImplementedError + + def compute_gradients( + self, postprocessed_batch: SampleBatch + ) -> Tuple[ModelGradients, Dict[str, TensorType]]: + """Computes gradients given a batch of experiences. + + Either this in combination with `apply_gradients()` or + `learn_on_batch()` must be implemented by subclasses. + + Args: + postprocessed_batch: The SampleBatch object to use + for calculating gradients. + + Returns: + grads: List of gradient output values. + grad_info: Extra policy-specific info values. + """ + raise NotImplementedError + + def apply_gradients(self, gradients: ModelGradients) -> None: + """Applies the (previously) computed gradients. + + Either this in combination with `compute_gradients()` or + `learn_on_batch()` must be implemented by subclasses. + + Args: + gradients: The already calculated gradients to apply to this + Policy. + """ + raise NotImplementedError + + def get_weights(self) -> ModelWeights: + """Returns model weights. + + Note: The return value of this method will reside under the "weights" + key in the return value of Policy.get_state(). Model weights are only + one part of a Policy's state. Other state information contains: + optimizer variables, exploration state, and global state vars such as + the sampling timestep. + + Returns: + Serializable copy or view of model weights. + """ + raise NotImplementedError + + def set_weights(self, weights: ModelWeights) -> None: + """Sets this Policy's model's weights. + + Note: Model weights are only one part of a Policy's state. Other + state information contains: optimizer variables, exploration state, + and global state vars such as the sampling timestep. + + Args: + weights: Serializable copy or view of model weights. + """ + raise NotImplementedError + + def get_exploration_state(self) -> Dict[str, TensorType]: + """Returns the state of this Policy's exploration component. + + Returns: + Serializable information on the `self.exploration` object. + """ + return self.exploration.get_state() + + def is_recurrent(self) -> bool: + """Whether this Policy holds a recurrent Model. + + Returns: + True if this Policy has-a RNN-based Model. + """ + return False + + def num_state_tensors(self) -> int: + """The number of internal states needed by the RNN-Model of the Policy. + + Returns: + int: The number of RNN internal states kept by this Policy's Model. + """ + return 0 + + def get_initial_state(self) -> List[TensorType]: + """Returns initial RNN state for the current policy. + + Returns: + List[TensorType]: Initial RNN state for the current policy. + """ + return [] + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def get_state(self) -> PolicyState: + """Returns the entire current state of this Policy. + + Note: Not to be confused with an RNN model's internal state. + State includes the Model(s)' weights, optimizer weights, + the exploration component's state, as well as global variables, such + as sampling timesteps. + + Note that the state may contain references to the original variables. + This means that you may need to deepcopy() the state before mutating it. + + Returns: + Serialized local state. + """ + state = { + # All the policy's weights. + "weights": self.get_weights(), + # The current global timestep. + "global_timestep": self.global_timestep, + # The current num_grad_updates counter. + "num_grad_updates": self.num_grad_updates, + } + + # Add this Policy's spec so it can be retreived w/o access to the original + # code. + policy_spec = PolicySpec( + policy_class=type(self), + observation_space=self.observation_space, + action_space=self.action_space, + config=self.config, + ) + state["policy_spec"] = policy_spec.serialize() + + # Checkpoint connectors state as well if enabled. + connector_configs = {} + if self.agent_connectors: + connector_configs["agent"] = self.agent_connectors.to_state() + if self.action_connectors: + connector_configs["action"] = self.action_connectors.to_state() + state["connector_configs"] = connector_configs + + return state + + def restore_connectors(self, state: PolicyState): + """Restore agent and action connectors if configs available. + + Args: + state: The new state to set this policy to. Can be + obtained by calling `self.get_state()`. + """ + # To avoid a circular dependency problem cause by SampleBatch. + from ray.rllib.connectors.util import restore_connectors_for_policy + + connector_configs = state.get("connector_configs", {}) + if "agent" in connector_configs: + self.agent_connectors = restore_connectors_for_policy( + self, connector_configs["agent"] + ) + logger.debug("restoring agent connectors:") + logger.debug(self.agent_connectors.__str__(indentation=4)) + if "action" in connector_configs: + self.action_connectors = restore_connectors_for_policy( + self, connector_configs["action"] + ) + logger.debug("restoring action connectors:") + logger.debug(self.action_connectors.__str__(indentation=4)) + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def set_state(self, state: PolicyState) -> None: + """Restores the entire current state of this Policy from `state`. + + Args: + state: The new state to set this policy to. Can be + obtained by calling `self.get_state()`. + """ + if "policy_spec" in state: + policy_spec = PolicySpec.deserialize(state["policy_spec"]) + # Assert spaces remained the same. + if ( + policy_spec.observation_space is not None + and policy_spec.observation_space != self.observation_space + ): + logger.warning( + "`observation_space` in given policy state (" + f"{policy_spec.observation_space}) does not match this Policy's " + f"observation space ({self.observation_space})." + ) + if ( + policy_spec.action_space is not None + and policy_spec.action_space != self.action_space + ): + logger.warning( + "`action_space` in given policy state (" + f"{policy_spec.action_space}) does not match this Policy's " + f"action space ({self.action_space})." + ) + # Override config, if part of the spec. + if policy_spec.config: + self.config = policy_spec.config + + # Override NN weights. + self.set_weights(state["weights"]) + self.restore_connectors(state) + + def apply( + self, + func: Callable[["Policy", Optional[Any], Optional[Any]], T], + *args, + **kwargs, + ) -> T: + """Calls the given function with this Policy instance. + + Useful for when the Policy class has been converted into a ActorHandle + and the user needs to execute some functionality (e.g. add a property) + on the underlying policy object. + + Args: + func: The function to call, with this Policy as first + argument, followed by args, and kwargs. + args: Optional additional args to pass to the function call. + kwargs: Optional additional kwargs to pass to the function call. + + Returns: + The return value of the function call. + """ + return func(self, *args, **kwargs) + + def on_global_var_update(self, global_vars: Dict[str, TensorType]) -> None: + """Called on an update to global vars. + + Args: + global_vars: Global variables by str key, broadcast from the + driver. + """ + # Store the current global time step (sum over all policies' sample + # steps). + # Make sure, we keep global_timestep as a Tensor for tf-eager + # (leads to memory leaks if not doing so). + if self.framework == "tf2": + self.global_timestep.assign(global_vars["timestep"]) + else: + self.global_timestep = global_vars["timestep"] + # Update our lifetime gradient update counter. + num_grad_updates = global_vars.get("num_grad_updates") + if num_grad_updates is not None: + self.num_grad_updates = num_grad_updates + + def export_checkpoint( + self, + export_dir: str, + filename_prefix=DEPRECATED_VALUE, + *, + policy_state: Optional[PolicyState] = None, + checkpoint_format: str = "cloudpickle", + ) -> None: + """Exports Policy checkpoint to a local directory and returns an AIR Checkpoint. + + Args: + export_dir: Local writable directory to store the AIR Checkpoint + information into. + policy_state: An optional PolicyState to write to disk. Used by + `Algorithm.save_checkpoint()` to save on the additional + `self.get_state()` calls of its different Policies. + checkpoint_format: Either one of 'cloudpickle' or 'msgpack'. + + .. testcode:: + :skipif: True + + from ray.rllib.algorithms.ppo import PPOTorchPolicy + policy = PPOTorchPolicy(...) + policy.export_checkpoint("/tmp/export_dir") + """ + # `filename_prefix` should not longer be used as new Policy checkpoints + # contain more than one file with a fixed filename structure. + if filename_prefix != DEPRECATED_VALUE: + deprecation_warning( + old="Policy.export_checkpoint(filename_prefix=...)", + error=True, + ) + if checkpoint_format not in ["cloudpickle", "msgpack"]: + raise ValueError( + f"Value of `checkpoint_format` ({checkpoint_format}) must either be " + "'cloudpickle' or 'msgpack'!" + ) + + if policy_state is None: + policy_state = self.get_state() + + # Write main policy state file. + os.makedirs(export_dir, exist_ok=True) + if checkpoint_format == "cloudpickle": + policy_state["checkpoint_version"] = CHECKPOINT_VERSION + state_file = "policy_state.pkl" + with open(os.path.join(export_dir, state_file), "w+b") as f: + pickle.dump(policy_state, f) + else: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + msgpack = try_import_msgpack(error=True) + policy_state["checkpoint_version"] = str(CHECKPOINT_VERSION) + # Serialize the config for msgpack dump'ing. + policy_state["policy_spec"]["config"] = AlgorithmConfig._serialize_dict( + policy_state["policy_spec"]["config"] + ) + state_file = "policy_state.msgpck" + with open(os.path.join(export_dir, state_file), "w+b") as f: + msgpack.dump(policy_state, f) + + # Write RLlib checkpoint json. + with open(os.path.join(export_dir, "rllib_checkpoint.json"), "w") as f: + json.dump( + { + "type": "Policy", + "checkpoint_version": str(policy_state["checkpoint_version"]), + "format": checkpoint_format, + "state_file": state_file, + "ray_version": ray.__version__, + "ray_commit": ray.__commit__, + }, + f, + ) + + # Add external model files, if required. + if self.config["export_native_model_files"]: + self.export_model(os.path.join(export_dir, "model")) + + def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None: + """Exports the Policy's Model to local directory for serving. + + Note: The file format will depend on the deep learning framework used. + See the child classed of Policy and their `export_model` + implementations for more details. + + Args: + export_dir: Local writable directory. + onnx: If given, will export model in ONNX format. The + value of this parameter set the ONNX OpSet version to use. + + Raises: + ValueError: If a native DL-framework based model (e.g. a keras Model) + cannot be saved to disk for various reasons. + """ + raise NotImplementedError + + def import_model_from_h5(self, import_file: str) -> None: + """Imports Policy from local file. + + Args: + import_file: Local readable file. + """ + raise NotImplementedError + + def get_session(self) -> Optional["tf1.Session"]: + """Returns tf.Session object to use for computing actions or None. + + Note: This method only applies to TFPolicy sub-classes. All other + sub-classes should expect a None to be returned from this method. + + Returns: + The tf Session to use for computing actions and losses with + this policy or None. + """ + return None + + def get_host(self) -> str: + """Returns the computer's network name. + + Returns: + The computer's networks name or an empty string, if the network + name could not be determined. + """ + return platform.node() + + def _get_num_gpus_for_policy(self) -> int: + """Decide on the number of CPU/GPU nodes this policy should run on. + + Return: + 0 if policy should run on CPU. >0 if policy should run on 1 or + more GPUs. + """ + worker_idx = self.config.get("worker_index", 0) + fake_gpus = self.config.get("_fake_gpus", False) + + if ( + ray._private.worker._mode() == ray._private.worker.LOCAL_MODE + and not fake_gpus + ): + # If in local debugging mode, and _fake_gpus is not on. + num_gpus = 0 + elif worker_idx == 0: + # If head node, take num_gpus. + num_gpus = self.config["num_gpus"] + else: + # If worker node, take `num_gpus_per_env_runner`. + num_gpus = self.config["num_gpus_per_env_runner"] + + if num_gpus == 0: + dev = "CPU" + else: + dev = "{} {}".format(num_gpus, "fake-GPUs" if fake_gpus else "GPUs") + + logger.info( + "Policy (worker={}) running on {}.".format( + worker_idx if worker_idx > 0 else "local", dev + ) + ) + + return num_gpus + + def _create_exploration(self) -> Exploration: + """Creates the Policy's Exploration object. + + This method only exists b/c some Algorithms do not use TfPolicy nor + TorchPolicy, but inherit directly from Policy. Others inherit from + TfPolicy w/o using DynamicTFPolicy. + + Returns: + Exploration: The Exploration object to be used by this Policy. + """ + if getattr(self, "exploration", None) is not None: + return self.exploration + + exploration = from_config( + Exploration, + self.config.get("exploration_config", {"type": "StochasticSampling"}), + action_space=self.action_space, + policy_config=self.config, + model=getattr(self, "model", None), + num_workers=self.config.get("num_env_runners", 0), + worker_index=self.config.get("worker_index", 0), + framework=getattr(self, "framework", self.config.get("framework", "tf")), + ) + return exploration + + def _get_default_view_requirements(self): + """Returns a default ViewRequirements dict. + + Note: This is the base/maximum requirement dict, from which later + some requirements will be subtracted again automatically to streamline + data collection, batch creation, and data transfer. + + Returns: + ViewReqDict: The default view requirements dict. + """ + + # Default view requirements (equal to those that we would use before + # the trajectory view API was introduced). + return { + SampleBatch.OBS: ViewRequirement(space=self.observation_space), + SampleBatch.NEXT_OBS: ViewRequirement( + data_col=SampleBatch.OBS, + shift=1, + space=self.observation_space, + used_for_compute_actions=False, + ), + SampleBatch.ACTIONS: ViewRequirement( + space=self.action_space, used_for_compute_actions=False + ), + # For backward compatibility with custom Models that don't specify + # these explicitly (will be removed by Policy if not used). + SampleBatch.PREV_ACTIONS: ViewRequirement( + data_col=SampleBatch.ACTIONS, shift=-1, space=self.action_space + ), + SampleBatch.REWARDS: ViewRequirement(), + # For backward compatibility with custom Models that don't specify + # these explicitly (will be removed by Policy if not used). + SampleBatch.PREV_REWARDS: ViewRequirement( + data_col=SampleBatch.REWARDS, shift=-1 + ), + SampleBatch.TERMINATEDS: ViewRequirement(), + SampleBatch.TRUNCATEDS: ViewRequirement(), + SampleBatch.INFOS: ViewRequirement(used_for_compute_actions=False), + SampleBatch.EPS_ID: ViewRequirement(), + SampleBatch.UNROLL_ID: ViewRequirement(), + SampleBatch.AGENT_INDEX: ViewRequirement(), + SampleBatch.T: ViewRequirement(), + } + + def _initialize_loss_from_dummy_batch( + self, + auto_remove_unneeded_view_reqs: bool = True, + stats_fn=None, + ) -> None: + """Performs test calls through policy's model and loss. + + NOTE: This base method should work for define-by-run Policies such as + torch and tf-eager policies. + + If required, will thereby detect automatically, which data views are + required by a) the forward pass, b) the postprocessing, and c) the loss + functions, and remove those from self.view_requirements that are not + necessary for these computations (to save data storage and transfer). + + Args: + auto_remove_unneeded_view_reqs: Whether to automatically + remove those ViewRequirements records from + self.view_requirements that are not needed. + stats_fn (Optional[Callable[[Policy, SampleBatch], Dict[str, + TensorType]]]): An optional stats function to be called after + the loss. + """ + + if self.config.get("_disable_initialize_loss_from_dummy_batch", False): + return + # Signal Policy that currently we do not like to eager/jit trace + # any function calls. This is to be able to track, which columns + # in the dummy batch are accessed by the different function (e.g. + # loss) such that we can then adjust our view requirements. + self._no_tracing = True + # Save for later so that loss init does not change global timestep + global_ts_before_init = int(convert_to_numpy(self.global_timestep)) + + sample_batch_size = min( + max(self.batch_divisibility_req * 4, 32), + self.config["train_batch_size"], # Don't go over the asked batch size. + ) + self._dummy_batch = self._get_dummy_batch_from_view_requirements( + sample_batch_size + ) + self._lazy_tensor_dict(self._dummy_batch) + explore = False + actions, state_outs, extra_outs = self.compute_actions_from_input_dict( + self._dummy_batch, explore=explore + ) + for key, view_req in self.view_requirements.items(): + if key not in self._dummy_batch.accessed_keys: + view_req.used_for_compute_actions = False + # Add all extra action outputs to view reqirements (these may be + # filtered out later again, if not needed for postprocessing or loss). + for key, value in extra_outs.items(): + self._dummy_batch[key] = value + if key not in self.view_requirements: + if isinstance(value, (dict, np.ndarray)): + # the assumption is that value is a nested_dict of np.arrays leaves + space = get_gym_space_from_struct_of_tensors(value) + self.view_requirements[key] = ViewRequirement( + space=space, used_for_compute_actions=False + ) + else: + raise ValueError( + "policy.compute_actions_from_input_dict() returns an " + "extra action output that is neither a numpy array nor a dict." + ) + + for key in self._dummy_batch.accessed_keys: + if key not in self.view_requirements: + self.view_requirements[key] = ViewRequirement() + self.view_requirements[key].used_for_compute_actions = False + # TODO (kourosh) Why did we use to make used_for_compute_actions True here? + new_batch = self._get_dummy_batch_from_view_requirements(sample_batch_size) + # Make sure the dummy_batch will return numpy arrays when accessed + self._dummy_batch.set_get_interceptor(None) + + # try to re-use the output of the previous run to avoid overriding things that + # would break (e.g. scale = 0 of Normal distribution cannot be zero) + for k in new_batch: + if k not in self._dummy_batch: + self._dummy_batch[k] = new_batch[k] + + # Make sure the book-keeping of dummy_batch keys are reset to correcly track + # what is accessed, what is added and what's deleted from now on. + self._dummy_batch.accessed_keys.clear() + self._dummy_batch.deleted_keys.clear() + self._dummy_batch.added_keys.clear() + + if self.exploration: + # Policies with RLModules don't have an exploration object. + self.exploration.postprocess_trajectory(self, self._dummy_batch) + + postprocessed_batch = self.postprocess_trajectory(self._dummy_batch) + seq_lens = None + if state_outs: + B = 4 # For RNNs, have B=4, T=[depends on sample_batch_size] + i = 0 + while "state_in_{}".format(i) in postprocessed_batch: + postprocessed_batch["state_in_{}".format(i)] = postprocessed_batch[ + "state_in_{}".format(i) + ][:B] + if "state_out_{}".format(i) in postprocessed_batch: + postprocessed_batch["state_out_{}".format(i)] = postprocessed_batch[ + "state_out_{}".format(i) + ][:B] + i += 1 + + seq_len = sample_batch_size // B + seq_lens = np.array([seq_len for _ in range(B)], dtype=np.int32) + postprocessed_batch[SampleBatch.SEQ_LENS] = seq_lens + + # Switch on lazy to-tensor conversion on `postprocessed_batch`. + train_batch = self._lazy_tensor_dict(postprocessed_batch) + # Calling loss, so set `is_training` to True. + train_batch.set_training(True) + if seq_lens is not None: + train_batch[SampleBatch.SEQ_LENS] = seq_lens + train_batch.count = self._dummy_batch.count + + # Call the loss function, if it exists. + # TODO(jungong) : clean up after all agents get migrated. + # We should simply do self.loss(...) here. + if self._loss is not None: + self._loss(self, self.model, self.dist_class, train_batch) + elif is_overridden(self.loss) and not self.config["in_evaluation"]: + self.loss(self.model, self.dist_class, train_batch) + # Call the stats fn, if given. + # TODO(jungong) : clean up after all agents get migrated. + # We should simply do self.stats_fn(train_batch) here. + if stats_fn is not None: + stats_fn(self, train_batch) + if hasattr(self, "stats_fn") and not self.config["in_evaluation"]: + self.stats_fn(train_batch) + + # Re-enable tracing. + self._no_tracing = False + + # Add new columns automatically to view-reqs. + if auto_remove_unneeded_view_reqs: + # Add those needed for postprocessing and training. + all_accessed_keys = ( + train_batch.accessed_keys + | self._dummy_batch.accessed_keys + | self._dummy_batch.added_keys + ) + for key in all_accessed_keys: + if key not in self.view_requirements and key != SampleBatch.SEQ_LENS: + self.view_requirements[key] = ViewRequirement( + used_for_compute_actions=False + ) + if self._loss or is_overridden(self.loss): + # Tag those only needed for post-processing (with some + # exceptions). + for key in self._dummy_batch.accessed_keys: + if ( + key not in train_batch.accessed_keys + and key in self.view_requirements + and key not in self.model.view_requirements + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + ] + ): + self.view_requirements[key].used_for_training = False + # Remove those not needed at all (leave those that are needed + # by Sampler to properly execute sample collection). Also always leave + # TERMINATEDS, TRUNCATEDS, REWARDS, INFOS, no matter what. + for key in list(self.view_requirements.keys()): + if ( + key not in all_accessed_keys + and key + not in [ + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.UNROLL_ID, + SampleBatch.TERMINATEDS, + SampleBatch.TRUNCATEDS, + SampleBatch.REWARDS, + SampleBatch.INFOS, + SampleBatch.T, + ] + and key not in self.model.view_requirements + ): + # If user deleted this key manually in postprocessing + # fn, warn about it and do not remove from + # view-requirements. + if key in self._dummy_batch.deleted_keys: + logger.warning( + "SampleBatch key '{}' was deleted manually in " + "postprocessing function! RLlib will " + "automatically remove non-used items from the " + "data stream. Remove the `del` from your " + "postprocessing function.".format(key) + ) + # If we are not writing output to disk, save to erase + # this key to save space in the sample batch. + elif self.config["output"] is None: + del self.view_requirements[key] + + if type(self.global_timestep) is int: + self.global_timestep = global_ts_before_init + elif isinstance(self.global_timestep, tf.Variable): + self.global_timestep.assign(global_ts_before_init) + else: + raise ValueError( + "Variable self.global_timestep of policy {} needs to be " + "either of type `int` or `tf.Variable`, " + "but is of type {}.".format(self, type(self.global_timestep)) + ) + + def maybe_remove_time_dimension(self, input_dict: Dict[str, TensorType]): + """Removes a time dimension for recurrent RLModules. + + Args: + input_dict: The input dict. + + Returns: + The input dict with a possibly removed time dimension. + """ + raise NotImplementedError + + def _get_dummy_batch_from_view_requirements( + self, batch_size: int = 1 + ) -> SampleBatch: + """Creates a numpy dummy batch based on the Policy's view requirements. + + Args: + batch_size: The size of the batch to create. + + Returns: + Dict[str, TensorType]: The dummy batch containing all zero values. + """ + ret = {} + for view_col, view_req in self.view_requirements.items(): + data_col = view_req.data_col or view_col + # Flattened dummy batch. + if (isinstance(view_req.space, (gym.spaces.Tuple, gym.spaces.Dict))) and ( + ( + data_col == SampleBatch.OBS + and not self.config["_disable_preprocessor_api"] + ) + or ( + data_col == SampleBatch.ACTIONS + and not self.config.get("_disable_action_flattening") + ) + ): + _, shape = ModelCatalog.get_action_shape( + view_req.space, framework=self.config["framework"] + ) + ret[view_col] = np.zeros((batch_size,) + shape[1:], np.float32) + # Non-flattened dummy batch. + else: + # Range of indices on time-axis, e.g. "-50:-1". + if isinstance(view_req.space, gym.spaces.Space): + time_size = ( + len(view_req.shift_arr) if len(view_req.shift_arr) > 1 else None + ) + ret[view_col] = get_dummy_batch_for_space( + view_req.space, batch_size=batch_size, time_size=time_size + ) + else: + ret[view_col] = [view_req.space for _ in range(batch_size)] + + # Due to different view requirements for the different columns, + # columns in the resulting batch may not all have the same batch size. + return SampleBatch(ret) + + def _update_model_view_requirements_from_init_state(self): + """Uses Model's (or this Policy's) init state to add needed ViewReqs. + + Can be called from within a Policy to make sure RNNs automatically + update their internal state-related view requirements. + Changes the `self.view_requirements` dict. + """ + self._model_init_state_automatically_added = True + model = getattr(self, "model", None) + + obj = model or self + if model and not hasattr(model, "view_requirements"): + model.view_requirements = { + SampleBatch.OBS: ViewRequirement(space=self.observation_space) + } + view_reqs = obj.view_requirements + # Add state-ins to this model's view. + init_state = [] + if hasattr(obj, "get_initial_state") and callable(obj.get_initial_state): + init_state = obj.get_initial_state() + else: + # Add this functionality automatically for new native model API. + if ( + tf + and isinstance(model, tf.keras.Model) + and "state_in_0" not in view_reqs + ): + obj.get_initial_state = lambda: [ + np.zeros_like(view_req.space.sample()) + for k, view_req in model.view_requirements.items() + if k.startswith("state_in_") + ] + else: + obj.get_initial_state = lambda: [] + if "state_in_0" in view_reqs: + self.is_recurrent = lambda: True + + # Make sure auto-generated init-state view requirements get added + # to both Policy and Model, no matter what. + view_reqs = [view_reqs] + ( + [self.view_requirements] if hasattr(self, "view_requirements") else [] + ) + + for i, state in enumerate(init_state): + # Allow `state` to be either a Space (use zeros as initial values) + # or any value (e.g. a dict or a non-zero tensor). + fw = ( + np + if isinstance(state, np.ndarray) + else torch + if torch and torch.is_tensor(state) + else None + ) + if fw: + space = ( + Box(-1.0, 1.0, shape=state.shape) if fw.all(state == 0.0) else state + ) + else: + space = state + for vr in view_reqs: + # Only override if user has not already provided + # custom view-requirements for state_in_n. + if "state_in_{}".format(i) not in vr: + vr["state_in_{}".format(i)] = ViewRequirement( + "state_out_{}".format(i), + shift=-1, + used_for_compute_actions=True, + batch_repeat_value=self.config.get("model", {}).get( + "max_seq_len", 1 + ), + space=space, + ) + # Only override if user has not already provided + # custom view-requirements for state_out_n. + if "state_out_{}".format(i) not in vr: + vr["state_out_{}".format(i)] = ViewRequirement( + space=space, used_for_training=True + ) + + def __repr__(self): + return type(self).__name__ + + +@OldAPIStack +def get_gym_space_from_struct_of_tensors( + value: Union[Dict, Tuple, List, TensorType], + batched_input=True, +) -> gym.Space: + start_idx = 1 if batched_input else 0 + struct = tree.map_structure( + lambda x: gym.spaces.Box( + -1.0, 1.0, shape=x.shape[start_idx:], dtype=get_np_dtype(x) + ), + value, + ) + space = get_gym_space_from_struct_of_spaces(struct) + return space + + +@OldAPIStack +def get_gym_space_from_struct_of_spaces(value: Union[Dict, Tuple]) -> gym.spaces.Dict: + if isinstance(value, dict): + return gym.spaces.Dict( + {k: get_gym_space_from_struct_of_spaces(v) for k, v in value.items()} + ) + elif isinstance(value, (tuple, list)): + return gym.spaces.Tuple([get_gym_space_from_struct_of_spaces(v) for v in value]) + else: + assert isinstance(value, gym.spaces.Space), ( + f"The struct of spaces should only contain dicts, tiples and primitive " + f"gym spaces. Space is of type {type(value)}" + ) + return value diff --git a/lib/python3.12/site-packages/ray/rllib/policy/policy_map.py b/lib/python3.12/site-packages/ray/rllib/policy/policy_map.py new file mode 100644 index 0000000000000000000000000000000000000000..4f64750339ae1d6836e080d1af3aa1a824e994d2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/policy_map.py @@ -0,0 +1,294 @@ +import logging +import threading +from collections import deque +from typing import Dict, Set + +import ray +from ray._common.deprecation import deprecation_warning +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.threading import with_lock +from ray.rllib.utils.typing import PolicyID + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +@OldAPIStack +class PolicyMap(dict): + """Maps policy IDs to Policy objects. + + Thereby, keeps n policies in memory and - when capacity is reached - + writes the least recently used to disk. This allows adding 100s of + policies to a Algorithm for league-based setups w/o running out of memory. + """ + + def __init__( + self, + *, + capacity: int = 100, + policy_states_are_swappable: bool = False, + # Deprecated args. + worker_index=None, + num_workers=None, + policy_config=None, + session_creator=None, + seed=None, + ): + """Initializes a PolicyMap instance. + + Args: + capacity: The size of the Policy object cache. This is the maximum number + of policies that are held in RAM memory. When reaching this capacity, + the least recently used Policy's state will be stored in the Ray object + store and recovered from there when being accessed again. + policy_states_are_swappable: Whether all Policy objects in this map can be + "swapped out" via a simple `state = A.get_state(); B.set_state(state)`, + where `A` and `B` are policy instances in this map. You should set + this to True for significantly speeding up the PolicyMap's cache lookup + times, iff your policies all share the same neural network + architecture and optimizer types. If True, the PolicyMap will not + have to garbage collect old, least recently used policies, but instead + keep them in memory and simply override their state with the state of + the most recently accessed one. + For example, in a league-based training setup, you might have 100s of + the same policies in your map (playing against each other in various + combinations), but all of them share the same state structure + (are "swappable"). + """ + if policy_config is not None: + deprecation_warning( + old="PolicyMap(policy_config=..)", + error=True, + ) + + super().__init__() + + self.capacity = capacity + + if any( + i is not None + for i in [policy_config, worker_index, num_workers, session_creator, seed] + ): + deprecation_warning( + old="PolicyMap([deprecated args]...)", + new="PolicyMap(capacity=..., policy_states_are_swappable=...)", + error=False, + ) + + self.policy_states_are_swappable = policy_states_are_swappable + + # The actual cache with the in-memory policy objects. + self.cache: Dict[str, Policy] = {} + + # Set of keys that may be looked up (cached or not). + self._valid_keys: Set[str] = set() + # The doubly-linked list holding the currently in-memory objects. + self._deque = deque() + + # Ray object store references to the stashed Policy states. + self._policy_state_refs = {} + + # Lock used for locking some methods on the object-level. + # This prevents possible race conditions when accessing the map + # and the underlying structures, like self._deque and others. + self._lock = threading.RLock() + + @with_lock + @override(dict) + def __getitem__(self, item: PolicyID): + # Never seen this key -> Error. + if item not in self._valid_keys: + raise KeyError( + f"PolicyID '{item}' not found in this PolicyMap! " + f"IDs stored in this map: {self._valid_keys}." + ) + + # Item already in cache -> Rearrange deque (promote `item` to + # "most recently used") and return it. + if item in self.cache: + self._deque.remove(item) + self._deque.append(item) + return self.cache[item] + + # Item not currently in cache -> Get from stash and - if at capacity - + # remove leftmost one. + if item not in self._policy_state_refs: + raise AssertionError( + f"PolicyID {item} not found in internal Ray object store cache!" + ) + policy_state = ray.get(self._policy_state_refs[item]) + + policy = None + # We are at capacity: Remove the oldest policy from deque as well as the + # cache and return it. + if len(self._deque) == self.capacity: + policy = self._stash_least_used_policy() + + # All our policies have same NN-architecture (are "swappable"). + # -> Load new policy's state into the one that just got removed from the cache. + # This way, we save the costly re-creation step. + if policy is not None and self.policy_states_are_swappable: + logger.debug(f"restoring policy: {item}") + policy.set_state(policy_state) + else: + logger.debug(f"creating new policy: {item}") + policy = Policy.from_state(policy_state) + + self.cache[item] = policy + # Promote the item to most recently one. + self._deque.append(item) + + return policy + + @with_lock + @override(dict) + def __setitem__(self, key: PolicyID, value: Policy): + # Item already in cache -> Rearrange deque. + if key in self.cache: + self._deque.remove(key) + + # Item not currently in cache -> store new value and - if at capacity - + # remove leftmost one. + else: + # Cache at capacity -> Drop leftmost item. + if len(self._deque) == self.capacity: + self._stash_least_used_policy() + + # Promote `key` to "most recently used". + self._deque.append(key) + + # Update our cache. + self.cache[key] = value + self._valid_keys.add(key) + + @with_lock + @override(dict) + def __delitem__(self, key: PolicyID): + # Make key invalid. + self._valid_keys.remove(key) + # Remove policy from deque if contained + if key in self._deque: + self._deque.remove(key) + # Remove policy from memory if currently cached. + if key in self.cache: + policy = self.cache[key] + self._close_session(policy) + del self.cache[key] + # Remove Ray object store reference (if this ID has already been stored + # there), so the item gets garbage collected. + if key in self._policy_state_refs: + del self._policy_state_refs[key] + + @override(dict) + def __iter__(self): + return iter(self.keys()) + + @override(dict) + def items(self): + """Iterates over all policies, even the stashed ones.""" + + def gen(): + for key in self._valid_keys: + yield (key, self[key]) + + return gen() + + @override(dict) + def keys(self): + """Returns all valid keys, even the stashed ones.""" + self._lock.acquire() + ks = list(self._valid_keys) + self._lock.release() + + def gen(): + for key in ks: + yield key + + return gen() + + @override(dict) + def values(self): + """Returns all valid values, even the stashed ones.""" + self._lock.acquire() + vs = [self[k] for k in self._valid_keys] + self._lock.release() + + def gen(): + for value in vs: + yield value + + return gen() + + @with_lock + @override(dict) + def update(self, __m, **kwargs): + """Updates the map with the given dict and/or kwargs.""" + for k, v in __m.items(): + self[k] = v + for k, v in kwargs.items(): + self[k] = v + + @with_lock + @override(dict) + def get(self, key: PolicyID): + """Returns the value for the given key or None if not found.""" + if key not in self._valid_keys: + return None + return self[key] + + @with_lock + @override(dict) + def __len__(self) -> int: + """Returns number of all policies, including the stashed-to-disk ones.""" + return len(self._valid_keys) + + @with_lock + @override(dict) + def __contains__(self, item: PolicyID): + return item in self._valid_keys + + @override(dict) + def __str__(self) -> str: + # Only print out our keys (policy IDs), not values as this could trigger + # the LRU caching. + return ( + f"" + ) + + def _stash_least_used_policy(self) -> Policy: + """Writes the least-recently used policy's state to the Ray object store. + + Also closes the session - if applicable - of the stashed policy. + + Returns: + The least-recently used policy, that just got removed from the cache. + """ + # Get policy's state for writing to object store. + dropped_policy_id = self._deque.popleft() + assert dropped_policy_id in self.cache + policy = self.cache[dropped_policy_id] + policy_state = policy.get_state() + + # If we don't simply swap out vs an existing policy: + # Close the tf session, if any. + if not self.policy_states_are_swappable: + self._close_session(policy) + + # Remove from memory. This will clear the tf Graph as well. + del self.cache[dropped_policy_id] + + # Store state in Ray object store. + self._policy_state_refs[dropped_policy_id] = ray.put(policy_state) + + # Return the just removed policy, in case it's needed by the caller. + return policy + + @staticmethod + def _close_session(policy: Policy): + sess = policy.get_session() + # Closes the tf session, if any. + if sess is not None: + sess.close() diff --git a/lib/python3.12/site-packages/ray/rllib/policy/policy_template.py b/lib/python3.12/site-packages/ray/rllib/policy/policy_template.py new file mode 100644 index 0000000000000000000000000000000000000000..94d284ec22ef04025d10fc523741671fabed6f99 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/policy_template.py @@ -0,0 +1,448 @@ +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +import gymnasium as gym + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_policy import TorchPolicy +from ray.rllib.utils import NullContextManager, add_mixins +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_jax, try_import_torch +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import AlgorithmConfigDict, ModelGradients, TensorType + +jax, _ = try_import_jax() +torch, _ = try_import_torch() + + +@OldAPIStack +def build_policy_class( + name: str, + framework: str, + *, + loss_fn: Optional[ + Callable[ + [Policy, ModelV2, Type[TorchDistributionWrapper], SampleBatch], + Union[TensorType, List[TensorType]], + ] + ], + get_default_config: Optional[Callable[[], AlgorithmConfigDict]] = None, + stats_fn: Optional[Callable[[Policy, SampleBatch], Dict[str, TensorType]]] = None, + postprocess_fn: Optional[ + Callable[ + [ + Policy, + SampleBatch, + Optional[Dict[Any, SampleBatch]], + Optional[Any], + ], + SampleBatch, + ] + ] = None, + extra_action_out_fn: Optional[ + Callable[ + [ + Policy, + Dict[str, TensorType], + List[TensorType], + ModelV2, + TorchDistributionWrapper, + ], + Dict[str, TensorType], + ] + ] = None, + extra_grad_process_fn: Optional[ + Callable[[Policy, "torch.optim.Optimizer", TensorType], Dict[str, TensorType]] + ] = None, + # TODO: (sven) Replace "fetches" with "process". + extra_learn_fetches_fn: Optional[Callable[[Policy], Dict[str, TensorType]]] = None, + optimizer_fn: Optional[ + Callable[[Policy, AlgorithmConfigDict], "torch.optim.Optimizer"] + ] = None, + validate_spaces: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + before_init: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + before_loss_init: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], None + ] + ] = None, + after_init: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + _after_loss_init: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], None + ] + ] = None, + action_sampler_fn: Optional[ + Callable[[TensorType, List[TensorType]], Tuple[TensorType, TensorType]] + ] = None, + action_distribution_fn: Optional[ + Callable[ + [Policy, ModelV2, TensorType, TensorType, TensorType], + Tuple[TensorType, type, List[TensorType]], + ] + ] = None, + make_model: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], ModelV2 + ] + ] = None, + make_model_and_action_dist: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], + Tuple[ModelV2, Type[TorchDistributionWrapper]], + ] + ] = None, + compute_gradients_fn: Optional[ + Callable[[Policy, SampleBatch], Tuple[ModelGradients, dict]] + ] = None, + apply_gradients_fn: Optional[ + Callable[[Policy, "torch.optim.Optimizer"], None] + ] = None, + mixins: Optional[List[type]] = None, + get_batch_divisibility_req: Optional[Callable[[Policy], int]] = None, +) -> Type[TorchPolicy]: + """Helper function for creating a new Policy class at runtime. + + Supports frameworks JAX and PyTorch. + + Args: + name: name of the policy (e.g., "PPOTorchPolicy") + framework: Either "jax" or "torch". + loss_fn (Optional[Callable[[Policy, ModelV2, + Type[TorchDistributionWrapper], SampleBatch], Union[TensorType, + List[TensorType]]]]): Callable that returns a loss tensor. + get_default_config (Optional[Callable[[None], AlgorithmConfigDict]]): + Optional callable that returns the default config to merge with any + overrides. If None, uses only(!) the user-provided + PartialAlgorithmConfigDict as dict for this Policy. + postprocess_fn (Optional[Callable[[Policy, SampleBatch, + Optional[Dict[Any, SampleBatch]], Optional[Any]], + SampleBatch]]): Optional callable for post-processing experience + batches (called after the super's `postprocess_trajectory` method). + stats_fn (Optional[Callable[[Policy, SampleBatch], + Dict[str, TensorType]]]): Optional callable that returns a dict of + values given the policy and training batch. If None, + will use `TorchPolicy.extra_grad_info()` instead. The stats dict is + used for logging (e.g. in TensorBoard). + extra_action_out_fn (Optional[Callable[[Policy, Dict[str, TensorType], + List[TensorType], ModelV2, TorchDistributionWrapper]], Dict[str, + TensorType]]]): Optional callable that returns a dict of extra + values to include in experiences. If None, no extra computations + will be performed. + extra_grad_process_fn (Optional[Callable[[Policy, + "torch.optim.Optimizer", TensorType], Dict[str, TensorType]]]): + Optional callable that is called after gradients are computed and + returns a processing info dict. If None, will call the + `TorchPolicy.extra_grad_process()` method instead. + # TODO: (sven) dissolve naming mismatch between "learn" and "compute.." + extra_learn_fetches_fn (Optional[Callable[[Policy], + Dict[str, TensorType]]]): Optional callable that returns a dict of + extra tensors from the policy after loss evaluation. If None, + will call the `TorchPolicy.extra_compute_grad_fetches()` method + instead. + optimizer_fn (Optional[Callable[[Policy, AlgorithmConfigDict], + "torch.optim.Optimizer"]]): Optional callable that returns a + torch optimizer given the policy and config. If None, will call + the `TorchPolicy.optimizer()` method instead (which returns a + torch Adam optimizer). + validate_spaces (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): Optional callable that takes the + Policy, observation_space, action_space, and config to check for + correctness. If None, no spaces checking will be done. + before_init (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): Optional callable to run at the + beginning of `Policy.__init__` that takes the same arguments as + the Policy constructor. If None, this step will be skipped. + before_loss_init (Optional[Callable[[Policy, gym.spaces.Space, + gym.spaces.Space, AlgorithmConfigDict], None]]): Optional callable to + run prior to loss init. If None, this step will be skipped. + after_init (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): DEPRECATED: Use `before_loss_init` + instead. + _after_loss_init (Optional[Callable[[Policy, gym.spaces.Space, + gym.spaces.Space, AlgorithmConfigDict], None]]): Optional callable to + run after the loss init. If None, this step will be skipped. + This will be deprecated at some point and renamed into `after_init` + to match `build_tf_policy()` behavior. + action_sampler_fn (Optional[Callable[[TensorType, List[TensorType]], + Tuple[TensorType, TensorType]]]): Optional callable returning a + sampled action and its log-likelihood given some (obs and state) + inputs. If None, will either use `action_distribution_fn` or + compute actions by calling self.model, then sampling from the + so parameterized action distribution. + action_distribution_fn (Optional[Callable[[Policy, ModelV2, TensorType, + TensorType, TensorType], Tuple[TensorType, + Type[TorchDistributionWrapper], List[TensorType]]]]): A callable + that takes the Policy, Model, the observation batch, an + explore-flag, a timestep, and an is_training flag and returns a + tuple of a) distribution inputs (parameters), b) a dist-class to + generate an action distribution object from, and c) internal-state + outputs (empty list if not applicable). If None, will either use + `action_sampler_fn` or compute actions by calling self.model, + then sampling from the parameterized action distribution. + make_model (Optional[Callable[[Policy, gym.spaces.Space, + gym.spaces.Space, AlgorithmConfigDict], ModelV2]]): Optional callable + that takes the same arguments as Policy.__init__ and returns a + model instance. The distribution class will be determined + automatically. Note: Only one of `make_model` or + `make_model_and_action_dist` should be provided. If both are None, + a default Model will be created. + make_model_and_action_dist (Optional[Callable[[Policy, + gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], + Tuple[ModelV2, Type[TorchDistributionWrapper]]]]): Optional + callable that takes the same arguments as Policy.__init__ and + returns a tuple of model instance and torch action distribution + class. + Note: Only one of `make_model` or `make_model_and_action_dist` + should be provided. If both are None, a default Model will be + created. + compute_gradients_fn (Optional[Callable[ + [Policy, SampleBatch], Tuple[ModelGradients, dict]]]): Optional + callable that the sampled batch an computes the gradients w.r. + to the loss function. + If None, will call the `TorchPolicy.compute_gradients()` method + instead. + apply_gradients_fn (Optional[Callable[[Policy, + "torch.optim.Optimizer"], None]]): Optional callable that + takes a grads list and applies these to the Model's parameters. + If None, will call the `TorchPolicy.apply_gradients()` method + instead. + mixins (Optional[List[type]]): Optional list of any class mixins for + the returned policy class. These mixins will be applied in order + and will have higher precedence than the TorchPolicy class. + get_batch_divisibility_req (Optional[Callable[[Policy], int]]): + Optional callable that returns the divisibility requirement for + sample batches. If None, will assume a value of 1. + + Returns: + Type[TorchPolicy]: TorchPolicy child class constructed from the + specified args. + """ + + original_kwargs = locals().copy() + parent_cls = TorchPolicy + base = add_mixins(parent_cls, mixins) + + class policy_cls(base): + def __init__(self, obs_space, action_space, config): + self.config = config + + # Set the DL framework for this Policy. + self.framework = self.config["framework"] = framework + + # Validate observation- and action-spaces. + if validate_spaces: + validate_spaces(self, obs_space, action_space, self.config) + + # Do some pre-initialization steps. + if before_init: + before_init(self, obs_space, action_space, self.config) + + # Model is customized (use default action dist class). + if make_model: + assert make_model_and_action_dist is None, ( + "Either `make_model` or `make_model_and_action_dist`" + " must be None!" + ) + self.model = make_model(self, obs_space, action_space, config) + dist_class, _ = ModelCatalog.get_action_dist( + action_space, self.config["model"], framework=framework + ) + # Model and action dist class are customized. + elif make_model_and_action_dist: + self.model, dist_class = make_model_and_action_dist( + self, obs_space, action_space, config + ) + # Use default model and default action dist. + else: + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"], framework=framework + ) + self.model = ModelCatalog.get_model_v2( + obs_space=obs_space, + action_space=action_space, + num_outputs=logit_dim, + model_config=self.config["model"], + framework=framework, + ) + + # Make sure, we passed in a correct Model factory. + model_cls = TorchModelV2 + assert isinstance( + self.model, model_cls + ), "ERROR: Generated Model must be a TorchModelV2 object!" + + # Call the framework-specific Policy constructor. + self.parent_cls = parent_cls + self.parent_cls.__init__( + self, + observation_space=obs_space, + action_space=action_space, + config=config, + model=self.model, + loss=None if self.config["in_evaluation"] else loss_fn, + action_distribution_class=dist_class, + action_sampler_fn=action_sampler_fn, + action_distribution_fn=action_distribution_fn, + max_seq_len=config["model"]["max_seq_len"], + get_batch_divisibility_req=get_batch_divisibility_req, + ) + + # Merge Model's view requirements into Policy's. + self.view_requirements.update(self.model.view_requirements) + + _before_loss_init = before_loss_init or after_init + if _before_loss_init: + _before_loss_init( + self, self.observation_space, self.action_space, config + ) + + # Perform test runs through postprocessing- and loss functions. + self._initialize_loss_from_dummy_batch( + auto_remove_unneeded_view_reqs=True, + stats_fn=None if self.config["in_evaluation"] else stats_fn, + ) + + if _after_loss_init: + _after_loss_init(self, obs_space, action_space, config) + + # Got to reset global_timestep again after this fake run-through. + self.global_timestep = 0 + + @override(Policy) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + # Do all post-processing always with no_grad(). + # Not using this here will introduce a memory leak + # in torch (issue #6962). + with self._no_grad_context(): + # Call super's postprocess_trajectory first. + sample_batch = super().postprocess_trajectory( + sample_batch, other_agent_batches, episode + ) + if postprocess_fn: + return postprocess_fn( + self, sample_batch, other_agent_batches, episode + ) + + return sample_batch + + @override(parent_cls) + def extra_grad_process(self, optimizer, loss): + """Called after optimizer.zero_grad() and loss.backward() calls. + + Allows for gradient processing before optimizer.step() is called. + E.g. for gradient clipping. + """ + if extra_grad_process_fn: + return extra_grad_process_fn(self, optimizer, loss) + else: + return parent_cls.extra_grad_process(self, optimizer, loss) + + @override(parent_cls) + def extra_compute_grad_fetches(self): + if extra_learn_fetches_fn: + fetches = convert_to_numpy(extra_learn_fetches_fn(self)) + # Auto-add empty learner stats dict if needed. + return dict({LEARNER_STATS_KEY: {}}, **fetches) + else: + return parent_cls.extra_compute_grad_fetches(self) + + @override(parent_cls) + def compute_gradients(self, batch): + if compute_gradients_fn: + return compute_gradients_fn(self, batch) + else: + return parent_cls.compute_gradients(self, batch) + + @override(parent_cls) + def apply_gradients(self, gradients): + if apply_gradients_fn: + apply_gradients_fn(self, gradients) + else: + parent_cls.apply_gradients(self, gradients) + + @override(parent_cls) + def extra_action_out(self, input_dict, state_batches, model, action_dist): + with self._no_grad_context(): + if extra_action_out_fn: + stats_dict = extra_action_out_fn( + self, input_dict, state_batches, model, action_dist + ) + else: + stats_dict = parent_cls.extra_action_out( + self, input_dict, state_batches, model, action_dist + ) + return self._convert_to_numpy(stats_dict) + + @override(parent_cls) + def optimizer(self): + if optimizer_fn: + optimizers = optimizer_fn(self, self.config) + else: + optimizers = parent_cls.optimizer(self) + return optimizers + + @override(parent_cls) + def extra_grad_info(self, train_batch): + with self._no_grad_context(): + if stats_fn: + stats_dict = stats_fn(self, train_batch) + else: + stats_dict = self.parent_cls.extra_grad_info(self, train_batch) + return self._convert_to_numpy(stats_dict) + + def _no_grad_context(self): + if self.framework == "torch": + return torch.no_grad() + return NullContextManager() + + def _convert_to_numpy(self, data): + if self.framework == "torch": + return convert_to_numpy(data) + return data + + def with_updates(**overrides): + """Creates a Torch|JAXPolicy cls based on settings of another one. + + Keyword Args: + **overrides: The settings (passed into `build_torch_policy`) that + should be different from the class that this method is called + on. + + Returns: + type: A new Torch|JAXPolicy sub-class. + + Examples: + >> MySpecialDQNPolicyClass = DQNTorchPolicy.with_updates( + .. name="MySpecialDQNPolicyClass", + .. loss_function=[some_new_loss_function], + .. ) + """ + return build_policy_class(**dict(original_kwargs, **overrides)) + + policy_cls.with_updates = staticmethod(with_updates) + policy_cls.__name__ = name + policy_cls.__qualname__ = name + return policy_cls diff --git a/lib/python3.12/site-packages/ray/rllib/policy/rnn_sequencing.py b/lib/python3.12/site-packages/ray/rllib/policy/rnn_sequencing.py new file mode 100644 index 0000000000000000000000000000000000000000..00e2aecffe3b8366e015348348ba9b935dab81d9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/rnn_sequencing.py @@ -0,0 +1,683 @@ +"""RNN utils for RLlib. + +The main trick here is that we add the time dimension at the last moment. +The non-LSTM layers of the model see their inputs as one flat batch. Before +the LSTM cell, we reshape the input to add the expected time dimension. During +postprocessing, we dynamically pad the experience batches so that this +reshaping is possible. + +Note that this padding strategy only works out if we assume zero inputs don't +meaningfully affect the loss function. This happens to be true for all the +current algorithms: https://github.com/ray-project/ray/issues/2992 +""" + +import functools +import logging +from typing import List, Optional + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.typing import SampleBatchType, TensorType, ViewRequirementsDict +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +def pad_batch_to_sequences_of_same_size( + batch: SampleBatch, + max_seq_len: int, + shuffle: bool = False, + batch_divisibility_req: int = 1, + feature_keys: Optional[List[str]] = None, + view_requirements: Optional[ViewRequirementsDict] = None, + _enable_new_api_stack: bool = False, + padding: str = "zero", +): + """Applies padding to `batch` so it's choppable into same-size sequences. + + Shuffles `batch` (if desired), makes sure divisibility requirement is met, + then pads the batch ([B, ...]) into same-size chunks ([B, ...]) w/o + adding a time dimension (yet). + Padding depends on episodes found in batch and `max_seq_len`. + + Args: + batch: The SampleBatch object. All values in here have + the shape [B, ...]. + max_seq_len: The max. sequence length to use for chopping. + shuffle: Whether to shuffle batch sequences. Shuffle may + be done in-place. This only makes sense if you're further + applying minibatch SGD after getting the outputs. + batch_divisibility_req: The int by which the batch dimension + must be dividable. + feature_keys: An optional list of keys to apply sequence-chopping + to. If None, use all keys in batch that are not + "state_in/out_"-type keys. + view_requirements: An optional Policy ViewRequirements dict to + be able to infer whether e.g. dynamic max'ing should be + applied over the seq_lens. + _enable_new_api_stack: This is a temporary flag to enable the new RLModule API. + After a complete rollout of the new API, this flag will be removed. + padding: Padding type to use. Either "zero" or "last". Zero padding + will pad with zeros, last padding will pad with the last value. + """ + # If already zero-padded, skip. + if batch.zero_padded: + return + + batch.zero_padded = True + + if batch_divisibility_req > 1: + meets_divisibility_reqs = ( + len(batch[SampleBatch.CUR_OBS]) % batch_divisibility_req == 0 + # not multiagent + and max(batch[SampleBatch.AGENT_INDEX]) == 0 + ) + else: + meets_divisibility_reqs = True + + states_already_reduced_to_init = False + + # RNN/attention net case. Figure out whether we should apply dynamic + # max'ing over the list of sequence lengths. + if _enable_new_api_stack and ("state_in" in batch or "state_out" in batch): + # TODO (Kourosh): This is a temporary fix to enable the new RLModule API. + # We should think of a more elegant solution once we have confirmed that other + # parts of the API are stable and user-friendly. + seq_lens = batch.get(SampleBatch.SEQ_LENS) + + # state_in is a nested dict of tensors of states. We need to retreive the + # length of the inner most tensor (which should be already the same as the + # length of other tensors) and compare it to len(seq_lens). + state_ins = tree.flatten(batch["state_in"]) + if state_ins: + assert all( + len(state_in) == len(state_ins[0]) for state_in in state_ins + ), "All state_in tensors should have the same batch_dim size." + + # if the batch dim of states is the same as the number of sequences + if len(state_ins[0]) == len(seq_lens): + states_already_reduced_to_init = True + + # TODO (Kourosh): What is the use-case of DynamicMax functionality? + dynamic_max = True + else: + dynamic_max = False + + elif not _enable_new_api_stack and ( + "state_in_0" in batch or "state_out_0" in batch + ): + # Check, whether the state inputs have already been reduced to their + # init values at the beginning of each max_seq_len chunk. + if batch.get(SampleBatch.SEQ_LENS) is not None and len( + batch["state_in_0"] + ) == len(batch[SampleBatch.SEQ_LENS]): + states_already_reduced_to_init = True + + # RNN (or single timestep state-in): Set the max dynamically. + if view_requirements and view_requirements["state_in_0"].shift_from is None: + dynamic_max = True + # Attention Nets (state inputs are over some range): No dynamic maxing + # possible. + else: + dynamic_max = False + # Multi-agent case. + elif not meets_divisibility_reqs: + max_seq_len = batch_divisibility_req + dynamic_max = False + batch.max_seq_len = max_seq_len + # Simple case: No RNN/attention net, nor do we need to pad. + else: + if shuffle: + batch.shuffle() + return + + # RNN, attention net, or multi-agent case. + state_keys = [] + feature_keys_ = feature_keys or [] + for k, v in batch.items(): + if k.startswith("state_in"): + state_keys.append(k) + elif ( + not feature_keys + and (not k.startswith("state_out") if not _enable_new_api_stack else True) + and k not in [SampleBatch.SEQ_LENS] + ): + feature_keys_.append(k) + feature_sequences, initial_states, seq_lens = chop_into_sequences( + feature_columns=[batch[k] for k in feature_keys_], + state_columns=[batch[k] for k in state_keys], + episode_ids=batch.get(SampleBatch.EPS_ID), + unroll_ids=batch.get(SampleBatch.UNROLL_ID), + agent_indices=batch.get(SampleBatch.AGENT_INDEX), + seq_lens=batch.get(SampleBatch.SEQ_LENS), + max_seq_len=max_seq_len, + dynamic_max=dynamic_max, + states_already_reduced_to_init=states_already_reduced_to_init, + shuffle=shuffle, + handle_nested_data=True, + padding=padding, + pad_infos_with_empty_dicts=_enable_new_api_stack, + ) + for i, k in enumerate(feature_keys_): + batch[k] = tree.unflatten_as(batch[k], feature_sequences[i]) + for i, k in enumerate(state_keys): + batch[k] = initial_states[i] + batch[SampleBatch.SEQ_LENS] = np.array(seq_lens) + if dynamic_max: + batch.max_seq_len = max(seq_lens) + + if log_once("rnn_ma_feed_dict"): + logger.info( + "Padded input for RNN/Attn.Nets/MA:\n\n{}\n".format( + summarize( + { + "features": feature_sequences, + "initial_states": initial_states, + "seq_lens": seq_lens, + "max_seq_len": max_seq_len, + } + ) + ) + ) + + +@OldAPIStack +def add_time_dimension( + padded_inputs: TensorType, + *, + seq_lens: TensorType, + framework: str = "tf", + time_major: bool = False, +): + """Adds a time dimension to padded inputs. + + Args: + padded_inputs: a padded batch of sequences. That is, + for seq_lens=[1, 2, 2], then inputs=[A, *, B, B, C, C], where + A, B, C are sequence elements and * denotes padding. + seq_lens: A 1D tensor of sequence lengths, denoting the non-padded length + in timesteps of each rollout in the batch. + framework: The framework string ("tf2", "tf", "torch"). + time_major: Whether data should be returned in time-major (TxB) + format or not (BxT). + + Returns: + TensorType: Reshaped tensor of shape [B, T, ...] or [T, B, ...]. + """ + + # Sequence lengths have to be specified for LSTM batch inputs. The + # input batch must be padded to the max seq length given here. That is, + # batch_size == len(seq_lens) * max(seq_lens) + if framework in ["tf2", "tf"]: + assert time_major is False, "time-major not supported yet for tf!" + padded_inputs = tf.convert_to_tensor(padded_inputs) + padded_batch_size = tf.shape(padded_inputs)[0] + # Dynamically reshape the padded batch to introduce a time dimension. + new_batch_size = tf.shape(seq_lens)[0] + time_size = padded_batch_size // new_batch_size + new_shape = tf.concat( + [ + tf.expand_dims(new_batch_size, axis=0), + tf.expand_dims(time_size, axis=0), + tf.shape(padded_inputs)[1:], + ], + axis=0, + ) + return tf.reshape(padded_inputs, new_shape) + elif framework == "torch": + padded_inputs = torch.as_tensor(padded_inputs) + padded_batch_size = padded_inputs.shape[0] + + # Dynamically reshape the padded batch to introduce a time dimension. + new_batch_size = seq_lens.shape[0] + time_size = padded_batch_size // new_batch_size + batch_major_shape = (new_batch_size, time_size) + padded_inputs.shape[1:] + padded_outputs = padded_inputs.view(batch_major_shape) + + if time_major: + # Swap the batch and time dimensions + padded_outputs = padded_outputs.transpose(0, 1) + return padded_outputs + else: + assert framework == "np", "Unknown framework: {}".format(framework) + padded_inputs = np.asarray(padded_inputs) + padded_batch_size = padded_inputs.shape[0] + + # Dynamically reshape the padded batch to introduce a time dimension. + new_batch_size = seq_lens.shape[0] + time_size = padded_batch_size // new_batch_size + batch_major_shape = (new_batch_size, time_size) + padded_inputs.shape[1:] + padded_outputs = padded_inputs.reshape(batch_major_shape) + + if time_major: + # Swap the batch and time dimensions + padded_outputs = padded_outputs.transpose(0, 1) + return padded_outputs + + +@OldAPIStack +def chop_into_sequences( + *, + feature_columns, + state_columns, + max_seq_len, + episode_ids=None, + unroll_ids=None, + agent_indices=None, + dynamic_max=True, + shuffle=False, + seq_lens=None, + states_already_reduced_to_init=False, + handle_nested_data=False, + _extra_padding=0, + padding: str = "zero", + pad_infos_with_empty_dicts: bool = False, +): + """Truncate and pad experiences into fixed-length sequences. + + Args: + feature_columns: List of arrays containing features. + state_columns: List of arrays containing LSTM state values. + max_seq_len: Max length of sequences. Sequences longer than max_seq_len + will be split into subsequences that span the batch dimension + and sum to max_seq_len. + episode_ids (List[EpisodeID]): List of episode ids for each step. + unroll_ids (List[UnrollID]): List of identifiers for the sample batch. + This is used to make sure sequences are cut between sample batches. + agent_indices (List[AgentID]): List of agent ids for each step. Note + that this has to be combined with episode_ids for uniqueness. + dynamic_max: Whether to dynamically shrink the max seq len. + For example, if max len is 20 and the actual max seq len in the + data is 7, it will be shrunk to 7. + shuffle: Whether to shuffle the sequence outputs. + handle_nested_data: If True, assume that the data in + `feature_columns` could be nested structures (of data). + If False, assumes that all items in `feature_columns` are + only np.ndarrays (no nested structured of np.ndarrays). + _extra_padding: Add extra padding to the end of sequences. + padding: Padding type to use. Either "zero" or "last". Zero padding + will pad with zeros, last padding will pad with the last value. + pad_infos_with_empty_dicts: If True, will zero-pad INFOs with empty + dicts (instead of None). Used by the new API stack in the meantime, + however, as soon as the new ConnectorV2 API will be activated (as + part of the new API stack), we will no longer use this utility function + anyway. + + Returns: + f_pad: Padded feature columns. These will be of shape + [NUM_SEQUENCES * MAX_SEQ_LEN, ...]. + s_init: Initial states for each sequence, of shape + [NUM_SEQUENCES, ...]. + seq_lens: List of sequence lengths, of shape [NUM_SEQUENCES]. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.rnn_sequencing import chop_into_sequences + f_pad, s_init, seq_lens = chop_into_sequences( + episode_ids=[1, 1, 5, 5, 5, 5], + unroll_ids=[4, 4, 4, 4, 4, 4], + agent_indices=[0, 0, 0, 0, 0, 0], + feature_columns=[[4, 4, 8, 8, 8, 8], + [1, 1, 0, 1, 1, 0]], + state_columns=[[4, 5, 4, 5, 5, 5]], + max_seq_len=3) + print(f_pad) + print(s_init) + print(seq_lens) + + + .. testoutput:: + + [[4, 4, 0, 8, 8, 8, 8, 0, 0], + [1, 1, 0, 0, 1, 1, 0, 0, 0]] + [[4, 4, 5]] + [2, 3, 1] + """ + + if seq_lens is None or len(seq_lens) == 0: + prev_id = None + seq_lens = [] + seq_len = 0 + unique_ids = np.add( + np.add(episode_ids, agent_indices), + np.array(unroll_ids, dtype=np.int64) << 32, + ) + for uid in unique_ids: + if (prev_id is not None and uid != prev_id) or seq_len >= max_seq_len: + seq_lens.append(seq_len) + seq_len = 0 + seq_len += 1 + prev_id = uid + if seq_len: + seq_lens.append(seq_len) + seq_lens = np.array(seq_lens, dtype=np.int32) + + # Dynamically shrink max len as needed to optimize memory usage + if dynamic_max: + max_seq_len = max(seq_lens) + _extra_padding + + length = len(seq_lens) * max_seq_len + + feature_sequences = [] + for col in feature_columns: + if isinstance(col, list): + col = np.array(col) + feature_sequences.append([]) + + for f in tree.flatten(col): + # Save unnecessary copy. + if not isinstance(f, np.ndarray): + f = np.array(f) + + # New stack behavior (temporarily until we move to ConnectorV2 API, where + # this (admitedly convoluted) function will no longer be used at all). + if ( + f.dtype == object + and pad_infos_with_empty_dicts + and isinstance(f[0], dict) + ): + f_pad = [{} for _ in range(length)] + # Old stack behavior: Pad INFOs with None. + elif f.dtype == object or f.dtype.type is np.str_: + f_pad = [None] * length + # Pad everything else with zeros. + else: + # Make sure type doesn't change. + f_pad = np.zeros((length,) + np.shape(f)[1:], dtype=f.dtype) + seq_base = 0 + i = 0 + for len_ in seq_lens: + for seq_offset in range(len_): + f_pad[seq_base + seq_offset] = f[i] + i += 1 + + if padding == "last": + for seq_offset in range(len_, max_seq_len): + f_pad[seq_base + seq_offset] = f[i - 1] + + seq_base += max_seq_len + + assert i == len(f), f + feature_sequences[-1].append(f_pad) + + if states_already_reduced_to_init: + initial_states = state_columns + else: + initial_states = [] + for state_column in state_columns: + if isinstance(state_column, list): + state_column = np.array(state_column) + initial_state_flat = [] + # state_column may have a nested structure (e.g. LSTM state). + for s in tree.flatten(state_column): + # Skip unnecessary copy. + if not isinstance(s, np.ndarray): + s = np.array(s) + s_init = [] + i = 0 + for len_ in seq_lens: + s_init.append(s[i]) + i += len_ + initial_state_flat.append(np.array(s_init)) + initial_states.append(tree.unflatten_as(state_column, initial_state_flat)) + + if shuffle: + permutation = np.random.permutation(len(seq_lens)) + for i, f in enumerate(tree.flatten(feature_sequences)): + orig_shape = f.shape + f = np.reshape(f, (len(seq_lens), -1) + f.shape[1:]) + f = f[permutation] + f = np.reshape(f, orig_shape) + feature_sequences[i] = f + for i, s in enumerate(initial_states): + s = s[permutation] + initial_states[i] = s + seq_lens = seq_lens[permutation] + + # Classic behavior: Don't assume data in feature_columns are nested + # structs. Don't return them as flattened lists, but as is (index 0). + if not handle_nested_data: + feature_sequences = [f[0] for f in feature_sequences] + + return feature_sequences, initial_states, seq_lens + + +@OldAPIStack +def timeslice_along_seq_lens_with_overlap( + sample_batch: SampleBatchType, + seq_lens: Optional[List[int]] = None, + zero_pad_max_seq_len: int = 0, + pre_overlap: int = 0, + zero_init_states: bool = True, +) -> List["SampleBatch"]: + """Slices batch along `seq_lens` (each seq-len item produces one batch). + + Args: + sample_batch: The SampleBatch to timeslice. + seq_lens (Optional[List[int]]): An optional list of seq_lens to slice + at. If None, use `sample_batch[SampleBatch.SEQ_LENS]`. + zero_pad_max_seq_len: If >0, already zero-pad the resulting + slices up to this length. NOTE: This max-len will include the + additional timesteps gained via setting pre_overlap (see Example). + pre_overlap: If >0, will overlap each two consecutive slices by + this many timesteps (toward the left side). This will cause + zero-padding at the very beginning of the batch. + zero_init_states: Whether initial states should always be + zero'd. If False, will use the state_outs of the batch to + populate state_in values. + + Returns: + List[SampleBatch]: The list of (new) SampleBatches. + + Examples: + assert seq_lens == [5, 5, 2] + assert sample_batch.count == 12 + # self = 0 1 2 3 4 | 5 6 7 8 9 | 10 11 <- timesteps + slices = timeslice_along_seq_lens_with_overlap( + sample_batch=sample_batch. + zero_pad_max_seq_len=10, + pre_overlap=3) + # Z = zero padding (at beginning or end). + # |pre (3)| seq | max-seq-len (up to 10) + # slices[0] = | Z Z Z | 0 1 2 3 4 | Z Z + # slices[1] = | 2 3 4 | 5 6 7 8 9 | Z Z + # slices[2] = | 7 8 9 | 10 11 Z Z Z | Z Z + # Note that `zero_pad_max_seq_len=10` includes the 3 pre-overlaps + # count (makes sure each slice has exactly length 10). + """ + if seq_lens is None: + seq_lens = sample_batch.get(SampleBatch.SEQ_LENS) + else: + if sample_batch.get(SampleBatch.SEQ_LENS) is not None and log_once( + "overriding_sequencing_information" + ): + logger.warning( + "Found sequencing information in a batch that will be " + "ignored when slicing. Ignore this warning if you know " + "what you are doing." + ) + + if seq_lens is None: + max_seq_len = zero_pad_max_seq_len - pre_overlap + if log_once("no_sequence_lengths_available_for_time_slicing"): + logger.warning( + "Trying to slice a batch along sequences without " + "sequence lengths being provided in the batch. Batch will " + "be sliced into slices of size " + "{} = {} - {} = zero_pad_max_seq_len - pre_overlap.".format( + max_seq_len, zero_pad_max_seq_len, pre_overlap + ) + ) + num_seq_lens, last_seq_len = divmod(len(sample_batch), max_seq_len) + seq_lens = [zero_pad_max_seq_len] * num_seq_lens + ( + [last_seq_len] if last_seq_len else [] + ) + + assert ( + seq_lens is not None and len(seq_lens) > 0 + ), "Cannot timeslice along `seq_lens` when `seq_lens` is empty or None!" + # Generate n slices based on seq_lens. + start = 0 + slices = [] + for seq_len in seq_lens: + pre_begin = start - pre_overlap + slice_begin = start + end = start + seq_len + slices.append((pre_begin, slice_begin, end)) + start += seq_len + + timeslices = [] + for begin, slice_begin, end in slices: + zero_length = None + data_begin = 0 + zero_init_states_ = zero_init_states + if begin < 0: + zero_length = pre_overlap + data_begin = slice_begin + zero_init_states_ = True + else: + eps_ids = sample_batch[SampleBatch.EPS_ID][begin if begin >= 0 else 0 : end] + is_last_episode_ids = eps_ids == eps_ids[-1] + if not is_last_episode_ids[0]: + zero_length = int(sum(1.0 - is_last_episode_ids)) + data_begin = begin + zero_length + zero_init_states_ = True + + if zero_length is not None: + data = { + k: np.concatenate( + [ + np.zeros(shape=(zero_length,) + v.shape[1:], dtype=v.dtype), + v[data_begin:end], + ] + ) + for k, v in sample_batch.items() + if k != SampleBatch.SEQ_LENS + } + else: + data = { + k: v[begin:end] + for k, v in sample_batch.items() + if k != SampleBatch.SEQ_LENS + } + + if zero_init_states_: + i = 0 + key = "state_in_{}".format(i) + while key in data: + data[key] = np.zeros_like(sample_batch[key][0:1]) + # Del state_out_n from data if exists. + data.pop("state_out_{}".format(i), None) + i += 1 + key = "state_in_{}".format(i) + # TODO: This will not work with attention nets as their state_outs are + # not compatible with state_ins. + else: + i = 0 + key = "state_in_{}".format(i) + while key in data: + data[key] = sample_batch["state_out_{}".format(i)][begin - 1 : begin] + del data["state_out_{}".format(i)] + i += 1 + key = "state_in_{}".format(i) + + timeslices.append(SampleBatch(data, seq_lens=[end - begin])) + + # Zero-pad each slice if necessary. + if zero_pad_max_seq_len > 0: + for ts in timeslices: + ts.right_zero_pad(max_seq_len=zero_pad_max_seq_len, exclude_states=True) + + return timeslices + + +@OldAPIStack +def get_fold_unfold_fns(b_dim: int, t_dim: int, framework: str): + """Produces two functions to fold/unfold any Tensors in a struct. + + Args: + b_dim: The batch dimension to use for folding. + t_dim: The time dimension to use for folding. + framework: The framework to use for folding. One of "tf2" or "torch". + + Returns: + fold: A function that takes a struct of torch.Tensors and reshapes + them to have a first dimension of `b_dim * t_dim`. + unfold: A function that takes a struct of torch.Tensors and reshapes + them to have a first dimension of `b_dim` and a second dimension + of `t_dim`. + """ + if framework in "tf2": + # TensorFlow traced eager complains if we don't convert these to tensors here + b_dim = tf.convert_to_tensor(b_dim) + t_dim = tf.convert_to_tensor(t_dim) + + def fold_mapping(item): + if item is None: + # Torch has no representation for `None`, so we return None + return item + item = tf.convert_to_tensor(item) + shape = tf.shape(item) + other_dims = shape[2:] + return tf.reshape(item, tf.concat([[b_dim * t_dim], other_dims], axis=0)) + + def unfold_mapping(item): + if item is None: + return item + item = tf.convert_to_tensor(item) + shape = item.shape + other_dims = shape[1:] + + return tf.reshape(item, tf.concat([[b_dim], [t_dim], other_dims], axis=0)) + + elif framework == "torch": + + def fold_mapping(item): + if item is None: + # Torch has no representation for `None`, so we return None + return item + item = torch.as_tensor(item) + size = list(item.size()) + current_b_dim, current_t_dim = list(size[:2]) + + assert (b_dim, t_dim) == (current_b_dim, current_t_dim), ( + "All tensors in the struct must have the same batch and time " + "dimensions. Got {} and {}.".format( + (b_dim, t_dim), (current_b_dim, current_t_dim) + ) + ) + + other_dims = size[2:] + return item.reshape([b_dim * t_dim] + other_dims) + + def unfold_mapping(item): + if item is None: + return item + item = torch.as_tensor(item) + size = list(item.size()) + current_b_dim = size[0] + other_dims = size[1:] + assert current_b_dim == b_dim * t_dim, ( + "The first dimension of the tensor must be equal to the product of " + "the desired batch and time dimensions. Got {} and {}.".format( + current_b_dim, b_dim * t_dim + ) + ) + return item.reshape([b_dim, t_dim] + other_dims) + + else: + raise ValueError(f"framework {framework} not implemented!") + + return functools.partial(tree.map_structure, fold_mapping), functools.partial( + tree.map_structure, unfold_mapping + ) diff --git a/lib/python3.12/site-packages/ray/rllib/policy/sample_batch.py b/lib/python3.12/site-packages/ray/rllib/policy/sample_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..253844e1d8e7f04b00b80c5ed59bee8e14ed76c1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/sample_batch.py @@ -0,0 +1,1843 @@ +import collections +import itertools +import sys +from functools import partial +from numbers import Number +from typing import Dict, Iterator, List, Optional, Set, Union + +import numpy as np +import tree # pip install dm_tree + +from ray._common.deprecation import Deprecated, deprecation_warning +from ray.rllib.core.columns import Columns +from ray.rllib.utils.annotations import DeveloperAPI, ExperimentalAPI, PublicAPI +from ray.rllib.utils.compression import is_compressed, pack, unpack +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import ( + ModuleID, + PolicyID, + SampleBatchType, + TensorType, + ViewRequirementsDict, +) +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +# Default policy id for single agent environments +DEFAULT_POLICY_ID = "default_policy" + + +@DeveloperAPI +def attempt_count_timesteps(tensor_dict: dict): + """Attempt to count timesteps based on dimensions of individual elements. + + Returns the first successfully counted number of timesteps. + We do not attempt to count on INFOS or any state_in_* and state_out_* keys. The + number of timesteps we count in cases where we are unable to count is zero. + + Args: + tensor_dict: A SampleBatch or another dict. + + Returns: + count: The inferred number of timesteps >= 0. + """ + # Try to infer the "length" of the SampleBatch by finding the first + # value that is actually a ndarray/tensor. + # Skip manual counting routine if we can directly infer count from sequence lengths + seq_lens = tensor_dict.get(SampleBatch.SEQ_LENS) + if ( + seq_lens is not None + and not (tf and tf.is_tensor(seq_lens) and not hasattr(seq_lens, "numpy")) + and len(seq_lens) > 0 + ): + if torch and torch.is_tensor(seq_lens): + return int(seq_lens.sum().item()) + else: + return int(sum(seq_lens)) + + for k, v in tensor_dict.items(): + if k == SampleBatch.SEQ_LENS: + continue + + assert isinstance(k, str), tensor_dict + + if ( + k == SampleBatch.INFOS + or k.startswith("state_in_") + or k.startswith("state_out_") + ): + # Don't attempt to count on infos since we make no assumptions + # about its content + # Don't attempt to count on state since nesting can potentially mess + # things up + continue + + # If this is a nested dict (for example a nested observation), + # try to flatten it, assert that all elements have the same length (batch + # dimension) + v_list = tree.flatten(v) if isinstance(v, (dict, tuple)) else [v] + # TODO: Drop support for lists and Numbers as values. + # If v_list contains lists or Numbers, convert them to arrays, too. + v_list = [ + np.array(_v) if isinstance(_v, (Number, list)) else _v for _v in v_list + ] + try: + # Add one of the elements' length, since they are all the same + _len = len(v_list[0]) + if _len: + return _len + except Exception: + pass + + # Return zero if we are unable to count + return 0 + + +@PublicAPI +class SampleBatch(dict): + """Wrapper around a dictionary with string keys and array-like values. + + For example, {"obs": [1, 2, 3], "reward": [0, -1, 1]} is a batch of three + samples, each with an "obs" and "reward" attribute. + """ + + # On rows in SampleBatch: + # Each comment signifies how values relate to each other within a given row. + # A row generally signifies one timestep. Most importantly, at t=0, SampleBatch.OBS + # will usually be the reset-observation, while SampleBatch.ACTIONS will be the + # action based on the reset-observation and so on. This scheme is derived from + # RLlib's sampling logic. + + # The following fields have all been moved to `Columns` and are only left here + # for backward compatibility. + OBS = Columns.OBS + ACTIONS = Columns.ACTIONS + REWARDS = Columns.REWARDS + TERMINATEDS = Columns.TERMINATEDS + TRUNCATEDS = Columns.TRUNCATEDS + INFOS = Columns.INFOS + SEQ_LENS = Columns.SEQ_LENS + T = Columns.T + ACTION_DIST_INPUTS = Columns.ACTION_DIST_INPUTS + ACTION_PROB = Columns.ACTION_PROB + ACTION_LOGP = Columns.ACTION_LOGP + VF_PREDS = Columns.VF_PREDS + VALUES_BOOTSTRAPPED = Columns.VALUES_BOOTSTRAPPED + EPS_ID = Columns.EPS_ID + NEXT_OBS = Columns.NEXT_OBS + + # Action distribution object. + ACTION_DIST = "action_dist" + # Action chosen before SampleBatch.ACTIONS. + PREV_ACTIONS = "prev_actions" + # Reward received before SampleBatch.REWARDS. + PREV_REWARDS = "prev_rewards" + ENV_ID = "env_id" # An env ID (e.g. the index for a vectorized sub-env). + AGENT_INDEX = "agent_index" # Uniquely identifies an agent within an episode. + # Uniquely identifies a sample batch. This is important to distinguish RNN + # sequences from the same episode when multiple sample batches are + # concatenated (fusing sequences across batches can be unsafe). + UNROLL_ID = "unroll_id" + + # RE 3 + # This is only computed and used when RE3 exploration strategy is enabled. + OBS_EMBEDS = "obs_embeds" + # Decision Transformer + RETURNS_TO_GO = "returns_to_go" + ATTENTION_MASKS = "attention_masks" + # Do not set this key directly. Instead, the values under this key are + # auto-computed via the values of the TERMINATEDS and TRUNCATEDS keys. + DONES = "dones" + # Use SampleBatch.OBS instead. + CUR_OBS = "obs" + + @PublicAPI + def __init__(self, *args, **kwargs): + """Constructs a sample batch (same params as dict constructor). + + Note: All args and those kwargs not listed below will be passed + as-is to the parent dict constructor. + + Args: + _time_major: Whether data in this sample batch + is time-major. This is False by default and only relevant + if the data contains sequences. + _max_seq_len: The max sequence chunk length + if the data contains sequences. + _zero_padded: Whether the data in this batch + contains sequences AND these sequences are right-zero-padded + according to the `_max_seq_len` setting. + _is_training: Whether this batch is used for + training. If False, batch may be used for e.g. action + computations (inference). + """ + + if SampleBatch.DONES in kwargs: + raise KeyError( + "SampleBatch cannot be constructed anymore with a `DONES` key! " + "Instead, set the new TERMINATEDS and TRUNCATEDS keys. The values under" + " DONES will then be automatically computed using terminated|truncated." + ) + + # Possible seq_lens (TxB or BxT) setup. + self.time_major = kwargs.pop("_time_major", None) + # Maximum seq len value. + self.max_seq_len = kwargs.pop("_max_seq_len", None) + # Is alredy right-zero-padded? + self.zero_padded = kwargs.pop("_zero_padded", False) + # Whether this batch is used for training (vs inference). + self._is_training = kwargs.pop("_is_training", None) + # Weighted average number of grad updates that have been performed on the + # policy/ies that were used to collect this batch. + # E.g.: Two rollout workers collect samples of 50ts each + # (rollout_fragment_length=50). One of them has a policy that has undergone + # 2 updates thus far, the other worker uses a policy that has undergone 3 + # updates thus far. The train batch size is 100, so we concatenate these 2 + # batches to a new one that's 100ts long. This new 100ts batch will have its + # `num_gradient_updates` property set to 2.5 as it's the weighted average + # (both original batches contribute 50%). + self.num_grad_updates: Optional[float] = kwargs.pop("_num_grad_updates", None) + + # Call super constructor. This will make the actual data accessible + # by column name (str) via e.g. self["some-col"]. + dict.__init__(self, *args, **kwargs) + + # Indicates whether, for this batch, sequence lengths should be slices by + # their index in the batch or by their index as a sequence. + # This is useful if a batch contains tensors of shape (B, T, ...), where each + # index of B indicates one sequence. In this case, when slicing the batch, + # we want one sequence to be slices out per index in B ( + # `_slice_seq_lens_by_batch_index=True`. However, if the padded batch + # contains tensors of shape (B*T, ...), where each index of B*T indicates + # one timestep, we want one sequence to be sliced per T steps in B*T ( + # `self._slice_seq_lens_in_B=False`). + # ._slice_seq_lens_in_B = True is only meant to be used for batches that we + # feed into Learner._update(), all other places in RLlib are not expected to + # need this. + self._slice_seq_lens_in_B = False + + self.accessed_keys = set() + self.added_keys = set() + self.deleted_keys = set() + self.intercepted_values = {} + self.get_interceptor = None + + # Clear out None seq-lens. + seq_lens_ = self.get(SampleBatch.SEQ_LENS) + if seq_lens_ is None or (isinstance(seq_lens_, list) and len(seq_lens_) == 0): + self.pop(SampleBatch.SEQ_LENS, None) + # Numpyfy seq_lens if list. + elif isinstance(seq_lens_, list): + self[SampleBatch.SEQ_LENS] = seq_lens_ = np.array(seq_lens_, dtype=np.int32) + elif (torch and torch.is_tensor(seq_lens_)) or (tf and tf.is_tensor(seq_lens_)): + self[SampleBatch.SEQ_LENS] = seq_lens_ + + if ( + self.max_seq_len is None + and seq_lens_ is not None + and not (tf and tf.is_tensor(seq_lens_)) + and len(seq_lens_) > 0 + ): + if torch and torch.is_tensor(seq_lens_): + self.max_seq_len = seq_lens_.max().item() + else: + self.max_seq_len = max(seq_lens_) + + if self._is_training is None: + self._is_training = self.pop("is_training", False) + + for k, v in self.items(): + # TODO: Drop support for lists and Numbers as values. + # Convert lists of int|float into numpy arrays make sure all data + # has same length. + if isinstance(v, (Number, list)) and not k == SampleBatch.INFOS: + self[k] = np.array(v) + + self.count = attempt_count_timesteps(self) + + # A convenience map for slicing this batch into sub-batches along + # the time axis. This helps reduce repeated iterations through the + # batch's seq_lens array to find good slicing points. Built lazily + # when needed. + self._slice_map = [] + + @PublicAPI + def __len__(self) -> int: + """Returns the amount of samples in the sample batch.""" + return self.count + + @PublicAPI + def agent_steps(self) -> int: + """Returns the same as len(self) (number of steps in this batch). + + To make this compatible with `MultiAgentBatch.agent_steps()`. + """ + return len(self) + + @PublicAPI + def env_steps(self) -> int: + """Returns the same as len(self) (number of steps in this batch). + + To make this compatible with `MultiAgentBatch.env_steps()`. + """ + return len(self) + + @DeveloperAPI + def enable_slicing_by_batch_id(self): + self._slice_seq_lens_in_B = True + + @DeveloperAPI + def disable_slicing_by_batch_id(self): + self._slice_seq_lens_in_B = False + + @ExperimentalAPI + def is_terminated_or_truncated(self) -> bool: + """Returns True if `self` is either terminated or truncated at idx -1.""" + return self[SampleBatch.TERMINATEDS][-1] or ( + SampleBatch.TRUNCATEDS in self and self[SampleBatch.TRUNCATEDS][-1] + ) + + @ExperimentalAPI + def is_single_trajectory(self) -> bool: + """Returns True if this SampleBatch only contains one trajectory. + + This is determined by checking all timesteps (except for the last) for being + not terminated AND (if applicable) not truncated. + """ + return not any(self[SampleBatch.TERMINATEDS][:-1]) and ( + SampleBatch.TRUNCATEDS not in self + or not any(self[SampleBatch.TRUNCATEDS][:-1]) + ) + + @staticmethod + @PublicAPI + @Deprecated(new="concat_samples() from rllib.policy.sample_batch", error=True) + def concat_samples(samples): + pass + + @PublicAPI + def concat(self, other: "SampleBatch") -> "SampleBatch": + """Concatenates `other` to this one and returns a new SampleBatch. + + Args: + other: The other SampleBatch object to concat to this one. + + Returns: + The new SampleBatch, resulting from concating `other` to `self`. + + .. testcode:: + :skipif: True + + import numpy as np + from ray.rllib.policy.sample_batch import SampleBatch + b1 = SampleBatch({"a": np.array([1, 2])}) + b2 = SampleBatch({"a": np.array([3, 4, 5])}) + print(b1.concat(b2)) + + .. testoutput:: + + {"a": np.array([1, 2, 3, 4, 5])} + """ + return concat_samples([self, other]) + + @PublicAPI + def copy(self, shallow: bool = False) -> "SampleBatch": + """Creates a deep or shallow copy of this SampleBatch and returns it. + + Args: + shallow: Whether the copying should be done shallowly. + + Returns: + A deep or shallow copy of this SampleBatch object. + """ + copy_ = dict(self) + data = tree.map_structure( + lambda v: ( + np.array(v, copy=not shallow) if isinstance(v, np.ndarray) else v + ), + copy_, + ) + copy_ = SampleBatch( + data, + _time_major=self.time_major, + _zero_padded=self.zero_padded, + _max_seq_len=self.max_seq_len, + _num_grad_updates=self.num_grad_updates, + ) + copy_.set_get_interceptor(self.get_interceptor) + copy_.added_keys = self.added_keys + copy_.deleted_keys = self.deleted_keys + copy_.accessed_keys = self.accessed_keys + return copy_ + + @PublicAPI + def rows(self) -> Iterator[Dict[str, TensorType]]: + """Returns an iterator over data rows, i.e. dicts with column values. + + Note that if `seq_lens` is set in self, we set it to 1 in the rows. + + Yields: + The column values of the row in this iteration. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.sample_batch import SampleBatch + batch = SampleBatch({ + "a": [1, 2, 3], + "b": [4, 5, 6], + "seq_lens": [1, 2] + }) + for row in batch.rows(): + print(row) + + .. testoutput:: + + {"a": 1, "b": 4, "seq_lens": 1} + {"a": 2, "b": 5, "seq_lens": 1} + {"a": 3, "b": 6, "seq_lens": 1} + """ + + seq_lens = None if self.get(SampleBatch.SEQ_LENS, 1) is None else 1 + + self_as_dict = dict(self) + + for i in range(self.count): + yield tree.map_structure_with_path( + lambda p, v, i=i: v[i] if p[0] != self.SEQ_LENS else seq_lens, + self_as_dict, + ) + + @PublicAPI + def columns(self, keys: List[str]) -> List[any]: + """Returns a list of the batch-data in the specified columns. + + Args: + keys: List of column names fo which to return the data. + + Returns: + The list of data items ordered by the order of column + names in `keys`. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.sample_batch import SampleBatch + batch = SampleBatch({"a": [1], "b": [2], "c": [3]}) + print(batch.columns(["a", "b"])) + + .. testoutput:: + + [[1], [2]] + """ + + # TODO: (sven) Make this work for nested data as well. + out = [] + for k in keys: + out.append(self[k]) + return out + + @PublicAPI + def shuffle(self) -> "SampleBatch": + """Shuffles the rows of this batch in-place. + + Returns: + This very (now shuffled) SampleBatch. + + Raises: + ValueError: If self[SampleBatch.SEQ_LENS] is defined. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.sample_batch import SampleBatch + batch = SampleBatch({"a": [1, 2, 3, 4]}) + print(batch.shuffle()) + + .. testoutput:: + + {"a": [4, 1, 3, 2]} + """ + has_time_rank = self.get(SampleBatch.SEQ_LENS) is not None + + # Shuffling the data when we have `seq_lens` defined is probably + # a bad idea! + if has_time_rank and not self.zero_padded: + raise ValueError( + "SampleBatch.shuffle not possible when your data has " + "`seq_lens` defined AND is not zero-padded yet!" + ) + + # Get a permutation over the single items once and use the same + # permutation for all the data (otherwise, data would become + # meaningless). + # - Shuffle by individual item. + if not has_time_rank: + permutation = np.random.permutation(self.count) + # - Shuffle along batch axis (leave axis=1/time-axis as-is). + else: + permutation = np.random.permutation(len(self[SampleBatch.SEQ_LENS])) + + self_as_dict = dict(self) + infos = self_as_dict.pop(Columns.INFOS, None) + shuffled = tree.map_structure(lambda v: v[permutation], self_as_dict) + if infos is not None: + self_as_dict[Columns.INFOS] = [infos[i] for i in permutation] + + self.update(shuffled) + + # Flush cache such that intercepted values are recalculated after the + # shuffling. + self.intercepted_values = {} + return self + + @PublicAPI + def split_by_episode(self, key: Optional[str] = None) -> List["SampleBatch"]: + """Splits by `eps_id` column and returns list of new batches. + If `eps_id` is not present, splits by `dones` instead. + + Args: + key: If specified, overwrite default and use key to split. + + Returns: + List of batches, one per distinct episode. + + Raises: + KeyError: If the `eps_id` AND `dones` columns are not present. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.sample_batch import SampleBatch + # "eps_id" is present + batch = SampleBatch( + {"a": [1, 2, 3], "eps_id": [0, 0, 1]}) + print(batch.split_by_episode()) + + # "eps_id" not present, split by "dones" instead + batch = SampleBatch( + {"a": [1, 2, 3, 4, 5], "dones": [0, 0, 1, 0, 1]}) + print(batch.split_by_episode()) + + # The last episode is appended even if it does not end with done + batch = SampleBatch( + {"a": [1, 2, 3, 4, 5], "dones": [0, 0, 1, 0, 0]}) + print(batch.split_by_episode()) + + batch = SampleBatch( + {"a": [1, 2, 3, 4, 5], "dones": [0, 0, 0, 0, 0]}) + print(batch.split_by_episode()) + + + .. testoutput:: + + [{"a": [1, 2], "eps_id": [0, 0]}, {"a": [3], "eps_id": [1]}] + [{"a": [1, 2, 3], "dones": [0, 0, 1]}, {"a": [4, 5], "dones": [0, 1]}] + [{"a": [1, 2, 3], "dones": [0, 0, 1]}, {"a": [4, 5], "dones": [0, 0]}] + [{"a": [1, 2, 3, 4, 5], "dones": [0, 0, 0, 0, 0]}] + + + """ + + assert key is None or key in [SampleBatch.EPS_ID, SampleBatch.DONES], ( + f"`SampleBatch.split_by_episode(key={key})` invalid! " + f"Must be [None|'dones'|'eps_id']." + ) + + def slice_by_eps_id(): + slices = [] + # Produce a new slice whenever we find a new episode ID. + cur_eps_id = self[SampleBatch.EPS_ID][0] + offset = 0 + for i in range(self.count): + next_eps_id = self[SampleBatch.EPS_ID][i] + if next_eps_id != cur_eps_id: + slices.append(self[offset:i]) + offset = i + cur_eps_id = next_eps_id + # Add final slice. + slices.append(self[offset : self.count]) + return slices + + def slice_by_terminateds_or_truncateds(): + slices = [] + offset = 0 + for i in range(self.count): + if self[SampleBatch.TERMINATEDS][i] or ( + SampleBatch.TRUNCATEDS in self and self[SampleBatch.TRUNCATEDS][i] + ): + # Since self[i] is the last timestep of the episode, + # append it to the batch, then set offset to the start + # of the next batch + slices.append(self[offset : i + 1]) + offset = i + 1 + # Add final slice. + if offset != self.count: + slices.append(self[offset:]) + return slices + + key_to_method = { + SampleBatch.EPS_ID: slice_by_eps_id, + SampleBatch.DONES: slice_by_terminateds_or_truncateds, + } + + # If key not specified, default to this order. + key_resolve_order = [SampleBatch.EPS_ID, SampleBatch.DONES] + + slices = None + if key is not None: + # If key specified, directly use it. + if key == SampleBatch.EPS_ID and key not in self: + raise KeyError(f"{self} does not have key `{key}`!") + slices = key_to_method[key]() + else: + # If key not specified, go in order. + for key in key_resolve_order: + if key == SampleBatch.DONES or key in self: + slices = key_to_method[key]() + break + if slices is None: + raise KeyError(f"{self} does not have keys {key_resolve_order}!") + + assert ( + sum(s.count for s in slices) == self.count + ), f"Calling split_by_episode on {self} returns {slices}" + f"which should in total have {self.count} timesteps!" + return slices + + def slice( + self, start: int, end: int, state_start=None, state_end=None + ) -> "SampleBatch": + """Returns a slice of the row data of this batch (w/o copying). + + Args: + start: Starting index. If < 0, will left-zero-pad. + end: Ending index. + + Returns: + A new SampleBatch, which has a slice of this batch's data. + """ + if ( + self.get(SampleBatch.SEQ_LENS) is not None + and len(self[SampleBatch.SEQ_LENS]) > 0 + ): + if start < 0: + data = { + k: np.concatenate( + [ + np.zeros(shape=(-start,) + v.shape[1:], dtype=v.dtype), + v[0:end], + ] + ) + for k, v in self.items() + if k != SampleBatch.SEQ_LENS and not k.startswith("state_in_") + } + else: + data = { + k: tree.map_structure(lambda s: s[start:end], v) + for k, v in self.items() + if k != SampleBatch.SEQ_LENS and not k.startswith("state_in_") + } + if state_start is not None: + assert state_end is not None + state_idx = 0 + state_key = "state_in_{}".format(state_idx) + while state_key in self: + data[state_key] = self[state_key][state_start:state_end] + state_idx += 1 + state_key = "state_in_{}".format(state_idx) + seq_lens = list(self[SampleBatch.SEQ_LENS][state_start:state_end]) + # Adjust seq_lens if necessary. + data_len = len(data[next(iter(data))]) + if sum(seq_lens) != data_len: + assert sum(seq_lens) > data_len + seq_lens[-1] = data_len - sum(seq_lens[:-1]) + else: + # Fix state_in_x data. + count = 0 + state_start = None + seq_lens = None + for i, seq_len in enumerate(self[SampleBatch.SEQ_LENS]): + count += seq_len + if count >= end: + state_idx = 0 + state_key = "state_in_{}".format(state_idx) + if state_start is None: + state_start = i + while state_key in self: + data[state_key] = self[state_key][state_start : i + 1] + state_idx += 1 + state_key = "state_in_{}".format(state_idx) + seq_lens = list(self[SampleBatch.SEQ_LENS][state_start:i]) + [ + seq_len - (count - end) + ] + if start < 0: + seq_lens[0] += -start + diff = sum(seq_lens) - (end - start) + if diff > 0: + seq_lens[0] -= diff + assert sum(seq_lens) == (end - start) + break + elif state_start is None and count > start: + state_start = i + + return SampleBatch( + data, + seq_lens=seq_lens, + _is_training=self.is_training, + _time_major=self.time_major, + _num_grad_updates=self.num_grad_updates, + ) + else: + return SampleBatch( + tree.map_structure(lambda value: value[start:end], self), + _is_training=self.is_training, + _time_major=self.time_major, + _num_grad_updates=self.num_grad_updates, + ) + + def _batch_slice(self, slice_: slice) -> "SampleBatch": + """Helper method to handle SampleBatch slicing using a slice object. + + The returned SampleBatch uses the same underlying data object as + `self`, so changing the slice will also change `self`. + + Note that only zero or positive bounds are allowed for both start + and stop values. The slice step must be 1 (or None, which is the + same). + + Args: + slice_: The python slice object to slice by. + + Returns: + A new SampleBatch, however "linking" into the same data + (sliced) as self. + """ + start = slice_.start or 0 + stop = slice_.stop or len(self[SampleBatch.SEQ_LENS]) + # If stop goes beyond the length of this batch -> Make it go till the + # end only (including last item). + # Analogous to `l = [0, 1, 2]; l[:100] -> [0, 1, 2];`. + if stop > len(self): + stop = len(self) + assert start >= 0 and stop >= 0 and slice_.step in [1, None] + + # Exclude INFOs from regular array slicing as the data under this column might + # be a list (not good for `tree.map_structure` call). + # Furthermore, slicing does not work when the data in the column is + # singular (not a list or array). + infos = self.pop(SampleBatch.INFOS, None) + data = tree.map_structure(lambda value: value[start:stop], self) + if infos is not None: + # Slice infos according to SEQ_LENS. + info_slice_start = int(sum(self[SampleBatch.SEQ_LENS][:start])) + info_slice_stop = int(sum(self[SampleBatch.SEQ_LENS][start:stop])) + data[SampleBatch.INFOS] = infos[info_slice_start:info_slice_stop] + # Put infos back into `self`. + self[Columns.INFOS] = infos + + return SampleBatch( + data, + _is_training=self.is_training, + _time_major=self.time_major, + _num_grad_updates=self.num_grad_updates, + ) + + @PublicAPI + def timeslices( + self, + size: Optional[int] = None, + num_slices: Optional[int] = None, + k: Optional[int] = None, + ) -> List["SampleBatch"]: + """Returns SampleBatches, each one representing a k-slice of this one. + + Will start from timestep 0 and produce slices of size=k. + + Args: + size: The size (in timesteps) of each returned SampleBatch. + num_slices: The number of slices to produce. + k: Deprecated: Use size or num_slices instead. The size + (in timesteps) of each returned SampleBatch. + + Returns: + The list of `num_slices` (new) SampleBatches or n (new) + SampleBatches each one of size `size`. + """ + if size is None and num_slices is None: + deprecation_warning("k", "size or num_slices") + assert k is not None + size = k + + if size is None: + assert isinstance(num_slices, int) + + slices = [] + left = len(self) + start = 0 + while left: + len_ = left // (num_slices - len(slices)) + stop = start + len_ + slices.append(self[start:stop]) + left -= len_ + start = stop + + return slices + + else: + assert isinstance(size, int) + + slices = [] + left = len(self) + start = 0 + while left: + stop = start + size + slices.append(self[start:stop]) + left -= size + start = stop + + return slices + + @Deprecated(new="SampleBatch.right_zero_pad", error=True) + def zero_pad(self, max_seq_len, exclude_states=True): + pass + + def right_zero_pad(self, max_seq_len: int, exclude_states: bool = True): + """Right (adding zeros at end) zero-pads this SampleBatch in-place. + + This will set the `self.zero_padded` flag to True and + `self.max_seq_len` to the given `max_seq_len` value. + + Args: + max_seq_len: The max (total) length to zero pad to. + exclude_states: If False, also right-zero-pad all + `state_in_x` data. If True, leave `state_in_x` keys + as-is. + + Returns: + This very (now right-zero-padded) SampleBatch. + + Raises: + ValueError: If self[SampleBatch.SEQ_LENS] is None (not defined). + + .. testcode:: + :skipif: True + + from ray.rllib.policy.sample_batch import SampleBatch + batch = SampleBatch( + {"a": [1, 2, 3], "seq_lens": [1, 2]}) + print(batch.right_zero_pad(max_seq_len=4)) + + batch = SampleBatch({"a": [1, 2, 3], + "state_in_0": [1.0, 3.0], + "seq_lens": [1, 2]}) + print(batch.right_zero_pad(max_seq_len=5)) + + .. testoutput:: + + {"a": [1, 0, 0, 0, 2, 3, 0, 0], "seq_lens": [1, 2]} + {"a": [1, 0, 0, 0, 0, 2, 3, 0, 0, 0], + "state_in_0": [1.0, 3.0], # <- all state-ins remain as-is + "seq_lens": [1, 2]} + + """ + seq_lens = self.get(SampleBatch.SEQ_LENS) + if seq_lens is None: + raise ValueError( + "Cannot right-zero-pad SampleBatch if no `seq_lens` field " + f"present! SampleBatch={self}" + ) + + length = len(seq_lens) * max_seq_len + + def _zero_pad_in_place(path, value): + # Skip "state_in_..." columns and "seq_lens". + if (exclude_states is True and path[0].startswith("state_in_")) or path[ + 0 + ] == SampleBatch.SEQ_LENS: + return + # Generate zero-filled primer of len=max_seq_len. + if value.dtype == object or value.dtype.type is np.str_: + f_pad = [None] * length + else: + # Make sure type doesn't change. + f_pad = np.zeros((length,) + np.shape(value)[1:], dtype=value.dtype) + # Fill primer with data. + f_pad_base = f_base = 0 + for len_ in self[SampleBatch.SEQ_LENS]: + f_pad[f_pad_base : f_pad_base + len_] = value[f_base : f_base + len_] + f_pad_base += max_seq_len + f_base += len_ + assert f_base == len(value), value + + # Update our data in-place. + curr = self + for i, p in enumerate(path): + if i == len(path) - 1: + curr[p] = f_pad + curr = curr[p] + + self_as_dict = dict(self) + tree.map_structure_with_path(_zero_pad_in_place, self_as_dict) + + # Set flags to indicate, we are now zero-padded (and to what extend). + self.zero_padded = True + self.max_seq_len = max_seq_len + + return self + + @ExperimentalAPI + def to_device( + self, + device, + framework: str = "torch", + pin_memory: bool = False, + use_stream: bool = False, + stream: Optional[Union["torch.cuda.Stream", "torch.cuda.Stream"]] = None, + ): + """TODO: transfer batch to given device as framework tensor.""" + if framework == "torch": + assert torch is not None + for k, v in self.items(): + self[k] = convert_to_torch_tensor( + v, + device, + pin_memory=pin_memory, + use_stream=use_stream, + stream=stream, + ) + else: + raise NotImplementedError + return self + + @PublicAPI + def size_bytes(self) -> int: + """Returns sum over number of bytes of all data buffers. + + For numpy arrays, we use ``.nbytes``. For all other value types, we use + sys.getsizeof(...). + + Returns: + The overall size in bytes of the data buffer (all columns). + """ + return sum( + v.nbytes if isinstance(v, np.ndarray) else sys.getsizeof(v) + for v in tree.flatten(self) + ) + + def get(self, key, default=None): + """Returns one column (by key) from the data or a default value.""" + try: + return self.__getitem__(key) + except KeyError: + return default + + @PublicAPI + def as_multi_agent(self, module_id: Optional[ModuleID] = None) -> "MultiAgentBatch": + """Returns the respective MultiAgentBatch + + Note, if `module_id` is not provided uses `DEFAULT_POLICY`_ID`. + + Args; + module_id: An optional module ID. If `None` the `DEFAULT_POLICY_ID` + is used. + + Returns: + The MultiAgentBatch (using DEFAULT_POLICY_ID) corresponding + to this SampleBatch. + """ + return MultiAgentBatch({module_id or DEFAULT_POLICY_ID: self}, self.count) + + @PublicAPI + def __getitem__(self, key: Union[str, slice]) -> TensorType: + """Returns one column (by key) from the data or a sliced new batch. + + Args: + key: The key (column name) to return or + a slice object for slicing this SampleBatch. + + Returns: + The data under the given key or a sliced version of this batch. + """ + if isinstance(key, slice): + return self._slice(key) + + # Special key DONES -> Translate to `TERMINATEDS | TRUNCATEDS` to reflect + # the old meaning of DONES. + if key == SampleBatch.DONES: + return self[SampleBatch.TERMINATEDS] + # Backward compatibility for when "input-dicts" were used. + elif key == "is_training": + if log_once("SampleBatch['is_training']"): + deprecation_warning( + old="SampleBatch['is_training']", + new="SampleBatch.is_training", + error=False, + ) + return self.is_training + + if not hasattr(self, key) and key in self: + self.accessed_keys.add(key) + + value = dict.__getitem__(self, key) + if self.get_interceptor is not None: + if key not in self.intercepted_values: + self.intercepted_values[key] = self.get_interceptor(value) + value = self.intercepted_values[key] + return value + + @PublicAPI + def __setitem__(self, key, item) -> None: + """Inserts (overrides) an entire column (by key) in the data buffer. + + Args: + key: The column name to set a value for. + item: The data to insert. + """ + # Disallow setting DONES key directly. + if key == SampleBatch.DONES: + raise KeyError( + "Cannot set `DONES` anymore in a SampleBatch! " + "Instead, set the new TERMINATEDS and TRUNCATEDS keys. The values under" + " DONES will then be automatically computed using terminated|truncated." + ) + # Defend against creating SampleBatch via pickle (no property + # `added_keys` and first item is already set). + elif not hasattr(self, "added_keys"): + dict.__setitem__(self, key, item) + return + + # Backward compatibility for when "input-dicts" were used. + if key == "is_training": + if log_once("SampleBatch['is_training']"): + deprecation_warning( + old="SampleBatch['is_training']", + new="SampleBatch.is_training", + error=False, + ) + self._is_training = item + return + + if key not in self: + self.added_keys.add(key) + + dict.__setitem__(self, key, item) + if key in self.intercepted_values: + self.intercepted_values[key] = item + + @property + def is_training(self): + if self.get_interceptor is not None and isinstance(self._is_training, bool): + if "_is_training" not in self.intercepted_values: + self.intercepted_values["_is_training"] = self.get_interceptor( + self._is_training + ) + return self.intercepted_values["_is_training"] + return self._is_training + + def set_training(self, training: Union[bool, "tf1.placeholder"] = True): + """Sets the `is_training` flag for this SampleBatch.""" + self._is_training = training + self.intercepted_values.pop("_is_training", None) + + @PublicAPI + def __delitem__(self, key): + self.deleted_keys.add(key) + dict.__delitem__(self, key) + + @DeveloperAPI + def compress( + self, bulk: bool = False, columns: Set[str] = frozenset(["obs", "new_obs"]) + ) -> "SampleBatch": + """Compresses the data buffers (by column) in place. + + Args: + bulk: Whether to compress across the batch dimension (0) + as well. If False will compress n separate list items, where n + is the batch size. + columns: The columns to compress. Default: Only + compress the obs and new_obs columns. + + Returns: + This very (now compressed) SampleBatch. + """ + + def _compress_in_place(path, value): + if path[0] not in columns: + return + curr = self + for i, p in enumerate(path): + if i == len(path) - 1: + if bulk: + curr[p] = pack(value) + else: + curr[p] = np.array([pack(o) for o in value]) + curr = curr[p] + + tree.map_structure_with_path(_compress_in_place, self) + + return self + + @DeveloperAPI + def decompress_if_needed( + self, columns: Set[str] = frozenset(["obs", "new_obs"]) + ) -> "SampleBatch": + """Decompresses data buffers (per column if not compressed) in place. + + Args: + columns: The columns to decompress. Default: Only + decompress the obs and new_obs columns. + + Returns: + This very (now uncompressed) SampleBatch. + """ + + def _decompress_in_place(path, value): + if path[0] not in columns: + return + curr = self + for p in path[:-1]: + curr = curr[p] + # Bulk compressed. + if is_compressed(value): + curr[path[-1]] = unpack(value) + # Non bulk compressed. + elif len(value) > 0 and is_compressed(value[0]): + curr[path[-1]] = np.array([unpack(o) for o in value]) + + tree.map_structure_with_path(_decompress_in_place, self) + + return self + + @DeveloperAPI + def set_get_interceptor(self, fn): + """Sets a function to be called on every getitem.""" + # If get-interceptor changes, must erase old intercepted values. + if fn is not self.get_interceptor: + self.intercepted_values = {} + self.get_interceptor = fn + + def __repr__(self): + keys = list(self.keys()) + if self.get(SampleBatch.SEQ_LENS) is None: + return f"SampleBatch({self.count}: {keys})" + else: + keys.remove(SampleBatch.SEQ_LENS) + return ( + f"SampleBatch({self.count} " f"(seqs={len(self['seq_lens'])}): {keys})" + ) + + def _slice(self, slice_: slice) -> "SampleBatch": + """Helper method to handle SampleBatch slicing using a slice object. + + The returned SampleBatch uses the same underlying data object as + `self`, so changing the slice will also change `self`. + + Note that only zero or positive bounds are allowed for both start + and stop values. The slice step must be 1 (or None, which is the + same). + + Args: + slice_: The python slice object to slice by. + + Returns: + A new SampleBatch, however "linking" into the same data + (sliced) as self. + """ + if self._slice_seq_lens_in_B: + return self._batch_slice(slice_) + + start = slice_.start or 0 + stop = slice_.stop or len(self) + # If stop goes beyond the length of this batch -> Make it go till the + # end only (including last item). + # Analogous to `l = [0, 1, 2]; l[:100] -> [0, 1, 2];`. + if stop > len(self): + stop = len(self) + + if ( + self.get(SampleBatch.SEQ_LENS) is not None + and len(self[SampleBatch.SEQ_LENS]) > 0 + ): + # Build our slice-map, if not done already. + if not self._slice_map: + sum_ = 0 + for i, l in enumerate(map(int, self[SampleBatch.SEQ_LENS])): + self._slice_map.extend([(i, sum_)] * l) + sum_ = sum_ + l + # In case `stop` points to the very end (lengths of this + # batch), return the last sequence (the -1 here makes sure we + # never go beyond it; would result in an index error below). + self._slice_map.append((len(self[SampleBatch.SEQ_LENS]), sum_)) + + start_seq_len, start_unpadded = self._slice_map[start] + stop_seq_len, stop_unpadded = self._slice_map[stop] + start_padded = start_unpadded + stop_padded = stop_unpadded + if self.zero_padded: + start_padded = start_seq_len * self.max_seq_len + stop_padded = stop_seq_len * self.max_seq_len + + def map_(path, value): + if path[0] != SampleBatch.SEQ_LENS and not path[0].startswith( + "state_in_" + ): + return value[start_padded:stop_padded] + else: + return value[start_seq_len:stop_seq_len] + + infos = self.pop(SampleBatch.INFOS, None) + data = tree.map_structure_with_path(map_, self) + if infos is not None and isinstance(infos, (list, np.ndarray)): + self[SampleBatch.INFOS] = infos + data[SampleBatch.INFOS] = infos[start_unpadded:stop_unpadded] + + return SampleBatch( + data, + _is_training=self.is_training, + _time_major=self.time_major, + _zero_padded=self.zero_padded, + _max_seq_len=self.max_seq_len if self.zero_padded else None, + _num_grad_updates=self.num_grad_updates, + ) + else: + infos = self.pop(SampleBatch.INFOS, None) + data = tree.map_structure(lambda s: s[start:stop], self) + if infos is not None and isinstance(infos, (list, np.ndarray)): + self[SampleBatch.INFOS] = infos + data[SampleBatch.INFOS] = infos[start:stop] + + return SampleBatch( + data, + _is_training=self.is_training, + _time_major=self.time_major, + _num_grad_updates=self.num_grad_updates, + ) + + @Deprecated(error=False) + def _get_slice_indices(self, slice_size): + data_slices = [] + data_slices_states = [] + if ( + self.get(SampleBatch.SEQ_LENS) is not None + and len(self[SampleBatch.SEQ_LENS]) > 0 + ): + assert np.all(self[SampleBatch.SEQ_LENS] < slice_size), ( + "ERROR: `slice_size` must be larger than the max. seq-len " + "in the batch!" + ) + start_pos = 0 + current_slize_size = 0 + actual_slice_idx = 0 + start_idx = 0 + idx = 0 + while idx < len(self[SampleBatch.SEQ_LENS]): + seq_len = self[SampleBatch.SEQ_LENS][idx] + current_slize_size += seq_len + actual_slice_idx += ( + seq_len if not self.zero_padded else self.max_seq_len + ) + # Complete minibatch -> Append to data_slices. + if current_slize_size >= slice_size: + end_idx = idx + 1 + # We are not zero-padded yet; all sequences are + # back-to-back. + if not self.zero_padded: + data_slices.append((start_pos, start_pos + slice_size)) + start_pos += slice_size + if current_slize_size > slice_size: + overhead = current_slize_size - slice_size + start_pos -= seq_len - overhead + idx -= 1 + # We are already zero-padded: Cut in chunks of max_seq_len. + else: + data_slices.append((start_pos, actual_slice_idx)) + start_pos = actual_slice_idx + + data_slices_states.append((start_idx, end_idx)) + current_slize_size = 0 + start_idx = idx + 1 + idx += 1 + else: + i = 0 + while i < self.count: + data_slices.append((i, i + slice_size)) + i += slice_size + return data_slices, data_slices_states + + @ExperimentalAPI + def get_single_step_input_dict( + self, + view_requirements: ViewRequirementsDict, + index: Union[str, int] = "last", + ) -> "SampleBatch": + """Creates single ts SampleBatch at given index from `self`. + + For usage as input-dict for model (action or value function) calls. + + Args: + view_requirements: A view requirements dict from the model for + which to produce the input_dict. + index: An integer index value indicating the + position in the trajectory for which to generate the + compute_actions input dict. Set to "last" to generate the dict + at the very end of the trajectory (e.g. for value estimation). + Note that "last" is different from -1, as "last" will use the + final NEXT_OBS as observation input. + + Returns: + The (single-timestep) input dict for ModelV2 calls. + """ + last_mappings = { + SampleBatch.OBS: SampleBatch.NEXT_OBS, + SampleBatch.PREV_ACTIONS: SampleBatch.ACTIONS, + SampleBatch.PREV_REWARDS: SampleBatch.REWARDS, + } + + input_dict = {} + for view_col, view_req in view_requirements.items(): + if view_req.used_for_compute_actions is False: + continue + + # Create batches of size 1 (single-agent input-dict). + data_col = view_req.data_col or view_col + if index == "last": + data_col = last_mappings.get(data_col, data_col) + # Range needed. + if view_req.shift_from is not None: + # Batch repeat value > 1: We have single frames in the + # batch at each timestep (for the `data_col`). + data = self[view_col][-1] + traj_len = len(self[data_col]) + missing_at_end = traj_len % view_req.batch_repeat_value + # Index into the observations column must be shifted by + # -1 b/c index=0 for observations means the current (last + # seen) observation (after having taken an action). + obs_shift = ( + -1 if data_col in [SampleBatch.OBS, SampleBatch.NEXT_OBS] else 0 + ) + from_ = view_req.shift_from + obs_shift + to_ = view_req.shift_to + obs_shift + 1 + if to_ == 0: + to_ = None + input_dict[view_col] = np.array( + [ + np.concatenate([data, self[data_col][-missing_at_end:]])[ + from_:to_ + ] + ] + ) + # Single index. + else: + input_dict[view_col] = tree.map_structure( + lambda v: v[-1:], # keep as array (w/ 1 element) + self[data_col], + ) + # Single index somewhere inside the trajectory (non-last). + else: + input_dict[view_col] = self[data_col][ + index : index + 1 if index != -1 else None + ] + + return SampleBatch(input_dict, seq_lens=np.array([1], dtype=np.int32)) + + +@PublicAPI +class MultiAgentBatch: + """A batch of experiences from multiple agents in the environment. + + Attributes: + policy_batches (Dict[PolicyID, SampleBatch]): Dict mapping policy IDs to + SampleBatches of experiences. + count: The number of env steps in this batch. + """ + + @PublicAPI + def __init__(self, policy_batches: Dict[PolicyID, SampleBatch], env_steps: int): + """Initialize a MultiAgentBatch instance. + + Args: + policy_batches: Dict mapping policy IDs to SampleBatches of experiences. + env_steps: The number of environment steps in the environment + this batch contains. This will be less than the number of + transitions this batch contains across all policies in total. + """ + + for v in policy_batches.values(): + assert isinstance(v, SampleBatch) + self.policy_batches = policy_batches + # Called "count" for uniformity with SampleBatch. + # Prefer to access this via the `env_steps()` method when possible + # for clarity. + self.count = env_steps + + @PublicAPI + def env_steps(self) -> int: + """The number of env steps (there are >= 1 agent steps per env step). + + Returns: + The number of environment steps contained in this batch. + """ + return self.count + + @PublicAPI + def __len__(self) -> int: + """Same as `self.env_steps()`.""" + return self.count + + @PublicAPI + def agent_steps(self) -> int: + """The number of agent steps (there are >= 1 agent steps per env step). + + Returns: + The number of agent steps total in this batch. + """ + ct = 0 + for batch in self.policy_batches.values(): + ct += batch.count + return ct + + @PublicAPI + def timeslices(self, k: int) -> List["MultiAgentBatch"]: + """Returns k-step batches holding data for each agent at those steps. + + For examples, suppose we have agent1 observations [a1t1, a1t2, a1t3], + for agent2, [a2t1, a2t3], and for agent3, [a3t3] only. + + Calling timeslices(1) would return three MultiAgentBatches containing + [a1t1, a2t1], [a1t2], and [a1t3, a2t3, a3t3]. + + Calling timeslices(2) would return two MultiAgentBatches containing + [a1t1, a1t2, a2t1], and [a1t3, a2t3, a3t3]. + + This method is used to implement "lockstep" replay mode. Note that this + method does not guarantee each batch contains only data from a single + unroll. Batches might contain data from multiple different envs. + """ + from ray.rllib.evaluation.sample_batch_builder import SampleBatchBuilder + + # Build a sorted set of (eps_id, t, policy_id, data...) + steps = [] + for policy_id, batch in self.policy_batches.items(): + for row in batch.rows(): + steps.append( + ( + row[SampleBatch.EPS_ID], + row[SampleBatch.T], + row[SampleBatch.AGENT_INDEX], + policy_id, + row, + ) + ) + steps.sort() + + finished_slices = [] + cur_slice = collections.defaultdict(SampleBatchBuilder) + cur_slice_size = 0 + + def finish_slice(): + nonlocal cur_slice_size + assert cur_slice_size > 0 + batch = MultiAgentBatch( + {k: v.build_and_reset() for k, v in cur_slice.items()}, cur_slice_size + ) + cur_slice_size = 0 + cur_slice.clear() + finished_slices.append(batch) + + # For each unique env timestep. + for _, group in itertools.groupby(steps, lambda x: x[:2]): + # Accumulate into the current slice. + for _, _, _, policy_id, row in group: + cur_slice[policy_id].add_values(**row) + cur_slice_size += 1 + # Slice has reached target number of env steps. + if cur_slice_size >= k: + finish_slice() + assert cur_slice_size == 0 + + if cur_slice_size > 0: + finish_slice() + + assert len(finished_slices) > 0, finished_slices + return finished_slices + + @staticmethod + @PublicAPI + def wrap_as_needed( + policy_batches: Dict[PolicyID, SampleBatch], env_steps: int + ) -> Union[SampleBatch, "MultiAgentBatch"]: + """Returns SampleBatch or MultiAgentBatch, depending on given policies. + If policy_batches is empty (i.e. {}) it returns an empty MultiAgentBatch. + + Args: + policy_batches: Mapping from policy ids to SampleBatch. + env_steps: Number of env steps in the batch. + + Returns: + The single default policy's SampleBatch or a MultiAgentBatch + (more than one policy). + """ + if len(policy_batches) == 1 and DEFAULT_POLICY_ID in policy_batches: + return policy_batches[DEFAULT_POLICY_ID] + return MultiAgentBatch(policy_batches=policy_batches, env_steps=env_steps) + + @staticmethod + @PublicAPI + @Deprecated(new="concat_samples() from rllib.policy.sample_batch", error=True) + def concat_samples(samples: List["MultiAgentBatch"]) -> "MultiAgentBatch": + return concat_samples_into_ma_batch(samples) + + @PublicAPI + def copy(self) -> "MultiAgentBatch": + """Deep-copies self into a new MultiAgentBatch. + + Returns: + The copy of self with deep-copied data. + """ + return MultiAgentBatch( + {k: v.copy() for (k, v) in self.policy_batches.items()}, self.count + ) + + @ExperimentalAPI + def to_device( + self, + device, + framework="torch", + pin_memory: bool = False, + use_stream: bool = False, + stream: Optional[Union["torch.cuda.Stream", "torch.cuda.Stream"]] = None, + ): + """TODO: transfer batch to given device as framework tensor.""" + if framework == "torch": + assert torch is not None + for pid, policy_batch in self.policy_batches.items(): + self.policy_batches[pid] = policy_batch.to_device( + device, + framework=framework, + pin_memory=pin_memory, + use_stream=use_stream, + stream=stream, + ) + else: + raise NotImplementedError + return self + + @PublicAPI + def size_bytes(self) -> int: + """ + Returns: + The overall size in bytes of all policy batches (all columns). + """ + return sum(b.size_bytes() for b in self.policy_batches.values()) + + @DeveloperAPI + def compress( + self, bulk: bool = False, columns: Set[str] = frozenset(["obs", "new_obs"]) + ) -> None: + """Compresses each policy batch (per column) in place. + + Args: + bulk: Whether to compress across the batch dimension (0) + as well. If False will compress n separate list items, where n + is the batch size. + columns: Set of column names to compress. + """ + for batch in self.policy_batches.values(): + batch.compress(bulk=bulk, columns=columns) + + @DeveloperAPI + def decompress_if_needed( + self, columns: Set[str] = frozenset(["obs", "new_obs"]) + ) -> "MultiAgentBatch": + """Decompresses each policy batch (per column), if already compressed. + + Args: + columns: Set of column names to decompress. + + Returns: + Self. + """ + for batch in self.policy_batches.values(): + batch.decompress_if_needed(columns) + return self + + @DeveloperAPI + def as_multi_agent(self) -> "MultiAgentBatch": + """Simply returns `self` (already a MultiAgentBatch). + + Returns: + This very instance of MultiAgentBatch. + """ + return self + + def __getitem__(self, key: str) -> SampleBatch: + """Returns the SampleBatch for the given policy id.""" + return self.policy_batches[key] + + def __str__(self): + return "MultiAgentBatch({}, env_steps={})".format( + str(self.policy_batches), self.count + ) + + def __repr__(self): + return "MultiAgentBatch({}, env_steps={})".format( + str(self.policy_batches), self.count + ) + + +@PublicAPI +def concat_samples(samples: List[SampleBatchType]) -> SampleBatchType: + """Concatenates a list of SampleBatches or MultiAgentBatches. + + If all items in the list are or SampleBatch typ4, the output will be + a SampleBatch type. Otherwise, the output will be a MultiAgentBatch type. + If input is a mixture of SampleBatch and MultiAgentBatch types, it will treat + SampleBatch objects as MultiAgentBatch types with 'default_policy' key and + concatenate it with th rest of MultiAgentBatch objects. + Empty samples are simply ignored. + + Args: + samples: List of SampleBatches or MultiAgentBatches to be + concatenated. + + Returns: + A new (concatenated) SampleBatch or MultiAgentBatch. + + .. testcode:: + :skipif: True + + import numpy as np + from ray.rllib.policy.sample_batch import SampleBatch + b1 = SampleBatch({"a": np.array([1, 2]), + "b": np.array([10, 11])}) + b2 = SampleBatch({"a": np.array([3]), + "b": np.array([12])}) + print(concat_samples([b1, b2])) + + + c1 = MultiAgentBatch({'default_policy': { + "a": np.array([1, 2]), + "b": np.array([10, 11]) + }}, env_steps=2) + c2 = SampleBatch({"a": np.array([3]), + "b": np.array([12])}) + print(concat_samples([b1, b2])) + + .. testoutput:: + + {"a": np.array([1, 2, 3]), "b": np.array([10, 11, 12])} + MultiAgentBatch = {'default_policy': {"a": np.array([1, 2, 3]), + "b": np.array([10, 11, 12])}} + + """ + + if any(isinstance(s, MultiAgentBatch) for s in samples): + return concat_samples_into_ma_batch(samples) + + # the output is a SampleBatch type + concatd_seq_lens = [] + concatd_num_grad_updates = [0, 0.0] # [0]=count; [1]=weighted sum values + concated_samples = [] + # Make sure these settings are consistent amongst all batches. + zero_padded = max_seq_len = time_major = None + for s in samples: + if s.count <= 0: + continue + + if max_seq_len is None: + zero_padded = s.zero_padded + max_seq_len = s.max_seq_len + time_major = s.time_major + + # Make sure these settings are consistent amongst all batches. + if s.zero_padded != zero_padded or s.time_major != time_major: + raise ValueError( + "All SampleBatches' `zero_padded` and `time_major` settings " + "must be consistent!" + ) + if ( + s.max_seq_len is None or max_seq_len is None + ) and s.max_seq_len != max_seq_len: + raise ValueError( + "Samples must consistently either provide or omit `max_seq_len`!" + ) + elif zero_padded and s.max_seq_len != max_seq_len: + raise ValueError( + "For `zero_padded` SampleBatches, the values of `max_seq_len` " + "must be consistent!" + ) + + if max_seq_len is not None: + max_seq_len = max(max_seq_len, s.max_seq_len) + if s.get(SampleBatch.SEQ_LENS) is not None: + concatd_seq_lens.extend(s[SampleBatch.SEQ_LENS]) + if s.num_grad_updates is not None: + concatd_num_grad_updates[0] += s.count + concatd_num_grad_updates[1] += s.num_grad_updates * s.count + + concated_samples.append(s) + + # If we don't have any samples (0 or only empty SampleBatches), + # return an empty SampleBatch here. + if len(concated_samples) == 0: + return SampleBatch() + + # Collect the concat'd data. + concatd_data = {} + + for k in concated_samples[0].keys(): + if k == SampleBatch.INFOS: + concatd_data[k] = _concat_values( + *[s[k] for s in concated_samples], + time_major=time_major, + ) + else: + values_to_concat = [c[k] for c in concated_samples] + _concat_values_w_time = partial(_concat_values, time_major=time_major) + concatd_data[k] = tree.map_structure( + _concat_values_w_time, *values_to_concat + ) + + if concatd_seq_lens != [] and torch and torch.is_tensor(concatd_seq_lens[0]): + concatd_seq_lens = torch.Tensor(concatd_seq_lens) + elif concatd_seq_lens != [] and tf and tf.is_tensor(concatd_seq_lens[0]): + concatd_seq_lens = tf.convert_to_tensor(concatd_seq_lens) + + # Return a new (concat'd) SampleBatch. + return SampleBatch( + concatd_data, + seq_lens=concatd_seq_lens, + _time_major=time_major, + _zero_padded=zero_padded, + _max_seq_len=max_seq_len, + # Compute weighted average of the num_grad_updates for the batches + # (assuming they all come from the same policy). + _num_grad_updates=( + concatd_num_grad_updates[1] / (concatd_num_grad_updates[0] or 1.0) + ), + ) + + +@PublicAPI +def concat_samples_into_ma_batch(samples: List[SampleBatchType]) -> "MultiAgentBatch": + """Concatenates a list of SampleBatchTypes to a single MultiAgentBatch type. + + This function, as opposed to concat_samples() forces the output to always be + MultiAgentBatch which is more generic than SampleBatch. + + Args: + samples: List of SampleBatches or MultiAgentBatches to be + concatenated. + + Returns: + A new (concatenated) MultiAgentBatch. + + .. testcode:: + :skipif: True + + import numpy as np + from ray.rllib.policy.sample_batch import SampleBatch + b1 = MultiAgentBatch({'default_policy': { + "a": np.array([1, 2]), + "b": np.array([10, 11]) + }}, env_steps=2) + b2 = SampleBatch({"a": np.array([3]), + "b": np.array([12])}) + print(concat_samples([b1, b2])) + + .. testoutput:: + + {'default_policy': {"a": np.array([1, 2, 3]), + "b": np.array([10, 11, 12])}} + + """ + + policy_batches = collections.defaultdict(list) + env_steps = 0 + for s in samples: + # Some batches in `samples` may be SampleBatch. + if isinstance(s, SampleBatch): + # If empty SampleBatch: ok (just ignore). + if len(s) <= 0: + continue + else: + # if non-empty: just convert to MA-batch and move forward + s = s.as_multi_agent() + elif not isinstance(s, MultiAgentBatch): + # Otherwise: Error. + raise ValueError( + "`concat_samples_into_ma_batch` can only concat " + "SampleBatch|MultiAgentBatch objects, not {}!".format(type(s).__name__) + ) + + for key, batch in s.policy_batches.items(): + policy_batches[key].append(batch) + env_steps += s.env_steps() + + out = {} + for key, batches in policy_batches.items(): + out[key] = concat_samples(batches) + + return MultiAgentBatch(out, env_steps) + + +def _concat_values(*values, time_major=None) -> TensorType: + """Concatenates a list of values. + + Args: + values: The values to concatenate. + time_major: Whether to concatenate along the first axis + (time_major=False) or the second axis (time_major=True). + """ + if torch and torch.is_tensor(values[0]): + return torch.cat(values, dim=1 if time_major else 0) + elif isinstance(values[0], np.ndarray): + return np.concatenate(values, axis=1 if time_major else 0) + elif tf and tf.is_tensor(values[0]): + return tf.concat(values, axis=1 if time_major else 0) + elif isinstance(values[0], list): + concatenated_list = [] + for sublist in values: + concatenated_list.extend(sublist) + return concatenated_list + else: + raise ValueError( + f"Unsupported type for concatenation: {type(values[0])} " + f"first element: {values[0]}" + ) + + +@DeveloperAPI +def convert_ma_batch_to_sample_batch(batch: SampleBatchType) -> SampleBatch: + """Converts a MultiAgentBatch to a SampleBatch if neccessary. + + Args: + batch: The SampleBatchType to convert. + + Returns: + batch: the converted SampleBatch + + Raises: + ValueError if the MultiAgentBatch has more than one policy_id + or if the policy_id is not `DEFAULT_POLICY_ID` + """ + if isinstance(batch, MultiAgentBatch): + policy_keys = batch.policy_batches.keys() + if len(policy_keys) == 1 and DEFAULT_POLICY_ID in policy_keys: + batch = batch.policy_batches[DEFAULT_POLICY_ID] + else: + raise ValueError( + "RLlib tried to convert a multi agent-batch with data from more " + "than one policy to a single-agent batch. This is not supported and " + "may be due to a number of issues. Here are two possible ones:" + "1) Off-Policy Estimation is not implemented for " + "multi-agent batches. You can set `off_policy_estimation_methods: {}` " + "to resolve this." + "2) Loading multi-agent data for offline training is not implemented." + "Load single-agent data instead to resolve this." + ) + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/policy/tf_mixins.py b/lib/python3.12/site-packages/ray/rllib/policy/tf_mixins.py new file mode 100644 index 0000000000000000000000000000000000000000..fa5d0862172945f1e11ce049f2534adf4ff4caa5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/tf_mixins.py @@ -0,0 +1,387 @@ +import logging +from typing import Dict, List + +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.eager_tf_policy import EagerTFPolicy +from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 +from ray.rllib.policy.policy import PolicyState +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import get_variable, try_import_tf +from ray.rllib.utils.schedules import PiecewiseSchedule +from ray.rllib.utils.tf_utils import make_tf_callable +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) + +logger = logging.getLogger(__name__) +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class LearningRateSchedule: + """Mixin for TFPolicy that adds a learning rate schedule.""" + + def __init__(self, lr, lr_schedule): + self._lr_schedule = None + if lr_schedule is None: + self.cur_lr = tf1.get_variable("lr", initializer=lr, trainable=False) + else: + self._lr_schedule = PiecewiseSchedule( + lr_schedule, outside_value=lr_schedule[-1][-1], framework=None + ) + self.cur_lr = tf1.get_variable( + "lr", initializer=self._lr_schedule.value(0), trainable=False + ) + if self.framework == "tf": + self._lr_placeholder = tf1.placeholder(dtype=tf.float32, name="lr") + self._lr_update = self.cur_lr.assign( + self._lr_placeholder, read_value=False + ) + + def on_global_var_update(self, global_vars): + super().on_global_var_update(global_vars) + if self._lr_schedule is not None: + new_val = self._lr_schedule.value(global_vars["timestep"]) + if self.framework == "tf": + self.get_session().run( + self._lr_update, feed_dict={self._lr_placeholder: new_val} + ) + else: + self.cur_lr.assign(new_val, read_value=False) + # This property (self._optimizer) is (still) accessible for + # both TFPolicy and any TFPolicy_eager. + self._optimizer.learning_rate.assign(self.cur_lr) + + def optimizer(self): + if self.framework == "tf": + return tf1.train.AdamOptimizer(learning_rate=self.cur_lr) + else: + return tf.keras.optimizers.Adam(self.cur_lr) + + +@OldAPIStack +class EntropyCoeffSchedule: + """Mixin for TFPolicy that adds entropy coeff decay.""" + + def __init__(self, entropy_coeff, entropy_coeff_schedule): + self._entropy_coeff_schedule = None + if entropy_coeff_schedule is None: + self.entropy_coeff = get_variable( + entropy_coeff, framework="tf", tf_name="entropy_coeff", trainable=False + ) + else: + # Allows for custom schedule similar to lr_schedule format + if isinstance(entropy_coeff_schedule, list): + self._entropy_coeff_schedule = PiecewiseSchedule( + entropy_coeff_schedule, + outside_value=entropy_coeff_schedule[-1][-1], + framework=None, + ) + else: + # Implements previous version but enforces outside_value + self._entropy_coeff_schedule = PiecewiseSchedule( + [[0, entropy_coeff], [entropy_coeff_schedule, 0.0]], + outside_value=0.0, + framework=None, + ) + + self.entropy_coeff = get_variable( + self._entropy_coeff_schedule.value(0), + framework="tf", + tf_name="entropy_coeff", + trainable=False, + ) + if self.framework == "tf": + self._entropy_coeff_placeholder = tf1.placeholder( + dtype=tf.float32, name="entropy_coeff" + ) + self._entropy_coeff_update = self.entropy_coeff.assign( + self._entropy_coeff_placeholder, read_value=False + ) + + def on_global_var_update(self, global_vars): + super().on_global_var_update(global_vars) + if self._entropy_coeff_schedule is not None: + new_val = self._entropy_coeff_schedule.value(global_vars["timestep"]) + if self.framework == "tf": + self.get_session().run( + self._entropy_coeff_update, + feed_dict={self._entropy_coeff_placeholder: new_val}, + ) + else: + self.entropy_coeff.assign(new_val, read_value=False) + + +@OldAPIStack +class KLCoeffMixin: + """Assigns the `update_kl()` and other KL-related methods to a TFPolicy. + + This is used in Algorithms to update the KL coefficient after each + learning step based on `config.kl_target` and the measured KL value + (from the train_batch). + """ + + def __init__(self, config: AlgorithmConfigDict): + # The current KL value (as python float). + self.kl_coeff_val = config["kl_coeff"] + # The current KL value (as tf Variable for in-graph operations). + self.kl_coeff = get_variable( + float(self.kl_coeff_val), + tf_name="kl_coeff", + trainable=False, + framework=config["framework"], + ) + # Constant target value. + self.kl_target = config["kl_target"] + if self.framework == "tf": + self._kl_coeff_placeholder = tf1.placeholder( + dtype=tf.float32, name="kl_coeff" + ) + self._kl_coeff_update = self.kl_coeff.assign( + self._kl_coeff_placeholder, read_value=False + ) + + def update_kl(self, sampled_kl): + # Update the current KL value based on the recently measured value. + # Increase. + if sampled_kl > 2.0 * self.kl_target: + self.kl_coeff_val *= 1.5 + # Decrease. + elif sampled_kl < 0.5 * self.kl_target: + self.kl_coeff_val *= 0.5 + # No change. + else: + return self.kl_coeff_val + + # Make sure, new value is also stored in graph/tf variable. + self._set_kl_coeff(self.kl_coeff_val) + + # Return the current KL value. + return self.kl_coeff_val + + def _set_kl_coeff(self, new_kl_coeff): + # Set the (off graph) value. + self.kl_coeff_val = new_kl_coeff + + # Update the tf/tf2 Variable (via session call for tf or `assign`). + if self.framework == "tf": + self.get_session().run( + self._kl_coeff_update, + feed_dict={self._kl_coeff_placeholder: self.kl_coeff_val}, + ) + else: + self.kl_coeff.assign(self.kl_coeff_val, read_value=False) + + def get_state(self) -> PolicyState: + state = super().get_state() + # Add current kl-coeff value. + state["current_kl_coeff"] = self.kl_coeff_val + return state + + def set_state(self, state: PolicyState) -> None: + # Set current kl-coeff value first. + self._set_kl_coeff(state.pop("current_kl_coeff", self.config["kl_coeff"])) + # Call super's set_state with rest of the state dict. + super().set_state(state) + + +@OldAPIStack +class TargetNetworkMixin: + """Assign the `update_target` method to the policy. + + The function is called every `target_network_update_freq` steps by the + master learner. + """ + + def __init__(self): + model_vars = self.model.trainable_variables() + target_model_vars = self.target_model.trainable_variables() + + @make_tf_callable(self.get_session()) + def update_target_fn(tau): + tau = tf.convert_to_tensor(tau, dtype=tf.float32) + update_target_expr = [] + assert len(model_vars) == len(target_model_vars), ( + model_vars, + target_model_vars, + ) + for var, var_target in zip(model_vars, target_model_vars): + update_target_expr.append( + var_target.assign(tau * var + (1.0 - tau) * var_target) + ) + logger.debug("Update target op {}".format(var_target)) + return tf.group(*update_target_expr) + + # Hard initial update. + self._do_update = update_target_fn + # TODO: The previous SAC implementation does an update(1.0) here. + # If this is changed to tau != 1.0 the sac_loss_function test fails. Why? + # Also the test is not very maintainable, we need to change that unittest + # anyway. + self.update_target(tau=1.0) # self.config.get("tau", 1.0)) + + @property + def q_func_vars(self): + if not hasattr(self, "_q_func_vars"): + self._q_func_vars = self.model.variables() + return self._q_func_vars + + @property + def target_q_func_vars(self): + if not hasattr(self, "_target_q_func_vars"): + self._target_q_func_vars = self.target_model.variables() + return self._target_q_func_vars + + # Support both hard and soft sync. + def update_target(self, tau: int = None) -> None: + self._do_update(np.float32(tau or self.config.get("tau", 1.0))) + + def variables(self) -> List[TensorType]: + return self.model.variables() + + def set_weights(self, weights): + if isinstance(self, TFPolicy): + TFPolicy.set_weights(self, weights) + elif isinstance(self, EagerTFPolicyV2): # Handle TF2V2 policies. + EagerTFPolicyV2.set_weights(self, weights) + elif isinstance(self, EagerTFPolicy): # Handle TF2 policies. + EagerTFPolicy.set_weights(self, weights) + self.update_target(self.config.get("tau", 1.0)) + + +@OldAPIStack +class ValueNetworkMixin: + """Assigns the `_value()` method to a TFPolicy. + + This way, Policy can call `_value()` to get the current VF estimate on a + single(!) observation (as done in `postprocess_trajectory_fn`). + Note: When doing this, an actual forward pass is being performed. + This is different from only calling `model.value_function()`, where + the result of the most recent forward pass is being used to return an + already calculated tensor. + """ + + def __init__(self, config): + # When doing GAE or vtrace, we need the value function estimate on the + # observation. + if config.get("use_gae") or config.get("vtrace"): + # Input dict is provided to us automatically via the Model's + # requirements. It's a single-timestep (last one in trajectory) + # input_dict. + @make_tf_callable(self.get_session()) + def value(**input_dict): + input_dict = SampleBatch(input_dict) + if isinstance(self.model, tf.keras.Model): + _, _, extra_outs = self.model(input_dict) + return extra_outs[SampleBatch.VF_PREDS][0] + else: + model_out, _ = self.model(input_dict) + # [0] = remove the batch dim. + return self.model.value_function()[0] + + # When not doing GAE, we do not require the value function's output. + else: + + @make_tf_callable(self.get_session()) + def value(*args, **kwargs): + return tf.constant(0.0) + + self._value = value + self._should_cache_extra_action = config["framework"] == "tf" + self._cached_extra_action_fetches = None + + def _extra_action_out_impl(self) -> Dict[str, TensorType]: + extra_action_out = super().extra_action_out_fn() + # Keras models return values for each call in third return argument + # (dict). + if isinstance(self.model, tf.keras.Model): + return extra_action_out + # Return value function outputs. VF estimates will hence be added to the + # SampleBatches produced by the sampler(s) to generate the train batches + # going into the loss function. + extra_action_out.update( + { + SampleBatch.VF_PREDS: self.model.value_function(), + } + ) + return extra_action_out + + def extra_action_out_fn(self) -> Dict[str, TensorType]: + if not self._should_cache_extra_action: + return self._extra_action_out_impl() + + # Note: there are 2 reasons we are caching the extra_action_fetches for + # TF1 static graph here. + # 1. for better performance, so we don't query base class and model for + # extra fetches every single time. + # 2. for correctness. TF1 is special because the static graph may contain + # two logical graphs. One created by DynamicTFPolicy for action + # computation, and one created by MultiGPUTower for GPU training. + # Depending on which logical graph ran last time, + # self.model.value_function() will point to the output tensor + # of the specific logical graph, causing problem if we try to + # fetch action (run inference) using the training output tensor. + # For that reason, we cache the action output tensor from the + # vanilla DynamicTFPolicy once and call it a day. + if self._cached_extra_action_fetches is not None: + return self._cached_extra_action_fetches + + self._cached_extra_action_fetches = self._extra_action_out_impl() + return self._cached_extra_action_fetches + + +@OldAPIStack +class GradStatsMixin: + def __init__(self): + pass + + def grad_stats_fn( + self, train_batch: SampleBatch, grads: ModelGradients + ) -> Dict[str, TensorType]: + # We have support for more than one loss (list of lists of grads). + if self.config.get("_tf_policy_handles_more_than_one_loss"): + grad_gnorm = [tf.linalg.global_norm(g) for g in grads] + # Old case: We have a single list of grads (only one loss term and + # optimizer). + else: + grad_gnorm = tf.linalg.global_norm(grads) + + return { + "grad_gnorm": grad_gnorm, + } + + +def compute_gradients( + policy, optimizer: LocalOptimizer, loss: TensorType +) -> ModelGradients: + # Compute the gradients. + variables = policy.model.trainable_variables + if isinstance(policy.model, ModelV2): + variables = variables() + grads_and_vars = optimizer.compute_gradients(loss, variables) + + # Clip by global norm, if necessary. + if policy.config.get("grad_clip") is not None: + # Defuse inf gradients (due to super large losses). + grads = [g for (g, v) in grads_and_vars] + grads, _ = tf.clip_by_global_norm(grads, policy.config["grad_clip"]) + # If the global_norm is inf -> All grads will be NaN. Stabilize this + # here by setting them to 0.0. This will simply ignore destructive loss + # calculations. + policy.grads = [] + for g in grads: + if g is not None: + policy.grads.append(tf.where(tf.math.is_nan(g), tf.zeros_like(g), g)) + else: + policy.grads.append(None) + clipped_grads_and_vars = list(zip(policy.grads, variables)) + return clipped_grads_and_vars + else: + return grads_and_vars diff --git a/lib/python3.12/site-packages/ray/rllib/policy/tf_policy.py b/lib/python3.12/site-packages/ray/rllib/policy/tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..15712bf84faf4c2271dcbd2cf8fe1c67c39973dd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/tf_policy.py @@ -0,0 +1,1197 @@ +import logging +import math +from typing import Dict, List, Optional, Tuple, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.policy import Policy, PolicySpec, PolicyState +from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.error import ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_AGENT_STEPS_TRAINED, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.spaces.space_utils import normalize_action +from ray.rllib.utils.tf_run_builder import _TFRunBuilder +from ray.rllib.utils.tf_utils import TensorFlowVariables, get_gpu_devices +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +@OldAPIStack +class TFPolicy(Policy): + """An agent policy and loss implemented in TensorFlow. + + Do not sub-class this class directly (neither should you sub-class + DynamicTFPolicy), but rather use + rllib.policy.tf_policy_template.build_tf_policy + to generate your custom tf (graph-mode or eager) Policy classes. + + Extending this class enables RLlib to perform TensorFlow specific + optimizations on the policy, e.g., parallelization across gpus or + fusing multiple graphs together in the multi-agent setting. + + Input tensors are typically shaped like [BATCH_SIZE, ...]. + + .. testcode:: + :skipif: True + + from ray.rllib.policy import TFPolicy + class TFPolicySubclass(TFPolicy): + ... + + sess, obs_input, sampled_action, loss, loss_inputs = ... + policy = TFPolicySubclass( + sess, obs_input, sampled_action, loss, loss_inputs) + print(policy.compute_actions([1, 0, 2])) + print(policy.postprocess_trajectory(SampleBatch({...}))) + + .. testoutput:: + + (array([0, 1, 1]), [], {}) + SampleBatch({"action": ..., "advantages": ..., ...}) + + """ + + # In order to create tf_policies from checkpoints, this class needs to separate + # variables into their own scopes. Normally, we would do this in the model + # catalog, but since Policy.from_state() can be called anywhere, we need to + # keep track of it here to not break the from_state API. + tf_var_creation_scope_counter = 0 + + @staticmethod + def next_tf_var_scope_name(): + # Tracks multiple instances that are spawned from this policy via .from_state() + TFPolicy.tf_var_creation_scope_counter += 1 + return f"var_scope_{TFPolicy.tf_var_creation_scope_counter}" + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + sess: "tf1.Session", + obs_input: TensorType, + sampled_action: TensorType, + loss: Union[TensorType, List[TensorType]], + loss_inputs: List[Tuple[str, TensorType]], + model: Optional[ModelV2] = None, + sampled_action_logp: Optional[TensorType] = None, + action_input: Optional[TensorType] = None, + log_likelihood: Optional[TensorType] = None, + dist_inputs: Optional[TensorType] = None, + dist_class: Optional[type] = None, + state_inputs: Optional[List[TensorType]] = None, + state_outputs: Optional[List[TensorType]] = None, + prev_action_input: Optional[TensorType] = None, + prev_reward_input: Optional[TensorType] = None, + seq_lens: Optional[TensorType] = None, + max_seq_len: int = 20, + batch_divisibility_req: int = 1, + update_ops: List[TensorType] = None, + explore: Optional[TensorType] = None, + timestep: Optional[TensorType] = None, + ): + """Initializes a Policy object. + + Args: + observation_space: Observation space of the policy. + action_space: Action space of the policy. + config: Policy-specific configuration data. + sess: The TensorFlow session to use. + obs_input: Input placeholder for observations, of shape + [BATCH_SIZE, obs...]. + sampled_action: Tensor for sampling an action, of shape + [BATCH_SIZE, action...] + loss: Scalar policy loss output tensor or a list thereof + (in case there is more than one loss). + loss_inputs: A (name, placeholder) tuple for each loss input + argument. Each placeholder name must + correspond to a SampleBatch column key returned by + postprocess_trajectory(), and has shape [BATCH_SIZE, data...]. + These keys will be read from postprocessed sample batches and + fed into the specified placeholders during loss computation. + model: The optional ModelV2 to use for calculating actions and + losses. If not None, TFPolicy will provide functionality for + getting variables, calling the model's custom loss (if + provided), and importing weights into the model. + sampled_action_logp: log probability of the sampled action. + action_input: Input placeholder for actions for + logp/log-likelihood calculations. + log_likelihood: Tensor to calculate the log_likelihood (given + action_input and obs_input). + dist_class: An optional ActionDistribution class to use for + generating a dist object from distribution inputs. + dist_inputs: Tensor to calculate the distribution + inputs/parameters. + state_inputs: List of RNN state input Tensors. + state_outputs: List of RNN state output Tensors. + prev_action_input: placeholder for previous actions. + prev_reward_input: placeholder for previous rewards. + seq_lens: Placeholder for RNN sequence lengths, of shape + [NUM_SEQUENCES]. + Note that NUM_SEQUENCES << BATCH_SIZE. See + policy/rnn_sequencing.py for more information. + max_seq_len: Max sequence length for LSTM training. + batch_divisibility_req: pad all agent experiences batches to + multiples of this value. This only has an effect if not using + a LSTM model. + update_ops: override the batchnorm update ops + to run when applying gradients. Otherwise we run all update + ops found in the current variable scope. + explore: Placeholder for `explore` parameter into call to + Exploration.get_exploration_action. Explicitly set this to + False for not creating any Exploration component. + timestep: Placeholder for the global sampling timestep. + """ + self.framework = "tf" + super().__init__(observation_space, action_space, config) + + # Get devices to build the graph on. + num_gpus = self._get_num_gpus_for_policy() + gpu_ids = get_gpu_devices() + logger.info(f"Found {len(gpu_ids)} visible cuda devices.") + + # Place on one or more CPU(s) when either: + # - Fake GPU mode. + # - num_gpus=0 (either set by user or we are in local_mode=True). + # - no GPUs available. + if config["_fake_gpus"] or num_gpus == 0 or not gpu_ids: + self.devices = ["/cpu:0" for _ in range(int(math.ceil(num_gpus)) or 1)] + # Place on one or more actual GPU(s), when: + # - num_gpus > 0 (set by user) AND + # - local_mode=False AND + # - actual GPUs available AND + # - non-fake GPU mode. + else: + # We are a remote worker (WORKER_MODE=1): + # GPUs should be assigned to us by ray. + if ray._private.worker._mode() == ray._private.worker.WORKER_MODE: + gpu_ids = ray.get_gpu_ids() + + if len(gpu_ids) < num_gpus: + raise ValueError( + "TFPolicy was not able to find enough GPU IDs! Found " + f"{gpu_ids}, but num_gpus={num_gpus}." + ) + + self.devices = [f"/gpu:{i}" for i, _ in enumerate(gpu_ids) if i < num_gpus] + + # Disable env-info placeholder. + if SampleBatch.INFOS in self.view_requirements: + self.view_requirements[SampleBatch.INFOS].used_for_compute_actions = False + self.view_requirements[SampleBatch.INFOS].used_for_training = False + # Optionally add `infos` to the output dataset + if self.config["output_config"].get("store_infos", False): + self.view_requirements[SampleBatch.INFOS].used_for_training = True + + assert model is None or isinstance(model, (ModelV2, tf.keras.Model)), ( + "Model classes for TFPolicy other than `ModelV2|tf.keras.Model` " + "not allowed! You passed in {}.".format(model) + ) + self.model = model + # Auto-update model's inference view requirements, if recurrent. + if self.model is not None: + self._update_model_view_requirements_from_init_state() + + # If `explore` is explicitly set to False, don't create an exploration + # component. + self.exploration = self._create_exploration() if explore is not False else None + + self._sess = sess + self._obs_input = obs_input + self._prev_action_input = prev_action_input + self._prev_reward_input = prev_reward_input + self._sampled_action = sampled_action + self._is_training = self._get_is_training_placeholder() + self._is_exploring = ( + explore + if explore is not None + else tf1.placeholder_with_default(True, (), name="is_exploring") + ) + self._sampled_action_logp = sampled_action_logp + self._sampled_action_prob = ( + tf.math.exp(self._sampled_action_logp) + if self._sampled_action_logp is not None + else None + ) + self._action_input = action_input # For logp calculations. + self._dist_inputs = dist_inputs + self.dist_class = dist_class + self._cached_extra_action_out = None + self._state_inputs = state_inputs or [] + self._state_outputs = state_outputs or [] + self._seq_lens = seq_lens + self._max_seq_len = max_seq_len + + if self._state_inputs and self._seq_lens is None: + raise ValueError( + "seq_lens tensor must be given if state inputs are defined" + ) + + self._batch_divisibility_req = batch_divisibility_req + self._update_ops = update_ops + self._apply_op = None + self._stats_fetches = {} + self._timestep = ( + timestep + if timestep is not None + else tf1.placeholder_with_default( + tf.zeros((), dtype=tf.int64), (), name="timestep" + ) + ) + + self._optimizers: List[LocalOptimizer] = [] + # Backward compatibility and for some code shared with tf-eager Policy. + self._optimizer = None + + self._grads_and_vars: Union[ModelGradients, List[ModelGradients]] = [] + self._grads: Union[ModelGradients, List[ModelGradients]] = [] + # Policy tf-variables (weights), whose values to get/set via + # get_weights/set_weights. + self._variables = None + # Local optimizer(s)' tf-variables (e.g. state vars for Adam). + # Will be stored alongside `self._variables` when checkpointing. + self._optimizer_variables: Optional[TensorFlowVariables] = None + + # The loss tf-op(s). Number of losses must match number of optimizers. + self._losses = [] + # Backward compatibility (in case custom child TFPolicies access this + # property). + self._loss = None + # A batch dict passed into loss function as input. + self._loss_input_dict = {} + losses = force_list(loss) + if len(losses) > 0: + self._initialize_loss(losses, loss_inputs) + + # The log-likelihood calculator op. + self._log_likelihood = log_likelihood + if ( + self._log_likelihood is None + and self._dist_inputs is not None + and self.dist_class is not None + ): + self._log_likelihood = self.dist_class(self._dist_inputs, self.model).logp( + self._action_input + ) + + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Union[SampleBatch, Dict[str, TensorType]], + explore: bool = None, + timestep: Optional[int] = None, + episode=None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep + + # Switch off is_training flag in our batch. + if isinstance(input_dict, SampleBatch): + input_dict.set_training(False) + else: + # Deprecated dict input. + input_dict["is_training"] = False + + builder = _TFRunBuilder(self.get_session(), "compute_actions_from_input_dict") + obs_batch = input_dict[SampleBatch.OBS] + to_fetch = self._build_compute_actions( + builder, input_dict=input_dict, explore=explore, timestep=timestep + ) + + # Execute session run to get action (and other fetches). + fetched = builder.get(to_fetch) + + # Update our global timestep by the batch size. + self.global_timestep += ( + len(obs_batch) + if isinstance(obs_batch, list) + else len(input_dict) + if isinstance(input_dict, SampleBatch) + else obs_batch.shape[0] + ) + + return fetched + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorType], TensorType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorType], TensorType] = None, + prev_reward_batch: Union[List[TensorType], TensorType] = None, + info_batch: Optional[Dict[str, list]] = None, + episodes=None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + **kwargs, + ): + explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep + + builder = _TFRunBuilder(self.get_session(), "compute_actions") + + input_dict = {SampleBatch.OBS: obs_batch, "is_training": False} + if state_batches: + for i, s in enumerate(state_batches): + input_dict[f"state_in_{i}"] = s + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + + to_fetch = self._build_compute_actions( + builder, input_dict=input_dict, explore=explore, timestep=timestep + ) + + # Execute session run to get action (and other fetches). + fetched = builder.get(to_fetch) + + # Update our global timestep by the batch size. + self.global_timestep += ( + len(obs_batch) + if isinstance(obs_batch, list) + else tree.flatten(obs_batch)[0].shape[0] + ) + + return fetched + + @override(Policy) + def compute_log_likelihoods( + self, + actions: Union[List[TensorType], TensorType], + obs_batch: Union[List[TensorType], TensorType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Optional[Union[List[TensorType], TensorType]] = None, + prev_reward_batch: Optional[Union[List[TensorType], TensorType]] = None, + actions_normalized: bool = True, + **kwargs, + ) -> TensorType: + if self._log_likelihood is None: + raise ValueError( + "Cannot compute log-prob/likelihood w/o a self._log_likelihood op!" + ) + + # Exploration hook before each forward pass. + self.exploration.before_compute_actions( + explore=False, tf_sess=self.get_session() + ) + + builder = _TFRunBuilder(self.get_session(), "compute_log_likelihoods") + + # Normalize actions if necessary. + if actions_normalized is False and self.config["normalize_actions"]: + actions = normalize_action(actions, self.action_space_struct) + + # Feed actions (for which we want logp values) into graph. + builder.add_feed_dict({self._action_input: actions}) + # Feed observations. + builder.add_feed_dict({self._obs_input: obs_batch}) + # Internal states. + state_batches = state_batches or [] + if len(self._state_inputs) != len(state_batches): + raise ValueError( + "Must pass in RNN state batches for placeholders {}, got {}".format( + self._state_inputs, state_batches + ) + ) + builder.add_feed_dict(dict(zip(self._state_inputs, state_batches))) + if state_batches: + builder.add_feed_dict({self._seq_lens: np.ones(len(obs_batch))}) + # Prev-a and r. + if self._prev_action_input is not None and prev_action_batch is not None: + builder.add_feed_dict({self._prev_action_input: prev_action_batch}) + if self._prev_reward_input is not None and prev_reward_batch is not None: + builder.add_feed_dict({self._prev_reward_input: prev_reward_batch}) + # Fetch the log_likelihoods output and return. + fetches = builder.add_fetches([self._log_likelihood]) + return builder.get(fetches)[0] + + @override(Policy) + def learn_on_batch(self, postprocessed_batch: SampleBatch) -> Dict[str, TensorType]: + assert self.loss_initialized() + + # Switch on is_training flag in our batch. + postprocessed_batch.set_training(True) + + builder = _TFRunBuilder(self.get_session(), "learn_on_batch") + + # Callback handling. + learn_stats = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=postprocessed_batch, result=learn_stats + ) + + fetches = self._build_learn_on_batch(builder, postprocessed_batch) + stats = builder.get(fetches) + self.num_grad_updates += 1 + + stats.update( + { + "custom_metrics": learn_stats, + NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count, + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates + - 1 + - (postprocessed_batch.num_grad_updates or 0) + ), + } + ) + + return stats + + @override(Policy) + def compute_gradients( + self, postprocessed_batch: SampleBatch + ) -> Tuple[ModelGradients, Dict[str, TensorType]]: + assert self.loss_initialized() + # Switch on is_training flag in our batch. + postprocessed_batch.set_training(True) + builder = _TFRunBuilder(self.get_session(), "compute_gradients") + fetches = self._build_compute_gradients(builder, postprocessed_batch) + return builder.get(fetches) + + @staticmethod + def _tf1_from_state_helper(state: PolicyState) -> "Policy": + """Recovers a TFPolicy from a state object. + + The `state` of an instantiated TFPolicy can be retrieved by calling its + `get_state` method. Is meant to be used by the Policy.from_state() method to + aid with tracking variable creation. + + Args: + state: The state to recover a new TFPolicy instance from. + + Returns: + A new TFPolicy instance. + """ + serialized_pol_spec: Optional[dict] = state.get("policy_spec") + if serialized_pol_spec is None: + raise ValueError( + "No `policy_spec` key was found in given `state`! " + "Cannot create new Policy." + ) + pol_spec = PolicySpec.deserialize(serialized_pol_spec) + + with tf1.variable_scope(TFPolicy.next_tf_var_scope_name()): + # Create the new policy. + new_policy = pol_spec.policy_class( + # Note(jungong) : we are intentionally not using keyward arguments here + # because some policies name the observation space parameter obs_space, + # and some others name it observation_space. + pol_spec.observation_space, + pol_spec.action_space, + pol_spec.config, + ) + + # Set the new policy's state (weights, optimizer vars, exploration state, + # etc..). + new_policy.set_state(state) + + # Return the new policy. + return new_policy + + @override(Policy) + def apply_gradients(self, gradients: ModelGradients) -> None: + assert self.loss_initialized() + builder = _TFRunBuilder(self.get_session(), "apply_gradients") + fetches = self._build_apply_gradients(builder, gradients) + builder.get(fetches) + + @override(Policy) + def get_weights(self) -> Union[Dict[str, TensorType], List[TensorType]]: + return self._variables.get_weights() + + @override(Policy) + def set_weights(self, weights) -> None: + return self._variables.set_weights(weights) + + @override(Policy) + def get_exploration_state(self) -> Dict[str, TensorType]: + return self.exploration.get_state(sess=self.get_session()) + + @Deprecated(new="get_exploration_state", error=True) + def get_exploration_info(self) -> Dict[str, TensorType]: + return self.get_exploration_state() + + @override(Policy) + def is_recurrent(self) -> bool: + return len(self._state_inputs) > 0 + + @override(Policy) + def num_state_tensors(self) -> int: + return len(self._state_inputs) + + @override(Policy) + def get_state(self) -> PolicyState: + # For tf Policies, return Policy weights and optimizer var values. + state = super().get_state() + + if len(self._optimizer_variables.variables) > 0: + state["_optimizer_variables"] = self.get_session().run( + self._optimizer_variables.variables + ) + # Add exploration state. + state["_exploration_state"] = self.exploration.get_state(self.get_session()) + return state + + @override(Policy) + def set_state(self, state: PolicyState) -> None: + # Set optimizer vars first. + optimizer_vars = state.get("_optimizer_variables", None) + if optimizer_vars is not None: + self._optimizer_variables.set_weights(optimizer_vars) + # Set exploration's state. + if hasattr(self, "exploration") and "_exploration_state" in state: + self.exploration.set_state( + state=state["_exploration_state"], sess=self.get_session() + ) + + # Restore global timestep. + self.global_timestep = state["global_timestep"] + + # Then the Policy's (NN) weights and connectors. + super().set_state(state) + + @override(Policy) + def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None: + """Export tensorflow graph to export_dir for serving.""" + if onnx: + try: + import tf2onnx + except ImportError as e: + raise RuntimeError( + "Converting a TensorFlow model to ONNX requires " + "`tf2onnx` to be installed. Install with " + "`pip install tf2onnx`." + ) from e + + with self.get_session().graph.as_default(): + signature_def_map = self._build_signature_def() + + sd = signature_def_map[ + tf1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY # noqa: E501 + ] + inputs = [v.name for k, v in sd.inputs.items()] + outputs = [v.name for k, v in sd.outputs.items()] + + from tf2onnx import tf_loader + + frozen_graph_def = tf_loader.freeze_session( + self.get_session(), input_names=inputs, output_names=outputs + ) + + with tf1.Session(graph=tf.Graph()) as session: + tf.import_graph_def(frozen_graph_def, name="") + + g = tf2onnx.tfonnx.process_tf_graph( + session.graph, + input_names=inputs, + output_names=outputs, + inputs_as_nchw=inputs, + ) + + model_proto = g.make_model("onnx_model") + tf2onnx.utils.save_onnx_model( + export_dir, "model", feed_dict={}, model_proto=model_proto + ) + # Save the tf.keras.Model (architecture and weights, so it can be retrieved + # w/o access to the original (custom) Model or Policy code). + elif ( + hasattr(self, "model") + and hasattr(self.model, "base_model") + and isinstance(self.model.base_model, tf.keras.Model) + ): + with self.get_session().graph.as_default(): + try: + self.model.base_model.save(filepath=export_dir, save_format="tf") + except Exception: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + else: + logger.warning(ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL) + + @override(Policy) + def import_model_from_h5(self, import_file: str) -> None: + """Imports weights into tf model.""" + if self.model is None: + raise NotImplementedError("No `self.model` to import into!") + + # Make sure the session is the right one (see issue #7046). + with self.get_session().graph.as_default(): + with self.get_session().as_default(): + return self.model.import_from_h5(import_file) + + @override(Policy) + def get_session(self) -> Optional["tf1.Session"]: + """Returns a reference to the TF session for this policy.""" + return self._sess + + def variables(self): + """Return the list of all savable variables for this policy.""" + if self.model is None: + raise NotImplementedError("No `self.model` to get variables for!") + elif isinstance(self.model, tf.keras.Model): + return self.model.variables + else: + return self.model.variables() + + def get_placeholder(self, name) -> "tf1.placeholder": + """Returns the given action or loss input placeholder by name. + + If the loss has not been initialized and a loss input placeholder is + requested, an error is raised. + + Args: + name: The name of the placeholder to return. One of + SampleBatch.CUR_OBS|PREV_ACTION/REWARD or a valid key from + `self._loss_input_dict`. + + Returns: + tf1.placeholder: The placeholder under the given str key. + """ + if name == SampleBatch.CUR_OBS: + return self._obs_input + elif name == SampleBatch.PREV_ACTIONS: + return self._prev_action_input + elif name == SampleBatch.PREV_REWARDS: + return self._prev_reward_input + + assert self._loss_input_dict, ( + "You need to populate `self._loss_input_dict` before " + "`get_placeholder()` can be called" + ) + return self._loss_input_dict[name] + + def loss_initialized(self) -> bool: + """Returns whether the loss term(s) have been initialized.""" + return len(self._losses) > 0 + + def _initialize_loss( + self, losses: List[TensorType], loss_inputs: List[Tuple[str, TensorType]] + ) -> None: + """Initializes the loss op from given loss tensor and placeholders. + + Args: + loss (List[TensorType]): The list of loss ops returned by some + loss function. + loss_inputs (List[Tuple[str, TensorType]]): The list of Tuples: + (name, tf1.placeholders) needed for calculating the loss. + """ + self._loss_input_dict = dict(loss_inputs) + self._loss_input_dict_no_rnn = { + k: v + for k, v in self._loss_input_dict.items() + if (v not in self._state_inputs and v != self._seq_lens) + } + for i, ph in enumerate(self._state_inputs): + self._loss_input_dict["state_in_{}".format(i)] = ph + + if self.model and not isinstance(self.model, tf.keras.Model): + self._losses = force_list( + self.model.custom_loss(losses, self._loss_input_dict) + ) + self._stats_fetches.update({"model": self.model.metrics()}) + else: + self._losses = losses + # Backward compatibility. + self._loss = self._losses[0] if self._losses is not None else None + + if not self._optimizers: + self._optimizers = force_list(self.optimizer()) + # Backward compatibility. + self._optimizer = self._optimizers[0] if self._optimizers else None + + # Supporting more than one loss/optimizer. + if self.config["_tf_policy_handles_more_than_one_loss"]: + self._grads_and_vars = [] + self._grads = [] + for group in self.gradients(self._optimizers, self._losses): + g_and_v = [(g, v) for (g, v) in group if g is not None] + self._grads_and_vars.append(g_and_v) + self._grads.append([g for (g, _) in g_and_v]) + # Only one optimizer and and loss term. + else: + self._grads_and_vars = [ + (g, v) + for (g, v) in self.gradients(self._optimizer, self._loss) + if g is not None + ] + self._grads = [g for (g, _) in self._grads_and_vars] + + if self.model: + self._variables = TensorFlowVariables( + [], self.get_session(), self.variables() + ) + + # Gather update ops for any batch norm layers. + if len(self.devices) <= 1: + if not self._update_ops: + self._update_ops = tf1.get_collection( + tf1.GraphKeys.UPDATE_OPS, scope=tf1.get_variable_scope().name + ) + if self._update_ops: + logger.info( + "Update ops to run on apply gradient: {}".format(self._update_ops) + ) + with tf1.control_dependencies(self._update_ops): + self._apply_op = self.build_apply_op( + optimizer=self._optimizers + if self.config["_tf_policy_handles_more_than_one_loss"] + else self._optimizer, + grads_and_vars=self._grads_and_vars, + ) + + if log_once("loss_used"): + logger.debug( + "These tensors were used in the loss functions:" + f"\n{summarize(self._loss_input_dict)}\n" + ) + + self.get_session().run(tf1.global_variables_initializer()) + + # TensorFlowVariables holding a flat list of all our optimizers' + # variables. + self._optimizer_variables = TensorFlowVariables( + [v for o in self._optimizers for v in o.variables()], self.get_session() + ) + + def copy(self, existing_inputs: List[Tuple[str, "tf1.placeholder"]]) -> "TFPolicy": + """Creates a copy of self using existing input placeholders. + + Optional: Only required to work with the multi-GPU optimizer. + + Args: + existing_inputs (List[Tuple[str, tf1.placeholder]]): Dict mapping + names (str) to tf1.placeholders to re-use (share) with the + returned copy of self. + + Returns: + TFPolicy: A copy of self. + """ + raise NotImplementedError + + def extra_compute_action_feed_dict(self) -> Dict[TensorType, TensorType]: + """Extra dict to pass to the compute actions session run. + + Returns: + Dict[TensorType, TensorType]: A feed dict to be added to the + feed_dict passed to the compute_actions session.run() call. + """ + return {} + + def extra_compute_action_fetches(self) -> Dict[str, TensorType]: + # Cache graph fetches for action computation for better + # performance. + # This function is called every time the static graph is run + # to compute actions. + if not self._cached_extra_action_out: + self._cached_extra_action_out = self.extra_action_out_fn() + return self._cached_extra_action_out + + def extra_action_out_fn(self) -> Dict[str, TensorType]: + """Extra values to fetch and return from compute_actions(). + + By default we return action probability/log-likelihood info + and action distribution inputs (if present). + + Returns: + Dict[str, TensorType]: An extra fetch-dict to be passed to and + returned from the compute_actions() call. + """ + extra_fetches = {} + # Action-logp and action-prob. + if self._sampled_action_logp is not None: + extra_fetches[SampleBatch.ACTION_PROB] = self._sampled_action_prob + extra_fetches[SampleBatch.ACTION_LOGP] = self._sampled_action_logp + # Action-dist inputs. + if self._dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = self._dist_inputs + return extra_fetches + + def extra_compute_grad_feed_dict(self) -> Dict[TensorType, TensorType]: + """Extra dict to pass to the compute gradients session run. + + Returns: + Dict[TensorType, TensorType]: Extra feed_dict to be passed to the + compute_gradients Session.run() call. + """ + return {} # e.g, kl_coeff + + def extra_compute_grad_fetches(self) -> Dict[str, any]: + """Extra values to fetch and return from compute_gradients(). + + Returns: + Dict[str, any]: Extra fetch dict to be added to the fetch dict + of the compute_gradients Session.run() call. + """ + return {LEARNER_STATS_KEY: {}} # e.g, stats, td error, etc. + + def optimizer(self) -> "tf.keras.optimizers.Optimizer": + """TF optimizer to use for policy optimization. + + Returns: + tf.keras.optimizers.Optimizer: The local optimizer to use for this + Policy's Model. + """ + if hasattr(self, "config") and "lr" in self.config: + return tf1.train.AdamOptimizer(learning_rate=self.config["lr"]) + else: + return tf1.train.AdamOptimizer() + + def gradients( + self, + optimizer: Union[LocalOptimizer, List[LocalOptimizer]], + loss: Union[TensorType, List[TensorType]], + ) -> Union[List[ModelGradients], List[List[ModelGradients]]]: + """Override this for a custom gradient computation behavior. + + Args: + optimizer (Union[LocalOptimizer, List[LocalOptimizer]]): A single + LocalOptimizer of a list thereof to use for gradient + calculations. If more than one optimizer given, the number of + optimizers must match the number of losses provided. + loss (Union[TensorType, List[TensorType]]): A single loss term + or a list thereof to use for gradient calculations. + If more than one loss given, the number of loss terms must + match the number of optimizers provided. + + Returns: + Union[List[ModelGradients], List[List[ModelGradients]]]: List of + ModelGradients (grads and vars OR just grads) OR List of List + of ModelGradients in case we have more than one + optimizer/loss. + """ + optimizers = force_list(optimizer) + losses = force_list(loss) + + # We have more than one optimizers and loss terms. + if self.config["_tf_policy_handles_more_than_one_loss"]: + grads = [] + for optim, loss_ in zip(optimizers, losses): + grads.append(optim.compute_gradients(loss_)) + # We have only one optimizer and one loss term. + else: + return optimizers[0].compute_gradients(losses[0]) + + def build_apply_op( + self, + optimizer: Union[LocalOptimizer, List[LocalOptimizer]], + grads_and_vars: Union[ModelGradients, List[ModelGradients]], + ) -> "tf.Operation": + """Override this for a custom gradient apply computation behavior. + + Args: + optimizer (Union[LocalOptimizer, List[LocalOptimizer]]): The local + tf optimizer to use for applying the grads and vars. + grads_and_vars (Union[ModelGradients, List[ModelGradients]]): List + of tuples with grad values and the grad-value's corresponding + tf.variable in it. + + Returns: + tf.Operation: The tf op that applies all computed gradients + (`grads_and_vars`) to the model(s) via the given optimizer(s). + """ + optimizers = force_list(optimizer) + + # We have more than one optimizers and loss terms. + if self.config["_tf_policy_handles_more_than_one_loss"]: + ops = [] + for i, optim in enumerate(optimizers): + # Specify global_step (e.g. for TD3 which needs to count the + # num updates that have happened). + ops.append( + optim.apply_gradients( + grads_and_vars[i], + global_step=tf1.train.get_or_create_global_step(), + ) + ) + return tf.group(ops) + # We have only one optimizer and one loss term. + else: + return optimizers[0].apply_gradients( + grads_and_vars, global_step=tf1.train.get_or_create_global_step() + ) + + def _get_is_training_placeholder(self): + """Get the placeholder for _is_training, i.e., for batch norm layers. + + This can be called safely before __init__ has run. + """ + if not hasattr(self, "_is_training"): + self._is_training = tf1.placeholder_with_default( + False, (), name="is_training" + ) + return self._is_training + + def _debug_vars(self): + if log_once("grad_vars"): + if self.config["_tf_policy_handles_more_than_one_loss"]: + for group in self._grads_and_vars: + for _, v in group: + logger.info("Optimizing variable {}".format(v)) + else: + for _, v in self._grads_and_vars: + logger.info("Optimizing variable {}".format(v)) + + def _extra_input_signature_def(self): + """Extra input signatures to add when exporting tf model. + Inferred from extra_compute_action_feed_dict() + """ + feed_dict = self.extra_compute_action_feed_dict() + return { + k.name: tf1.saved_model.utils.build_tensor_info(k) for k in feed_dict.keys() + } + + def _extra_output_signature_def(self): + """Extra output signatures to add when exporting tf model. + Inferred from extra_compute_action_fetches() + """ + fetches = self.extra_compute_action_fetches() + return { + k: tf1.saved_model.utils.build_tensor_info(fetches[k]) + for k in fetches.keys() + } + + def _build_signature_def(self): + """Build signature def map for tensorflow SavedModelBuilder.""" + # build input signatures + input_signature = self._extra_input_signature_def() + input_signature["observations"] = tf1.saved_model.utils.build_tensor_info( + self._obs_input + ) + + if self._seq_lens is not None: + input_signature[ + SampleBatch.SEQ_LENS + ] = tf1.saved_model.utils.build_tensor_info(self._seq_lens) + if self._prev_action_input is not None: + input_signature["prev_action"] = tf1.saved_model.utils.build_tensor_info( + self._prev_action_input + ) + if self._prev_reward_input is not None: + input_signature["prev_reward"] = tf1.saved_model.utils.build_tensor_info( + self._prev_reward_input + ) + + input_signature["is_training"] = tf1.saved_model.utils.build_tensor_info( + self._is_training + ) + + if self._timestep is not None: + input_signature["timestep"] = tf1.saved_model.utils.build_tensor_info( + self._timestep + ) + + for state_input in self._state_inputs: + input_signature[state_input.name] = tf1.saved_model.utils.build_tensor_info( + state_input + ) + + # build output signatures + output_signature = self._extra_output_signature_def() + for i, a in enumerate(tf.nest.flatten(self._sampled_action)): + output_signature[ + "actions_{}".format(i) + ] = tf1.saved_model.utils.build_tensor_info(a) + + for state_output in self._state_outputs: + output_signature[ + state_output.name + ] = tf1.saved_model.utils.build_tensor_info(state_output) + signature_def = tf1.saved_model.signature_def_utils.build_signature_def( + input_signature, + output_signature, + tf1.saved_model.signature_constants.PREDICT_METHOD_NAME, + ) + signature_def_key = ( + tf1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + ) + signature_def_map = {signature_def_key: signature_def} + return signature_def_map + + def _build_compute_actions( + self, + builder, + *, + input_dict=None, + obs_batch=None, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + episodes=None, + explore=None, + timestep=None, + ): + explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep + + # Call the exploration before_compute_actions hook. + self.exploration.before_compute_actions( + timestep=timestep, explore=explore, tf_sess=self.get_session() + ) + + builder.add_feed_dict(self.extra_compute_action_feed_dict()) + + # `input_dict` given: Simply build what's in that dict. + if hasattr(self, "_input_dict"): + for key, value in input_dict.items(): + if key in self._input_dict: + # Handle complex/nested spaces as well. + tree.map_structure( + lambda k, v: builder.add_feed_dict({k: v}), + self._input_dict[key], + value, + ) + # For policies that inherit directly from TFPolicy. + else: + builder.add_feed_dict({self._obs_input: input_dict[SampleBatch.OBS]}) + if SampleBatch.PREV_ACTIONS in input_dict: + builder.add_feed_dict( + {self._prev_action_input: input_dict[SampleBatch.PREV_ACTIONS]} + ) + if SampleBatch.PREV_REWARDS in input_dict: + builder.add_feed_dict( + {self._prev_reward_input: input_dict[SampleBatch.PREV_REWARDS]} + ) + state_batches = [] + i = 0 + while "state_in_{}".format(i) in input_dict: + state_batches.append(input_dict["state_in_{}".format(i)]) + i += 1 + builder.add_feed_dict(dict(zip(self._state_inputs, state_batches))) + + if "state_in_0" in input_dict and SampleBatch.SEQ_LENS not in input_dict: + builder.add_feed_dict( + {self._seq_lens: np.ones(len(input_dict["state_in_0"]))} + ) + + builder.add_feed_dict({self._is_exploring: explore}) + if timestep is not None: + builder.add_feed_dict({self._timestep: timestep}) + + # Determine, what exactly to fetch from the graph. + to_fetch = ( + [self._sampled_action] + + self._state_outputs + + [self.extra_compute_action_fetches()] + ) + + # Add the ops to fetch for the upcoming session call. + fetches = builder.add_fetches(to_fetch) + return fetches[0], fetches[1:-1], fetches[-1] + + def _build_compute_gradients(self, builder, postprocessed_batch): + self._debug_vars() + builder.add_feed_dict(self.extra_compute_grad_feed_dict()) + builder.add_feed_dict( + self._get_loss_inputs_dict(postprocessed_batch, shuffle=False) + ) + fetches = builder.add_fetches([self._grads, self._get_grad_and_stats_fetches()]) + return fetches[0], fetches[1] + + def _build_apply_gradients(self, builder, gradients): + if len(gradients) != len(self._grads): + raise ValueError( + "Unexpected number of gradients to apply, got {} for {}".format( + gradients, self._grads + ) + ) + builder.add_feed_dict({self._is_training: True}) + builder.add_feed_dict(dict(zip(self._grads, gradients))) + fetches = builder.add_fetches([self._apply_op]) + return fetches[0] + + def _build_learn_on_batch(self, builder, postprocessed_batch): + self._debug_vars() + + builder.add_feed_dict(self.extra_compute_grad_feed_dict()) + builder.add_feed_dict( + self._get_loss_inputs_dict(postprocessed_batch, shuffle=False) + ) + fetches = builder.add_fetches( + [ + self._apply_op, + self._get_grad_and_stats_fetches(), + ] + ) + return fetches[1] + + def _get_grad_and_stats_fetches(self): + fetches = self.extra_compute_grad_fetches() + if LEARNER_STATS_KEY not in fetches: + raise ValueError("Grad fetches should contain 'stats': {...} entry") + if self._stats_fetches: + fetches[LEARNER_STATS_KEY] = dict( + self._stats_fetches, **fetches[LEARNER_STATS_KEY] + ) + return fetches + + def _get_loss_inputs_dict(self, train_batch: SampleBatch, shuffle: bool): + """Return a feed dict from a batch. + + Args: + train_batch: batch of data to derive inputs from. + shuffle: whether to shuffle batch sequences. Shuffle may + be done in-place. This only makes sense if you're further + applying minibatch SGD after getting the outputs. + + Returns: + Feed dict of data. + """ + + # Get batch ready for RNNs, if applicable. + if not isinstance(train_batch, SampleBatch) or not train_batch.zero_padded: + pad_batch_to_sequences_of_same_size( + train_batch, + max_seq_len=self._max_seq_len, + shuffle=shuffle, + batch_divisibility_req=self._batch_divisibility_req, + feature_keys=list(self._loss_input_dict_no_rnn.keys()), + view_requirements=self.view_requirements, + ) + + # Mark the batch as "is_training" so the Model can use this + # information. + train_batch.set_training(True) + + # Build the feed dict from the batch. + feed_dict = {} + for key, placeholders in self._loss_input_dict.items(): + a = tree.map_structure( + lambda ph, v: feed_dict.__setitem__(ph, v), + placeholders, + train_batch[key], + ) + del a + + state_keys = ["state_in_{}".format(i) for i in range(len(self._state_inputs))] + for key in state_keys: + feed_dict[self._loss_input_dict[key]] = train_batch[key] + if state_keys: + feed_dict[self._seq_lens] = train_batch[SampleBatch.SEQ_LENS] + + return feed_dict diff --git a/lib/python3.12/site-packages/ray/rllib/policy/tf_policy_template.py b/lib/python3.12/site-packages/ray/rllib/policy/tf_policy_template.py new file mode 100644 index 0000000000000000000000000000000000000000..0b4f1be1d0fb682765b961d3c9a22d6977ca6b36 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/tf_policy_template.py @@ -0,0 +1,366 @@ +from typing import Callable, Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy import eager_tf_policy +from ray.rllib.policy.dynamic_tf_policy import DynamicTFPolicy +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.utils import add_mixins, force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + ModelGradients, + TensorType, +) + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +def build_tf_policy( + name: str, + *, + loss_fn: Callable[ + [Policy, ModelV2, Type[TFActionDistribution], SampleBatch], + Union[TensorType, List[TensorType]], + ], + get_default_config: Optional[Callable[[None], AlgorithmConfigDict]] = None, + postprocess_fn=None, + stats_fn: Optional[Callable[[Policy, SampleBatch], Dict[str, TensorType]]] = None, + optimizer_fn: Optional[ + Callable[[Policy, AlgorithmConfigDict], "tf.keras.optimizers.Optimizer"] + ] = None, + compute_gradients_fn: Optional[ + Callable[[Policy, "tf.keras.optimizers.Optimizer", TensorType], ModelGradients] + ] = None, + apply_gradients_fn: Optional[ + Callable[ + [Policy, "tf.keras.optimizers.Optimizer", ModelGradients], "tf.Operation" + ] + ] = None, + grad_stats_fn: Optional[ + Callable[[Policy, SampleBatch, ModelGradients], Dict[str, TensorType]] + ] = None, + extra_action_out_fn: Optional[Callable[[Policy], Dict[str, TensorType]]] = None, + extra_learn_fetches_fn: Optional[Callable[[Policy], Dict[str, TensorType]]] = None, + validate_spaces: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + before_init: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + before_loss_init: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], None + ] + ] = None, + after_init: Optional[ + Callable[[Policy, gym.Space, gym.Space, AlgorithmConfigDict], None] + ] = None, + make_model: Optional[ + Callable[ + [Policy, gym.spaces.Space, gym.spaces.Space, AlgorithmConfigDict], ModelV2 + ] + ] = None, + action_sampler_fn: Optional[ + Callable[[TensorType, List[TensorType]], Tuple[TensorType, TensorType]] + ] = None, + action_distribution_fn: Optional[ + Callable[ + [Policy, ModelV2, TensorType, TensorType, TensorType], + Tuple[TensorType, type, List[TensorType]], + ] + ] = None, + mixins: Optional[List[type]] = None, + get_batch_divisibility_req: Optional[Callable[[Policy], int]] = None, + # Deprecated args. + obs_include_prev_action_reward=DEPRECATED_VALUE, + extra_action_fetches_fn=None, # Use `extra_action_out_fn`. + gradients_fn=None, # Use `compute_gradients_fn`. +) -> Type[DynamicTFPolicy]: + """Helper function for creating a dynamic tf policy at runtime. + + Functions will be run in this order to initialize the policy: + 1. Placeholder setup: postprocess_fn + 2. Loss init: loss_fn, stats_fn + 3. Optimizer init: optimizer_fn, gradients_fn, apply_gradients_fn, + grad_stats_fn + + This means that you can e.g., depend on any policy attributes created in + the running of `loss_fn` in later functions such as `stats_fn`. + + In eager mode, the following functions will be run repeatedly on each + eager execution: loss_fn, stats_fn, gradients_fn, apply_gradients_fn, + and grad_stats_fn. + + This means that these functions should not define any variables internally, + otherwise they will fail in eager mode execution. Variable should only + be created in make_model (if defined). + + Args: + name: Name of the policy (e.g., "PPOTFPolicy"). + loss_fn (Callable[[ + Policy, ModelV2, Type[TFActionDistribution], SampleBatch], + Union[TensorType, List[TensorType]]]): Callable for calculating a + loss tensor. + get_default_config (Optional[Callable[[None], AlgorithmConfigDict]]): + Optional callable that returns the default config to merge with any + overrides. If None, uses only(!) the user-provided + PartialAlgorithmConfigDict as dict for this Policy. + postprocess_fn (Optional[Callable[[Policy, SampleBatch, + Optional[Dict[AgentID, SampleBatch]], Episode], None]]): + Optional callable for post-processing experience batches (called + after the parent class' `postprocess_trajectory` method). + stats_fn (Optional[Callable[[Policy, SampleBatch], + Dict[str, TensorType]]]): Optional callable that returns a dict of + TF tensors to fetch given the policy and batch input tensors. If + None, will not compute any stats. + optimizer_fn (Optional[Callable[[Policy, AlgorithmConfigDict], + "tf.keras.optimizers.Optimizer"]]): Optional callable that returns + a tf.Optimizer given the policy and config. If None, will call + the base class' `optimizer()` method instead (which returns a + tf1.train.AdamOptimizer). + compute_gradients_fn (Optional[Callable[[Policy, + "tf.keras.optimizers.Optimizer", TensorType], ModelGradients]]): + Optional callable that returns a list of gradients. If None, + this defaults to optimizer.compute_gradients([loss]). + apply_gradients_fn (Optional[Callable[[Policy, + "tf.keras.optimizers.Optimizer", ModelGradients], + "tf.Operation"]]): Optional callable that returns an apply + gradients op given policy, tf-optimizer, and grads_and_vars. If + None, will call the base class' `build_apply_op()` method instead. + grad_stats_fn (Optional[Callable[[Policy, SampleBatch, ModelGradients], + Dict[str, TensorType]]]): Optional callable that returns a dict of + TF fetches given the policy, batch input, and gradient tensors. If + None, will not collect any gradient stats. + extra_action_out_fn (Optional[Callable[[Policy], + Dict[str, TensorType]]]): Optional callable that returns + a dict of TF fetches given the policy object. If None, will not + perform any extra fetches. + extra_learn_fetches_fn (Optional[Callable[[Policy], + Dict[str, TensorType]]]): Optional callable that returns a dict of + extra values to fetch and return when learning on a batch. If None, + will call the base class' `extra_compute_grad_fetches()` method + instead. + validate_spaces (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): Optional callable that takes the + Policy, observation_space, action_space, and config to check + the spaces for correctness. If None, no spaces checking will be + done. + before_init (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): Optional callable to run at the + beginning of policy init that takes the same arguments as the + policy constructor. If None, this step will be skipped. + before_loss_init (Optional[Callable[[Policy, gym.spaces.Space, + gym.spaces.Space, AlgorithmConfigDict], None]]): Optional callable to + run prior to loss init. If None, this step will be skipped. + after_init (Optional[Callable[[Policy, gym.Space, gym.Space, + AlgorithmConfigDict], None]]): Optional callable to run at the end of + policy init. If None, this step will be skipped. + make_model (Optional[Callable[[Policy, gym.spaces.Space, + gym.spaces.Space, AlgorithmConfigDict], ModelV2]]): Optional callable + that returns a ModelV2 object. + All policy variables should be created in this function. If None, + a default ModelV2 object will be created. + action_sampler_fn (Optional[Callable[[TensorType, List[TensorType]], + Tuple[TensorType, TensorType]]]): A callable returning a sampled + action and its log-likelihood given observation and state inputs. + If None, will either use `action_distribution_fn` or + compute actions by calling self.model, then sampling from the + so parameterized action distribution. + action_distribution_fn (Optional[Callable[[Policy, ModelV2, TensorType, + TensorType, TensorType], + Tuple[TensorType, type, List[TensorType]]]]): Optional callable + returning distribution inputs (parameters), a dist-class to + generate an action distribution object from, and internal-state + outputs (or an empty list if not applicable). If None, will either + use `action_sampler_fn` or compute actions by calling self.model, + then sampling from the so parameterized action distribution. + mixins (Optional[List[type]]): Optional list of any class mixins for + the returned policy class. These mixins will be applied in order + and will have higher precedence than the DynamicTFPolicy class. + get_batch_divisibility_req (Optional[Callable[[Policy], int]]): + Optional callable that returns the divisibility requirement for + sample batches. If None, will assume a value of 1. + + Returns: + Type[DynamicTFPolicy]: A child class of DynamicTFPolicy based on the + specified args. + """ + original_kwargs = locals().copy() + base = add_mixins(DynamicTFPolicy, mixins) + + if obs_include_prev_action_reward != DEPRECATED_VALUE: + deprecation_warning(old="obs_include_prev_action_reward", error=True) + + if extra_action_fetches_fn is not None: + deprecation_warning( + old="extra_action_fetches_fn", new="extra_action_out_fn", error=True + ) + + if gradients_fn is not None: + deprecation_warning(old="gradients_fn", new="compute_gradients_fn", error=True) + + class policy_cls(base): + def __init__( + self, + obs_space, + action_space, + config, + existing_model=None, + existing_inputs=None, + ): + if validate_spaces: + validate_spaces(self, obs_space, action_space, config) + + if before_init: + before_init(self, obs_space, action_space, config) + + def before_loss_init_wrapper(policy, obs_space, action_space, config): + if before_loss_init: + before_loss_init(policy, obs_space, action_space, config) + + if extra_action_out_fn is None or policy._is_tower: + extra_action_fetches = {} + else: + extra_action_fetches = extra_action_out_fn(policy) + + if hasattr(policy, "_extra_action_fetches"): + policy._extra_action_fetches.update(extra_action_fetches) + else: + policy._extra_action_fetches = extra_action_fetches + + DynamicTFPolicy.__init__( + self, + obs_space=obs_space, + action_space=action_space, + config=config, + loss_fn=loss_fn, + stats_fn=stats_fn, + grad_stats_fn=grad_stats_fn, + before_loss_init=before_loss_init_wrapper, + make_model=make_model, + action_sampler_fn=action_sampler_fn, + action_distribution_fn=action_distribution_fn, + existing_inputs=existing_inputs, + existing_model=existing_model, + get_batch_divisibility_req=get_batch_divisibility_req, + ) + + if after_init: + after_init(self, obs_space, action_space, config) + + # Got to reset global_timestep again after this fake run-through. + self.global_timestep = 0 + + @override(Policy) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + # Call super's postprocess_trajectory first. + sample_batch = Policy.postprocess_trajectory(self, sample_batch) + if postprocess_fn: + return postprocess_fn(self, sample_batch, other_agent_batches, episode) + return sample_batch + + @override(TFPolicy) + def optimizer(self): + if optimizer_fn: + optimizers = optimizer_fn(self, self.config) + else: + optimizers = base.optimizer(self) + optimizers = force_list(optimizers) + if self.exploration: + optimizers = self.exploration.get_exploration_optimizer(optimizers) + + # No optimizers produced -> Return None. + if not optimizers: + return None + # New API: Allow more than one optimizer to be returned. + # -> Return list. + elif self.config["_tf_policy_handles_more_than_one_loss"]: + return optimizers + # Old API: Return a single LocalOptimizer. + else: + return optimizers[0] + + @override(TFPolicy) + def gradients(self, optimizer, loss): + optimizers = force_list(optimizer) + losses = force_list(loss) + + if compute_gradients_fn: + # New API: Allow more than one optimizer -> Return a list of + # lists of gradients. + if self.config["_tf_policy_handles_more_than_one_loss"]: + return compute_gradients_fn(self, optimizers, losses) + # Old API: Return a single List of gradients. + else: + return compute_gradients_fn(self, optimizers[0], losses[0]) + else: + return base.gradients(self, optimizers, losses) + + @override(TFPolicy) + def build_apply_op(self, optimizer, grads_and_vars): + if apply_gradients_fn: + return apply_gradients_fn(self, optimizer, grads_and_vars) + else: + return base.build_apply_op(self, optimizer, grads_and_vars) + + @override(TFPolicy) + def extra_compute_action_fetches(self): + return dict( + base.extra_compute_action_fetches(self), **self._extra_action_fetches + ) + + @override(TFPolicy) + def extra_compute_grad_fetches(self): + if extra_learn_fetches_fn: + # TODO: (sven) in torch, extra_learn_fetches do not exist. + # Hence, things like td_error are returned by the stats_fn + # and end up under the LEARNER_STATS_KEY. We should + # change tf to do this as well. However, this will confilct + # the handling of LEARNER_STATS_KEY inside the multi-GPU + # train op. + # Auto-add empty learner stats dict if needed. + return dict({LEARNER_STATS_KEY: {}}, **extra_learn_fetches_fn(self)) + else: + return base.extra_compute_grad_fetches(self) + + def with_updates(**overrides): + """Allows creating a TFPolicy cls based on settings of another one. + + Keyword Args: + **overrides: The settings (passed into `build_tf_policy`) that + should be different from the class that this method is called + on. + + Returns: + type: A new TFPolicy sub-class. + + Examples: + >> MySpecialDQNPolicyClass = DQNTFPolicy.with_updates( + .. name="MySpecialDQNPolicyClass", + .. loss_function=[some_new_loss_function], + .. ) + """ + return build_tf_policy(**dict(original_kwargs, **overrides)) + + def as_eager(): + return eager_tf_policy._build_eager_tf_policy(**original_kwargs) + + policy_cls.with_updates = staticmethod(with_updates) + policy_cls.as_eager = staticmethod(as_eager) + policy_cls.__name__ = name + policy_cls.__qualname__ = name + return policy_cls diff --git a/lib/python3.12/site-packages/ray/rllib/policy/torch_mixins.py b/lib/python3.12/site-packages/ray/rllib/policy/torch_mixins.py new file mode 100644 index 0000000000000000000000000000000000000000..a255f4e7a577cdc9bb4d0870457649b1f2e958e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/torch_mixins.py @@ -0,0 +1,221 @@ +from ray.rllib.policy.policy import PolicyState +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_policy import TorchPolicy +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.schedules import PiecewiseSchedule + +torch, nn = try_import_torch() + + +@OldAPIStack +class LearningRateSchedule: + """Mixin for TorchPolicy that adds a learning rate schedule.""" + + def __init__(self, lr, lr_schedule, lr2=None, lr2_schedule=None): + self._lr_schedule = None + self._lr2_schedule = None + # Disable any scheduling behavior related to learning if Learner API is active. + # Schedules are handled by Learner class. + if lr_schedule is None: + self.cur_lr = lr + else: + self._lr_schedule = PiecewiseSchedule( + lr_schedule, outside_value=lr_schedule[-1][-1], framework=None + ) + self.cur_lr = self._lr_schedule.value(0) + if lr2_schedule is None: + self.cur_lr2 = lr2 + else: + self._lr2_schedule = PiecewiseSchedule( + lr2_schedule, outside_value=lr2_schedule[-1][-1], framework=None + ) + self.cur_lr2 = self._lr2_schedule.value(0) + + def on_global_var_update(self, global_vars): + super().on_global_var_update(global_vars) + if self._lr_schedule: + self.cur_lr = self._lr_schedule.value(global_vars["timestep"]) + for opt in self._optimizers: + for p in opt.param_groups: + p["lr"] = self.cur_lr + if self._lr2_schedule: + assert len(self._optimizers) == 2 + self.cur_lr2 = self._lr2_schedule.value(global_vars["timestep"]) + opt = self._optimizers[1] + for p in opt.param_groups: + p["lr"] = self.cur_lr2 + + +@OldAPIStack +class EntropyCoeffSchedule: + """Mixin for TorchPolicy that adds entropy coeff decay.""" + + def __init__(self, entropy_coeff, entropy_coeff_schedule): + self._entropy_coeff_schedule = None + # Disable any scheduling behavior related to learning if Learner API is active. + # Schedules are handled by Learner class. + if entropy_coeff_schedule is None: + self.entropy_coeff = entropy_coeff + else: + # Allows for custom schedule similar to lr_schedule format + if isinstance(entropy_coeff_schedule, list): + self._entropy_coeff_schedule = PiecewiseSchedule( + entropy_coeff_schedule, + outside_value=entropy_coeff_schedule[-1][-1], + framework=None, + ) + else: + # Implements previous version but enforces outside_value + self._entropy_coeff_schedule = PiecewiseSchedule( + [[0, entropy_coeff], [entropy_coeff_schedule, 0.0]], + outside_value=0.0, + framework=None, + ) + self.entropy_coeff = self._entropy_coeff_schedule.value(0) + + def on_global_var_update(self, global_vars): + super(EntropyCoeffSchedule, self).on_global_var_update(global_vars) + if self._entropy_coeff_schedule is not None: + self.entropy_coeff = self._entropy_coeff_schedule.value( + global_vars["timestep"] + ) + + +@OldAPIStack +class KLCoeffMixin: + """Assigns the `update_kl()` method to a TorchPolicy. + + This is used by Algorithms to update the KL coefficient + after each learning step based on `config.kl_target` and + the measured KL value (from the train_batch). + """ + + def __init__(self, config): + # The current KL value (as python float). + self.kl_coeff = config["kl_coeff"] + # Constant target value. + self.kl_target = config["kl_target"] + + def update_kl(self, sampled_kl): + # Update the current KL value based on the recently measured value. + if sampled_kl > 2.0 * self.kl_target: + self.kl_coeff *= 1.5 + elif sampled_kl < 0.5 * self.kl_target: + self.kl_coeff *= 0.5 + # Return the current KL value. + return self.kl_coeff + + def get_state(self) -> PolicyState: + state = super().get_state() + # Add current kl-coeff value. + state["current_kl_coeff"] = self.kl_coeff + return state + + def set_state(self, state: PolicyState) -> None: + # Set current kl-coeff value first. + self.kl_coeff = state.pop("current_kl_coeff", self.config["kl_coeff"]) + # Call super's set_state with rest of the state dict. + super().set_state(state) + + +@OldAPIStack +class ValueNetworkMixin: + """Assigns the `_value()` method to a TorchPolicy. + + This way, Policy can call `_value()` to get the current VF estimate on a + single(!) observation (as done in `postprocess_trajectory_fn`). + Note: When doing this, an actual forward pass is being performed. + This is different from only calling `model.value_function()`, where + the result of the most recent forward pass is being used to return an + already calculated tensor. + """ + + def __init__(self, config): + # When doing GAE, we need the value function estimate on the + # observation. + if config.get("use_gae") or config.get("vtrace"): + # Input dict is provided to us automatically via the Model's + # requirements. It's a single-timestep (last one in trajectory) + # input_dict. + + def value(**input_dict): + input_dict = SampleBatch(input_dict) + input_dict = self._lazy_tensor_dict(input_dict) + model_out, _ = self.model(input_dict) + # [0] = remove the batch dim. + return self.model.value_function()[0].item() + + # When not doing GAE, we do not require the value function's output. + else: + + def value(*args, **kwargs): + return 0.0 + + self._value = value + + def extra_action_out(self, input_dict, state_batches, model, action_dist): + """Defines extra fetches per action computation. + + Args: + input_dict (Dict[str, TensorType]): The input dict used for the action + computing forward pass. + state_batches (List[TensorType]): List of state tensors (empty for + non-RNNs). + model (ModelV2): The Model object of the Policy. + action_dist: The instantiated distribution + object, resulting from the model's outputs and the given + distribution class. + + Returns: + Dict[str, TensorType]: Dict with extra tf fetches to perform per + action computation. + """ + # Return value function outputs. VF estimates will hence be added to + # the SampleBatches produced by the sampler(s) to generate the train + # batches going into the loss function. + return { + SampleBatch.VF_PREDS: model.value_function(), + } + + +@OldAPIStack +class TargetNetworkMixin: + """Mixin class adding a method for (soft) target net(s) synchronizations. + + - Adds the `update_target` method to the policy. + Calling `update_target` updates all target Q-networks' weights from their + respective "main" Q-networks, based on tau (smooth, partial updating). + """ + + def __init__(self): + # Hard initial update from Q-net(s) to target Q-net(s). + tau = self.config.get("tau", 1.0) + self.update_target(tau=tau) + + def update_target(self, tau=None): + # Update_target_fn will be called periodically to copy Q network to + # target Q network, using (soft) tau-synching. + tau = tau or self.config.get("tau", 1.0) + + model_state_dict = self.model.state_dict() + + # Support partial (soft) synching. + # If tau == 1.0: Full sync from Q-model to target Q-model. + # Support partial (soft) synching. + # If tau == 1.0: Full sync from Q-model to target Q-model. + target_state_dict = next(iter(self.target_models.values())).state_dict() + model_state_dict = { + k: tau * model_state_dict[k] + (1 - tau) * v + for k, v in target_state_dict.items() + } + + for target in self.target_models.values(): + target.load_state_dict(model_state_dict) + + def set_weights(self, weights): + # Makes sure that whenever we restore weights for this policy's + # model, we sync the target network (from the main model) + # at the same time. + TorchPolicy.set_weights(self, weights) + self.update_target() diff --git a/lib/python3.12/site-packages/ray/rllib/policy/torch_policy.py b/lib/python3.12/site-packages/ray/rllib/policy/torch_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..64eeb8374001983a4c9b6bb0782a203ad42199c3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/torch_policy.py @@ -0,0 +1,1201 @@ +import copy +import functools +import logging +import math +import os +import threading +import time +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Set, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +import ray +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.policy import Policy, PolicyState +from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import NullContextManager, force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.error import ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_AGENT_STEPS_TRAINED, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.spaces.space_utils import normalize_action +from ray.rllib.utils.threading import with_lock +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + GradInfoDict, + ModelGradients, + ModelWeights, + TensorStructType, + TensorType, +) + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class TorchPolicy(Policy): + """PyTorch specific Policy class to use with RLlib.""" + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + *, + model: Optional[TorchModelV2] = None, + loss: Optional[ + Callable[ + [Policy, ModelV2, Type[TorchDistributionWrapper], SampleBatch], + Union[TensorType, List[TensorType]], + ] + ] = None, + action_distribution_class: Optional[Type[TorchDistributionWrapper]] = None, + action_sampler_fn: Optional[ + Callable[ + [TensorType, List[TensorType]], + Union[ + Tuple[TensorType, TensorType, List[TensorType]], + Tuple[TensorType, TensorType, TensorType, List[TensorType]], + ], + ] + ] = None, + action_distribution_fn: Optional[ + Callable[ + [Policy, ModelV2, TensorType, TensorType, TensorType], + Tuple[TensorType, Type[TorchDistributionWrapper], List[TensorType]], + ] + ] = None, + max_seq_len: int = 20, + get_batch_divisibility_req: Optional[Callable[[Policy], int]] = None, + ): + """Initializes a TorchPolicy instance. + + Args: + observation_space: Observation space of the policy. + action_space: Action space of the policy. + config: The Policy's config dict. + model: PyTorch policy module. Given observations as + input, this module must return a list of outputs where the + first item is action logits, and the rest can be any value. + loss: Callable that returns one or more (a list of) scalar loss + terms. + action_distribution_class: Class for a torch action distribution. + action_sampler_fn: A callable returning either a sampled action, + its log-likelihood and updated state or a sampled action, its + log-likelihood, updated state and action distribution inputs + given Policy, ModelV2, input_dict, state batches (optional), + explore, and timestep. Provide `action_sampler_fn` if you would + like to have full control over the action computation step, + including the model forward pass, possible sampling from a + distribution, and exploration logic. + Note: If `action_sampler_fn` is given, `action_distribution_fn` + must be None. If both `action_sampler_fn` and + `action_distribution_fn` are None, RLlib will simply pass + inputs through `self.model` to get distribution inputs, create + the distribution object, sample from it, and apply some + exploration logic to the results. + The callable takes as inputs: Policy, ModelV2, input_dict + (SampleBatch), state_batches (optional), explore, and timestep. + action_distribution_fn: A callable returning distribution inputs + (parameters), a dist-class to generate an action distribution + object from, and internal-state outputs (or an empty list if + not applicable). + Provide `action_distribution_fn` if you would like to only + customize the model forward pass call. The resulting + distribution parameters are then used by RLlib to create a + distribution object, sample from it, and execute any + exploration logic. + Note: If `action_distribution_fn` is given, `action_sampler_fn` + must be None. If both `action_sampler_fn` and + `action_distribution_fn` are None, RLlib will simply pass + inputs through `self.model` to get distribution inputs, create + the distribution object, sample from it, and apply some + exploration logic to the results. + The callable takes as inputs: Policy, ModelV2, ModelInputDict, + explore, timestep, is_training. + max_seq_len: Max sequence length for LSTM training. + get_batch_divisibility_req: Optional callable that returns the + divisibility requirement for sample batches given the Policy. + """ + self.framework = config["framework"] = "torch" + self._loss_initialized = False + super().__init__(observation_space, action_space, config) + + # Create multi-GPU model towers, if necessary. + # - The central main model will be stored under self.model, residing + # on self.device (normally, a CPU). + # - Each GPU will have a copy of that model under + # self.model_gpu_towers, matching the devices in self.devices. + # - Parallelization is done by splitting the train batch and passing + # it through the model copies in parallel, then averaging over the + # resulting gradients, applying these averages on the main model and + # updating all towers' weights from the main model. + # - In case of just one device (1 (fake or real) GPU or 1 CPU), no + # parallelization will be done. + + # If no Model is provided, build a default one here. + if model is None: + dist_class, logit_dim = ModelCatalog.get_action_dist( + action_space, self.config["model"], framework=self.framework + ) + model = ModelCatalog.get_model_v2( + obs_space=self.observation_space, + action_space=self.action_space, + num_outputs=logit_dim, + model_config=self.config["model"], + framework=self.framework, + ) + if action_distribution_class is None: + action_distribution_class = dist_class + + # Get devices to build the graph on. + num_gpus = self._get_num_gpus_for_policy() + gpu_ids = list(range(torch.cuda.device_count())) + logger.info(f"Found {len(gpu_ids)} visible cuda devices.") + + # Place on one or more CPU(s) when either: + # - Fake GPU mode. + # - num_gpus=0 (either set by user or we are in local_mode=True). + # - No GPUs available. + if config["_fake_gpus"] or num_gpus == 0 or not gpu_ids: + self.device = torch.device("cpu") + self.devices = [self.device for _ in range(int(math.ceil(num_gpus)) or 1)] + self.model_gpu_towers = [ + model if i == 0 else copy.deepcopy(model) + for i in range(int(math.ceil(num_gpus)) or 1) + ] + if hasattr(self, "target_model"): + self.target_models = { + m: self.target_model for m in self.model_gpu_towers + } + self.model = model + # Place on one or more actual GPU(s), when: + # - num_gpus > 0 (set by user) AND + # - local_mode=False AND + # - actual GPUs available AND + # - non-fake GPU mode. + else: + # We are a remote worker (WORKER_MODE=1): + # GPUs should be assigned to us by ray. + if ray._private.worker._mode() == ray._private.worker.WORKER_MODE: + gpu_ids = ray.get_gpu_ids() + + if len(gpu_ids) < num_gpus: + raise ValueError( + "TorchPolicy was not able to find enough GPU IDs! Found " + f"{gpu_ids}, but num_gpus={num_gpus}." + ) + + self.devices = [ + torch.device("cuda:{}".format(i)) + for i, id_ in enumerate(gpu_ids) + if i < num_gpus + ] + self.device = self.devices[0] + ids = [id_ for i, id_ in enumerate(gpu_ids) if i < num_gpus] + self.model_gpu_towers = [] + for i, _ in enumerate(ids): + model_copy = copy.deepcopy(model) + self.model_gpu_towers.append(model_copy.to(self.devices[i])) + if hasattr(self, "target_model"): + self.target_models = { + m: copy.deepcopy(self.target_model).to(self.devices[i]) + for i, m in enumerate(self.model_gpu_towers) + } + self.model = self.model_gpu_towers[0] + + # Lock used for locking some methods on the object-level. + # This prevents possible race conditions when calling the model + # first, then its value function (e.g. in a loss function), in + # between of which another model call is made (e.g. to compute an + # action). + self._lock = threading.RLock() + + self._state_inputs = self.model.get_initial_state() + self._is_recurrent = len(self._state_inputs) > 0 + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + # Combine view_requirements for Model and Policy. + self.view_requirements.update(self.model.view_requirements) + + self.exploration = self._create_exploration() + self.unwrapped_model = model # used to support DistributedDataParallel + # To ensure backward compatibility: + # Old way: If `loss` provided here, use as-is (as a function). + if loss is not None: + self._loss = loss + # New way: Convert the overridden `self.loss` into a plain function, + # so it can be called the same way as `loss` would be, ensuring + # backward compatibility. + elif self.loss.__func__.__qualname__ != "Policy.loss": + self._loss = self.loss.__func__ + # `loss` not provided nor overridden from Policy -> Set to None. + else: + self._loss = None + self._optimizers = force_list(self.optimizer()) + # Store, which params (by index within the model's list of + # parameters) should be updated per optimizer. + # Maps optimizer idx to set or param indices. + self.multi_gpu_param_groups: List[Set[int]] = [] + main_params = {p: i for i, p in enumerate(self.model.parameters())} + for o in self._optimizers: + param_indices = [] + for pg_idx, pg in enumerate(o.param_groups): + for p in pg["params"]: + param_indices.append(main_params[p]) + self.multi_gpu_param_groups.append(set(param_indices)) + + # Create n sample-batch buffers (num_multi_gpu_tower_stacks), each + # one with m towers (num_gpus). + num_buffers = self.config.get("num_multi_gpu_tower_stacks", 1) + self._loaded_batches = [[] for _ in range(num_buffers)] + + self.dist_class = action_distribution_class + self.action_sampler_fn = action_sampler_fn + self.action_distribution_fn = action_distribution_fn + + # If set, means we are using distributed allreduce during learning. + self.distributed_world_size = None + + self.max_seq_len = max_seq_len + self.batch_divisibility_req = ( + get_batch_divisibility_req(self) + if callable(get_batch_divisibility_req) + else (get_batch_divisibility_req or 1) + ) + + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Dict[str, TensorType], + explore: bool = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + with torch.no_grad(): + # Pass lazy (torch) tensor dict to Model as `input_dict`. + input_dict = self._lazy_tensor_dict(input_dict) + input_dict.set_training(True) + # Pack internal state inputs into (separate) list. + state_batches = [ + input_dict[k] for k in input_dict.keys() if "state_in" in k[:8] + ] + # Calculate RNN sequence lengths. + seq_lens = ( + torch.tensor( + [1] * len(state_batches[0]), + dtype=torch.long, + device=state_batches[0].device, + ) + if state_batches + else None + ) + + return self._compute_action_helper( + input_dict, state_batches, seq_lens, explore, timestep + ) + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorStructType], TensorStructType] = None, + prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None, + info_batch: Optional[Dict[str, list]] = None, + episodes=None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorStructType, List[TensorType], Dict[str, TensorType]]: + with torch.no_grad(): + seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) + input_dict = self._lazy_tensor_dict( + { + SampleBatch.CUR_OBS: obs_batch, + "is_training": False, + } + ) + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = np.asarray(prev_action_batch) + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = np.asarray(prev_reward_batch) + state_batches = [ + convert_to_torch_tensor(s, self.device) for s in (state_batches or []) + ] + return self._compute_action_helper( + input_dict, state_batches, seq_lens, explore, timestep + ) + + @with_lock + @override(Policy) + def compute_log_likelihoods( + self, + actions: Union[List[TensorStructType], TensorStructType], + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Optional[ + Union[List[TensorStructType], TensorStructType] + ] = None, + prev_reward_batch: Optional[ + Union[List[TensorStructType], TensorStructType] + ] = None, + actions_normalized: bool = True, + **kwargs, + ) -> TensorType: + if self.action_sampler_fn and self.action_distribution_fn is None: + raise ValueError( + "Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!" + ) + + with torch.no_grad(): + input_dict = self._lazy_tensor_dict( + {SampleBatch.CUR_OBS: obs_batch, SampleBatch.ACTIONS: actions} + ) + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) + state_batches = [ + convert_to_torch_tensor(s, self.device) for s in (state_batches or []) + ] + + # Exploration hook before each forward pass. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if self.action_distribution_fn: + # Try new action_distribution_fn signature, supporting + # state_batches and seq_lens. + try: + dist_inputs, dist_class, state_out = self.action_distribution_fn( + self, + self.model, + input_dict=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=False, + is_training=False, + ) + # Trying the old way (to stay backward compatible). + # TODO: Remove in future. + except TypeError as e: + if ( + "positional argument" in e.args[0] + or "unexpected keyword argument" in e.args[0] + ): + dist_inputs, dist_class, _ = self.action_distribution_fn( + policy=self, + model=self.model, + obs_batch=input_dict[SampleBatch.CUR_OBS], + explore=False, + is_training=False, + ) + else: + raise e + + # Default action-dist inputs calculation. + else: + dist_class = self.dist_class + dist_inputs, _ = self.model(input_dict, state_batches, seq_lens) + + action_dist = dist_class(dist_inputs, self.model) + + # Normalize actions if necessary. + actions = input_dict[SampleBatch.ACTIONS] + if not actions_normalized and self.config["normalize_actions"]: + actions = normalize_action(actions, self.action_space_struct) + + log_likelihoods = action_dist.logp(actions) + + return log_likelihoods + + @with_lock + @override(Policy) + def learn_on_batch(self, postprocessed_batch: SampleBatch) -> Dict[str, TensorType]: + # Set Model to train mode. + if self.model: + self.model.train() + # Callback handling. + learn_stats = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=postprocessed_batch, result=learn_stats + ) + + # Compute gradients (will calculate all losses and `backward()` + # them to get the grads). + grads, fetches = self.compute_gradients(postprocessed_batch) + + # Step the optimizers. + self.apply_gradients(_directStepOptimizerSingleton) + + self.num_grad_updates += 1 + + if self.model: + fetches["model"] = self.model.metrics() + + fetches.update( + { + "custom_metrics": learn_stats, + NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count, + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates + - 1 + - (postprocessed_batch.num_grad_updates or 0) + ), + } + ) + + return fetches + + @override(Policy) + def load_batch_into_buffer( + self, + batch: SampleBatch, + buffer_index: int = 0, + ) -> int: + # Set the is_training flag of the batch. + batch.set_training(True) + + # Shortcut for 1 CPU only: Store batch in `self._loaded_batches`. + if len(self.devices) == 1 and self.devices[0].type == "cpu": + assert buffer_index == 0 + pad_batch_to_sequences_of_same_size( + batch=batch, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + self._lazy_tensor_dict(batch) + self._loaded_batches[0] = [batch] + return len(batch) + + # Batch (len=28, seq-lens=[4, 7, 4, 10, 3]): + # 0123 0123456 0123 0123456789ABC + + # 1) split into n per-GPU sub batches (n=2). + # [0123 0123456] [012] [3 0123456789 ABC] + # (len=14, 14 seq-lens=[4, 7, 3] [1, 10, 3]) + slices = batch.timeslices(num_slices=len(self.devices)) + + # 2) zero-padding (max-seq-len=10). + # - [0123000000 0123456000 0120000000] + # - [3000000000 0123456789 ABC0000000] + for slice in slices: + pad_batch_to_sequences_of_same_size( + batch=slice, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + # 3) Load splits into the given buffer (consisting of n GPUs). + slices = [slice.to_device(self.devices[i]) for i, slice in enumerate(slices)] + self._loaded_batches[buffer_index] = slices + + # Return loaded samples per-device. + return len(slices[0]) + + @override(Policy) + def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int: + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + return sum(len(b) for b in self._loaded_batches[buffer_index]) + + @override(Policy) + def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0): + if not self._loaded_batches[buffer_index]: + raise ValueError( + "Must call Policy.load_batch_into_buffer() before " + "Policy.learn_on_loaded_batch()!" + ) + + # Get the correct slice of the already loaded batch to use, + # based on offset and batch size. + device_batch_size = self.config.get("minibatch_size") + if device_batch_size is None: + device_batch_size = self.config.get( + "sgd_minibatch_size", + self.config["train_batch_size"], + ) + device_batch_size //= len(self.devices) + + # Set Model to train mode. + if self.model_gpu_towers: + for t in self.model_gpu_towers: + t.train() + + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_batches`. + if len(self.devices) == 1 and self.devices[0].type == "cpu": + assert buffer_index == 0 + if device_batch_size >= len(self._loaded_batches[0][0]): + batch = self._loaded_batches[0][0] + else: + batch = self._loaded_batches[0][0][offset : offset + device_batch_size] + return self.learn_on_batch(batch) + + if len(self.devices) > 1: + # Copy weights of main model (tower-0) to all other towers. + state_dict = self.model.state_dict() + # Just making sure tower-0 is really the same as self.model. + assert self.model_gpu_towers[0] is self.model + for tower in self.model_gpu_towers[1:]: + tower.load_state_dict(state_dict) + + if device_batch_size >= sum(len(s) for s in self._loaded_batches[buffer_index]): + device_batches = self._loaded_batches[buffer_index] + else: + device_batches = [ + b[offset : offset + device_batch_size] + for b in self._loaded_batches[buffer_index] + ] + + # Callback handling. + batch_fetches = {} + for i, batch in enumerate(device_batches): + custom_metrics = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=batch, result=custom_metrics + ) + batch_fetches[f"tower_{i}"] = {"custom_metrics": custom_metrics} + + # Do the (maybe parallelized) gradient calculation step. + tower_outputs = self._multi_gpu_parallel_grad_calc(device_batches) + + # Mean-reduce gradients over GPU-towers (do this on CPU: self.device). + all_grads = [] + for i in range(len(tower_outputs[0][0])): + if tower_outputs[0][0][i] is not None: + all_grads.append( + torch.mean( + torch.stack([t[0][i].to(self.device) for t in tower_outputs]), + dim=0, + ) + ) + else: + all_grads.append(None) + # Set main model's grads to mean-reduced values. + for i, p in enumerate(self.model.parameters()): + p.grad = all_grads[i] + + self.apply_gradients(_directStepOptimizerSingleton) + + self.num_grad_updates += 1 + + for i, (model, batch) in enumerate(zip(self.model_gpu_towers, device_batches)): + batch_fetches[f"tower_{i}"].update( + { + LEARNER_STATS_KEY: self.extra_grad_info(batch), + "model": model.metrics(), + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update + # above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates - 1 - (batch.num_grad_updates or 0) + ), + } + ) + batch_fetches.update(self.extra_compute_grad_fetches()) + + return batch_fetches + + @with_lock + @override(Policy) + def compute_gradients(self, postprocessed_batch: SampleBatch) -> ModelGradients: + assert len(self.devices) == 1 + + # If not done yet, see whether we have to zero-pad this batch. + if not postprocessed_batch.zero_padded: + pad_batch_to_sequences_of_same_size( + batch=postprocessed_batch, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + ) + + postprocessed_batch.set_training(True) + self._lazy_tensor_dict(postprocessed_batch, device=self.devices[0]) + + # Do the (maybe parallelized) gradient calculation step. + tower_outputs = self._multi_gpu_parallel_grad_calc([postprocessed_batch]) + + all_grads, grad_info = tower_outputs[0] + + grad_info["allreduce_latency"] /= len(self._optimizers) + grad_info.update(self.extra_grad_info(postprocessed_batch)) + + fetches = self.extra_compute_grad_fetches() + + return all_grads, dict(fetches, **{LEARNER_STATS_KEY: grad_info}) + + @override(Policy) + def apply_gradients(self, gradients: ModelGradients) -> None: + if gradients == _directStepOptimizerSingleton: + for i, opt in enumerate(self._optimizers): + opt.step() + else: + # TODO(sven): Not supported for multiple optimizers yet. + assert len(self._optimizers) == 1 + for g, p in zip(gradients, self.model.parameters()): + if g is not None: + if torch.is_tensor(g): + p.grad = g.to(self.device) + else: + p.grad = torch.from_numpy(g).to(self.device) + + self._optimizers[0].step() + + def get_tower_stats(self, stats_name: str) -> List[TensorStructType]: + """Returns list of per-tower stats, copied to this Policy's device. + + Args: + stats_name: The name of the stats to average over (this str + must exist as a key inside each tower's `tower_stats` dict). + + Returns: + The list of stats tensor (structs) of all towers, copied to this + Policy's device. + + Raises: + AssertionError: If the `stats_name` cannot be found in any one + of the tower's `tower_stats` dicts. + """ + data = [] + for tower in self.model_gpu_towers: + if stats_name in tower.tower_stats: + data.append( + tree.map_structure( + lambda s: s.to(self.device), tower.tower_stats[stats_name] + ) + ) + assert len(data) > 0, ( + f"Stats `{stats_name}` not found in any of the towers (you have " + f"{len(self.model_gpu_towers)} towers in total)! Make " + "sure you call the loss function on at least one of the towers." + ) + return data + + @override(Policy) + def get_weights(self) -> ModelWeights: + return {k: v.cpu().detach().numpy() for k, v in self.model.state_dict().items()} + + @override(Policy) + def set_weights(self, weights: ModelWeights) -> None: + weights = convert_to_torch_tensor(weights, device=self.device) + self.model.load_state_dict(weights) + + @override(Policy) + def is_recurrent(self) -> bool: + return self._is_recurrent + + @override(Policy) + def num_state_tensors(self) -> int: + return len(self.model.get_initial_state()) + + @override(Policy) + def get_initial_state(self) -> List[TensorType]: + return [s.detach().cpu().numpy() for s in self.model.get_initial_state()] + + @override(Policy) + def get_state(self) -> PolicyState: + state = super().get_state() + + state["_optimizer_variables"] = [] + for i, o in enumerate(self._optimizers): + optim_state_dict = convert_to_numpy(o.state_dict()) + state["_optimizer_variables"].append(optim_state_dict) + # Add exploration state. + if self.exploration: + # This is not compatible with RLModules, which have a method + # `forward_exploration` to specify custom exploration behavior. + state["_exploration_state"] = self.exploration.get_state() + return state + + @override(Policy) + def set_state(self, state: PolicyState) -> None: + # Set optimizer vars first. + optimizer_vars = state.get("_optimizer_variables", None) + if optimizer_vars: + assert len(optimizer_vars) == len(self._optimizers) + for o, s in zip(self._optimizers, optimizer_vars): + # Torch optimizer param_groups include things like beta, etc. These + # parameters should be left as scalar and not converted to tensors. + # otherwise, torch.optim.step() will start to complain. + optim_state_dict = {"param_groups": s["param_groups"]} + optim_state_dict["state"] = convert_to_torch_tensor( + s["state"], device=self.device + ) + o.load_state_dict(optim_state_dict) + # Set exploration's state. + if hasattr(self, "exploration") and "_exploration_state" in state: + self.exploration.set_state(state=state["_exploration_state"]) + + # Restore global timestep. + self.global_timestep = state["global_timestep"] + + # Then the Policy's (NN) weights and connectors. + super().set_state(state) + + def extra_grad_process( + self, optimizer: "torch.optim.Optimizer", loss: TensorType + ) -> Dict[str, TensorType]: + """Called after each optimizer.zero_grad() + loss.backward() call. + + Called for each self._optimizers/loss-value pair. + Allows for gradient processing before optimizer.step() is called. + E.g. for gradient clipping. + + Args: + optimizer: A torch optimizer object. + loss: The loss tensor associated with the optimizer. + + Returns: + An dict with information on the gradient processing step. + """ + return {} + + def extra_compute_grad_fetches(self) -> Dict[str, Any]: + """Extra values to fetch and return from compute_gradients(). + + Returns: + Extra fetch dict to be added to the fetch dict of the + `compute_gradients` call. + """ + return {LEARNER_STATS_KEY: {}} # e.g, stats, td error, etc. + + def extra_action_out( + self, + input_dict: Dict[str, TensorType], + state_batches: List[TensorType], + model: TorchModelV2, + action_dist: TorchDistributionWrapper, + ) -> Dict[str, TensorType]: + """Returns dict of extra info to include in experience batch. + + Args: + input_dict: Dict of model input tensors. + state_batches: List of state tensors. + model: Reference to the model object. + action_dist: Torch action dist object + to get log-probs (e.g. for already sampled actions). + + Returns: + Extra outputs to return in a `compute_actions_from_input_dict()` + call (3rd return value). + """ + return {} + + def extra_grad_info(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Return dict of extra grad info. + + Args: + train_batch: The training batch for which to produce + extra grad info for. + + Returns: + The info dict carrying grad info per str key. + """ + return {} + + def optimizer( + self, + ) -> Union[List["torch.optim.Optimizer"], "torch.optim.Optimizer"]: + """Custom the local PyTorch optimizer(s) to use. + + Returns: + The local PyTorch optimizer(s) to use for this Policy. + """ + if hasattr(self, "config"): + optimizers = [ + torch.optim.Adam(self.model.parameters(), lr=self.config["lr"]) + ] + else: + optimizers = [torch.optim.Adam(self.model.parameters())] + if self.exploration: + optimizers = self.exploration.get_exploration_optimizer(optimizers) + return optimizers + + @override(Policy) + def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None: + """Exports the Policy's Model to local directory for serving. + + Creates a TorchScript model and saves it. + + Args: + export_dir: Local writable directory or filename. + onnx: If given, will export model in ONNX format. The + value of this parameter set the ONNX OpSet version to use. + """ + os.makedirs(export_dir, exist_ok=True) + + if onnx: + self._lazy_tensor_dict(self._dummy_batch) + # Provide dummy state inputs if not an RNN (torch cannot jit with + # returned empty internal states list). + if "state_in_0" not in self._dummy_batch: + self._dummy_batch["state_in_0"] = self._dummy_batch[ + SampleBatch.SEQ_LENS + ] = np.array([1.0]) + seq_lens = self._dummy_batch[SampleBatch.SEQ_LENS] + + state_ins = [] + i = 0 + while "state_in_{}".format(i) in self._dummy_batch: + state_ins.append(self._dummy_batch["state_in_{}".format(i)]) + i += 1 + dummy_inputs = { + k: self._dummy_batch[k] + for k in self._dummy_batch.keys() + if k != "is_training" + } + + file_name = os.path.join(export_dir, "model.onnx") + torch.onnx.export( + self.model, + (dummy_inputs, state_ins, seq_lens), + file_name, + export_params=True, + opset_version=onnx, + do_constant_folding=True, + input_names=list(dummy_inputs.keys()) + + ["state_ins", SampleBatch.SEQ_LENS], + output_names=["output", "state_outs"], + dynamic_axes={ + k: {0: "batch_size"} + for k in list(dummy_inputs.keys()) + + ["state_ins", SampleBatch.SEQ_LENS] + }, + ) + # Save the torch.Model (architecture and weights, so it can be retrieved + # w/o access to the original (custom) Model or Policy code). + else: + filename = os.path.join(export_dir, "model.pt") + try: + torch.save(self.model, f=filename) + except Exception: + if os.path.exists(filename): + os.remove(filename) + logger.warning(ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL) + + @override(Policy) + def import_model_from_h5(self, import_file: str) -> None: + """Imports weights into torch model.""" + return self.model.import_from_h5(import_file) + + @with_lock + def _compute_action_helper( + self, input_dict, state_batches, seq_lens, explore, timestep + ): + """Shared forward pass logic (w/ and w/o trajectory view API). + + Returns: + A tuple consisting of a) actions, b) state_out, c) extra_fetches. + """ + explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep + self._is_recurrent = state_batches is not None and state_batches != [] + + # Switch to eval mode. + if self.model: + self.model.eval() + + if self.action_sampler_fn: + action_dist = dist_inputs = None + action_sampler_outputs = self.action_sampler_fn( + self, + self.model, + input_dict, + state_batches, + explore=explore, + timestep=timestep, + ) + if len(action_sampler_outputs) == 4: + actions, logp, dist_inputs, state_out = action_sampler_outputs + else: + actions, logp, state_out = action_sampler_outputs + else: + # Call the exploration before_compute_actions hook. + self.exploration.before_compute_actions(explore=explore, timestep=timestep) + if self.action_distribution_fn: + # Try new action_distribution_fn signature, supporting + # state_batches and seq_lens. + try: + dist_inputs, dist_class, state_out = self.action_distribution_fn( + self, + self.model, + input_dict=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=explore, + timestep=timestep, + is_training=False, + ) + # Trying the old way (to stay backward compatible). + # TODO: Remove in future. + except TypeError as e: + if ( + "positional argument" in e.args[0] + or "unexpected keyword argument" in e.args[0] + ): + ( + dist_inputs, + dist_class, + state_out, + ) = self.action_distribution_fn( + self, + self.model, + input_dict[SampleBatch.CUR_OBS], + explore=explore, + timestep=timestep, + is_training=False, + ) + else: + raise e + else: + dist_class = self.dist_class + dist_inputs, state_out = self.model(input_dict, state_batches, seq_lens) + + if not ( + isinstance(dist_class, functools.partial) + or issubclass(dist_class, TorchDistributionWrapper) + ): + raise ValueError( + "`dist_class` ({}) not a TorchDistributionWrapper " + "subclass! Make sure your `action_distribution_fn` or " + "`make_model_and_action_dist` return a correct " + "distribution class.".format(dist_class.__name__) + ) + action_dist = dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = self.exploration.get_exploration_action( + action_distribution=action_dist, timestep=timestep, explore=explore + ) + + input_dict[SampleBatch.ACTIONS] = actions + + # Add default and custom fetches. + extra_fetches = self.extra_action_out( + input_dict, state_batches, self.model, action_dist + ) + + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + + # Action-logp and action-prob. + if logp is not None: + extra_fetches[SampleBatch.ACTION_PROB] = torch.exp(logp.float()) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + + # Update our global timestep by the batch size. + self.global_timestep += len(input_dict[SampleBatch.CUR_OBS]) + + return convert_to_numpy((actions, state_out, extra_fetches)) + + def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch, device=None): + # TODO: (sven): Keep for a while to ensure backward compatibility. + if not isinstance(postprocessed_batch, SampleBatch): + postprocessed_batch = SampleBatch(postprocessed_batch) + postprocessed_batch.set_get_interceptor( + functools.partial(convert_to_torch_tensor, device=device or self.device) + ) + return postprocessed_batch + + def _multi_gpu_parallel_grad_calc( + self, sample_batches: List[SampleBatch] + ) -> List[Tuple[List[TensorType], GradInfoDict]]: + """Performs a parallelized loss and gradient calculation over the batch. + + Splits up the given train batch into n shards (n=number of this + Policy's devices) and passes each data shard (in parallel) through + the loss function using the individual devices' models + (self.model_gpu_towers). Then returns each tower's outputs. + + Args: + sample_batches: A list of SampleBatch shards to + calculate loss and gradients for. + + Returns: + A list (one item per device) of 2-tuples, each with 1) gradient + list and 2) grad info dict. + """ + assert len(self.model_gpu_towers) == len(sample_batches) + lock = threading.Lock() + results = {} + grad_enabled = torch.is_grad_enabled() + + def _worker(shard_idx, model, sample_batch, device): + torch.set_grad_enabled(grad_enabled) + try: + with NullContextManager() if device.type == "cpu" else torch.cuda.device( # noqa: E501 + device + ): + loss_out = force_list( + self._loss(self, model, self.dist_class, sample_batch) + ) + + # Call Model's custom-loss with Policy loss outputs and + # train_batch. + loss_out = model.custom_loss(loss_out, sample_batch) + + assert len(loss_out) == len(self._optimizers) + + # Loop through all optimizers. + grad_info = {"allreduce_latency": 0.0} + + parameters = list(model.parameters()) + all_grads = [None for _ in range(len(parameters))] + for opt_idx, opt in enumerate(self._optimizers): + # Erase gradients in all vars of the tower that this + # optimizer would affect. + param_indices = self.multi_gpu_param_groups[opt_idx] + for param_idx, param in enumerate(parameters): + if param_idx in param_indices and param.grad is not None: + param.grad.data.zero_() + # Recompute gradients of loss over all variables. + loss_out[opt_idx].backward(retain_graph=True) + grad_info.update( + self.extra_grad_process(opt, loss_out[opt_idx]) + ) + + grads = [] + # Note that return values are just references; + # Calling zero_grad would modify the values. + for param_idx, param in enumerate(parameters): + if param_idx in param_indices: + if param.grad is not None: + grads.append(param.grad) + all_grads[param_idx] = param.grad + + if self.distributed_world_size: + start = time.time() + if torch.cuda.is_available(): + # Sadly, allreduce_coalesced does not work with + # CUDA yet. + for g in grads: + torch.distributed.all_reduce( + g, op=torch.distributed.ReduceOp.SUM + ) + else: + torch.distributed.all_reduce_coalesced( + grads, op=torch.distributed.ReduceOp.SUM + ) + + for param_group in opt.param_groups: + for p in param_group["params"]: + if p.grad is not None: + p.grad /= self.distributed_world_size + + grad_info["allreduce_latency"] += time.time() - start + + with lock: + results[shard_idx] = (all_grads, grad_info) + except Exception as e: + import traceback + + with lock: + results[shard_idx] = ( + ValueError( + f"Error In tower {shard_idx} on device " + f"{device} during multi GPU parallel gradient " + f"calculation:" + f": {e}\n" + f"Traceback: \n" + f"{traceback.format_exc()}\n" + ), + e, + ) + + # Single device (GPU) or fake-GPU case (serialize for better + # debugging). + if len(self.devices) == 1 or self.config["_fake_gpus"]: + for shard_idx, (model, sample_batch, device) in enumerate( + zip(self.model_gpu_towers, sample_batches, self.devices) + ): + _worker(shard_idx, model, sample_batch, device) + # Raise errors right away for better debugging. + last_result = results[len(results) - 1] + if isinstance(last_result[0], ValueError): + raise last_result[0] from last_result[1] + # Multi device (GPU) case: Parallelize via threads. + else: + threads = [ + threading.Thread( + target=_worker, args=(shard_idx, model, sample_batch, device) + ) + for shard_idx, (model, sample_batch, device) in enumerate( + zip(self.model_gpu_towers, sample_batches, self.devices) + ) + ] + + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # Gather all threads' outputs and return. + outputs = [] + for shard_idx in range(len(sample_batches)): + output = results[shard_idx] + if isinstance(output[0], Exception): + raise output[0] from output[1] + outputs.append(results[shard_idx]) + return outputs + + +@OldAPIStack +class DirectStepOptimizer: + """Typesafe method for indicating `apply_gradients` can directly step the + optimizers with in-place gradients. + """ + + _instance = None + + def __new__(cls): + if DirectStepOptimizer._instance is None: + DirectStepOptimizer._instance = super().__new__(cls) + return DirectStepOptimizer._instance + + def __eq__(self, other): + return type(self) is type(other) + + def __repr__(self): + return "DirectStepOptimizer" + + +_directStepOptimizerSingleton = DirectStepOptimizer() diff --git a/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py b/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..365aa2cfc7d875c27285c39c2f6e937cc7285b78 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py @@ -0,0 +1,1260 @@ +import copy +import functools +import logging +import math +import os +import threading +import time +from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from packaging import version + +import ray +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_policy import _directStepOptimizerSingleton +from ray.rllib.utils import NullContextManager, force_list +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + is_overridden, + override, +) +from ray.rllib.utils.error import ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + NUM_AGENT_STEPS_TRAINED, + NUM_GRAD_UPDATES_LIFETIME, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.spaces.space_utils import normalize_action +from ray.rllib.utils.threading import with_lock +from ray.rllib.utils.torch_utils import ( + TORCH_COMPILE_REQUIRED_VERSION, + convert_to_torch_tensor, +) +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + GradInfoDict, + ModelGradients, + ModelWeights, + PolicyState, + TensorStructType, + TensorType, +) + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class TorchPolicyV2(Policy): + """PyTorch specific Policy class to use with RLlib.""" + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, + *, + max_seq_len: int = 20, + ): + """Initializes a TorchPolicy instance. + + Args: + observation_space: Observation space of the policy. + action_space: Action space of the policy. + config: The Policy's config dict. + max_seq_len: Max sequence length for LSTM training. + """ + self.framework = config["framework"] = "torch" + + self._loss_initialized = False + super().__init__(observation_space, action_space, config) + + # Create model. + model, dist_class = self._init_model_and_dist_class() + + # Create multi-GPU model towers, if necessary. + # - The central main model will be stored under self.model, residing + # on self.device (normally, a CPU). + # - Each GPU will have a copy of that model under + # self.model_gpu_towers, matching the devices in self.devices. + # - Parallelization is done by splitting the train batch and passing + # it through the model copies in parallel, then averaging over the + # resulting gradients, applying these averages on the main model and + # updating all towers' weights from the main model. + # - In case of just one device (1 (fake or real) GPU or 1 CPU), no + # parallelization will be done. + + # Get devices to build the graph on. + num_gpus = self._get_num_gpus_for_policy() + gpu_ids = list(range(torch.cuda.device_count())) + logger.info(f"Found {len(gpu_ids)} visible cuda devices.") + + # Place on one or more CPU(s) when either: + # - Fake GPU mode. + # - num_gpus=0 (either set by user or we are in local_mode=True). + # - No GPUs available. + if config["_fake_gpus"] or num_gpus == 0 or not gpu_ids: + self.device = torch.device("cpu") + self.devices = [self.device for _ in range(int(math.ceil(num_gpus)) or 1)] + self.model_gpu_towers = [ + model if i == 0 else copy.deepcopy(model) + for i in range(int(math.ceil(num_gpus)) or 1) + ] + if hasattr(self, "target_model"): + self.target_models = { + m: self.target_model for m in self.model_gpu_towers + } + self.model = model + # Place on one or more actual GPU(s), when: + # - num_gpus > 0 (set by user) AND + # - local_mode=False AND + # - actual GPUs available AND + # - non-fake GPU mode. + else: + # We are a remote worker (WORKER_MODE=1): + # GPUs should be assigned to us by ray. + if ray._private.worker._mode() == ray._private.worker.WORKER_MODE: + gpu_ids = ray.get_gpu_ids() + + if len(gpu_ids) < num_gpus: + raise ValueError( + "TorchPolicy was not able to find enough GPU IDs! Found " + f"{gpu_ids}, but num_gpus={num_gpus}." + ) + + self.devices = [ + torch.device("cuda:{}".format(i)) + for i, id_ in enumerate(gpu_ids) + if i < num_gpus + ] + self.device = self.devices[0] + ids = [id_ for i, id_ in enumerate(gpu_ids) if i < num_gpus] + self.model_gpu_towers = [] + for i, _ in enumerate(ids): + model_copy = copy.deepcopy(model) + self.model_gpu_towers.append(model_copy.to(self.devices[i])) + if hasattr(self, "target_model"): + self.target_models = { + m: copy.deepcopy(self.target_model).to(self.devices[i]) + for i, m in enumerate(self.model_gpu_towers) + } + self.model = self.model_gpu_towers[0] + + self.dist_class = dist_class + self.unwrapped_model = model # used to support DistributedDataParallel + + # Lock used for locking some methods on the object-level. + # This prevents possible race conditions when calling the model + # first, then its value function (e.g. in a loss function), in + # between of which another model call is made (e.g. to compute an + # action). + self._lock = threading.RLock() + + self._state_inputs = self.model.get_initial_state() + self._is_recurrent = len(tree.flatten(self._state_inputs)) > 0 + # Auto-update model's inference view requirements, if recurrent. + self._update_model_view_requirements_from_init_state() + # Combine view_requirements for Model and Policy. + self.view_requirements.update(self.model.view_requirements) + + self.exploration = self._create_exploration() + self._optimizers = force_list(self.optimizer()) + + # Backward compatibility workaround so Policy will call self.loss() + # directly. + # TODO (jungong): clean up after all policies are migrated to new sub-class + # implementation. + self._loss = None + + # Store, which params (by index within the model's list of + # parameters) should be updated per optimizer. + # Maps optimizer idx to set or param indices. + self.multi_gpu_param_groups: List[Set[int]] = [] + main_params = {p: i for i, p in enumerate(self.model.parameters())} + for o in self._optimizers: + param_indices = [] + for pg_idx, pg in enumerate(o.param_groups): + for p in pg["params"]: + param_indices.append(main_params[p]) + self.multi_gpu_param_groups.append(set(param_indices)) + + # Create n sample-batch buffers (num_multi_gpu_tower_stacks), each + # one with m towers (num_gpus). + num_buffers = self.config.get("num_multi_gpu_tower_stacks", 1) + self._loaded_batches = [[] for _ in range(num_buffers)] + + # If set, means we are using distributed allreduce during learning. + self.distributed_world_size = None + + self.batch_divisibility_req = self.get_batch_divisibility_req() + self.max_seq_len = max_seq_len + + # If model is an RLModule it won't have tower_stats instead there will be a + # self.tower_state[model] -> dict for each tower. + self.tower_stats = {} + if not hasattr(self.model, "tower_stats"): + for model in self.model_gpu_towers: + self.tower_stats[model] = {} + + def loss_initialized(self): + return self._loss_initialized + + @OverrideToImplementCustomLogic + @override(Policy) + def loss( + self, + model: ModelV2, + dist_class: Type[TorchDistributionWrapper], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + """Constructs the loss function. + + Args: + model: The Model to calculate the loss for. + dist_class: The action distr. class. + train_batch: The training data. + + Returns: + Loss tensor given the input batch. + """ + raise NotImplementedError + + @OverrideToImplementCustomLogic + def action_sampler_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, TensorType, TensorType, List[TensorType]]: + """Custom function for sampling new actions given policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Sampled action + Log-likelihood + Action distribution inputs + Updated state + """ + return None, None, None, None + + @OverrideToImplementCustomLogic + def action_distribution_fn( + self, + model: ModelV2, + *, + obs_batch: TensorType, + state_batches: TensorType, + **kwargs, + ) -> Tuple[TensorType, type, List[TensorType]]: + """Action distribution function for this Policy. + + Args: + model: Underlying model. + obs_batch: Observation tensor batch. + state_batches: Action sampling state batch. + + Returns: + Distribution input. + ActionDistribution class. + State outs. + """ + return None, None, None + + @OverrideToImplementCustomLogic + def make_model(self) -> ModelV2: + """Create model. + + Note: only one of make_model or make_model_and_action_dist + can be overridden. + + Returns: + ModelV2 model. + """ + return None + + @OverrideToImplementCustomLogic + def make_model_and_action_dist( + self, + ) -> Tuple[ModelV2, Type[TorchDistributionWrapper]]: + """Create model and action distribution function. + + Returns: + ModelV2 model. + ActionDistribution class. + """ + return None, None + + @OverrideToImplementCustomLogic + def get_batch_divisibility_req(self) -> int: + """Get batch divisibility request. + + Returns: + Size N. A sample batch must be of size K*N. + """ + # By default, any sized batch is ok, so simply return 1. + return 1 + + @OverrideToImplementCustomLogic + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function. Returns a dict of statistics. + + Args: + train_batch: The SampleBatch (already) used for training. + + Returns: + The stats dict. + """ + return {} + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_grad_process( + self, optimizer: "torch.optim.Optimizer", loss: TensorType + ) -> Dict[str, TensorType]: + """Called after each optimizer.zero_grad() + loss.backward() call. + + Called for each self._optimizers/loss-value pair. + Allows for gradient processing before optimizer.step() is called. + E.g. for gradient clipping. + + Args: + optimizer: A torch optimizer object. + loss: The loss tensor associated with the optimizer. + + Returns: + An dict with information on the gradient processing step. + """ + return {} + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_compute_grad_fetches(self) -> Dict[str, Any]: + """Extra values to fetch and return from compute_gradients(). + + Returns: + Extra fetch dict to be added to the fetch dict of the + `compute_gradients` call. + """ + return {LEARNER_STATS_KEY: {}} # e.g, stats, td error, etc. + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def extra_action_out( + self, + input_dict: Dict[str, TensorType], + state_batches: List[TensorType], + model: TorchModelV2, + action_dist: TorchDistributionWrapper, + ) -> Dict[str, TensorType]: + """Returns dict of extra info to include in experience batch. + + Args: + input_dict: Dict of model input tensors. + state_batches: List of state tensors. + model: Reference to the model object. + action_dist: Torch action dist object + to get log-probs (e.g. for already sampled actions). + + Returns: + Extra outputs to return in a `compute_actions_from_input_dict()` + call (3rd return value). + """ + return {} + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[Dict[Any, SampleBatch]] = None, + episode=None, + ) -> SampleBatch: + """Postprocesses a trajectory and returns the processed trajectory. + + The trajectory contains only data from one episode and from one agent. + - If `config.batch_mode=truncate_episodes` (default), sample_batch may + contain a truncated (at-the-end) episode, in case the + `config.rollout_fragment_length` was reached by the sampler. + - If `config.batch_mode=complete_episodes`, sample_batch will contain + exactly one episode (no matter how long). + New columns can be added to sample_batch and existing ones may be altered. + + Args: + sample_batch: The SampleBatch to postprocess. + other_agent_batches (Optional[Dict[PolicyID, SampleBatch]]): Optional + dict of AgentIDs mapping to other agents' trajectory data (from the + same episode). NOTE: The other agents use the same policy. + episode (Optional[Episode]): Optional multi-agent episode + object in which the agents operated. + + Returns: + SampleBatch: The postprocessed, modified SampleBatch (or a new one). + """ + return sample_batch + + @OverrideToImplementCustomLogic + def optimizer( + self, + ) -> Union[List["torch.optim.Optimizer"], "torch.optim.Optimizer"]: + """Custom the local PyTorch optimizer(s) to use. + + Returns: + The local PyTorch optimizer(s) to use for this Policy. + """ + if hasattr(self, "config"): + optimizers = [ + torch.optim.Adam(self.model.parameters(), lr=self.config["lr"]) + ] + else: + optimizers = [torch.optim.Adam(self.model.parameters())] + if self.exploration: + optimizers = self.exploration.get_exploration_optimizer(optimizers) + return optimizers + + def _init_model_and_dist_class(self): + if is_overridden(self.make_model) and is_overridden( + self.make_model_and_action_dist + ): + raise ValueError( + "Only one of make_model or make_model_and_action_dist " + "can be overridden." + ) + + if is_overridden(self.make_model): + model = self.make_model() + dist_class, _ = ModelCatalog.get_action_dist( + self.action_space, self.config["model"], framework=self.framework + ) + elif is_overridden(self.make_model_and_action_dist): + model, dist_class = self.make_model_and_action_dist() + else: + dist_class, logit_dim = ModelCatalog.get_action_dist( + self.action_space, self.config["model"], framework=self.framework + ) + model = ModelCatalog.get_model_v2( + obs_space=self.observation_space, + action_space=self.action_space, + num_outputs=logit_dim, + model_config=self.config["model"], + framework=self.framework, + ) + + # Compile the model, if requested by the user. + if self.config.get("torch_compile_learner"): + if ( + torch is not None + and version.parse(torch.__version__) < TORCH_COMPILE_REQUIRED_VERSION + ): + raise ValueError("`torch.compile` is not supported for torch < 2.0.0!") + + lw = "learner" if self.config.get("worker_index") else "worker" + model = torch.compile( + model, + backend=self.config.get( + f"torch_compile_{lw}_dynamo_backend", "inductor" + ), + dynamic=False, + mode=self.config.get(f"torch_compile_{lw}_dynamo_mode"), + ) + return model, dist_class + + @override(Policy) + def compute_actions_from_input_dict( + self, + input_dict: Dict[str, TensorType], + explore: bool = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + + seq_lens = None + with torch.no_grad(): + # Pass lazy (torch) tensor dict to Model as `input_dict`. + input_dict = self._lazy_tensor_dict(input_dict) + input_dict.set_training(True) + # Pack internal state inputs into (separate) list. + state_batches = [ + input_dict[k] for k in input_dict.keys() if "state_in" in k[:8] + ] + # Calculate RNN sequence lengths. + if state_batches: + seq_lens = torch.tensor( + [1] * len(state_batches[0]), + dtype=torch.long, + device=state_batches[0].device, + ) + + return self._compute_action_helper( + input_dict, state_batches, seq_lens, explore, timestep + ) + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorStructType], TensorStructType] = None, + prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None, + info_batch: Optional[Dict[str, list]] = None, + episodes=None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + **kwargs, + ) -> Tuple[TensorStructType, List[TensorType], Dict[str, TensorType]]: + + with torch.no_grad(): + seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) + input_dict = self._lazy_tensor_dict( + { + SampleBatch.CUR_OBS: obs_batch, + "is_training": False, + } + ) + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = np.asarray(prev_action_batch) + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = np.asarray(prev_reward_batch) + state_batches = [ + convert_to_torch_tensor(s, self.device) for s in (state_batches or []) + ] + return self._compute_action_helper( + input_dict, state_batches, seq_lens, explore, timestep + ) + + @with_lock + @override(Policy) + def compute_log_likelihoods( + self, + actions: Union[List[TensorStructType], TensorStructType], + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Optional[ + Union[List[TensorStructType], TensorStructType] + ] = None, + prev_reward_batch: Optional[ + Union[List[TensorStructType], TensorStructType] + ] = None, + actions_normalized: bool = True, + in_training: bool = True, + ) -> TensorType: + + if is_overridden(self.action_sampler_fn) and not is_overridden( + self.action_distribution_fn + ): + raise ValueError( + "Cannot compute log-prob/likelihood w/o an " + "`action_distribution_fn` and a provided " + "`action_sampler_fn`!" + ) + + with torch.no_grad(): + input_dict = self._lazy_tensor_dict( + {SampleBatch.CUR_OBS: obs_batch, SampleBatch.ACTIONS: actions} + ) + if prev_action_batch is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch + if prev_reward_batch is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch + seq_lens = torch.ones(len(obs_batch), dtype=torch.int32) + state_batches = [ + convert_to_torch_tensor(s, self.device) for s in (state_batches or []) + ] + + if self.exploration: + # Exploration hook before each forward pass. + self.exploration.before_compute_actions(explore=False) + + # Action dist class and inputs are generated via custom function. + if is_overridden(self.action_distribution_fn): + dist_inputs, dist_class, state_out = self.action_distribution_fn( + self.model, + obs_batch=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=False, + is_training=False, + ) + action_dist = dist_class(dist_inputs, self.model) + # Default action-dist inputs calculation. + else: + dist_class = self.dist_class + dist_inputs, _ = self.model(input_dict, state_batches, seq_lens) + + action_dist = dist_class(dist_inputs, self.model) + + # Normalize actions if necessary. + actions = input_dict[SampleBatch.ACTIONS] + if not actions_normalized and self.config["normalize_actions"]: + actions = normalize_action(actions, self.action_space_struct) + + log_likelihoods = action_dist.logp(actions) + + return log_likelihoods + + @with_lock + @override(Policy) + def learn_on_batch(self, postprocessed_batch: SampleBatch) -> Dict[str, TensorType]: + + # Set Model to train mode. + if self.model: + self.model.train() + # Callback handling. + learn_stats = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=postprocessed_batch, result=learn_stats + ) + + # Compute gradients (will calculate all losses and `backward()` + # them to get the grads). + grads, fetches = self.compute_gradients(postprocessed_batch) + + # Step the optimizers. + self.apply_gradients(_directStepOptimizerSingleton) + + self.num_grad_updates += 1 + if self.model and hasattr(self.model, "metrics"): + fetches["model"] = self.model.metrics() + else: + fetches["model"] = {} + + fetches.update( + { + "custom_metrics": learn_stats, + NUM_AGENT_STEPS_TRAINED: postprocessed_batch.count, + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates + - 1 + - (postprocessed_batch.num_grad_updates or 0) + ), + } + ) + + return fetches + + @override(Policy) + def load_batch_into_buffer( + self, + batch: SampleBatch, + buffer_index: int = 0, + ) -> int: + # Set the is_training flag of the batch. + batch.set_training(True) + + # Shortcut for 1 CPU only: Store batch in `self._loaded_batches`. + if len(self.devices) == 1 and self.devices[0].type == "cpu": + assert buffer_index == 0 + pad_batch_to_sequences_of_same_size( + batch=batch, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + _enable_new_api_stack=False, + padding="zero", + ) + self._lazy_tensor_dict(batch) + self._loaded_batches[0] = [batch] + return len(batch) + + # Batch (len=28, seq-lens=[4, 7, 4, 10, 3]): + # 0123 0123456 0123 0123456789ABC + + # 1) split into n per-GPU sub batches (n=2). + # [0123 0123456] [012] [3 0123456789 ABC] + # (len=14, 14 seq-lens=[4, 7, 3] [1, 10, 3]) + slices = batch.timeslices(num_slices=len(self.devices)) + + # 2) zero-padding (max-seq-len=10). + # - [0123000000 0123456000 0120000000] + # - [3000000000 0123456789 ABC0000000] + for slice in slices: + pad_batch_to_sequences_of_same_size( + batch=slice, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + _enable_new_api_stack=False, + padding="zero", + ) + + # 3) Load splits into the given buffer (consisting of n GPUs). + slices = [slice.to_device(self.devices[i]) for i, slice in enumerate(slices)] + self._loaded_batches[buffer_index] = slices + + # Return loaded samples per-device. + return len(slices[0]) + + @override(Policy) + def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int: + if len(self.devices) == 1 and self.devices[0] == "/cpu:0": + assert buffer_index == 0 + return sum(len(b) for b in self._loaded_batches[buffer_index]) + + @override(Policy) + def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0): + if not self._loaded_batches[buffer_index]: + raise ValueError( + "Must call Policy.load_batch_into_buffer() before " + "Policy.learn_on_loaded_batch()!" + ) + + # Get the correct slice of the already loaded batch to use, + # based on offset and batch size. + device_batch_size = self.config.get("minibatch_size") + if device_batch_size is None: + device_batch_size = self.config.get( + "sgd_minibatch_size", + self.config["train_batch_size"], + ) + device_batch_size //= len(self.devices) + + # Set Model to train mode. + if self.model_gpu_towers: + for t in self.model_gpu_towers: + t.train() + + # Shortcut for 1 CPU only: Batch should already be stored in + # `self._loaded_batches`. + if len(self.devices) == 1 and self.devices[0].type == "cpu": + assert buffer_index == 0 + if device_batch_size >= len(self._loaded_batches[0][0]): + batch = self._loaded_batches[0][0] + else: + batch = self._loaded_batches[0][0][offset : offset + device_batch_size] + + return self.learn_on_batch(batch) + + if len(self.devices) > 1: + # Copy weights of main model (tower-0) to all other towers. + state_dict = self.model.state_dict() + # Just making sure tower-0 is really the same as self.model. + assert self.model_gpu_towers[0] is self.model + for tower in self.model_gpu_towers[1:]: + tower.load_state_dict(state_dict) + + if device_batch_size >= sum(len(s) for s in self._loaded_batches[buffer_index]): + device_batches = self._loaded_batches[buffer_index] + else: + device_batches = [ + b[offset : offset + device_batch_size] + for b in self._loaded_batches[buffer_index] + ] + + # Callback handling. + batch_fetches = {} + for i, batch in enumerate(device_batches): + custom_metrics = {} + self.callbacks.on_learn_on_batch( + policy=self, train_batch=batch, result=custom_metrics + ) + batch_fetches[f"tower_{i}"] = {"custom_metrics": custom_metrics} + + # Do the (maybe parallelized) gradient calculation step. + tower_outputs = self._multi_gpu_parallel_grad_calc(device_batches) + + # Mean-reduce gradients over GPU-towers (do this on CPU: self.device). + all_grads = [] + for i in range(len(tower_outputs[0][0])): + if tower_outputs[0][0][i] is not None: + all_grads.append( + torch.mean( + torch.stack([t[0][i].to(self.device) for t in tower_outputs]), + dim=0, + ) + ) + else: + all_grads.append(None) + # Set main model's grads to mean-reduced values. + for i, p in enumerate(self.model.parameters()): + p.grad = all_grads[i] + + self.apply_gradients(_directStepOptimizerSingleton) + + self.num_grad_updates += 1 + + for i, (model, batch) in enumerate(zip(self.model_gpu_towers, device_batches)): + batch_fetches[f"tower_{i}"].update( + { + LEARNER_STATS_KEY: self.stats_fn(batch), + "model": model.metrics(), + NUM_GRAD_UPDATES_LIFETIME: self.num_grad_updates, + # -1, b/c we have to measure this diff before we do the update + # above. + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY: ( + self.num_grad_updates - 1 - (batch.num_grad_updates or 0) + ), + } + ) + batch_fetches.update(self.extra_compute_grad_fetches()) + + return batch_fetches + + @with_lock + @override(Policy) + def compute_gradients(self, postprocessed_batch: SampleBatch) -> ModelGradients: + + assert len(self.devices) == 1 + + # If not done yet, see whether we have to zero-pad this batch. + if not postprocessed_batch.zero_padded: + pad_batch_to_sequences_of_same_size( + batch=postprocessed_batch, + max_seq_len=self.max_seq_len, + shuffle=False, + batch_divisibility_req=self.batch_divisibility_req, + view_requirements=self.view_requirements, + _enable_new_api_stack=False, + padding="zero", + ) + + postprocessed_batch.set_training(True) + self._lazy_tensor_dict(postprocessed_batch, device=self.devices[0]) + + # Do the (maybe parallelized) gradient calculation step. + tower_outputs = self._multi_gpu_parallel_grad_calc([postprocessed_batch]) + + all_grads, grad_info = tower_outputs[0] + + grad_info["allreduce_latency"] /= len(self._optimizers) + grad_info.update(self.stats_fn(postprocessed_batch)) + + fetches = self.extra_compute_grad_fetches() + + return all_grads, dict(fetches, **{LEARNER_STATS_KEY: grad_info}) + + @override(Policy) + def apply_gradients(self, gradients: ModelGradients) -> None: + if gradients == _directStepOptimizerSingleton: + for i, opt in enumerate(self._optimizers): + opt.step() + else: + # TODO(sven): Not supported for multiple optimizers yet. + assert len(self._optimizers) == 1 + for g, p in zip(gradients, self.model.parameters()): + if g is not None: + if torch.is_tensor(g): + p.grad = g.to(self.device) + else: + p.grad = torch.from_numpy(g).to(self.device) + + self._optimizers[0].step() + + def get_tower_stats(self, stats_name: str) -> List[TensorStructType]: + """Returns list of per-tower stats, copied to this Policy's device. + + Args: + stats_name: The name of the stats to average over (this str + must exist as a key inside each tower's `tower_stats` dict). + + Returns: + The list of stats tensor (structs) of all towers, copied to this + Policy's device. + + Raises: + AssertionError: If the `stats_name` cannot be found in any one + of the tower's `tower_stats` dicts. + """ + data = [] + for model in self.model_gpu_towers: + if self.tower_stats: + tower_stats = self.tower_stats[model] + else: + tower_stats = model.tower_stats + + if stats_name in tower_stats: + data.append( + tree.map_structure( + lambda s: s.to(self.device), tower_stats[stats_name] + ) + ) + + assert len(data) > 0, ( + f"Stats `{stats_name}` not found in any of the towers (you have " + f"{len(self.model_gpu_towers)} towers in total)! Make " + "sure you call the loss function on at least one of the towers." + ) + return data + + @override(Policy) + def get_weights(self) -> ModelWeights: + return {k: v.cpu().detach().numpy() for k, v in self.model.state_dict().items()} + + @override(Policy) + def set_weights(self, weights: ModelWeights) -> None: + weights = convert_to_torch_tensor(weights, device=self.device) + self.model.load_state_dict(weights) + + @override(Policy) + def is_recurrent(self) -> bool: + return self._is_recurrent + + @override(Policy) + def num_state_tensors(self) -> int: + return len(self.model.get_initial_state()) + + @override(Policy) + def get_initial_state(self) -> List[TensorType]: + return [s.detach().cpu().numpy() for s in self.model.get_initial_state()] + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def get_state(self) -> PolicyState: + # Legacy Policy state (w/o torch.nn.Module and w/o PolicySpec). + state = super().get_state() + + state["_optimizer_variables"] = [] + for i, o in enumerate(self._optimizers): + optim_state_dict = convert_to_numpy(o.state_dict()) + state["_optimizer_variables"].append(optim_state_dict) + # Add exploration state. + if self.exploration: + # This is not compatible with RLModules, which have a method + # `forward_exploration` to specify custom exploration behavior. + state["_exploration_state"] = self.exploration.get_state() + return state + + @override(Policy) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def set_state(self, state: PolicyState) -> None: + # Set optimizer vars first. + optimizer_vars = state.get("_optimizer_variables", None) + if optimizer_vars: + assert len(optimizer_vars) == len(self._optimizers) + for o, s in zip(self._optimizers, optimizer_vars): + # Torch optimizer param_groups include things like beta, etc. These + # parameters should be left as scalar and not converted to tensors. + # otherwise, torch.optim.step() will start to complain. + optim_state_dict = {"param_groups": s["param_groups"]} + optim_state_dict["state"] = convert_to_torch_tensor( + s["state"], device=self.device + ) + o.load_state_dict(optim_state_dict) + # Set exploration's state. + if hasattr(self, "exploration") and "_exploration_state" in state: + self.exploration.set_state(state=state["_exploration_state"]) + + # Restore global timestep. + self.global_timestep = state["global_timestep"] + + # Then the Policy's (NN) weights and connectors. + super().set_state(state) + + @override(Policy) + def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None: + """Exports the Policy's Model to local directory for serving. + + Creates a TorchScript model and saves it. + + Args: + export_dir: Local writable directory or filename. + onnx: If given, will export model in ONNX format. The + value of this parameter set the ONNX OpSet version to use. + """ + + os.makedirs(export_dir, exist_ok=True) + + if onnx: + self._lazy_tensor_dict(self._dummy_batch) + # Provide dummy state inputs if not an RNN (torch cannot jit with + # returned empty internal states list). + if "state_in_0" not in self._dummy_batch: + self._dummy_batch["state_in_0"] = self._dummy_batch[ + SampleBatch.SEQ_LENS + ] = np.array([1.0]) + seq_lens = self._dummy_batch[SampleBatch.SEQ_LENS] + + state_ins = [] + i = 0 + while "state_in_{}".format(i) in self._dummy_batch: + state_ins.append(self._dummy_batch["state_in_{}".format(i)]) + i += 1 + dummy_inputs = { + k: self._dummy_batch[k] + for k in self._dummy_batch.keys() + if k != "is_training" + } + + file_name = os.path.join(export_dir, "model.onnx") + torch.onnx.export( + self.model, + (dummy_inputs, state_ins, seq_lens), + file_name, + export_params=True, + opset_version=onnx, + do_constant_folding=True, + input_names=list(dummy_inputs.keys()) + + ["state_ins", SampleBatch.SEQ_LENS], + output_names=["output", "state_outs"], + dynamic_axes={ + k: {0: "batch_size"} + for k in list(dummy_inputs.keys()) + + ["state_ins", SampleBatch.SEQ_LENS] + }, + ) + # Save the torch.Model (architecture and weights, so it can be retrieved + # w/o access to the original (custom) Model or Policy code). + else: + filename = os.path.join(export_dir, "model.pt") + try: + torch.save(self.model, f=filename) + except Exception: + if os.path.exists(filename): + os.remove(filename) + logger.warning(ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL) + + @override(Policy) + def import_model_from_h5(self, import_file: str) -> None: + """Imports weights into torch model.""" + return self.model.import_from_h5(import_file) + + @with_lock + def _compute_action_helper( + self, input_dict, state_batches, seq_lens, explore, timestep + ): + """Shared forward pass logic (w/ and w/o trajectory view API). + + Returns: + A tuple consisting of a) actions, b) state_out, c) extra_fetches. + The input_dict is modified in-place to include a numpy copy of the computed + actions under `SampleBatch.ACTIONS`. + """ + explore = explore if explore is not None else self.config["explore"] + timestep = timestep if timestep is not None else self.global_timestep + + # Switch to eval mode. + if self.model: + self.model.eval() + + extra_fetches = dist_inputs = logp = None + + if is_overridden(self.action_sampler_fn): + action_dist = None + actions, logp, dist_inputs, state_out = self.action_sampler_fn( + self.model, + obs_batch=input_dict, + state_batches=state_batches, + explore=explore, + timestep=timestep, + ) + else: + # Call the exploration before_compute_actions hook. + self.exploration.before_compute_actions(explore=explore, timestep=timestep) + if is_overridden(self.action_distribution_fn): + dist_inputs, dist_class, state_out = self.action_distribution_fn( + self.model, + obs_batch=input_dict, + state_batches=state_batches, + seq_lens=seq_lens, + explore=explore, + timestep=timestep, + is_training=False, + ) + else: + dist_class = self.dist_class + dist_inputs, state_out = self.model(input_dict, state_batches, seq_lens) + + if not ( + isinstance(dist_class, functools.partial) + or issubclass(dist_class, TorchDistributionWrapper) + ): + raise ValueError( + "`dist_class` ({}) not a TorchDistributionWrapper " + "subclass! Make sure your `action_distribution_fn` or " + "`make_model_and_action_dist` return a correct " + "distribution class.".format(dist_class.__name__) + ) + action_dist = dist_class(dist_inputs, self.model) + + # Get the exploration action from the forward results. + actions, logp = self.exploration.get_exploration_action( + action_distribution=action_dist, timestep=timestep, explore=explore + ) + + # Add default and custom fetches. + if extra_fetches is None: + extra_fetches = self.extra_action_out( + input_dict, state_batches, self.model, action_dist + ) + + # Action-dist inputs. + if dist_inputs is not None: + extra_fetches[SampleBatch.ACTION_DIST_INPUTS] = dist_inputs + + # Action-logp and action-prob. + if logp is not None: + extra_fetches[SampleBatch.ACTION_PROB] = torch.exp(logp.float()) + extra_fetches[SampleBatch.ACTION_LOGP] = logp + + # Update our global timestep by the batch size. + self.global_timestep += len(input_dict[SampleBatch.CUR_OBS]) + return convert_to_numpy((actions, state_out, extra_fetches)) + + def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch, device=None): + if not isinstance(postprocessed_batch, SampleBatch): + postprocessed_batch = SampleBatch(postprocessed_batch) + postprocessed_batch.set_get_interceptor( + functools.partial(convert_to_torch_tensor, device=device or self.device) + ) + return postprocessed_batch + + def _multi_gpu_parallel_grad_calc( + self, sample_batches: List[SampleBatch] + ) -> List[Tuple[List[TensorType], GradInfoDict]]: + """Performs a parallelized loss and gradient calculation over the batch. + + Splits up the given train batch into n shards (n=number of this + Policy's devices) and passes each data shard (in parallel) through + the loss function using the individual devices' models + (self.model_gpu_towers). Then returns each tower's outputs. + + Args: + sample_batches: A list of SampleBatch shards to + calculate loss and gradients for. + + Returns: + A list (one item per device) of 2-tuples, each with 1) gradient + list and 2) grad info dict. + """ + assert len(self.model_gpu_towers) == len(sample_batches) + lock = threading.Lock() + results = {} + grad_enabled = torch.is_grad_enabled() + + def _worker(shard_idx, model, sample_batch, device): + torch.set_grad_enabled(grad_enabled) + try: + with NullContextManager() if device.type == "cpu" else torch.cuda.device( # noqa: E501 + device + ): + loss_out = force_list( + self.loss(model, self.dist_class, sample_batch) + ) + + # Call Model's custom-loss with Policy loss outputs and + # train_batch. + if hasattr(model, "custom_loss"): + loss_out = model.custom_loss(loss_out, sample_batch) + + assert len(loss_out) == len(self._optimizers) + + # Loop through all optimizers. + grad_info = {"allreduce_latency": 0.0} + + parameters = list(model.parameters()) + all_grads = [None for _ in range(len(parameters))] + for opt_idx, opt in enumerate(self._optimizers): + # Erase gradients in all vars of the tower that this + # optimizer would affect. + param_indices = self.multi_gpu_param_groups[opt_idx] + for param_idx, param in enumerate(parameters): + if param_idx in param_indices and param.grad is not None: + param.grad.data.zero_() + # Recompute gradients of loss over all variables. + loss_out[opt_idx].backward(retain_graph=True) + grad_info.update( + self.extra_grad_process(opt, loss_out[opt_idx]) + ) + + grads = [] + # Note that return values are just references; + # Calling zero_grad would modify the values. + for param_idx, param in enumerate(parameters): + if param_idx in param_indices: + if param.grad is not None: + grads.append(param.grad) + all_grads[param_idx] = param.grad + + if self.distributed_world_size: + start = time.time() + if torch.cuda.is_available(): + # Sadly, allreduce_coalesced does not work with + # CUDA yet. + for g in grads: + torch.distributed.all_reduce( + g, op=torch.distributed.ReduceOp.SUM + ) + else: + torch.distributed.all_reduce_coalesced( + grads, op=torch.distributed.ReduceOp.SUM + ) + + for param_group in opt.param_groups: + for p in param_group["params"]: + if p.grad is not None: + p.grad /= self.distributed_world_size + + grad_info["allreduce_latency"] += time.time() - start + + with lock: + results[shard_idx] = (all_grads, grad_info) + except Exception as e: + import traceback + + with lock: + results[shard_idx] = ( + ValueError( + e.args[0] + + "\n traceback" + + traceback.format_exc() + + "\n" + + "In tower {} on device {}".format(shard_idx, device) + ), + e, + ) + + # Single device (GPU) or fake-GPU case (serialize for better + # debugging). + if len(self.devices) == 1 or self.config["_fake_gpus"]: + for shard_idx, (model, sample_batch, device) in enumerate( + zip(self.model_gpu_towers, sample_batches, self.devices) + ): + _worker(shard_idx, model, sample_batch, device) + # Raise errors right away for better debugging. + last_result = results[len(results) - 1] + if isinstance(last_result[0], ValueError): + raise last_result[0] from last_result[1] + # Multi device (GPU) case: Parallelize via threads. + else: + threads = [ + threading.Thread( + target=_worker, args=(shard_idx, model, sample_batch, device) + ) + for shard_idx, (model, sample_batch, device) in enumerate( + zip(self.model_gpu_towers, sample_batches, self.devices) + ) + ] + + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # Gather all threads' outputs and return. + outputs = [] + for shard_idx in range(len(sample_batches)): + output = results[shard_idx] + if isinstance(output[0], Exception): + raise output[0] from output[1] + outputs.append(results[shard_idx]) + return outputs diff --git a/lib/python3.12/site-packages/ray/rllib/policy/view_requirement.py b/lib/python3.12/site-packages/ray/rllib/policy/view_requirement.py new file mode 100644 index 0000000000000000000000000000000000000000..c3ff12ff40da6fab45d82a86f8844b4b2a7e342f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/policy/view_requirement.py @@ -0,0 +1,153 @@ +import dataclasses +from typing import Dict, List, Optional, Union + +import gymnasium as gym +import numpy as np + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.serialization import ( + gym_space_from_dict, + gym_space_to_dict, +) + +torch, _ = try_import_torch() + + +@OldAPIStack +@dataclasses.dataclass +class ViewRequirement: + """Single view requirement (for one column in an SampleBatch/input_dict). + + Policies and ModelV2s return a Dict[str, ViewRequirement] upon calling + their `[train|inference]_view_requirements()` methods, where the str key + represents the column name (C) under which the view is available in the + input_dict/SampleBatch and ViewRequirement specifies the actual underlying + column names (in the original data buffer), timestep shifts, and other + options to build the view. + + .. testcode:: + :skipif: True + + from ray.rllib.models.modelv2 import ModelV2 + # The default ViewRequirement for a Model is: + req = ModelV2(...).view_requirements + print(req) + + .. testoutput:: + + {"obs": ViewRequirement(shift=0)} + + Args: + data_col: The data column name from the SampleBatch + (str key). If None, use the dict key under which this + ViewRequirement resides. + space: The gym Space used in case we need to pad data + in inaccessible areas of the trajectory (t<0 or t>H). + Default: Simple box space, e.g. rewards. + shift: Single shift value or + list of relative positions to use (relative to the underlying + `data_col`). + Example: For a view column "prev_actions", you can set + `data_col="actions"` and `shift=-1`. + Example: For a view column "obs" in an Atari framestacking + fashion, you can set `data_col="obs"` and + `shift=[-3, -2, -1, 0]`. + Example: For the obs input to an attention net, you can specify + a range via a str: `shift="-100:0"`, which will pass in + the past 100 observations plus the current one. + index: An optional absolute position arg, + used e.g. for the location of a requested inference dict within + the trajectory. Negative values refer to counting from the end + of a trajectory. (#TODO: Is this still used?) + batch_repeat_value: determines how many time steps we should skip + before we repeat the view indexing for the next timestep. For RNNs this + number is usually the sequence length that we will rollout over. + Example: + view_col = "state_in_0", data_col = "state_out_0" + batch_repeat_value = 5, shift = -1 + buffer["state_out_0"] = [-1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + output["state_in_0"] = [-1, 4, 9] + Explanation: For t=0, we output buffer["state_out_0"][-1]. We then skip 5 + time steps and repeat the view. for t=5, we output buffer["state_out_0"][4] + . Continuing on this pattern, for t=10, we output buffer["state_out_0"][9]. + used_for_compute_actions: Whether the data will be used for + creating input_dicts for `Policy.compute_actions()` calls (or + `Policy.compute_actions_from_input_dict()`). + used_for_training: Whether the data will be used for + training. If False, the column will not be copied into the + final train batch. + """ + + data_col: Optional[str] = None + space: gym.Space = None + shift: Union[int, str, List[int]] = 0 + index: Optional[int] = None + batch_repeat_value: int = 1 + used_for_compute_actions: bool = True + used_for_training: bool = True + shift_arr: Optional[np.ndarray] = dataclasses.field(init=False) + + def __post_init__(self): + """Initializes a ViewRequirement object. + + shift_arr is infered from the shift value. + + For example: + - if shift is -1, then shift_arr is np.array([-1]). + - if shift is [-1, -2], then shift_arr is np.array([-2, -1]). + - if shift is "-2:2", then shift_arr is np.array([-2, -1, 0, 1, 2]). + """ + + if self.space is None: + self.space = gym.spaces.Box(float("-inf"), float("inf"), shape=()) + + # TODO: ideally we won't need shift_from and shift_to, and shift_step. + # all of them should be captured within shift_arr. + # Special case: Providing a (probably larger) range of indices, e.g. + # "-100:0" (past 100 timesteps plus current one). + self.shift_from = self.shift_to = self.shift_step = None + if isinstance(self.shift, str): + split = self.shift.split(":") + assert len(split) in [2, 3], f"Invalid shift str format: {self.shift}" + if len(split) == 2: + f, t = split + self.shift_step = 1 + else: + f, t, s = split + self.shift_step = int(s) + + self.shift_from = int(f) + self.shift_to = int(t) + + shift = self.shift + self.shfit_arr = None + if self.shift_from: + self.shift_arr = np.arange( + self.shift_from, self.shift_to + 1, self.shift_step + ) + else: + if isinstance(shift, int): + self.shift_arr = np.array([shift]) + elif isinstance(shift, list): + self.shift_arr = np.array(shift) + else: + ValueError(f'unrecognized shift type: "{shift}"') + + def to_dict(self) -> Dict: + """Return a dict for this ViewRequirement that can be JSON serialized.""" + return { + "data_col": self.data_col, + "space": gym_space_to_dict(self.space), + "shift": self.shift, + "index": self.index, + "batch_repeat_value": self.batch_repeat_value, + "used_for_training": self.used_for_training, + "used_for_compute_actions": self.used_for_compute_actions, + } + + @classmethod + def from_dict(cls, d: Dict): + """Construct a ViewRequirement instance from JSON deserialized dict.""" + d["space"] = gym_space_from_dict(d["space"]) + return cls(**d) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/INSTALLER b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/METADATA b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..a1b5c57543005b310ac2efb02243a6b6ab6d8932 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/METADATA @@ -0,0 +1,441 @@ +Metadata-Version: 2.1 +Name: colorama +Version: 0.4.6 +Summary: Cross-platform colored terminal text. +Project-URL: Homepage, https://github.com/tartley/colorama +Author-email: Jonathan Hartley +License-File: LICENSE.txt +Keywords: ansi,color,colour,crossplatform,terminal,text,windows,xplatform +Classifier: Development Status :: 5 - Production/Stable +Classifier: Environment :: Console +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: BSD License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 2 +Classifier: Programming Language :: Python :: 2.7 +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Terminals +Requires-Python: !=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7 +Description-Content-Type: text/x-rst + +.. image:: https://img.shields.io/pypi/v/colorama.svg + :target: https://pypi.org/project/colorama/ + :alt: Latest Version + +.. image:: https://img.shields.io/pypi/pyversions/colorama.svg + :target: https://pypi.org/project/colorama/ + :alt: Supported Python versions + +.. image:: https://github.com/tartley/colorama/actions/workflows/test.yml/badge.svg + :target: https://github.com/tartley/colorama/actions/workflows/test.yml + :alt: Build Status + +Colorama +======== + +Makes ANSI escape character sequences (for producing colored terminal text and +cursor positioning) work under MS Windows. + +.. |donate| image:: https://www.paypalobjects.com/en_US/i/btn/btn_donate_SM.gif + :target: https://www.paypal.com/cgi-bin/webscr?cmd=_donations&business=2MZ9D2GMLYCUJ&item_name=Colorama¤cy_code=USD + :alt: Donate with Paypal + +`PyPI for releases `_ | +`Github for source `_ | +`Colorama for enterprise on Tidelift `_ + +If you find Colorama useful, please |donate| to the authors. Thank you! + +Installation +------------ + +Tested on CPython 2.7, 3.7, 3.8, 3.9 and 3.10 and Pypy 2.7 and 3.8. + +No requirements other than the standard library. + +.. code-block:: bash + + pip install colorama + # or + conda install -c anaconda colorama + +Description +----------- + +ANSI escape character sequences have long been used to produce colored terminal +text and cursor positioning on Unix and Macs. Colorama makes this work on +Windows, too, by wrapping ``stdout``, stripping ANSI sequences it finds (which +would appear as gobbledygook in the output), and converting them into the +appropriate win32 calls to modify the state of the terminal. On other platforms, +Colorama does nothing. + +This has the upshot of providing a simple cross-platform API for printing +colored terminal text from Python, and has the happy side-effect that existing +applications or libraries which use ANSI sequences to produce colored output on +Linux or Macs can now also work on Windows, simply by calling +``colorama.just_fix_windows_console()`` (since v0.4.6) or ``colorama.init()`` +(all versions, but may have other side-effects – see below). + +An alternative approach is to install ``ansi.sys`` on Windows machines, which +provides the same behaviour for all applications running in terminals. Colorama +is intended for situations where that isn't easy (e.g., maybe your app doesn't +have an installer.) + +Demo scripts in the source code repository print some colored text using +ANSI sequences. Compare their output under Gnome-terminal's built in ANSI +handling, versus on Windows Command-Prompt using Colorama: + +.. image:: https://github.com/tartley/colorama/raw/master/screenshots/ubuntu-demo.png + :width: 661 + :height: 357 + :alt: ANSI sequences on Ubuntu under gnome-terminal. + +.. image:: https://github.com/tartley/colorama/raw/master/screenshots/windows-demo.png + :width: 668 + :height: 325 + :alt: Same ANSI sequences on Windows, using Colorama. + +These screenshots show that, on Windows, Colorama does not support ANSI 'dim +text'; it looks the same as 'normal text'. + +Usage +----- + +Initialisation +.............. + +If the only thing you want from Colorama is to get ANSI escapes to work on +Windows, then run: + +.. code-block:: python + + from colorama import just_fix_windows_console + just_fix_windows_console() + +If you're on a recent version of Windows 10 or better, and your stdout/stderr +are pointing to a Windows console, then this will flip the magic configuration +switch to enable Windows' built-in ANSI support. + +If you're on an older version of Windows, and your stdout/stderr are pointing to +a Windows console, then this will wrap ``sys.stdout`` and/or ``sys.stderr`` in a +magic file object that intercepts ANSI escape sequences and issues the +appropriate Win32 calls to emulate them. + +In all other circumstances, it does nothing whatsoever. Basically the idea is +that this makes Windows act like Unix with respect to ANSI escape handling. + +It's safe to call this function multiple times. It's safe to call this function +on non-Windows platforms, but it won't do anything. It's safe to call this +function when one or both of your stdout/stderr are redirected to a file – it +won't do anything to those streams. + +Alternatively, you can use the older interface with more features (but also more +potential footguns): + +.. code-block:: python + + from colorama import init + init() + +This does the same thing as ``just_fix_windows_console``, except for the +following differences: + +- It's not safe to call ``init`` multiple times; you can end up with multiple + layers of wrapping and broken ANSI support. + +- Colorama will apply a heuristic to guess whether stdout/stderr support ANSI, + and if it thinks they don't, then it will wrap ``sys.stdout`` and + ``sys.stderr`` in a magic file object that strips out ANSI escape sequences + before printing them. This happens on all platforms, and can be convenient if + you want to write your code to emit ANSI escape sequences unconditionally, and + let Colorama decide whether they should actually be output. But note that + Colorama's heuristic is not particularly clever. + +- ``init`` also accepts explicit keyword args to enable/disable various + functionality – see below. + +To stop using Colorama before your program exits, simply call ``deinit()``. +This will restore ``stdout`` and ``stderr`` to their original values, so that +Colorama is disabled. To resume using Colorama again, call ``reinit()``; it is +cheaper than calling ``init()`` again (but does the same thing). + +Most users should depend on ``colorama >= 0.4.6``, and use +``just_fix_windows_console``. The old ``init`` interface will be supported +indefinitely for backwards compatibility, but we don't plan to fix any issues +with it, also for backwards compatibility. + +Colored Output +.............. + +Cross-platform printing of colored text can then be done using Colorama's +constant shorthand for ANSI escape sequences. These are deliberately +rudimentary, see below. + +.. code-block:: python + + from colorama import Fore, Back, Style + print(Fore.RED + 'some red text') + print(Back.GREEN + 'and with a green background') + print(Style.DIM + 'and in dim text') + print(Style.RESET_ALL) + print('back to normal now') + +...or simply by manually printing ANSI sequences from your own code: + +.. code-block:: python + + print('\033[31m' + 'some red text') + print('\033[39m') # and reset to default color + +...or, Colorama can be used in conjunction with existing ANSI libraries +such as the venerable `Termcolor `_ +the fabulous `Blessings `_, +or the incredible `_Rich `_. + +If you wish Colorama's Fore, Back and Style constants were more capable, +then consider using one of the above highly capable libraries to generate +colors, etc, and use Colorama just for its primary purpose: to convert +those ANSI sequences to also work on Windows: + +SIMILARLY, do not send PRs adding the generation of new ANSI types to Colorama. +We are only interested in converting ANSI codes to win32 API calls, not +shortcuts like the above to generate ANSI characters. + +.. code-block:: python + + from colorama import just_fix_windows_console + from termcolor import colored + + # use Colorama to make Termcolor work on Windows too + just_fix_windows_console() + + # then use Termcolor for all colored text output + print(colored('Hello, World!', 'green', 'on_red')) + +Available formatting constants are:: + + Fore: BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE, RESET. + Back: BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE, RESET. + Style: DIM, NORMAL, BRIGHT, RESET_ALL + +``Style.RESET_ALL`` resets foreground, background, and brightness. Colorama will +perform this reset automatically on program exit. + +These are fairly well supported, but not part of the standard:: + + Fore: LIGHTBLACK_EX, LIGHTRED_EX, LIGHTGREEN_EX, LIGHTYELLOW_EX, LIGHTBLUE_EX, LIGHTMAGENTA_EX, LIGHTCYAN_EX, LIGHTWHITE_EX + Back: LIGHTBLACK_EX, LIGHTRED_EX, LIGHTGREEN_EX, LIGHTYELLOW_EX, LIGHTBLUE_EX, LIGHTMAGENTA_EX, LIGHTCYAN_EX, LIGHTWHITE_EX + +Cursor Positioning +.................. + +ANSI codes to reposition the cursor are supported. See ``demos/demo06.py`` for +an example of how to generate them. + +Init Keyword Args +................. + +``init()`` accepts some ``**kwargs`` to override default behaviour. + +init(autoreset=False): + If you find yourself repeatedly sending reset sequences to turn off color + changes at the end of every print, then ``init(autoreset=True)`` will + automate that: + + .. code-block:: python + + from colorama import init + init(autoreset=True) + print(Fore.RED + 'some red text') + print('automatically back to default color again') + +init(strip=None): + Pass ``True`` or ``False`` to override whether ANSI codes should be + stripped from the output. The default behaviour is to strip if on Windows + or if output is redirected (not a tty). + +init(convert=None): + Pass ``True`` or ``False`` to override whether to convert ANSI codes in the + output into win32 calls. The default behaviour is to convert if on Windows + and output is to a tty (terminal). + +init(wrap=True): + On Windows, Colorama works by replacing ``sys.stdout`` and ``sys.stderr`` + with proxy objects, which override the ``.write()`` method to do their work. + If this wrapping causes you problems, then this can be disabled by passing + ``init(wrap=False)``. The default behaviour is to wrap if ``autoreset`` or + ``strip`` or ``convert`` are True. + + When wrapping is disabled, colored printing on non-Windows platforms will + continue to work as normal. To do cross-platform colored output, you can + use Colorama's ``AnsiToWin32`` proxy directly: + + .. code-block:: python + + import sys + from colorama import init, AnsiToWin32 + init(wrap=False) + stream = AnsiToWin32(sys.stderr).stream + + # Python 2 + print >>stream, Fore.BLUE + 'blue text on stderr' + + # Python 3 + print(Fore.BLUE + 'blue text on stderr', file=stream) + +Recognised ANSI Sequences +......................... + +ANSI sequences generally take the form:: + + ESC [ ; ... + +Where ```` is an integer, and ```` is a single letter. Zero or +more params are passed to a ````. If no params are passed, it is +generally synonymous with passing a single zero. No spaces exist in the +sequence; they have been inserted here simply to read more easily. + +The only ANSI sequences that Colorama converts into win32 calls are:: + + ESC [ 0 m # reset all (colors and brightness) + ESC [ 1 m # bright + ESC [ 2 m # dim (looks same as normal brightness) + ESC [ 22 m # normal brightness + + # FOREGROUND: + ESC [ 30 m # black + ESC [ 31 m # red + ESC [ 32 m # green + ESC [ 33 m # yellow + ESC [ 34 m # blue + ESC [ 35 m # magenta + ESC [ 36 m # cyan + ESC [ 37 m # white + ESC [ 39 m # reset + + # BACKGROUND + ESC [ 40 m # black + ESC [ 41 m # red + ESC [ 42 m # green + ESC [ 43 m # yellow + ESC [ 44 m # blue + ESC [ 45 m # magenta + ESC [ 46 m # cyan + ESC [ 47 m # white + ESC [ 49 m # reset + + # cursor positioning + ESC [ y;x H # position cursor at x across, y down + ESC [ y;x f # position cursor at x across, y down + ESC [ n A # move cursor n lines up + ESC [ n B # move cursor n lines down + ESC [ n C # move cursor n characters forward + ESC [ n D # move cursor n characters backward + + # clear the screen + ESC [ mode J # clear the screen + + # clear the line + ESC [ mode K # clear the line + +Multiple numeric params to the ``'m'`` command can be combined into a single +sequence:: + + ESC [ 36 ; 45 ; 1 m # bright cyan text on magenta background + +All other ANSI sequences of the form ``ESC [ ; ... `` +are silently stripped from the output on Windows. + +Any other form of ANSI sequence, such as single-character codes or alternative +initial characters, are not recognised or stripped. It would be cool to add +them though. Let me know if it would be useful for you, via the Issues on +GitHub. + +Status & Known Problems +----------------------- + +I've personally only tested it on Windows XP (CMD, Console2), Ubuntu +(gnome-terminal, xterm), and OS X. + +Some valid ANSI sequences aren't recognised. + +If you're hacking on the code, see `README-hacking.md`_. ESPECIALLY, see the +explanation there of why we do not want PRs that allow Colorama to generate new +types of ANSI codes. + +See outstanding issues and wish-list: +https://github.com/tartley/colorama/issues + +If anything doesn't work for you, or doesn't do what you expected or hoped for, +I'd love to hear about it on that issues list, would be delighted by patches, +and would be happy to grant commit access to anyone who submits a working patch +or two. + +.. _README-hacking.md: README-hacking.md + +License +------- + +Copyright Jonathan Hartley & Arnon Yaari, 2013-2020. BSD 3-Clause license; see +LICENSE file. + +Professional support +-------------------- + +.. |tideliftlogo| image:: https://cdn2.hubspot.net/hubfs/4008838/website/logos/logos_for_download/Tidelift_primary-shorthand-logo.png + :alt: Tidelift + :target: https://tidelift.com/subscription/pkg/pypi-colorama?utm_source=pypi-colorama&utm_medium=referral&utm_campaign=readme + +.. list-table:: + :widths: 10 100 + + * - |tideliftlogo| + - Professional support for colorama is available as part of the + `Tidelift Subscription`_. + Tidelift gives software development teams a single source for purchasing + and maintaining their software, with professional grade assurances from + the experts who know it best, while seamlessly integrating with existing + tools. + +.. _Tidelift Subscription: https://tidelift.com/subscription/pkg/pypi-colorama?utm_source=pypi-colorama&utm_medium=referral&utm_campaign=readme + +Thanks +------ + +See the CHANGELOG for more thanks! + +* Marc Schlaich (schlamar) for a ``setup.py`` fix for Python2.5. +* Marc Abramowitz, reported & fixed a crash on exit with closed ``stdout``, + providing a solution to issue #7's setuptools/distutils debate, + and other fixes. +* User 'eryksun', for guidance on correctly instantiating ``ctypes.windll``. +* Matthew McCormick for politely pointing out a longstanding crash on non-Win. +* Ben Hoyt, for a magnificent fix under 64-bit Windows. +* Jesse at Empty Square for submitting a fix for examples in the README. +* User 'jamessp', an observant documentation fix for cursor positioning. +* User 'vaal1239', Dave Mckee & Lackner Kristof for a tiny but much-needed Win7 + fix. +* Julien Stuyck, for wisely suggesting Python3 compatible updates to README. +* Daniel Griffith for multiple fabulous patches. +* Oscar Lesta for a valuable fix to stop ANSI chars being sent to non-tty + output. +* Roger Binns, for many suggestions, valuable feedback, & bug reports. +* Tim Golden for thought and much appreciated feedback on the initial idea. +* User 'Zearin' for updates to the README file. +* John Szakmeister for adding support for light colors +* Charles Merriam for adding documentation to demos +* Jurko for a fix on 64-bit Windows CPython2.5 w/o ctypes +* Florian Bruhin for a fix when stdout or stderr are None +* Thomas Weininger for fixing ValueError on Windows +* Remi Rampin for better Github integration and fixes to the README file +* Simeon Visser for closing a file handle using 'with' and updating classifiers + to include Python 3.3 and 3.4 +* Andy Neff for fixing RESET of LIGHT_EX colors. +* Jonathan Hartley for the initial idea and implementation. diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/RECORD b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..f71babf13786131e6ba7735b9e4ebbc4e6860262 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/RECORD @@ -0,0 +1,32 @@ +colorama-0.4.6.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +colorama-0.4.6.dist-info/METADATA,sha256=e67SnrUMOym9sz_4TjF3vxvAV4T3aF7NyqRHHH3YEMw,17158 +colorama-0.4.6.dist-info/RECORD,, +colorama-0.4.6.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +colorama-0.4.6.dist-info/WHEEL,sha256=cdcF4Fbd0FPtw2EMIOwH-3rSOTUdTCeOSXRMD1iLUb8,105 +colorama-0.4.6.dist-info/licenses/LICENSE.txt,sha256=ysNcAmhuXQSlpxQL-zs25zrtSWZW6JEQLkKIhteTAxg,1491 +colorama/__init__.py,sha256=wePQA4U20tKgYARySLEC047ucNX-g8pRLpYBuiHlLb8,266 +colorama/__pycache__/__init__.cpython-312.pyc,, +colorama/__pycache__/ansi.cpython-312.pyc,, +colorama/__pycache__/ansitowin32.cpython-312.pyc,, +colorama/__pycache__/initialise.cpython-312.pyc,, +colorama/__pycache__/win32.cpython-312.pyc,, +colorama/__pycache__/winterm.cpython-312.pyc,, +colorama/ansi.py,sha256=Top4EeEuaQdBWdteKMEcGOTeKeF19Q-Wo_6_Cj5kOzQ,2522 +colorama/ansitowin32.py,sha256=vPNYa3OZbxjbuFyaVo0Tmhmy1FZ1lKMWCnT7odXpItk,11128 +colorama/initialise.py,sha256=-hIny86ClXo39ixh5iSCfUIa2f_h_bgKRDW7gqs-KLU,3325 +colorama/tests/__init__.py,sha256=MkgPAEzGQd-Rq0w0PZXSX2LadRWhUECcisJY8lSrm4Q,75 +colorama/tests/__pycache__/__init__.cpython-312.pyc,, +colorama/tests/__pycache__/ansi_test.cpython-312.pyc,, +colorama/tests/__pycache__/ansitowin32_test.cpython-312.pyc,, +colorama/tests/__pycache__/initialise_test.cpython-312.pyc,, +colorama/tests/__pycache__/isatty_test.cpython-312.pyc,, +colorama/tests/__pycache__/utils.cpython-312.pyc,, +colorama/tests/__pycache__/winterm_test.cpython-312.pyc,, +colorama/tests/ansi_test.py,sha256=FeViDrUINIZcr505PAxvU4AjXz1asEiALs9GXMhwRaE,2839 +colorama/tests/ansitowin32_test.py,sha256=RN7AIhMJ5EqDsYaCjVo-o4u8JzDD4ukJbmevWKS70rY,10678 +colorama/tests/initialise_test.py,sha256=BbPy-XfyHwJ6zKozuQOvNvQZzsx9vdb_0bYXn7hsBTc,6741 +colorama/tests/isatty_test.py,sha256=Pg26LRpv0yQDB5Ac-sxgVXG7hsA1NYvapFgApZfYzZg,1866 +colorama/tests/utils.py,sha256=1IIRylG39z5-dzq09R_ngufxyPZxgldNbrxKxUGwGKE,1079 +colorama/tests/winterm_test.py,sha256=qoWFPEjym5gm2RuMwpf3pOis3a5r_PJZFCzK254JL8A,3709 +colorama/win32.py,sha256=YQOKwMTwtGBbsY4dL5HYTvwTeP9wIQra5MvPNddpxZs,6181 +colorama/winterm.py,sha256=XCQFDHjPi6AHYNdZwy0tA02H-Jh48Jp-HvCjeLeLp3U,7134 diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/REQUESTED b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/WHEEL b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..d79189fda3251187de18c3998f23ae6fec11b20f --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: hatchling 1.11.1 +Root-Is-Purelib: true +Tag: py2-none-any +Tag: py3-none-any diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/licenses/LICENSE.txt b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/licenses/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..3105888ec149d10cad51c11d332779e94b548661 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama-0.4.6.dist-info/licenses/LICENSE.txt @@ -0,0 +1,27 @@ +Copyright (c) 2010 Jonathan Hartley +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holders, nor those of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__init__.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..383101cdb38706c305449674044e9288b92b7d75 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__init__.py @@ -0,0 +1,7 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. +from .initialise import init, deinit, reinit, colorama_text, just_fix_windows_console +from .ansi import Fore, Back, Style, Cursor +from .ansitowin32 import AnsiToWin32 + +__version__ = '0.4.6' + diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08bb5929264f4bb0d8781c1ee5894bc0d1875f6d Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansi.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansi.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b63001e994a0ed73efae5100174794de1c123953 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansi.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansitowin32.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansitowin32.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54835615436f5807f6a82e19eb21b40eab369f62 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/ansitowin32.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/initialise.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/initialise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be0c5867d08f7b29dc2e6c1da708ed4e530284fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/initialise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/win32.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/win32.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..02dd172f32329303df6e7c2f1165bd66a24f4d1d Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/win32.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/winterm.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/winterm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..814ab51007ffa8789f39559645513bb87f043828 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/__pycache__/winterm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansi.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansi.py new file mode 100644 index 0000000000000000000000000000000000000000..11ec695ff79627463a0282d25079527562de9e42 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansi.py @@ -0,0 +1,102 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. +''' +This module generates ANSI character codes to printing colors to terminals. +See: http://en.wikipedia.org/wiki/ANSI_escape_code +''' + +CSI = '\033[' +OSC = '\033]' +BEL = '\a' + + +def code_to_chars(code): + return CSI + str(code) + 'm' + +def set_title(title): + return OSC + '2;' + title + BEL + +def clear_screen(mode=2): + return CSI + str(mode) + 'J' + +def clear_line(mode=2): + return CSI + str(mode) + 'K' + + +class AnsiCodes(object): + def __init__(self): + # the subclasses declare class attributes which are numbers. + # Upon instantiation we define instance attributes, which are the same + # as the class attributes but wrapped with the ANSI escape sequence + for name in dir(self): + if not name.startswith('_'): + value = getattr(self, name) + setattr(self, name, code_to_chars(value)) + + +class AnsiCursor(object): + def UP(self, n=1): + return CSI + str(n) + 'A' + def DOWN(self, n=1): + return CSI + str(n) + 'B' + def FORWARD(self, n=1): + return CSI + str(n) + 'C' + def BACK(self, n=1): + return CSI + str(n) + 'D' + def POS(self, x=1, y=1): + return CSI + str(y) + ';' + str(x) + 'H' + + +class AnsiFore(AnsiCodes): + BLACK = 30 + RED = 31 + GREEN = 32 + YELLOW = 33 + BLUE = 34 + MAGENTA = 35 + CYAN = 36 + WHITE = 37 + RESET = 39 + + # These are fairly well supported, but not part of the standard. + LIGHTBLACK_EX = 90 + LIGHTRED_EX = 91 + LIGHTGREEN_EX = 92 + LIGHTYELLOW_EX = 93 + LIGHTBLUE_EX = 94 + LIGHTMAGENTA_EX = 95 + LIGHTCYAN_EX = 96 + LIGHTWHITE_EX = 97 + + +class AnsiBack(AnsiCodes): + BLACK = 40 + RED = 41 + GREEN = 42 + YELLOW = 43 + BLUE = 44 + MAGENTA = 45 + CYAN = 46 + WHITE = 47 + RESET = 49 + + # These are fairly well supported, but not part of the standard. + LIGHTBLACK_EX = 100 + LIGHTRED_EX = 101 + LIGHTGREEN_EX = 102 + LIGHTYELLOW_EX = 103 + LIGHTBLUE_EX = 104 + LIGHTMAGENTA_EX = 105 + LIGHTCYAN_EX = 106 + LIGHTWHITE_EX = 107 + + +class AnsiStyle(AnsiCodes): + BRIGHT = 1 + DIM = 2 + NORMAL = 22 + RESET_ALL = 0 + +Fore = AnsiFore() +Back = AnsiBack() +Style = AnsiStyle() +Cursor = AnsiCursor() diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansitowin32.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansitowin32.py new file mode 100644 index 0000000000000000000000000000000000000000..abf209e60c7c4a9b1ae57452e36b383969848c2e --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/ansitowin32.py @@ -0,0 +1,277 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. +import re +import sys +import os + +from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style, BEL +from .winterm import enable_vt_processing, WinTerm, WinColor, WinStyle +from .win32 import windll, winapi_test + + +winterm = None +if windll is not None: + winterm = WinTerm() + + +class StreamWrapper(object): + ''' + Wraps a stream (such as stdout), acting as a transparent proxy for all + attribute access apart from method 'write()', which is delegated to our + Converter instance. + ''' + def __init__(self, wrapped, converter): + # double-underscore everything to prevent clashes with names of + # attributes on the wrapped stream object. + self.__wrapped = wrapped + self.__convertor = converter + + def __getattr__(self, name): + return getattr(self.__wrapped, name) + + def __enter__(self, *args, **kwargs): + # special method lookup bypasses __getattr__/__getattribute__, see + # https://stackoverflow.com/questions/12632894/why-doesnt-getattr-work-with-exit + # thus, contextlib magic methods are not proxied via __getattr__ + return self.__wrapped.__enter__(*args, **kwargs) + + def __exit__(self, *args, **kwargs): + return self.__wrapped.__exit__(*args, **kwargs) + + def __setstate__(self, state): + self.__dict__ = state + + def __getstate__(self): + return self.__dict__ + + def write(self, text): + self.__convertor.write(text) + + def isatty(self): + stream = self.__wrapped + if 'PYCHARM_HOSTED' in os.environ: + if stream is not None and (stream is sys.__stdout__ or stream is sys.__stderr__): + return True + try: + stream_isatty = stream.isatty + except AttributeError: + return False + else: + return stream_isatty() + + @property + def closed(self): + stream = self.__wrapped + try: + return stream.closed + # AttributeError in the case that the stream doesn't support being closed + # ValueError for the case that the stream has already been detached when atexit runs + except (AttributeError, ValueError): + return True + + +class AnsiToWin32(object): + ''' + Implements a 'write()' method which, on Windows, will strip ANSI character + sequences from the text, and if outputting to a tty, will convert them into + win32 function calls. + ''' + ANSI_CSI_RE = re.compile('\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?') # Control Sequence Introducer + ANSI_OSC_RE = re.compile('\001?\033\\]([^\a]*)(\a)\002?') # Operating System Command + + def __init__(self, wrapped, convert=None, strip=None, autoreset=False): + # The wrapped stream (normally sys.stdout or sys.stderr) + self.wrapped = wrapped + + # should we reset colors to defaults after every .write() + self.autoreset = autoreset + + # create the proxy wrapping our output stream + self.stream = StreamWrapper(wrapped, self) + + on_windows = os.name == 'nt' + # We test if the WinAPI works, because even if we are on Windows + # we may be using a terminal that doesn't support the WinAPI + # (e.g. Cygwin Terminal). In this case it's up to the terminal + # to support the ANSI codes. + conversion_supported = on_windows and winapi_test() + try: + fd = wrapped.fileno() + except Exception: + fd = -1 + system_has_native_ansi = not on_windows or enable_vt_processing(fd) + have_tty = not self.stream.closed and self.stream.isatty() + need_conversion = conversion_supported and not system_has_native_ansi + + # should we strip ANSI sequences from our output? + if strip is None: + strip = need_conversion or not have_tty + self.strip = strip + + # should we should convert ANSI sequences into win32 calls? + if convert is None: + convert = need_conversion and have_tty + self.convert = convert + + # dict of ansi codes to win32 functions and parameters + self.win32_calls = self.get_win32_calls() + + # are we wrapping stderr? + self.on_stderr = self.wrapped is sys.stderr + + def should_wrap(self): + ''' + True if this class is actually needed. If false, then the output + stream will not be affected, nor will win32 calls be issued, so + wrapping stdout is not actually required. This will generally be + False on non-Windows platforms, unless optional functionality like + autoreset has been requested using kwargs to init() + ''' + return self.convert or self.strip or self.autoreset + + def get_win32_calls(self): + if self.convert and winterm: + return { + AnsiStyle.RESET_ALL: (winterm.reset_all, ), + AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT), + AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL), + AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL), + AnsiFore.BLACK: (winterm.fore, WinColor.BLACK), + AnsiFore.RED: (winterm.fore, WinColor.RED), + AnsiFore.GREEN: (winterm.fore, WinColor.GREEN), + AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW), + AnsiFore.BLUE: (winterm.fore, WinColor.BLUE), + AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA), + AnsiFore.CYAN: (winterm.fore, WinColor.CYAN), + AnsiFore.WHITE: (winterm.fore, WinColor.GREY), + AnsiFore.RESET: (winterm.fore, ), + AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True), + AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True), + AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True), + AnsiFore.LIGHTYELLOW_EX: (winterm.fore, WinColor.YELLOW, True), + AnsiFore.LIGHTBLUE_EX: (winterm.fore, WinColor.BLUE, True), + AnsiFore.LIGHTMAGENTA_EX: (winterm.fore, WinColor.MAGENTA, True), + AnsiFore.LIGHTCYAN_EX: (winterm.fore, WinColor.CYAN, True), + AnsiFore.LIGHTWHITE_EX: (winterm.fore, WinColor.GREY, True), + AnsiBack.BLACK: (winterm.back, WinColor.BLACK), + AnsiBack.RED: (winterm.back, WinColor.RED), + AnsiBack.GREEN: (winterm.back, WinColor.GREEN), + AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW), + AnsiBack.BLUE: (winterm.back, WinColor.BLUE), + AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA), + AnsiBack.CYAN: (winterm.back, WinColor.CYAN), + AnsiBack.WHITE: (winterm.back, WinColor.GREY), + AnsiBack.RESET: (winterm.back, ), + AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True), + AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True), + AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True), + AnsiBack.LIGHTYELLOW_EX: (winterm.back, WinColor.YELLOW, True), + AnsiBack.LIGHTBLUE_EX: (winterm.back, WinColor.BLUE, True), + AnsiBack.LIGHTMAGENTA_EX: (winterm.back, WinColor.MAGENTA, True), + AnsiBack.LIGHTCYAN_EX: (winterm.back, WinColor.CYAN, True), + AnsiBack.LIGHTWHITE_EX: (winterm.back, WinColor.GREY, True), + } + return dict() + + def write(self, text): + if self.strip or self.convert: + self.write_and_convert(text) + else: + self.wrapped.write(text) + self.wrapped.flush() + if self.autoreset: + self.reset_all() + + + def reset_all(self): + if self.convert: + self.call_win32('m', (0,)) + elif not self.strip and not self.stream.closed: + self.wrapped.write(Style.RESET_ALL) + + + def write_and_convert(self, text): + ''' + Write the given text to our wrapped stream, stripping any ANSI + sequences from the text, and optionally converting them into win32 + calls. + ''' + cursor = 0 + text = self.convert_osc(text) + for match in self.ANSI_CSI_RE.finditer(text): + start, end = match.span() + self.write_plain_text(text, cursor, start) + self.convert_ansi(*match.groups()) + cursor = end + self.write_plain_text(text, cursor, len(text)) + + + def write_plain_text(self, text, start, end): + if start < end: + self.wrapped.write(text[start:end]) + self.wrapped.flush() + + + def convert_ansi(self, paramstring, command): + if self.convert: + params = self.extract_params(command, paramstring) + self.call_win32(command, params) + + + def extract_params(self, command, paramstring): + if command in 'Hf': + params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';')) + while len(params) < 2: + # defaults: + params = params + (1,) + else: + params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0) + if len(params) == 0: + # defaults: + if command in 'JKm': + params = (0,) + elif command in 'ABCD': + params = (1,) + + return params + + + def call_win32(self, command, params): + if command == 'm': + for param in params: + if param in self.win32_calls: + func_args = self.win32_calls[param] + func = func_args[0] + args = func_args[1:] + kwargs = dict(on_stderr=self.on_stderr) + func(*args, **kwargs) + elif command in 'J': + winterm.erase_screen(params[0], on_stderr=self.on_stderr) + elif command in 'K': + winterm.erase_line(params[0], on_stderr=self.on_stderr) + elif command in 'Hf': # cursor position - absolute + winterm.set_cursor_position(params, on_stderr=self.on_stderr) + elif command in 'ABCD': # cursor position - relative + n = params[0] + # A - up, B - down, C - forward, D - back + x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command] + winterm.cursor_adjust(x, y, on_stderr=self.on_stderr) + + + def convert_osc(self, text): + for match in self.ANSI_OSC_RE.finditer(text): + start, end = match.span() + text = text[:start] + text[end:] + paramstring, command = match.groups() + if command == BEL: + if paramstring.count(";") == 1: + params = paramstring.split(";") + # 0 - change title and icon (we will only change title) + # 1 - change icon (we don't support this) + # 2 - change title + if params[0] in '02': + winterm.set_title(params[1]) + return text + + + def flush(self): + self.wrapped.flush() diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/initialise.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/initialise.py new file mode 100644 index 0000000000000000000000000000000000000000..d5fd4b71fed1bb4871717f978f0c470280f099c1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/initialise.py @@ -0,0 +1,121 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. +import atexit +import contextlib +import sys + +from .ansitowin32 import AnsiToWin32 + + +def _wipe_internal_state_for_tests(): + global orig_stdout, orig_stderr + orig_stdout = None + orig_stderr = None + + global wrapped_stdout, wrapped_stderr + wrapped_stdout = None + wrapped_stderr = None + + global atexit_done + atexit_done = False + + global fixed_windows_console + fixed_windows_console = False + + try: + # no-op if it wasn't registered + atexit.unregister(reset_all) + except AttributeError: + # python 2: no atexit.unregister. Oh well, we did our best. + pass + + +def reset_all(): + if AnsiToWin32 is not None: # Issue #74: objects might become None at exit + AnsiToWin32(orig_stdout).reset_all() + + +def init(autoreset=False, convert=None, strip=None, wrap=True): + + if not wrap and any([autoreset, convert, strip]): + raise ValueError('wrap=False conflicts with any other arg=True') + + global wrapped_stdout, wrapped_stderr + global orig_stdout, orig_stderr + + orig_stdout = sys.stdout + orig_stderr = sys.stderr + + if sys.stdout is None: + wrapped_stdout = None + else: + sys.stdout = wrapped_stdout = \ + wrap_stream(orig_stdout, convert, strip, autoreset, wrap) + if sys.stderr is None: + wrapped_stderr = None + else: + sys.stderr = wrapped_stderr = \ + wrap_stream(orig_stderr, convert, strip, autoreset, wrap) + + global atexit_done + if not atexit_done: + atexit.register(reset_all) + atexit_done = True + + +def deinit(): + if orig_stdout is not None: + sys.stdout = orig_stdout + if orig_stderr is not None: + sys.stderr = orig_stderr + + +def just_fix_windows_console(): + global fixed_windows_console + + if sys.platform != "win32": + return + if fixed_windows_console: + return + if wrapped_stdout is not None or wrapped_stderr is not None: + # Someone already ran init() and it did stuff, so we won't second-guess them + return + + # On newer versions of Windows, AnsiToWin32.__init__ will implicitly enable the + # native ANSI support in the console as a side-effect. We only need to actually + # replace sys.stdout/stderr if we're in the old-style conversion mode. + new_stdout = AnsiToWin32(sys.stdout, convert=None, strip=None, autoreset=False) + if new_stdout.convert: + sys.stdout = new_stdout + new_stderr = AnsiToWin32(sys.stderr, convert=None, strip=None, autoreset=False) + if new_stderr.convert: + sys.stderr = new_stderr + + fixed_windows_console = True + +@contextlib.contextmanager +def colorama_text(*args, **kwargs): + init(*args, **kwargs) + try: + yield + finally: + deinit() + + +def reinit(): + if wrapped_stdout is not None: + sys.stdout = wrapped_stdout + if wrapped_stderr is not None: + sys.stderr = wrapped_stderr + + +def wrap_stream(stream, convert, strip, autoreset, wrap): + if wrap: + wrapper = AnsiToWin32(stream, + convert=convert, strip=strip, autoreset=autoreset) + if wrapper.should_wrap(): + stream = wrapper.stream + return stream + + +# Use this for initial setup as well, to reduce code duplication +_wipe_internal_state_for_tests() diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/win32.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/win32.py new file mode 100644 index 0000000000000000000000000000000000000000..841b0e270a381cdfaca544a9be976d7276d83b1e --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/win32.py @@ -0,0 +1,180 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. + +# from winbase.h +STDOUT = -11 +STDERR = -12 + +ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 + +try: + import ctypes + from ctypes import LibraryLoader + windll = LibraryLoader(ctypes.WinDLL) + from ctypes import wintypes +except (AttributeError, ImportError): + windll = None + SetConsoleTextAttribute = lambda *_: None + winapi_test = lambda *_: None +else: + from ctypes import byref, Structure, c_char, POINTER + + COORD = wintypes._COORD + + class CONSOLE_SCREEN_BUFFER_INFO(Structure): + """struct in wincon.h.""" + _fields_ = [ + ("dwSize", COORD), + ("dwCursorPosition", COORD), + ("wAttributes", wintypes.WORD), + ("srWindow", wintypes.SMALL_RECT), + ("dwMaximumWindowSize", COORD), + ] + def __str__(self): + return '(%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d)' % ( + self.dwSize.Y, self.dwSize.X + , self.dwCursorPosition.Y, self.dwCursorPosition.X + , self.wAttributes + , self.srWindow.Top, self.srWindow.Left, self.srWindow.Bottom, self.srWindow.Right + , self.dwMaximumWindowSize.Y, self.dwMaximumWindowSize.X + ) + + _GetStdHandle = windll.kernel32.GetStdHandle + _GetStdHandle.argtypes = [ + wintypes.DWORD, + ] + _GetStdHandle.restype = wintypes.HANDLE + + _GetConsoleScreenBufferInfo = windll.kernel32.GetConsoleScreenBufferInfo + _GetConsoleScreenBufferInfo.argtypes = [ + wintypes.HANDLE, + POINTER(CONSOLE_SCREEN_BUFFER_INFO), + ] + _GetConsoleScreenBufferInfo.restype = wintypes.BOOL + + _SetConsoleTextAttribute = windll.kernel32.SetConsoleTextAttribute + _SetConsoleTextAttribute.argtypes = [ + wintypes.HANDLE, + wintypes.WORD, + ] + _SetConsoleTextAttribute.restype = wintypes.BOOL + + _SetConsoleCursorPosition = windll.kernel32.SetConsoleCursorPosition + _SetConsoleCursorPosition.argtypes = [ + wintypes.HANDLE, + COORD, + ] + _SetConsoleCursorPosition.restype = wintypes.BOOL + + _FillConsoleOutputCharacterA = windll.kernel32.FillConsoleOutputCharacterA + _FillConsoleOutputCharacterA.argtypes = [ + wintypes.HANDLE, + c_char, + wintypes.DWORD, + COORD, + POINTER(wintypes.DWORD), + ] + _FillConsoleOutputCharacterA.restype = wintypes.BOOL + + _FillConsoleOutputAttribute = windll.kernel32.FillConsoleOutputAttribute + _FillConsoleOutputAttribute.argtypes = [ + wintypes.HANDLE, + wintypes.WORD, + wintypes.DWORD, + COORD, + POINTER(wintypes.DWORD), + ] + _FillConsoleOutputAttribute.restype = wintypes.BOOL + + _SetConsoleTitleW = windll.kernel32.SetConsoleTitleW + _SetConsoleTitleW.argtypes = [ + wintypes.LPCWSTR + ] + _SetConsoleTitleW.restype = wintypes.BOOL + + _GetConsoleMode = windll.kernel32.GetConsoleMode + _GetConsoleMode.argtypes = [ + wintypes.HANDLE, + POINTER(wintypes.DWORD) + ] + _GetConsoleMode.restype = wintypes.BOOL + + _SetConsoleMode = windll.kernel32.SetConsoleMode + _SetConsoleMode.argtypes = [ + wintypes.HANDLE, + wintypes.DWORD + ] + _SetConsoleMode.restype = wintypes.BOOL + + def _winapi_test(handle): + csbi = CONSOLE_SCREEN_BUFFER_INFO() + success = _GetConsoleScreenBufferInfo( + handle, byref(csbi)) + return bool(success) + + def winapi_test(): + return any(_winapi_test(h) for h in + (_GetStdHandle(STDOUT), _GetStdHandle(STDERR))) + + def GetConsoleScreenBufferInfo(stream_id=STDOUT): + handle = _GetStdHandle(stream_id) + csbi = CONSOLE_SCREEN_BUFFER_INFO() + success = _GetConsoleScreenBufferInfo( + handle, byref(csbi)) + return csbi + + def SetConsoleTextAttribute(stream_id, attrs): + handle = _GetStdHandle(stream_id) + return _SetConsoleTextAttribute(handle, attrs) + + def SetConsoleCursorPosition(stream_id, position, adjust=True): + position = COORD(*position) + # If the position is out of range, do nothing. + if position.Y <= 0 or position.X <= 0: + return + # Adjust for Windows' SetConsoleCursorPosition: + # 1. being 0-based, while ANSI is 1-based. + # 2. expecting (x,y), while ANSI uses (y,x). + adjusted_position = COORD(position.Y - 1, position.X - 1) + if adjust: + # Adjust for viewport's scroll position + sr = GetConsoleScreenBufferInfo(STDOUT).srWindow + adjusted_position.Y += sr.Top + adjusted_position.X += sr.Left + # Resume normal processing + handle = _GetStdHandle(stream_id) + return _SetConsoleCursorPosition(handle, adjusted_position) + + def FillConsoleOutputCharacter(stream_id, char, length, start): + handle = _GetStdHandle(stream_id) + char = c_char(char.encode()) + length = wintypes.DWORD(length) + num_written = wintypes.DWORD(0) + # Note that this is hard-coded for ANSI (vs wide) bytes. + success = _FillConsoleOutputCharacterA( + handle, char, length, start, byref(num_written)) + return num_written.value + + def FillConsoleOutputAttribute(stream_id, attr, length, start): + ''' FillConsoleOutputAttribute( hConsole, csbi.wAttributes, dwConSize, coordScreen, &cCharsWritten )''' + handle = _GetStdHandle(stream_id) + attribute = wintypes.WORD(attr) + length = wintypes.DWORD(length) + num_written = wintypes.DWORD(0) + # Note that this is hard-coded for ANSI (vs wide) bytes. + return _FillConsoleOutputAttribute( + handle, attribute, length, start, byref(num_written)) + + def SetConsoleTitle(title): + return _SetConsoleTitleW(title) + + def GetConsoleMode(handle): + mode = wintypes.DWORD() + success = _GetConsoleMode(handle, byref(mode)) + if not success: + raise ctypes.WinError() + return mode.value + + def SetConsoleMode(handle, mode): + success = _SetConsoleMode(handle, mode) + if not success: + raise ctypes.WinError() diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/colorama/winterm.py b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/winterm.py new file mode 100644 index 0000000000000000000000000000000000000000..aad867e8c80b826bf6a060116f17fa08a8eb0765 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/colorama/winterm.py @@ -0,0 +1,195 @@ +# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. +try: + from msvcrt import get_osfhandle +except ImportError: + def get_osfhandle(_): + raise OSError("This isn't windows!") + + +from . import win32 + +# from wincon.h +class WinColor(object): + BLACK = 0 + BLUE = 1 + GREEN = 2 + CYAN = 3 + RED = 4 + MAGENTA = 5 + YELLOW = 6 + GREY = 7 + +# from wincon.h +class WinStyle(object): + NORMAL = 0x00 # dim text, dim background + BRIGHT = 0x08 # bright text, dim background + BRIGHT_BACKGROUND = 0x80 # dim text, bright background + +class WinTerm(object): + + def __init__(self): + self._default = win32.GetConsoleScreenBufferInfo(win32.STDOUT).wAttributes + self.set_attrs(self._default) + self._default_fore = self._fore + self._default_back = self._back + self._default_style = self._style + # In order to emulate LIGHT_EX in windows, we borrow the BRIGHT style. + # So that LIGHT_EX colors and BRIGHT style do not clobber each other, + # we track them separately, since LIGHT_EX is overwritten by Fore/Back + # and BRIGHT is overwritten by Style codes. + self._light = 0 + + def get_attrs(self): + return self._fore + self._back * 16 + (self._style | self._light) + + def set_attrs(self, value): + self._fore = value & 7 + self._back = (value >> 4) & 7 + self._style = value & (WinStyle.BRIGHT | WinStyle.BRIGHT_BACKGROUND) + + def reset_all(self, on_stderr=None): + self.set_attrs(self._default) + self.set_console(attrs=self._default) + self._light = 0 + + def fore(self, fore=None, light=False, on_stderr=False): + if fore is None: + fore = self._default_fore + self._fore = fore + # Emulate LIGHT_EX with BRIGHT Style + if light: + self._light |= WinStyle.BRIGHT + else: + self._light &= ~WinStyle.BRIGHT + self.set_console(on_stderr=on_stderr) + + def back(self, back=None, light=False, on_stderr=False): + if back is None: + back = self._default_back + self._back = back + # Emulate LIGHT_EX with BRIGHT_BACKGROUND Style + if light: + self._light |= WinStyle.BRIGHT_BACKGROUND + else: + self._light &= ~WinStyle.BRIGHT_BACKGROUND + self.set_console(on_stderr=on_stderr) + + def style(self, style=None, on_stderr=False): + if style is None: + style = self._default_style + self._style = style + self.set_console(on_stderr=on_stderr) + + def set_console(self, attrs=None, on_stderr=False): + if attrs is None: + attrs = self.get_attrs() + handle = win32.STDOUT + if on_stderr: + handle = win32.STDERR + win32.SetConsoleTextAttribute(handle, attrs) + + def get_position(self, handle): + position = win32.GetConsoleScreenBufferInfo(handle).dwCursorPosition + # Because Windows coordinates are 0-based, + # and win32.SetConsoleCursorPosition expects 1-based. + position.X += 1 + position.Y += 1 + return position + + def set_cursor_position(self, position=None, on_stderr=False): + if position is None: + # I'm not currently tracking the position, so there is no default. + # position = self.get_position() + return + handle = win32.STDOUT + if on_stderr: + handle = win32.STDERR + win32.SetConsoleCursorPosition(handle, position) + + def cursor_adjust(self, x, y, on_stderr=False): + handle = win32.STDOUT + if on_stderr: + handle = win32.STDERR + position = self.get_position(handle) + adjusted_position = (position.Y + y, position.X + x) + win32.SetConsoleCursorPosition(handle, adjusted_position, adjust=False) + + def erase_screen(self, mode=0, on_stderr=False): + # 0 should clear from the cursor to the end of the screen. + # 1 should clear from the cursor to the beginning of the screen. + # 2 should clear the entire screen, and move cursor to (1,1) + handle = win32.STDOUT + if on_stderr: + handle = win32.STDERR + csbi = win32.GetConsoleScreenBufferInfo(handle) + # get the number of character cells in the current buffer + cells_in_screen = csbi.dwSize.X * csbi.dwSize.Y + # get number of character cells before current cursor position + cells_before_cursor = csbi.dwSize.X * csbi.dwCursorPosition.Y + csbi.dwCursorPosition.X + if mode == 0: + from_coord = csbi.dwCursorPosition + cells_to_erase = cells_in_screen - cells_before_cursor + elif mode == 1: + from_coord = win32.COORD(0, 0) + cells_to_erase = cells_before_cursor + elif mode == 2: + from_coord = win32.COORD(0, 0) + cells_to_erase = cells_in_screen + else: + # invalid mode + return + # fill the entire screen with blanks + win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord) + # now set the buffer's attributes accordingly + win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord) + if mode == 2: + # put the cursor where needed + win32.SetConsoleCursorPosition(handle, (1, 1)) + + def erase_line(self, mode=0, on_stderr=False): + # 0 should clear from the cursor to the end of the line. + # 1 should clear from the cursor to the beginning of the line. + # 2 should clear the entire line. + handle = win32.STDOUT + if on_stderr: + handle = win32.STDERR + csbi = win32.GetConsoleScreenBufferInfo(handle) + if mode == 0: + from_coord = csbi.dwCursorPosition + cells_to_erase = csbi.dwSize.X - csbi.dwCursorPosition.X + elif mode == 1: + from_coord = win32.COORD(0, csbi.dwCursorPosition.Y) + cells_to_erase = csbi.dwCursorPosition.X + elif mode == 2: + from_coord = win32.COORD(0, csbi.dwCursorPosition.Y) + cells_to_erase = csbi.dwSize.X + else: + # invalid mode + return + # fill the entire screen with blanks + win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord) + # now set the buffer's attributes accordingly + win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord) + + def set_title(self, title): + win32.SetConsoleTitle(title) + + +def enable_vt_processing(fd): + if win32.windll is None or not win32.winapi_test(): + return False + + try: + handle = get_osfhandle(fd) + mode = win32.GetConsoleMode(handle) + win32.SetConsoleMode( + handle, + mode | win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING, + ) + + mode = win32.GetConsoleMode(handle) + if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING: + return True + # Can get TypeError in testsuite where 'fd' is a Mock() + except (OSError, TypeError): + return False diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/INSTALLER b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/LICENSE b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..cff5eb74e1badd1c5237ed2654b349530179ad1d --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2009, Jay Loden, Dave Daeschler, Giampaolo Rodola +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name of the psutil authors nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/METADATA b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0c7c91f55fb7dec6d2fd9bd07084f9968a26840b --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/METADATA @@ -0,0 +1,553 @@ +Metadata-Version: 2.1 +Name: psutil +Version: 7.1.3 +Summary: Cross-platform lib for process and system monitoring. +Home-page: https://github.com/giampaolo/psutil +Author: Giampaolo Rodola +Author-email: g.rodola@gmail.com +License: BSD-3-Clause +Keywords: ps,top,kill,free,lsof,netstat,nice,tty,ionice,uptime,taskmgr,process,df,iotop,iostat,ifconfig,taskset,who,pidof,pmap,smem,pstree,monitoring,ulimit,prlimit,smem,performance,metrics,agent,observability +Platform: Platform Independent +Classifier: Development Status :: 5 - Production/Stable +Classifier: Environment :: Console +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Information Technology +Classifier: Intended Audience :: System Administrators +Classifier: Operating System :: MacOS :: MacOS X +Classifier: Operating System :: Microsoft :: Windows :: Windows 10 +Classifier: Operating System :: Microsoft :: Windows :: Windows 11 +Classifier: Operating System :: Microsoft :: Windows :: Windows 7 +Classifier: Operating System :: Microsoft :: Windows :: Windows 8 +Classifier: Operating System :: Microsoft :: Windows :: Windows 8.1 +Classifier: Operating System :: Microsoft :: Windows :: Windows Server 2003 +Classifier: Operating System :: Microsoft :: Windows :: Windows Server 2008 +Classifier: Operating System :: Microsoft :: Windows :: Windows Vista +Classifier: Operating System :: Microsoft :: Windows +Classifier: Operating System :: Microsoft +Classifier: Operating System :: OS Independent +Classifier: Operating System :: POSIX :: AIX +Classifier: Operating System :: POSIX :: BSD :: FreeBSD +Classifier: Operating System :: POSIX :: BSD :: NetBSD +Classifier: Operating System :: POSIX :: BSD :: OpenBSD +Classifier: Operating System :: POSIX :: BSD +Classifier: Operating System :: POSIX :: Linux +Classifier: Operating System :: POSIX :: SunOS/Solaris +Classifier: Operating System :: POSIX +Classifier: Programming Language :: C +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Programming Language :: Python +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Topic :: Software Development :: Libraries +Classifier: Topic :: System :: Benchmark +Classifier: Topic :: System :: Hardware +Classifier: Topic :: System :: Monitoring +Classifier: Topic :: System :: Networking :: Monitoring :: Hardware Watchdog +Classifier: Topic :: System :: Networking :: Monitoring +Classifier: Topic :: System :: Networking +Classifier: Topic :: System :: Operating System +Classifier: Topic :: System :: Systems Administration +Classifier: Topic :: Utilities +Requires-Python: >=3.6 +Description-Content-Type: text/x-rst +License-File: LICENSE +Provides-Extra: dev +Requires-Dist: pytest ; extra == 'dev' +Requires-Dist: pytest-instafail ; extra == 'dev' +Requires-Dist: pytest-subtests ; extra == 'dev' +Requires-Dist: pytest-xdist ; extra == 'dev' +Requires-Dist: setuptools ; extra == 'dev' +Requires-Dist: abi3audit ; extra == 'dev' +Requires-Dist: black ; extra == 'dev' +Requires-Dist: check-manifest ; extra == 'dev' +Requires-Dist: coverage ; extra == 'dev' +Requires-Dist: packaging ; extra == 'dev' +Requires-Dist: pylint ; extra == 'dev' +Requires-Dist: pyperf ; extra == 'dev' +Requires-Dist: pypinfo ; extra == 'dev' +Requires-Dist: pytest-cov ; extra == 'dev' +Requires-Dist: requests ; extra == 'dev' +Requires-Dist: rstcheck ; extra == 'dev' +Requires-Dist: ruff ; extra == 'dev' +Requires-Dist: sphinx ; extra == 'dev' +Requires-Dist: sphinx-rtd-theme ; extra == 'dev' +Requires-Dist: toml-sort ; extra == 'dev' +Requires-Dist: twine ; extra == 'dev' +Requires-Dist: validate-pyproject[all] ; extra == 'dev' +Requires-Dist: virtualenv ; extra == 'dev' +Requires-Dist: vulture ; extra == 'dev' +Requires-Dist: wheel ; extra == 'dev' +Requires-Dist: colorama ; (os_name == "nt") and extra == 'dev' +Requires-Dist: pyreadline ; (os_name == "nt") and extra == 'dev' +Requires-Dist: pywin32 ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'dev' +Requires-Dist: wheel ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'dev' +Requires-Dist: wmi ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'dev' +Provides-Extra: test +Requires-Dist: pytest ; extra == 'test' +Requires-Dist: pytest-instafail ; extra == 'test' +Requires-Dist: pytest-subtests ; extra == 'test' +Requires-Dist: pytest-xdist ; extra == 'test' +Requires-Dist: setuptools ; extra == 'test' +Requires-Dist: pywin32 ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'test' +Requires-Dist: wheel ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'test' +Requires-Dist: wmi ; (os_name == "nt" and platform_python_implementation != "PyPy") and extra == 'test' + +| |downloads| |stars| |forks| |contributors| |coverage| +| |version| |packages| |license| +| |github-actions-wheels| |github-actions-bsd| |doc| |twitter| |tidelift| + +.. |downloads| image:: https://img.shields.io/pypi/dm/psutil.svg + :target: https://clickpy.clickhouse.com/dashboard/psutil + :alt: Downloads + +.. |stars| image:: https://img.shields.io/github/stars/giampaolo/psutil.svg + :target: https://github.com/giampaolo/psutil/stargazers + :alt: Github stars + +.. |forks| image:: https://img.shields.io/github/forks/giampaolo/psutil.svg + :target: https://github.com/giampaolo/psutil/network/members + :alt: Github forks + +.. |contributors| image:: https://img.shields.io/github/contributors/giampaolo/psutil.svg + :target: https://github.com/giampaolo/psutil/graphs/contributors + :alt: Contributors + +.. |github-actions-wheels| image:: https://img.shields.io/github/actions/workflow/status/giampaolo/psutil/.github/workflows/build.yml.svg?label=Linux%2C%20macOS%2C%20Windows + :target: https://github.com/giampaolo/psutil/actions?query=workflow%3Abuild + :alt: Linux, macOS, Windows + +.. |github-actions-bsd| image:: https://img.shields.io/github/actions/workflow/status/giampaolo/psutil/.github/workflows/bsd.yml.svg?label=FreeBSD,%20NetBSD,%20OpenBSD + :target: https://github.com/giampaolo/psutil/actions?query=workflow%3Absd-tests + :alt: FreeBSD, NetBSD, OpenBSD + +.. |coverage| image:: https://coveralls.io/repos/github/giampaolo/psutil/badge.svg?branch=master + :target: https://coveralls.io/github/giampaolo/psutil?branch=master + :alt: Test coverage (coverall.io) + +.. |doc| image:: https://readthedocs.org/projects/psutil/badge/?version=latest + :target: https://psutil.readthedocs.io/en/latest/ + :alt: Documentation Status + +.. |version| image:: https://img.shields.io/pypi/v/psutil.svg?label=pypi + :target: https://pypi.org/project/psutil + :alt: Latest version + +.. |packages| image:: https://repology.org/badge/tiny-repos/python:psutil.svg + :target: https://repology.org/metapackage/python:psutil/versions + :alt: Binary packages + +.. |license| image:: https://img.shields.io/pypi/l/psutil.svg + :target: https://github.com/giampaolo/psutil/blob/master/LICENSE + :alt: License + +.. |twitter| image:: https://img.shields.io/twitter/follow/grodola.svg?label=follow&style=flat&logo=twitter&logoColor=4FADFF + :target: https://twitter.com/grodola + :alt: Twitter Follow + +.. |tidelift| image:: https://tidelift.com/badges/github/giampaolo/psutil?style=flat + :target: https://tidelift.com/subscription/pkg/pypi-psutil?utm_source=pypi-psutil&utm_medium=referral&utm_campaign=readme + :alt: Tidelift + +----- + +Quick links +=========== + +- `Home page `_ +- `Install `_ +- `Documentation `_ +- `Download `_ +- `Forum `_ +- `StackOverflow `_ +- `Blog `_ +- `What's new `_ + + +Summary +======= + +psutil (process and system utilities) is a cross-platform library for +retrieving information on **running processes** and **system utilization** +(CPU, memory, disks, network, sensors) in Python. +It is useful mainly for **system monitoring**, **profiling and limiting process +resources** and **management of running processes**. +It implements many functionalities offered by classic UNIX command line tools +such as *ps, top, iotop, lsof, netstat, ifconfig, free* and others. +psutil currently supports the following platforms: + +- **Linux** +- **Windows** +- **macOS** +- **FreeBSD, OpenBSD**, **NetBSD** +- **Sun Solaris** +- **AIX** + +Supported Python versions are cPython 3.6+ and `PyPy `__. +Latest psutil version supporting Python 2.7 is +`psutil 6.1.1 `__. + +Funding +======= + +While psutil is free software and will always be, the project would benefit +immensely from some funding. +Keeping up with bug reports and maintenance has become hardly sustainable for +me alone in terms of time. +If you're a company that's making significant use of psutil you can consider +becoming a sponsor via `GitHub Sponsors `__, +`Open Collective `__ or +`PayPal `__ +and have your logo displayed in here and psutil `doc `__. + +Sponsors +======== + +.. image:: https://github.com/giampaolo/psutil/raw/master/docs/_static/tidelift-logo.png + :width: 200 + :alt: Alternative text + +`Add your logo `__. + +Example usages +============== + +This represents pretty much the whole psutil API. + +CPU +--- + +.. code-block:: python + + >>> import psutil + >>> + >>> psutil.cpu_times() + scputimes(user=3961.46, nice=169.729, system=2150.659, idle=16900.540, iowait=629.59, irq=0.0, softirq=19.42, steal=0.0, guest=0, guest_nice=0.0) + >>> + >>> for x in range(3): + ... psutil.cpu_percent(interval=1) + ... + 4.0 + 5.9 + 3.8 + >>> + >>> for x in range(3): + ... psutil.cpu_percent(interval=1, percpu=True) + ... + [4.0, 6.9, 3.7, 9.2] + [7.0, 8.5, 2.4, 2.1] + [1.2, 9.0, 9.9, 7.2] + >>> + >>> for x in range(3): + ... psutil.cpu_times_percent(interval=1, percpu=False) + ... + scputimes(user=1.5, nice=0.0, system=0.5, idle=96.5, iowait=1.5, irq=0.0, softirq=0.0, steal=0.0, guest=0.0, guest_nice=0.0) + scputimes(user=1.0, nice=0.0, system=0.0, idle=99.0, iowait=0.0, irq=0.0, softirq=0.0, steal=0.0, guest=0.0, guest_nice=0.0) + scputimes(user=2.0, nice=0.0, system=0.0, idle=98.0, iowait=0.0, irq=0.0, softirq=0.0, steal=0.0, guest=0.0, guest_nice=0.0) + >>> + >>> psutil.cpu_count() + 4 + >>> psutil.cpu_count(logical=False) + 2 + >>> + >>> psutil.cpu_stats() + scpustats(ctx_switches=20455687, interrupts=6598984, soft_interrupts=2134212, syscalls=0) + >>> + >>> psutil.cpu_freq() + scpufreq(current=931.42925, min=800.0, max=3500.0) + >>> + >>> psutil.getloadavg() # also on Windows (emulated) + (3.14, 3.89, 4.67) + +Memory +------ + +.. code-block:: python + + >>> psutil.virtual_memory() + svmem(total=10367352832, available=6472179712, percent=37.6, used=8186245120, free=2181107712, active=4748992512, inactive=2758115328, buffers=790724608, cached=3500347392, shared=787554304) + >>> psutil.swap_memory() + sswap(total=2097147904, used=296128512, free=1801019392, percent=14.1, sin=304193536, sout=677842944) + >>> + +Disks +----- + +.. code-block:: python + + >>> psutil.disk_partitions() + [sdiskpart(device='/dev/sda1', mountpoint='/', fstype='ext4', opts='rw,nosuid'), + sdiskpart(device='/dev/sda2', mountpoint='/home', fstype='ext', opts='rw')] + >>> + >>> psutil.disk_usage('/') + sdiskusage(total=21378641920, used=4809781248, free=15482871808, percent=22.5) + >>> + >>> psutil.disk_io_counters(perdisk=False) + sdiskio(read_count=719566, write_count=1082197, read_bytes=18626220032, write_bytes=24081764352, read_time=5023392, write_time=63199568, read_merged_count=619166, write_merged_count=812396, busy_time=4523412) + >>> + +Network +------- + +.. code-block:: python + + >>> psutil.net_io_counters(pernic=True) + {'eth0': netio(bytes_sent=485291293, bytes_recv=6004858642, packets_sent=3251564, packets_recv=4787798, errin=0, errout=0, dropin=0, dropout=0), + 'lo': netio(bytes_sent=2838627, bytes_recv=2838627, packets_sent=30567, packets_recv=30567, errin=0, errout=0, dropin=0, dropout=0)} + >>> + >>> psutil.net_connections(kind='tcp') + [sconn(fd=115, family=, type=, laddr=addr(ip='10.0.0.1', port=48776), raddr=addr(ip='93.186.135.91', port=80), status='ESTABLISHED', pid=1254), + sconn(fd=117, family=, type=, laddr=addr(ip='10.0.0.1', port=43761), raddr=addr(ip='72.14.234.100', port=80), status='CLOSING', pid=2987), + ...] + >>> + >>> psutil.net_if_addrs() + {'lo': [snicaddr(family=, address='127.0.0.1', netmask='255.0.0.0', broadcast='127.0.0.1', ptp=None), + snicaddr(family=, address='::1', netmask='ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff', broadcast=None, ptp=None), + snicaddr(family=, address='00:00:00:00:00:00', netmask=None, broadcast='00:00:00:00:00:00', ptp=None)], + 'wlan0': [snicaddr(family=, address='192.168.1.3', netmask='255.255.255.0', broadcast='192.168.1.255', ptp=None), + snicaddr(family=, address='fe80::c685:8ff:fe45:641%wlan0', netmask='ffff:ffff:ffff:ffff::', broadcast=None, ptp=None), + snicaddr(family=, address='c4:85:08:45:06:41', netmask=None, broadcast='ff:ff:ff:ff:ff:ff', ptp=None)]} + >>> + >>> psutil.net_if_stats() + {'lo': snicstats(isup=True, duplex=, speed=0, mtu=65536, flags='up,loopback,running'), + 'wlan0': snicstats(isup=True, duplex=, speed=100, mtu=1500, flags='up,broadcast,running,multicast')} + >>> + +Sensors +------- + +.. code-block:: python + + >>> import psutil + >>> psutil.sensors_temperatures() + {'acpitz': [shwtemp(label='', current=47.0, high=103.0, critical=103.0)], + 'asus': [shwtemp(label='', current=47.0, high=None, critical=None)], + 'coretemp': [shwtemp(label='Physical id 0', current=52.0, high=100.0, critical=100.0), + shwtemp(label='Core 0', current=45.0, high=100.0, critical=100.0)]} + >>> + >>> psutil.sensors_fans() + {'asus': [sfan(label='cpu_fan', current=3200)]} + >>> + >>> psutil.sensors_battery() + sbattery(percent=93, secsleft=16628, power_plugged=False) + >>> + +Other system info +----------------- + +.. code-block:: python + + >>> import psutil + >>> psutil.users() + [suser(name='giampaolo', terminal='pts/2', host='localhost', started=1340737536.0, pid=1352), + suser(name='giampaolo', terminal='pts/3', host='localhost', started=1340737792.0, pid=1788)] + >>> + >>> psutil.boot_time() + 1365519115.0 + >>> + +Process management +------------------ + +.. code-block:: python + + >>> import psutil + >>> psutil.pids() + [1, 2, 3, 4, 5, 6, 7, 46, 48, 50, 51, 178, 182, 222, 223, 224, 268, 1215, + 1216, 1220, 1221, 1243, 1244, 1301, 1601, 2237, 2355, 2637, 2774, 3932, + 4176, 4177, 4185, 4187, 4189, 4225, 4243, 4245, 4263, 4282, 4306, 4311, + 4312, 4313, 4314, 4337, 4339, 4357, 4358, 4363, 4383, 4395, 4408, 4433, + 4443, 4445, 4446, 5167, 5234, 5235, 5252, 5318, 5424, 5644, 6987, 7054, + 7055, 7071] + >>> + >>> p = psutil.Process(7055) + >>> p + psutil.Process(pid=7055, name='python3', status='running', started='09:04:44') + >>> p.pid + 7055 + >>> p.name() + 'python3' + >>> p.exe() + '/usr/bin/python3' + >>> p.cwd() + '/home/giampaolo' + >>> p.cmdline() + ['/usr/bin/python3', 'main.py'] + >>> + >>> p.ppid() + 7054 + >>> p.parent() + psutil.Process(pid=4699, name='bash', status='sleeping', started='09:06:44') + >>> p.parents() + [psutil.Process(pid=4699, name='bash', started='09:06:44'), + psutil.Process(pid=4689, name='gnome-terminal-server', status='sleeping', started='0:06:44'), + psutil.Process(pid=1, name='systemd', status='sleeping', started='05:56:55')] + >>> p.children(recursive=True) + [psutil.Process(pid=29835, name='python3', status='sleeping', started='11:45:38'), + psutil.Process(pid=29836, name='python3', status='waking', started='11:43:39')] + >>> + >>> p.status() + 'running' + >>> p.create_time() + 1267551141.5019531 + >>> p.terminal() + '/dev/pts/0' + >>> + >>> p.username() + 'giampaolo' + >>> p.uids() + puids(real=1000, effective=1000, saved=1000) + >>> p.gids() + pgids(real=1000, effective=1000, saved=1000) + >>> + >>> p.cpu_times() + pcputimes(user=1.02, system=0.31, children_user=0.32, children_system=0.1, iowait=0.0) + >>> p.cpu_percent(interval=1.0) + 12.1 + >>> p.cpu_affinity() + [0, 1, 2, 3] + >>> p.cpu_affinity([0, 1]) # set + >>> p.cpu_num() + 1 + >>> + >>> p.memory_info() + pmem(rss=10915840, vms=67608576, shared=3313664, text=2310144, lib=0, data=7262208, dirty=0) + >>> p.memory_full_info() # "real" USS memory usage (Linux, macOS, Win only) + pfullmem(rss=10199040, vms=52133888, shared=3887104, text=2867200, lib=0, data=5967872, dirty=0, uss=6545408, pss=6872064, swap=0) + >>> p.memory_percent() + 0.7823 + >>> p.memory_maps() + [pmmap_grouped(path='/lib/x8664-linux-gnu/libutil-2.15.so', rss=32768, size=2125824, pss=32768, shared_clean=0, shared_dirty=0, private_clean=20480, private_dirty=12288, referenced=32768, anonymous=12288, swap=0), + pmmap_grouped(path='/lib/x8664-linux-gnu/libc-2.15.so', rss=3821568, size=3842048, pss=3821568, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=3821568, referenced=3575808, anonymous=3821568, swap=0), + pmmap_grouped(path='[heap]', rss=32768, size=139264, pss=32768, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=32768, referenced=32768, anonymous=32768, swap=0), + pmmap_grouped(path='[stack]', rss=2465792, size=2494464, pss=2465792, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=2465792, referenced=2277376, anonymous=2465792, swap=0), + ...] + >>> + >>> p.io_counters() + pio(read_count=478001, write_count=59371, read_bytes=700416, write_bytes=69632, read_chars=456232, write_chars=517543) + >>> + >>> p.open_files() + [popenfile(path='/home/giampaolo/monit.py', fd=3, position=0, mode='r', flags=32768), + popenfile(path='/var/log/monit.log', fd=4, position=235542, mode='a', flags=33793)] + >>> + >>> p.net_connections(kind='tcp') + [pconn(fd=115, family=, type=, laddr=addr(ip='10.0.0.1', port=48776), raddr=addr(ip='93.186.135.91', port=80), status='ESTABLISHED'), + pconn(fd=117, family=, type=, laddr=addr(ip='10.0.0.1', port=43761), raddr=addr(ip='72.14.234.100', port=80), status='CLOSING')] + >>> + >>> p.threads() + [pthread(id=5234, user_time=22.5, system_time=9.2891), + pthread(id=5237, user_time=0.0707, system_time=1.1)] + >>> + >>> p.num_threads() + 4 + >>> p.num_fds() + 8 + >>> p.num_ctx_switches() + pctxsw(voluntary=78, involuntary=19) + >>> + >>> p.nice() + 0 + >>> p.nice(10) # set + >>> + >>> p.ionice(psutil.IOPRIO_CLASS_IDLE) # IO priority (Win and Linux only) + >>> p.ionice() + pionice(ioclass=, value=0) + >>> + >>> p.rlimit(psutil.RLIMIT_NOFILE, (5, 5)) # set resource limits (Linux only) + >>> p.rlimit(psutil.RLIMIT_NOFILE) + (5, 5) + >>> + >>> p.environ() + {'LC_PAPER': 'it_IT.UTF-8', 'SHELL': '/bin/bash', 'GREP_OPTIONS': '--color=auto', + 'XDG_CONFIG_DIRS': '/etc/xdg/xdg-ubuntu:/usr/share/upstart/xdg:/etc/xdg', + ...} + >>> + >>> p.as_dict() + {'status': 'running', 'num_ctx_switches': pctxsw(voluntary=63, involuntary=1), 'pid': 5457, ...} + >>> p.is_running() + True + >>> p.suspend() + >>> p.resume() + >>> + >>> p.terminate() + >>> p.kill() + >>> p.wait(timeout=3) + + >>> + >>> psutil.test() + USER PID %CPU %MEM VSZ RSS TTY START TIME COMMAND + root 1 0.0 0.0 24584 2240 Jun17 00:00 init + root 2 0.0 0.0 0 0 Jun17 00:00 kthreadd + ... + giampaolo 31475 0.0 0.0 20760 3024 /dev/pts/0 Jun19 00:00 python2.4 + giampaolo 31721 0.0 2.2 773060 181896 00:04 10:30 chrome + root 31763 0.0 0.0 0 0 00:05 00:00 kworker/0:1 + >>> + +Further process APIs +-------------------- + +.. code-block:: python + + >>> import psutil + >>> for proc in psutil.process_iter(['pid', 'name']): + ... print(proc.info) + ... + {'pid': 1, 'name': 'systemd'} + {'pid': 2, 'name': 'kthreadd'} + {'pid': 3, 'name': 'ksoftirqd/0'} + ... + >>> + >>> psutil.pid_exists(3) + True + >>> + >>> def on_terminate(proc): + ... print("process {} terminated".format(proc)) + ... + >>> # waits for multiple processes to terminate + >>> gone, alive = psutil.wait_procs(procs_list, timeout=3, callback=on_terminate) + >>> + +Windows services +---------------- + +.. code-block:: python + + >>> list(psutil.win_service_iter()) + [, + , + , + , + ...] + >>> s = psutil.win_service_get('alg') + >>> s.as_dict() + {'binpath': 'C:\\Windows\\System32\\alg.exe', + 'description': 'Provides support for 3rd party protocol plug-ins for Internet Connection Sharing', + 'display_name': 'Application Layer Gateway Service', + 'name': 'alg', + 'pid': None, + 'start_type': 'manual', + 'status': 'stopped', + 'username': 'NT AUTHORITY\\LocalService'} + +Projects using psutil +===================== + +Here's some I find particularly interesting: + +- https://github.com/google/grr +- https://github.com/facebook/osquery/ +- https://github.com/nicolargo/glances +- https://github.com/aristocratos/bpytop +- https://github.com/Jahaja/psdash +- https://github.com/ajenti/ajenti +- https://github.com/home-assistant/home-assistant/ + +Portings +======== + +- Go: https://github.com/shirou/gopsutil +- C: https://github.com/hamon-in/cpslib +- Rust: https://github.com/rust-psutil/rust-psutil +- Nim: https://github.com/johnscillieri/psutil-nim + diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/RECORD b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..8722b2eb61f52debe817e013257fae0d79b0fd98 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/RECORD @@ -0,0 +1,66 @@ +psutil-7.1.3.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +psutil-7.1.3.dist-info/LICENSE,sha256=uJwGOzeG4o4MCjjxkx22H-015p3SopZvvs_-4PRsjRA,1548 +psutil-7.1.3.dist-info/METADATA,sha256=NXjcTuX0PlQtKYOXQbnj2HYMPA5djDiOwrC5gL2MZqE,23082 +psutil-7.1.3.dist-info/RECORD,, +psutil-7.1.3.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +psutil-7.1.3.dist-info/WHEEL,sha256=k0riQKE5_0GXActL4egeNSDuQ6nsFSKiWV91AD_j7xg,184 +psutil-7.1.3.dist-info/top_level.txt,sha256=gCNhn57wzksDjSAISmgMJ0aiXzQulk0GJhb2-BAyYgw,7 +psutil/__init__.py,sha256=7BTYRH1NsMvPBQyeRKqecd5d-JcxaeNFwnXtPYckDzo,87789 +psutil/__pycache__/__init__.cpython-312.pyc,, +psutil/__pycache__/_common.cpython-312.pyc,, +psutil/__pycache__/_psaix.cpython-312.pyc,, +psutil/__pycache__/_psbsd.cpython-312.pyc,, +psutil/__pycache__/_pslinux.cpython-312.pyc,, +psutil/__pycache__/_psosx.cpython-312.pyc,, +psutil/__pycache__/_psposix.cpython-312.pyc,, +psutil/__pycache__/_pssunos.cpython-312.pyc,, +psutil/__pycache__/_pswindows.cpython-312.pyc,, +psutil/_common.py,sha256=kDPOMktstdNqB0MbHgCvjR2wl8DQmwfIauBZXOsYLsI,28584 +psutil/_psaix.py,sha256=Pcos7TWAKSoU-eNFQ-u8wpMKlOS5xnA1f1TMgtmpC2s,18179 +psutil/_psbsd.py,sha256=K_Dn2NU6XsgW2ZTrVB5sH3AaQ6TGc3Q09a4rUcX-a6w,30701 +psutil/_pslinux.py,sha256=FFjrlyH8_vZndau5DBLpXFxePZ0AozDIKtF0Tn0_BgM,86321 +psutil/_psosx.py,sha256=ixb1t2ketUSzt8_AzVkixT8cly8qf7J5qet8KLZsZIo,16541 +psutil/_psposix.py,sha256=5pBE2Mk8LCWLthOpgKM5_S6AJOSankF0dHrozwtDi-Y,7141 +psutil/_pssunos.py,sha256=K99FU9-0k1W08ynXxX8LFaUaz4e59Iu-4vgp6gYSNCE,24848 +psutil/_psutil_linux.abi3.so,sha256=EiNyuoEmue7geZAxz5owrp6iFwHUfWsGEjtrDtgLJ_E,146080 +psutil/_pswindows.py,sha256=QInpXnEXWc_FyN0cHkgb53p-USMbWEgRqQG95iGnH64,36532 +psutil/tests/__init__.py,sha256=Q3__qjPt0V50DSfzMqjFEPyvDN2Tp8SmbZ-QRhPwPSM,62903 +psutil/tests/__main__.py,sha256=S1hzDPelhUBG_XfWAR0vh0bCeFFBOGpt_j1TAI5fnaM,305 +psutil/tests/__pycache__/__init__.cpython-312.pyc,, +psutil/tests/__pycache__/__main__.cpython-312.pyc,, +psutil/tests/__pycache__/test_aix.cpython-312.pyc,, +psutil/tests/__pycache__/test_bsd.cpython-312.pyc,, +psutil/tests/__pycache__/test_connections.cpython-312.pyc,, +psutil/tests/__pycache__/test_contracts.cpython-312.pyc,, +psutil/tests/__pycache__/test_linux.cpython-312.pyc,, +psutil/tests/__pycache__/test_memleaks.cpython-312.pyc,, +psutil/tests/__pycache__/test_misc.cpython-312.pyc,, +psutil/tests/__pycache__/test_osx.cpython-312.pyc,, +psutil/tests/__pycache__/test_posix.cpython-312.pyc,, +psutil/tests/__pycache__/test_process.cpython-312.pyc,, +psutil/tests/__pycache__/test_process_all.cpython-312.pyc,, +psutil/tests/__pycache__/test_scripts.cpython-312.pyc,, +psutil/tests/__pycache__/test_sudo.cpython-312.pyc,, +psutil/tests/__pycache__/test_sunos.cpython-312.pyc,, +psutil/tests/__pycache__/test_system.cpython-312.pyc,, +psutil/tests/__pycache__/test_testutils.cpython-312.pyc,, +psutil/tests/__pycache__/test_unicode.cpython-312.pyc,, +psutil/tests/__pycache__/test_windows.cpython-312.pyc,, +psutil/tests/test_aix.py,sha256=O5IqMAU3qw1NXvI-nhNb9v1LbNHb1q3iFRe5OkqgpoI,4408 +psutil/tests/test_bsd.py,sha256=Q2WmwU97yKuWYBE5qug_-MJdKW_AEKGoYs2RSYYWe7I,19997 +psutil/tests/test_connections.py,sha256=E8rOBe-jumVQ21ZWSo__JYepa9QNS1sdgQrQw75y47Q,21240 +psutil/tests/test_contracts.py,sha256=MQZN5981o_uXzmJnZMfzZo6NJGojekXDDCCbrZme7LQ,12183 +psutil/tests/test_linux.py,sha256=ELiN1bnDXBEArrAkQpvdOIG8YLxsG20NDkVxYWraYxA,88562 +psutil/tests/test_memleaks.py,sha256=MRc75HtM_LyK_Sq4OOtlHrsS0ykjx1CJkPwvpHn6wfs,15076 +psutil/tests/test_misc.py,sha256=ZrJCALG_JF2Wfe9cjJ7SgXx9EbPDZL0GGIhGiV3qS7s,29414 +psutil/tests/test_osx.py,sha256=-n81JFHAar8Ish7HDwyPxEChpsQZCS4d3egP1epkdH8,6668 +psutil/tests/test_posix.py,sha256=TREpzTFB02SpkDQLiFPdlyCtHjxgzk8y6LjvyX4PWe8,17391 +psutil/tests/test_process.py,sha256=PUfqA69CNKfwGSjeDQk0lciabpldFx_nxycZPIlQL5c,60114 +psutil/tests/test_process_all.py,sha256=EHZ4zJar6xHH-eAMyV92q0E4D__7LoVpB3wPJfDwk8A,18636 +psutil/tests/test_scripts.py,sha256=lIOH0ASZdsfVKN-SvefCSb4iZjlmEMT8pZ-XtTiOQ2w,7877 +psutil/tests/test_sudo.py,sha256=DdTBSh3l5pelXZ2O2plWcT2NZHF-29U2vK007S0OTfo,3837 +psutil/tests/test_sunos.py,sha256=78nUq_4I2ALidWLiOGRjgP7ykFiGm8eGV3JNAiIoOVs,1190 +psutil/tests/test_system.py,sha256=Juuk9qZ14yS4en5pteH5nxh2Xrbgu29MZklVqz4B8Qc,35533 +psutil/tests/test_testutils.py,sha256=8kf_h7Nezpbbqs7v2rNFw8bikhVpYnc-fS4hv8tFgE0,18788 +psutil/tests/test_unicode.py,sha256=m2fAfhg_2WzVHoKnkUu1dpPkA6TZ5-wmXJ1U7H3rcoo,10723 +psutil/tests/test_windows.py,sha256=P10qHGtzY0YLLPOJjuwU5B0eICy4sRpYes95DKnUtUk,34421 diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/REQUESTED b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/WHEEL b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..0c46bf3a9bf2a70caf599120106bbc63c220763a --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/WHEEL @@ -0,0 +1,7 @@ +Wheel-Version: 1.0 +Generator: setuptools (75.3.2) +Root-Is-Purelib: false +Tag: cp36-abi3-manylinux_2_12_x86_64 +Tag: cp36-abi3-manylinux2010_x86_64 +Tag: cp36-abi3-manylinux_2_28_x86_64 + diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/top_level.txt b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..a4d92cc08db6a0d8bfedbbbd620d1fb11f84677b --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil-7.1.3.dist-info/top_level.txt @@ -0,0 +1 @@ +psutil diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__init__.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92517709113b26e16d5f1def1c58c678972523d6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__init__.py @@ -0,0 +1,2425 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""psutil is a cross-platform library for retrieving information on +running processes and system utilization (CPU, memory, disks, network, +sensors) in Python. Supported platforms: + + - Linux + - Windows + - macOS + - FreeBSD + - OpenBSD + - NetBSD + - Sun Solaris + - AIX + +Supported Python versions are cPython 3.6+ and PyPy. +""" + +import collections +import contextlib +import datetime +import functools +import os +import signal +import socket +import subprocess +import sys +import threading +import time + +try: + import pwd +except ImportError: + pwd = None + +from . import _common +from ._common import AIX +from ._common import BSD +from ._common import CONN_CLOSE +from ._common import CONN_CLOSE_WAIT +from ._common import CONN_CLOSING +from ._common import CONN_ESTABLISHED +from ._common import CONN_FIN_WAIT1 +from ._common import CONN_FIN_WAIT2 +from ._common import CONN_LAST_ACK +from ._common import CONN_LISTEN +from ._common import CONN_NONE +from ._common import CONN_SYN_RECV +from ._common import CONN_SYN_SENT +from ._common import CONN_TIME_WAIT +from ._common import FREEBSD +from ._common import LINUX +from ._common import MACOS +from ._common import NETBSD +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import OPENBSD +from ._common import OSX # deprecated alias +from ._common import POSIX +from ._common import POWER_TIME_UNKNOWN +from ._common import POWER_TIME_UNLIMITED +from ._common import STATUS_DEAD +from ._common import STATUS_DISK_SLEEP +from ._common import STATUS_IDLE +from ._common import STATUS_LOCKED +from ._common import STATUS_PARKED +from ._common import STATUS_RUNNING +from ._common import STATUS_SLEEPING +from ._common import STATUS_STOPPED +from ._common import STATUS_TRACING_STOP +from ._common import STATUS_WAITING +from ._common import STATUS_WAKING +from ._common import STATUS_ZOMBIE +from ._common import SUNOS +from ._common import WINDOWS +from ._common import AccessDenied +from ._common import Error +from ._common import NoSuchProcess +from ._common import TimeoutExpired +from ._common import ZombieProcess +from ._common import debug +from ._common import memoize_when_activated +from ._common import wrap_numbers as _wrap_numbers + +if LINUX: + # This is public API and it will be retrieved from _pslinux.py + # via sys.modules. + PROCFS_PATH = "/proc" + + from . import _pslinux as _psplatform + from ._pslinux import IOPRIO_CLASS_BE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_IDLE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_NONE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_RT # noqa: F401 + +elif WINDOWS: + from . import _pswindows as _psplatform + from ._psutil_windows import ABOVE_NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import BELOW_NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import HIGH_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import IDLE_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import REALTIME_PRIORITY_CLASS # noqa: F401 + from ._pswindows import CONN_DELETE_TCB # noqa: F401 + from ._pswindows import IOPRIO_HIGH # noqa: F401 + from ._pswindows import IOPRIO_LOW # noqa: F401 + from ._pswindows import IOPRIO_NORMAL # noqa: F401 + from ._pswindows import IOPRIO_VERYLOW # noqa: F401 + +elif MACOS: + from . import _psosx as _psplatform + +elif BSD: + from . import _psbsd as _psplatform + +elif SUNOS: + from . import _pssunos as _psplatform + from ._pssunos import CONN_BOUND # noqa: F401 + from ._pssunos import CONN_IDLE # noqa: F401 + + # This is public writable API which is read from _pslinux.py and + # _pssunos.py via sys.modules. + PROCFS_PATH = "/proc" + +elif AIX: + from . import _psaix as _psplatform + + # This is public API and it will be retrieved from _pslinux.py + # via sys.modules. + PROCFS_PATH = "/proc" + +else: # pragma: no cover + msg = f"platform {sys.platform} is not supported" + raise NotImplementedError(msg) + + +# fmt: off +__all__ = [ + # exceptions + "Error", "NoSuchProcess", "ZombieProcess", "AccessDenied", + "TimeoutExpired", + + # constants + "version_info", "__version__", + + "STATUS_RUNNING", "STATUS_IDLE", "STATUS_SLEEPING", "STATUS_DISK_SLEEP", + "STATUS_STOPPED", "STATUS_TRACING_STOP", "STATUS_ZOMBIE", "STATUS_DEAD", + "STATUS_WAKING", "STATUS_LOCKED", "STATUS_WAITING", "STATUS_LOCKED", + "STATUS_PARKED", + + "CONN_ESTABLISHED", "CONN_SYN_SENT", "CONN_SYN_RECV", "CONN_FIN_WAIT1", + "CONN_FIN_WAIT2", "CONN_TIME_WAIT", "CONN_CLOSE", "CONN_CLOSE_WAIT", + "CONN_LAST_ACK", "CONN_LISTEN", "CONN_CLOSING", "CONN_NONE", + # "CONN_IDLE", "CONN_BOUND", + + "AF_LINK", + + "NIC_DUPLEX_FULL", "NIC_DUPLEX_HALF", "NIC_DUPLEX_UNKNOWN", + + "POWER_TIME_UNKNOWN", "POWER_TIME_UNLIMITED", + + "BSD", "FREEBSD", "LINUX", "NETBSD", "OPENBSD", "MACOS", "OSX", "POSIX", + "SUNOS", "WINDOWS", "AIX", + + # "RLIM_INFINITY", "RLIMIT_AS", "RLIMIT_CORE", "RLIMIT_CPU", "RLIMIT_DATA", + # "RLIMIT_FSIZE", "RLIMIT_LOCKS", "RLIMIT_MEMLOCK", "RLIMIT_NOFILE", + # "RLIMIT_NPROC", "RLIMIT_RSS", "RLIMIT_STACK", "RLIMIT_MSGQUEUE", + # "RLIMIT_NICE", "RLIMIT_RTPRIO", "RLIMIT_RTTIME", "RLIMIT_SIGPENDING", + + # classes + "Process", "Popen", + + # functions + "pid_exists", "pids", "process_iter", "wait_procs", # proc + "virtual_memory", "swap_memory", # memory + "cpu_times", "cpu_percent", "cpu_times_percent", "cpu_count", # cpu + "cpu_stats", # "cpu_freq", "getloadavg" + "net_io_counters", "net_connections", "net_if_addrs", # network + "net_if_stats", + "disk_io_counters", "disk_partitions", "disk_usage", # disk + # "sensors_temperatures", "sensors_battery", "sensors_fans" # sensors + "users", "boot_time", # others +] +# fmt: on + + +__all__.extend(_psplatform.__extra__all__) + +# Linux, FreeBSD +if hasattr(_psplatform.Process, "rlimit"): + # Populate global namespace with RLIM* constants. + _globals = globals() + _name = None + for _name in dir(_psplatform.cext): + if _name.startswith('RLIM') and _name.isupper(): + _globals[_name] = getattr(_psplatform.cext, _name) + __all__.append(_name) + del _globals, _name + +AF_LINK = _psplatform.AF_LINK + +__author__ = "Giampaolo Rodola'" +__version__ = "7.1.3" +version_info = tuple(int(num) for num in __version__.split('.')) + +_timer = getattr(time, 'monotonic', time.time) +_TOTAL_PHYMEM = None +_LOWEST_PID = None +_SENTINEL = object() + +# Sanity check in case the user messed up with psutil installation +# or did something weird with sys.path. In this case we might end +# up importing a python module using a C extension module which +# was compiled for a different version of psutil. +# We want to prevent that by failing sooner rather than later. +# See: https://github.com/giampaolo/psutil/issues/564 +if int(__version__.replace('.', '')) != getattr( + _psplatform.cext, 'version', None +): + msg = f"version conflict: {_psplatform.cext.__file__!r} C extension " + msg += "module was built for another version of psutil" + if hasattr(_psplatform.cext, 'version'): + v = ".".join(list(str(_psplatform.cext.version))) + msg += f" ({v} instead of {__version__})" + else: + msg += f" (different than {__version__})" + what = getattr( + _psplatform.cext, + "__file__", + "the existing psutil install directory", + ) + msg += f"; you may try to 'pip uninstall psutil', manually remove {what}" + msg += " or clean the virtual env somehow, then reinstall" + raise ImportError(msg) + + +# ===================================================================== +# --- Utils +# ===================================================================== + + +if hasattr(_psplatform, 'ppid_map'): + # Faster version (Windows and Linux). + _ppid_map = _psplatform.ppid_map +else: # pragma: no cover + + def _ppid_map(): + """Return a {pid: ppid, ...} dict for all running processes in + one shot. Used to speed up Process.children(). + """ + ret = {} + for pid in pids(): + try: + ret[pid] = _psplatform.Process(pid).ppid() + except (NoSuchProcess, ZombieProcess): + pass + return ret + + +def _pprint_secs(secs): + """Format seconds in a human readable form.""" + now = time.time() + secs_ago = int(now - secs) + fmt = "%H:%M:%S" if secs_ago < 60 * 60 * 24 else "%Y-%m-%d %H:%M:%S" + return datetime.datetime.fromtimestamp(secs).strftime(fmt) + + +def _check_conn_kind(kind): + """Check net_connections()'s `kind` parameter.""" + kinds = tuple(_common.conn_tmap) + if kind not in kinds: + msg = f"invalid kind argument {kind!r}; valid ones are: {kinds}" + raise ValueError(msg) + + +# ===================================================================== +# --- Process class +# ===================================================================== + + +class Process: + """Represents an OS process with the given PID. + If PID is omitted current process PID (os.getpid()) is used. + Raise NoSuchProcess if PID does not exist. + + Note that most of the methods of this class do not make sure that + the PID of the process being queried has been reused. That means + that you may end up retrieving information for another process. + + The only exceptions for which process identity is pre-emptively + checked and guaranteed are: + + - parent() + - children() + - nice() (set) + - ionice() (set) + - rlimit() (set) + - cpu_affinity (set) + - suspend() + - resume() + - send_signal() + - terminate() + - kill() + + To prevent this problem for all other methods you can use + is_running() before querying the process. + """ + + def __init__(self, pid=None): + self._init(pid) + + def _init(self, pid, _ignore_nsp=False): + if pid is None: + pid = os.getpid() + else: + if pid < 0: + msg = f"pid must be a positive integer (got {pid})" + raise ValueError(msg) + try: + _psplatform.cext.check_pid_range(pid) + except OverflowError as err: + msg = "process PID out of range" + raise NoSuchProcess(pid, msg=msg) from err + + self._pid = pid + self._name = None + self._exe = None + self._create_time = None + self._gone = False + self._pid_reused = False + self._hash = None + self._lock = threading.RLock() + # used for caching on Windows only (on POSIX ppid may change) + self._ppid = None + # platform-specific modules define an _psplatform.Process + # implementation class + self._proc = _psplatform.Process(pid) + self._last_sys_cpu_times = None + self._last_proc_cpu_times = None + self._exitcode = _SENTINEL + self._ident = (self.pid, None) + try: + self._ident = self._get_ident() + except AccessDenied: + # This should happen on Windows only, since we use the fast + # create time method. AFAIK, on all other platforms we are + # able to get create time for all PIDs. + pass + except ZombieProcess: + # Zombies can still be queried by this class (although + # not always) and pids() return them so just go on. + pass + except NoSuchProcess: + if not _ignore_nsp: + msg = "process PID not found" + raise NoSuchProcess(pid, msg=msg) from None + self._gone = True + + def _get_ident(self): + """Return a (pid, uid) tuple which is supposed to identify a + Process instance univocally over time. The PID alone is not + enough, as it can be assigned to a new process after this one + terminates, so we add process creation time to the mix. We need + this in order to prevent killing the wrong process later on. + This is also known as PID reuse or PID recycling problem. + + The reliability of this strategy mostly depends on + create_time() precision, which is 0.01 secs on Linux. The + assumption is that, after a process terminates, the kernel + won't reuse the same PID after such a short period of time + (0.01 secs). Technically this is inherently racy, but + practically it should be good enough. + + NOTE: unreliable on FreeBSD and OpenBSD as ctime is subject to + system clock updates. + """ + + if WINDOWS: + # Use create_time() fast method in order to speedup + # `process_iter()`. This means we'll get AccessDenied for + # most ADMIN processes, but that's fine since it means + # we'll also get AccessDenied on kill(). + # https://github.com/giampaolo/psutil/issues/2366#issuecomment-2381646555 + self._create_time = self._proc.create_time(fast_only=True) + return (self.pid, self._create_time) + elif LINUX or NETBSD or OSX: + # Use 'monotonic' process starttime since boot to form unique + # process identity, since it is stable over changes to system + # time. + return (self.pid, self._proc.create_time(monotonic=True)) + else: + return (self.pid, self.create_time()) + + def __str__(self): + info = collections.OrderedDict() + info["pid"] = self.pid + if self._name: + info['name'] = self._name + with self.oneshot(): + if self._pid_reused: + info["status"] = "terminated + PID reused" + else: + try: + info["name"] = self.name() + info["status"] = self.status() + except ZombieProcess: + info["status"] = "zombie" + except NoSuchProcess: + info["status"] = "terminated" + except AccessDenied: + pass + + if self._exitcode not in {_SENTINEL, None}: + info["exitcode"] = self._exitcode + if self._create_time is not None: + info['started'] = _pprint_secs(self._create_time) + + return "{}.{}({})".format( + self.__class__.__module__, + self.__class__.__name__, + ", ".join([f"{k}={v!r}" for k, v in info.items()]), + ) + + __repr__ = __str__ + + def __eq__(self, other): + # Test for equality with another Process object based + # on PID and creation time. + if not isinstance(other, Process): + return NotImplemented + if OPENBSD or NETBSD or SUNOS: # pragma: no cover + # Zombie processes on Open/NetBSD/illumos/Solaris have a + # creation time of 0.0. This covers the case when a process + # started normally (so it has a ctime), then it turned into a + # zombie. It's important to do this because is_running() + # depends on __eq__. + pid1, ident1 = self._ident + pid2, ident2 = other._ident + if pid1 == pid2: + if ident1 and not ident2: + try: + return self.status() == STATUS_ZOMBIE + except Error: + pass + return self._ident == other._ident + + def __ne__(self, other): + return not self == other + + def __hash__(self): + if self._hash is None: + self._hash = hash(self._ident) + return self._hash + + def _raise_if_pid_reused(self): + """Raises NoSuchProcess in case process PID has been reused.""" + if self._pid_reused or (not self.is_running() and self._pid_reused): + # We may directly raise NSP in here already if PID is just + # not running, but I prefer NSP to be raised naturally by + # the actual Process API call. This way unit tests will tell + # us if the API is broken (aka don't raise NSP when it + # should). We also remain consistent with all other "get" + # APIs which don't use _raise_if_pid_reused(). + msg = "process no longer exists and its PID has been reused" + raise NoSuchProcess(self.pid, self._name, msg=msg) + + @property + def pid(self): + """The process PID.""" + return self._pid + + # --- utility methods + + @contextlib.contextmanager + def oneshot(self): + """Utility context manager which considerably speeds up the + retrieval of multiple process information at the same time. + + Internally different process info (e.g. name, ppid, uids, + gids, ...) may be fetched by using the same routine, but + only one information is returned and the others are discarded. + When using this context manager the internal routine is + executed once (in the example below on name()) and the + other info are cached. + + The cache is cleared when exiting the context manager block. + The advice is to use this every time you retrieve more than + one information about the process. If you're lucky, you'll + get a hell of a speedup. + + >>> import psutil + >>> p = psutil.Process() + >>> with p.oneshot(): + ... p.name() # collect multiple info + ... p.cpu_times() # return cached value + ... p.cpu_percent() # return cached value + ... p.create_time() # return cached value + ... + >>> + """ + with self._lock: + if hasattr(self, "_cache"): + # NOOP: this covers the use case where the user enters the + # context twice: + # + # >>> with p.oneshot(): + # ... with p.oneshot(): + # ... + # + # Also, since as_dict() internally uses oneshot() + # I expect that the code below will be a pretty common + # "mistake" that the user will make, so let's guard + # against that: + # + # >>> with p.oneshot(): + # ... p.as_dict() + # ... + yield + else: + try: + # cached in case cpu_percent() is used + self.cpu_times.cache_activate(self) + # cached in case memory_percent() is used + self.memory_info.cache_activate(self) + # cached in case parent() is used + self.ppid.cache_activate(self) + # cached in case username() is used + if POSIX: + self.uids.cache_activate(self) + # specific implementation cache + self._proc.oneshot_enter() + yield + finally: + self.cpu_times.cache_deactivate(self) + self.memory_info.cache_deactivate(self) + self.ppid.cache_deactivate(self) + if POSIX: + self.uids.cache_deactivate(self) + self._proc.oneshot_exit() + + def as_dict(self, attrs=None, ad_value=None): + """Utility method returning process information as a + hashable dictionary. + If *attrs* is specified it must be a list of strings + reflecting available Process class' attribute names + (e.g. ['cpu_times', 'name']) else all public (read + only) attributes are assumed. + *ad_value* is the value which gets assigned in case + AccessDenied or ZombieProcess exception is raised when + retrieving that particular process information. + """ + valid_names = _as_dict_attrnames + if attrs is not None: + if not isinstance(attrs, (list, tuple, set, frozenset)): + msg = f"invalid attrs type {type(attrs)}" + raise TypeError(msg) + attrs = set(attrs) + invalid_names = attrs - valid_names + if invalid_names: + msg = "invalid attr name{} {}".format( + "s" if len(invalid_names) > 1 else "", + ", ".join(map(repr, invalid_names)), + ) + raise ValueError(msg) + + retdict = {} + ls = attrs or valid_names + with self.oneshot(): + for name in ls: + try: + if name == 'pid': + ret = self.pid + else: + meth = getattr(self, name) + ret = meth() + except (AccessDenied, ZombieProcess): + ret = ad_value + except NotImplementedError: + # in case of not implemented functionality (may happen + # on old or exotic systems) we want to crash only if + # the user explicitly asked for that particular attr + if attrs: + raise + continue + retdict[name] = ret + return retdict + + def parent(self): + """Return the parent process as a Process object pre-emptively + checking whether PID has been reused. + If no parent is known return None. + """ + lowest_pid = _LOWEST_PID if _LOWEST_PID is not None else pids()[0] + if self.pid == lowest_pid: + return None + ppid = self.ppid() + if ppid is not None: + # Get a fresh (non-cached) ctime in case the system clock + # was updated. TODO: use a monotonic ctime on platforms + # where it's supported. + proc_ctime = Process(self.pid).create_time() + try: + parent = Process(ppid) + if parent.create_time() <= proc_ctime: + return parent + # ...else ppid has been reused by another process + except NoSuchProcess: + pass + + def parents(self): + """Return the parents of this process as a list of Process + instances. If no parents are known return an empty list. + """ + parents = [] + proc = self.parent() + while proc is not None: + parents.append(proc) + proc = proc.parent() + return parents + + def is_running(self): + """Return whether this process is running. + + It also checks if PID has been reused by another process, in + which case it will remove the process from `process_iter()` + internal cache and return False. + """ + if self._gone or self._pid_reused: + return False + try: + # Checking if PID is alive is not enough as the PID might + # have been reused by another process. Process identity / + # uniqueness over time is guaranteed by (PID + creation + # time) and that is verified in __eq__. + self._pid_reused = self != Process(self.pid) + if self._pid_reused: + _pids_reused.add(self.pid) + raise NoSuchProcess(self.pid) + return True + except ZombieProcess: + # We should never get here as it's already handled in + # Process.__init__; here just for extra safety. + return True + except NoSuchProcess: + self._gone = True + return False + + # --- actual API + + @memoize_when_activated + def ppid(self): + """The process parent PID. + On Windows the return value is cached after first call. + """ + # On POSIX we don't want to cache the ppid as it may unexpectedly + # change to 1 (init) in case this process turns into a zombie: + # https://github.com/giampaolo/psutil/issues/321 + # http://stackoverflow.com/questions/356722/ + + # XXX should we check creation time here rather than in + # Process.parent()? + self._raise_if_pid_reused() + if POSIX: + return self._proc.ppid() + else: # pragma: no cover + self._ppid = self._ppid or self._proc.ppid() + return self._ppid + + def name(self): + """The process name. The return value is cached after first call.""" + # Process name is only cached on Windows as on POSIX it may + # change, see: + # https://github.com/giampaolo/psutil/issues/692 + if WINDOWS and self._name is not None: + return self._name + name = self._proc.name() + if POSIX and len(name) >= 15: + # On UNIX the name gets truncated to the first 15 characters. + # If it matches the first part of the cmdline we return that + # one instead because it's usually more explicative. + # Examples are "gnome-keyring-d" vs. "gnome-keyring-daemon". + try: + cmdline = self.cmdline() + except (AccessDenied, ZombieProcess): + # Just pass and return the truncated name: it's better + # than nothing. Note: there are actual cases where a + # zombie process can return a name() but not a + # cmdline(), see: + # https://github.com/giampaolo/psutil/issues/2239 + pass + else: + if cmdline: + extended_name = os.path.basename(cmdline[0]) + if extended_name.startswith(name): + name = extended_name + self._name = name + self._proc._name = name + return name + + def exe(self): + """The process executable as an absolute path. + May also be an empty string. + The return value is cached after first call. + """ + + def guess_it(fallback): + # try to guess exe from cmdline[0] in absence of a native + # exe representation + cmdline = self.cmdline() + if cmdline and hasattr(os, 'access') and hasattr(os, 'X_OK'): + exe = cmdline[0] # the possible exe + # Attempt to guess only in case of an absolute path. + # It is not safe otherwise as the process might have + # changed cwd. + if ( + os.path.isabs(exe) + and os.path.isfile(exe) + and os.access(exe, os.X_OK) + ): + return exe + if isinstance(fallback, AccessDenied): + raise fallback + return fallback + + if self._exe is None: + try: + exe = self._proc.exe() + except AccessDenied as err: + return guess_it(fallback=err) + else: + if not exe: + # underlying implementation can legitimately return an + # empty string; if that's the case we don't want to + # raise AD while guessing from the cmdline + try: + exe = guess_it(fallback=exe) + except AccessDenied: + pass + self._exe = exe + return self._exe + + def cmdline(self): + """The command line this process has been called with.""" + return self._proc.cmdline() + + def status(self): + """The process current status as a STATUS_* constant.""" + try: + return self._proc.status() + except ZombieProcess: + return STATUS_ZOMBIE + + def username(self): + """The name of the user that owns the process. + On UNIX this is calculated by using *real* process uid. + """ + if POSIX: + if pwd is None: + # might happen if python was installed from sources + msg = "requires pwd module shipped with standard python" + raise ImportError(msg) + real_uid = self.uids().real + try: + return pwd.getpwuid(real_uid).pw_name + except KeyError: + # the uid can't be resolved by the system + return str(real_uid) + else: + return self._proc.username() + + def create_time(self): + """The process creation time as a floating point number + expressed in seconds since the epoch (seconds since January 1, + 1970, at midnight UTC). The return value, which is cached after + first call, is based on the system clock, which means it may be + affected by changes such as manual adjustments or time + synchronization (e.g. NTP). + """ + if self._create_time is None: + self._create_time = self._proc.create_time() + return self._create_time + + def cwd(self): + """Process current working directory as an absolute path.""" + return self._proc.cwd() + + def nice(self, value=None): + """Get or set process niceness (priority).""" + if value is None: + return self._proc.nice_get() + else: + self._raise_if_pid_reused() + self._proc.nice_set(value) + + if POSIX: + + @memoize_when_activated + def uids(self): + """Return process UIDs as a (real, effective, saved) + namedtuple. + """ + return self._proc.uids() + + def gids(self): + """Return process GIDs as a (real, effective, saved) + namedtuple. + """ + return self._proc.gids() + + def terminal(self): + """The terminal associated with this process, if any, + else None. + """ + return self._proc.terminal() + + def num_fds(self): + """Return the number of file descriptors opened by this + process (POSIX only). + """ + return self._proc.num_fds() + + # Linux, BSD, AIX and Windows only + if hasattr(_psplatform.Process, "io_counters"): + + def io_counters(self): + """Return process I/O statistics as a + (read_count, write_count, read_bytes, write_bytes) + namedtuple. + Those are the number of read/write calls performed and the + amount of bytes read and written by the process. + """ + return self._proc.io_counters() + + # Linux and Windows + if hasattr(_psplatform.Process, "ionice_get"): + + def ionice(self, ioclass=None, value=None): + """Get or set process I/O niceness (priority). + + On Linux *ioclass* is one of the IOPRIO_CLASS_* constants. + *value* is a number which goes from 0 to 7. The higher the + value, the lower the I/O priority of the process. + + On Windows only *ioclass* is used and it can be set to 2 + (normal), 1 (low) or 0 (very low). + + Available on Linux and Windows > Vista only. + """ + if ioclass is None: + if value is not None: + msg = "'ioclass' argument must be specified" + raise ValueError(msg) + return self._proc.ionice_get() + else: + self._raise_if_pid_reused() + return self._proc.ionice_set(ioclass, value) + + # Linux / FreeBSD only + if hasattr(_psplatform.Process, "rlimit"): + + def rlimit(self, resource, limits=None): + """Get or set process resource limits as a (soft, hard) + tuple. + + *resource* is one of the RLIMIT_* constants. + *limits* is supposed to be a (soft, hard) tuple. + + See "man prlimit" for further info. + Available on Linux and FreeBSD only. + """ + if limits is not None: + self._raise_if_pid_reused() + return self._proc.rlimit(resource, limits) + + # Windows, Linux and FreeBSD only + if hasattr(_psplatform.Process, "cpu_affinity_get"): + + def cpu_affinity(self, cpus=None): + """Get or set process CPU affinity. + If specified, *cpus* must be a list of CPUs for which you + want to set the affinity (e.g. [0, 1]). + If an empty list is passed, all egible CPUs are assumed + (and set). + (Windows, Linux and BSD only). + """ + if cpus is None: + return sorted(set(self._proc.cpu_affinity_get())) + else: + self._raise_if_pid_reused() + if not cpus: + if hasattr(self._proc, "_get_eligible_cpus"): + cpus = self._proc._get_eligible_cpus() + else: + cpus = tuple(range(len(cpu_times(percpu=True)))) + self._proc.cpu_affinity_set(list(set(cpus))) + + # Linux, FreeBSD, SunOS + if hasattr(_psplatform.Process, "cpu_num"): + + def cpu_num(self): + """Return what CPU this process is currently running on. + The returned number should be <= psutil.cpu_count() + and <= len(psutil.cpu_percent(percpu=True)). + It may be used in conjunction with + psutil.cpu_percent(percpu=True) to observe the system + workload distributed across CPUs. + """ + return self._proc.cpu_num() + + # All platforms has it, but maybe not in the future. + if hasattr(_psplatform.Process, "environ"): + + def environ(self): + """The environment variables of the process as a dict. Note: this + might not reflect changes made after the process started. + """ + return self._proc.environ() + + if WINDOWS: + + def num_handles(self): + """Return the number of handles opened by this process + (Windows only). + """ + return self._proc.num_handles() + + def num_ctx_switches(self): + """Return the number of voluntary and involuntary context + switches performed by this process. + """ + return self._proc.num_ctx_switches() + + def num_threads(self): + """Return the number of threads used by this process.""" + return self._proc.num_threads() + + if hasattr(_psplatform.Process, "threads"): + + def threads(self): + """Return threads opened by process as a list of + (id, user_time, system_time) namedtuples representing + thread id and thread CPU times (user/system). + On OpenBSD this method requires root access. + """ + return self._proc.threads() + + def children(self, recursive=False): + """Return the children of this process as a list of Process + instances, pre-emptively checking whether PID has been reused. + If *recursive* is True return all the parent descendants. + + Example (A == this process): + + A ─┐ + │ + ├─ B (child) ─┐ + │ └─ X (grandchild) ─┐ + │ └─ Y (great grandchild) + ├─ C (child) + └─ D (child) + + >>> import psutil + >>> p = psutil.Process() + >>> p.children() + B, C, D + >>> p.children(recursive=True) + B, X, Y, C, D + + Note that in the example above if process X disappears + process Y won't be listed as the reference to process A + is lost. + """ + self._raise_if_pid_reused() + ppid_map = _ppid_map() + # Get a fresh (non-cached) ctime in case the system clock was + # updated. TODO: use a monotonic ctime on platforms where it's + # supported. + proc_ctime = Process(self.pid).create_time() + ret = [] + if not recursive: + for pid, ppid in ppid_map.items(): + if ppid == self.pid: + try: + child = Process(pid) + # if child happens to be older than its parent + # (self) it means child's PID has been reused + if proc_ctime <= child.create_time(): + ret.append(child) + except (NoSuchProcess, ZombieProcess): + pass + else: + # Construct a {pid: [child pids]} dict + reverse_ppid_map = collections.defaultdict(list) + for pid, ppid in ppid_map.items(): + reverse_ppid_map[ppid].append(pid) + # Recursively traverse that dict, starting from self.pid, + # such that we only call Process() on actual children + seen = set() + stack = [self.pid] + while stack: + pid = stack.pop() + if pid in seen: + # Since pids can be reused while the ppid_map is + # constructed, there may be rare instances where + # there's a cycle in the recorded process "tree". + continue + seen.add(pid) + for child_pid in reverse_ppid_map[pid]: + try: + child = Process(child_pid) + # if child happens to be older than its parent + # (self) it means child's PID has been reused + intime = proc_ctime <= child.create_time() + if intime: + ret.append(child) + stack.append(child_pid) + except (NoSuchProcess, ZombieProcess): + pass + return ret + + def cpu_percent(self, interval=None): + """Return a float representing the current process CPU + utilization as a percentage. + + When *interval* is 0.0 or None (default) compares process times + to system CPU times elapsed since last call, returning + immediately (non-blocking). That means that the first time + this is called it will return a meaningful 0.0 value. + + When *interval* is > 0.0 compares process times to system CPU + times elapsed before and after the interval (blocking). + + In this case is recommended for accuracy that this function + be called with at least 0.1 seconds between calls. + + A value > 100.0 can be returned in case of processes running + multiple threads on different CPU cores. + + The returned value is explicitly NOT split evenly between + all available logical CPUs. This means that a busy loop process + running on a system with 2 logical CPUs will be reported as + having 100% CPU utilization instead of 50%. + + Examples: + + >>> import psutil + >>> p = psutil.Process(os.getpid()) + >>> # blocking + >>> p.cpu_percent(interval=1) + 2.0 + >>> # non-blocking (percentage since last call) + >>> p.cpu_percent(interval=None) + 2.9 + >>> + """ + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval!r})" + raise ValueError(msg) + num_cpus = cpu_count() or 1 + + def timer(): + return _timer() * num_cpus + + if blocking: + st1 = timer() + pt1 = self._proc.cpu_times() + time.sleep(interval) + st2 = timer() + pt2 = self._proc.cpu_times() + else: + st1 = self._last_sys_cpu_times + pt1 = self._last_proc_cpu_times + st2 = timer() + pt2 = self._proc.cpu_times() + if st1 is None or pt1 is None: + self._last_sys_cpu_times = st2 + self._last_proc_cpu_times = pt2 + return 0.0 + + delta_proc = (pt2.user - pt1.user) + (pt2.system - pt1.system) + delta_time = st2 - st1 + # reset values for next call in case of interval == None + self._last_sys_cpu_times = st2 + self._last_proc_cpu_times = pt2 + + try: + # This is the utilization split evenly between all CPUs. + # E.g. a busy loop process on a 2-CPU-cores system at this + # point is reported as 50% instead of 100%. + overall_cpus_percent = (delta_proc / delta_time) * 100 + except ZeroDivisionError: + # interval was too low + return 0.0 + else: + # Note 1: + # in order to emulate "top" we multiply the value for the num + # of CPU cores. This way the busy process will be reported as + # having 100% (or more) usage. + # + # Note 2: + # taskmgr.exe on Windows differs in that it will show 50% + # instead. + # + # Note 3: + # a percentage > 100 is legitimate as it can result from a + # process with multiple threads running on different CPU + # cores (top does the same), see: + # http://stackoverflow.com/questions/1032357 + # https://github.com/giampaolo/psutil/issues/474 + single_cpu_percent = overall_cpus_percent * num_cpus + return round(single_cpu_percent, 1) + + @memoize_when_activated + def cpu_times(self): + """Return a (user, system, children_user, children_system) + namedtuple representing the accumulated process time, in + seconds. + This is similar to os.times() but per-process. + On macOS and Windows children_user and children_system are + always set to 0. + """ + return self._proc.cpu_times() + + @memoize_when_activated + def memory_info(self): + """Return a namedtuple with variable fields depending on the + platform, representing memory information about the process. + + The "portable" fields available on all platforms are `rss` and `vms`. + + All numbers are expressed in bytes. + """ + return self._proc.memory_info() + + def memory_full_info(self): + """This method returns the same information as memory_info(), + plus, on some platform (Linux, macOS, Windows), also provides + additional metrics (USS, PSS and swap). + The additional metrics provide a better representation of actual + process memory usage. + + Namely USS is the memory which is unique to a process and which + would be freed if the process was terminated right now. + + It does so by passing through the whole process address. + As such it usually requires higher user privileges than + memory_info() and is considerably slower. + """ + return self._proc.memory_full_info() + + def memory_percent(self, memtype="rss"): + """Compare process memory to total physical system memory and + calculate process memory utilization as a percentage. + *memtype* argument is a string that dictates what type of + process memory you want to compare against (defaults to "rss"). + The list of available strings can be obtained like this: + + >>> psutil.Process().memory_info()._fields + ('rss', 'vms', 'shared', 'text', 'lib', 'data', 'dirty', 'uss', 'pss') + """ + valid_types = list(_psplatform.pfullmem._fields) + if memtype not in valid_types: + msg = ( + f"invalid memtype {memtype!r}; valid types are" + f" {tuple(valid_types)!r}" + ) + raise ValueError(msg) + fun = ( + self.memory_info + if memtype in _psplatform.pmem._fields + else self.memory_full_info + ) + metrics = fun() + value = getattr(metrics, memtype) + + # use cached value if available + total_phymem = _TOTAL_PHYMEM or virtual_memory().total + if not total_phymem > 0: + # we should never get here + msg = ( + "can't calculate process memory percent because total physical" + f" system memory is not positive ({total_phymem!r})" + ) + raise ValueError(msg) + return (value / float(total_phymem)) * 100 + + if hasattr(_psplatform.Process, "memory_maps"): + + def memory_maps(self, grouped=True): + """Return process' mapped memory regions as a list of namedtuples + whose fields are variable depending on the platform. + + If *grouped* is True the mapped regions with the same 'path' + are grouped together and the different memory fields are summed. + + If *grouped* is False every mapped region is shown as a single + entity and the namedtuple will also include the mapped region's + address space ('addr') and permission set ('perms'). + """ + it = self._proc.memory_maps() + if grouped: + d = {} + for tupl in it: + path = tupl[2] + nums = tupl[3:] + try: + d[path] = list(map(lambda x, y: x + y, d[path], nums)) + except KeyError: + d[path] = nums + nt = _psplatform.pmmap_grouped + return [nt(path, *d[path]) for path in d] + else: + nt = _psplatform.pmmap_ext + return [nt(*x) for x in it] + + def open_files(self): + """Return files opened by process as a list of + (path, fd) namedtuples including the absolute file name + and file descriptor number. + """ + return self._proc.open_files() + + def net_connections(self, kind='inet'): + """Return socket connections opened by process as a list of + (fd, family, type, laddr, raddr, status) namedtuples. + The *kind* parameter filters for connections that match the + following criteria: + + +------------+----------------------------------------------------+ + | Kind Value | Connections using | + +------------+----------------------------------------------------+ + | inet | IPv4 and IPv6 | + | inet4 | IPv4 | + | inet6 | IPv6 | + | tcp | TCP | + | tcp4 | TCP over IPv4 | + | tcp6 | TCP over IPv6 | + | udp | UDP | + | udp4 | UDP over IPv4 | + | udp6 | UDP over IPv6 | + | unix | UNIX socket (both UDP and TCP protocols) | + | all | the sum of all the possible families and protocols | + +------------+----------------------------------------------------+ + """ + _check_conn_kind(kind) + return self._proc.net_connections(kind) + + @_common.deprecated_method(replacement="net_connections") + def connections(self, kind="inet"): + return self.net_connections(kind=kind) + + # --- signals + + if POSIX: + + def _send_signal(self, sig): + assert not self.pid < 0, self.pid + self._raise_if_pid_reused() + + pid, ppid, name = self.pid, self._ppid, self._name + if pid == 0: + # see "man 2 kill" + msg = ( + "preventing sending signal to process with PID 0 as it " + "would affect every process in the process group of the " + "calling process (os.getpid()) instead of PID 0" + ) + raise ValueError(msg) + try: + os.kill(pid, sig) + except ProcessLookupError as err: + if OPENBSD and pid_exists(pid): + # We do this because os.kill() lies in case of + # zombie processes. + raise ZombieProcess(pid, name, ppid) from err + self._gone = True + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + def send_signal(self, sig): + """Send a signal *sig* to process pre-emptively checking + whether PID has been reused (see signal module constants) . + On Windows only SIGTERM is valid and is treated as an alias + for kill(). + """ + if POSIX: + self._send_signal(sig) + else: # pragma: no cover + self._raise_if_pid_reused() + if sig != signal.SIGTERM and not self.is_running(): + msg = "process no longer exists" + raise NoSuchProcess(self.pid, self._name, msg=msg) + self._proc.send_signal(sig) + + def suspend(self): + """Suspend process execution with SIGSTOP pre-emptively checking + whether PID has been reused. + On Windows this has the effect of suspending all process threads. + """ + if POSIX: + self._send_signal(signal.SIGSTOP) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.suspend() + + def resume(self): + """Resume process execution with SIGCONT pre-emptively checking + whether PID has been reused. + On Windows this has the effect of resuming all process threads. + """ + if POSIX: + self._send_signal(signal.SIGCONT) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.resume() + + def terminate(self): + """Terminate the process with SIGTERM pre-emptively checking + whether PID has been reused. + On Windows this is an alias for kill(). + """ + if POSIX: + self._send_signal(signal.SIGTERM) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.kill() + + def kill(self): + """Kill the current process with SIGKILL pre-emptively checking + whether PID has been reused. + """ + if POSIX: + self._send_signal(signal.SIGKILL) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.kill() + + def wait(self, timeout=None): + """Wait for process to terminate and, if process is a children + of os.getpid(), also return its exit code, else None. + On Windows there's no such limitation (exit code is always + returned). + + If the process is already terminated immediately return None + instead of raising NoSuchProcess. + + If *timeout* (in seconds) is specified and process is still + alive raise TimeoutExpired. + + To wait for multiple Process(es) use psutil.wait_procs(). + """ + if timeout is not None and not timeout >= 0: + msg = "timeout must be a positive integer" + raise ValueError(msg) + if self._exitcode is not _SENTINEL: + return self._exitcode + self._exitcode = self._proc.wait(timeout) + return self._exitcode + + +# The valid attr names which can be processed by Process.as_dict(). +# fmt: off +_as_dict_attrnames = { + x for x in dir(Process) if not x.startswith("_") and x not in + {'send_signal', 'suspend', 'resume', 'terminate', 'kill', 'wait', + 'is_running', 'as_dict', 'parent', 'parents', 'children', 'rlimit', + 'connections', 'oneshot'} +} +# fmt: on + + +# ===================================================================== +# --- Popen class +# ===================================================================== + + +class Popen(Process): + """Same as subprocess.Popen, but in addition it provides all + psutil.Process methods in a single class. + For the following methods which are common to both classes, psutil + implementation takes precedence: + + * send_signal() + * terminate() + * kill() + + This is done in order to avoid killing another process in case its + PID has been reused, fixing BPO-6973. + + >>> import psutil + >>> from subprocess import PIPE + >>> p = psutil.Popen(["python", "-c", "print 'hi'"], stdout=PIPE) + >>> p.name() + 'python' + >>> p.uids() + user(real=1000, effective=1000, saved=1000) + >>> p.username() + 'giampaolo' + >>> p.communicate() + ('hi', None) + >>> p.terminate() + >>> p.wait(timeout=2) + 0 + >>> + """ + + def __init__(self, *args, **kwargs): + # Explicitly avoid to raise NoSuchProcess in case the process + # spawned by subprocess.Popen terminates too quickly, see: + # https://github.com/giampaolo/psutil/issues/193 + self.__subproc = subprocess.Popen(*args, **kwargs) + self._init(self.__subproc.pid, _ignore_nsp=True) + + def __dir__(self): + return sorted(set(dir(Popen) + dir(subprocess.Popen))) + + def __enter__(self): + if hasattr(self.__subproc, '__enter__'): + self.__subproc.__enter__() + return self + + def __exit__(self, *args, **kwargs): + if hasattr(self.__subproc, '__exit__'): + return self.__subproc.__exit__(*args, **kwargs) + else: + if self.stdout: + self.stdout.close() + if self.stderr: + self.stderr.close() + try: + # Flushing a BufferedWriter may raise an error. + if self.stdin: + self.stdin.close() + finally: + # Wait for the process to terminate, to avoid zombies. + self.wait() + + def __getattribute__(self, name): + try: + return object.__getattribute__(self, name) + except AttributeError: + try: + return object.__getattribute__(self.__subproc, name) + except AttributeError: + msg = f"{self.__class__!r} has no attribute {name!r}" + raise AttributeError(msg) from None + + def wait(self, timeout=None): + if self.__subproc.returncode is not None: + return self.__subproc.returncode + ret = super().wait(timeout) + self.__subproc.returncode = ret + return ret + + +# ===================================================================== +# --- system processes related functions +# ===================================================================== + + +def pids(): + """Return a list of current running PIDs.""" + global _LOWEST_PID + ret = sorted(_psplatform.pids()) + _LOWEST_PID = ret[0] + return ret + + +def pid_exists(pid): + """Return True if given PID exists in the current process list. + This is faster than doing "pid in psutil.pids()" and + should be preferred. + """ + if pid < 0: + return False + elif pid == 0 and POSIX: + # On POSIX we use os.kill() to determine PID existence. + # According to "man 2 kill" PID 0 has a special meaning + # though: it refers to <> and that is not we want + # to do here. + return pid in pids() + else: + return _psplatform.pid_exists(pid) + + +_pmap = {} +_pids_reused = set() + + +def process_iter(attrs=None, ad_value=None): + """Return a generator yielding a Process instance for all + running processes. + + Every new Process instance is only created once and then cached + into an internal table which is updated every time this is used. + Cache can optionally be cleared via `process_iter.cache_clear()`. + + The sorting order in which processes are yielded is based on + their PIDs. + + *attrs* and *ad_value* have the same meaning as in + Process.as_dict(). If *attrs* is specified as_dict() is called + and the resulting dict is stored as a 'info' attribute attached + to returned Process instance. + If *attrs* is an empty list it will retrieve all process info + (slow). + """ + global _pmap + + def add(pid): + proc = Process(pid) + pmap[proc.pid] = proc + return proc + + def remove(pid): + pmap.pop(pid, None) + + pmap = _pmap.copy() + a = set(pids()) + b = set(pmap.keys()) + new_pids = a - b + gone_pids = b - a + for pid in gone_pids: + remove(pid) + while _pids_reused: + pid = _pids_reused.pop() + debug(f"refreshing Process instance for reused PID {pid}") + remove(pid) + try: + ls = sorted(list(pmap.items()) + list(dict.fromkeys(new_pids).items())) + for pid, proc in ls: + try: + if proc is None: # new process + proc = add(pid) + if attrs is not None: + proc.info = proc.as_dict(attrs=attrs, ad_value=ad_value) + yield proc + except NoSuchProcess: + remove(pid) + finally: + _pmap = pmap + + +process_iter.cache_clear = lambda: _pmap.clear() # noqa: PLW0108 +process_iter.cache_clear.__doc__ = "Clear process_iter() internal cache." + + +def wait_procs(procs, timeout=None, callback=None): + """Convenience function which waits for a list of processes to + terminate. + + Return a (gone, alive) tuple indicating which processes + are gone and which ones are still alive. + + The gone ones will have a new *returncode* attribute indicating + process exit status (may be None). + + *callback* is a function which gets called every time a process + terminates (a Process instance is passed as callback argument). + + Function will return as soon as all processes terminate or when + *timeout* occurs. + Differently from Process.wait() it will not raise TimeoutExpired if + *timeout* occurs. + + Typical use case is: + + - send SIGTERM to a list of processes + - give them some time to terminate + - send SIGKILL to those ones which are still alive + + Example: + + >>> def on_terminate(proc): + ... print("process {} terminated".format(proc)) + ... + >>> for p in procs: + ... p.terminate() + ... + >>> gone, alive = wait_procs(procs, timeout=3, callback=on_terminate) + >>> for p in alive: + ... p.kill() + """ + + def check_gone(proc, timeout): + try: + returncode = proc.wait(timeout=timeout) + except (TimeoutExpired, subprocess.TimeoutExpired): + pass + else: + if returncode is not None or not proc.is_running(): + # Set new Process instance attribute. + proc.returncode = returncode + gone.add(proc) + if callback is not None: + callback(proc) + + if timeout is not None and not timeout >= 0: + msg = f"timeout must be a positive integer, got {timeout}" + raise ValueError(msg) + gone = set() + alive = set(procs) + if callback is not None and not callable(callback): + msg = f"callback {callback!r} is not a callable" + raise TypeError(msg) + if timeout is not None: + deadline = _timer() + timeout + + while alive: + if timeout is not None and timeout <= 0: + break + for proc in alive: + # Make sure that every complete iteration (all processes) + # will last max 1 sec. + # We do this because we don't want to wait too long on a + # single process: in case it terminates too late other + # processes may disappear in the meantime and their PID + # reused. + max_timeout = 1.0 / len(alive) + if timeout is not None: + timeout = min((deadline - _timer()), max_timeout) + if timeout <= 0: + break + check_gone(proc, timeout) + else: + check_gone(proc, max_timeout) + alive = alive - gone # noqa: PLR6104 + + if alive: + # Last attempt over processes survived so far. + # timeout == 0 won't make this function wait any further. + for proc in alive: + check_gone(proc, 0) + alive = alive - gone # noqa: PLR6104 + + return (list(gone), list(alive)) + + +# ===================================================================== +# --- CPU related functions +# ===================================================================== + + +def cpu_count(logical=True): + """Return the number of logical CPUs in the system (same as + os.cpu_count()). + + If *logical* is False return the number of physical cores only + (e.g. hyper thread CPUs are excluded). + + Return None if undetermined. + + The return value is cached after first call. + If desired cache can be cleared like this: + + >>> psutil.cpu_count.cache_clear() + """ + if logical: + ret = _psplatform.cpu_count_logical() + else: + ret = _psplatform.cpu_count_cores() + if ret is not None and ret < 1: + ret = None + return ret + + +def cpu_times(percpu=False): + """Return system-wide CPU times as a namedtuple. + Every CPU time represents the seconds the CPU has spent in the + given mode. The namedtuple's fields availability varies depending on the + platform: + + - user + - system + - idle + - nice (UNIX) + - iowait (Linux) + - irq (Linux, FreeBSD) + - softirq (Linux) + - steal (Linux >= 2.6.11) + - guest (Linux >= 2.6.24) + - guest_nice (Linux >= 3.2.0) + + When *percpu* is True return a list of namedtuples for each CPU. + First element of the list refers to first CPU, second element + to second CPU and so on. + The order of the list is consistent across calls. + """ + if not percpu: + return _psplatform.cpu_times() + else: + return _psplatform.per_cpu_times() + + +try: + _last_cpu_times = {threading.current_thread().ident: cpu_times()} +except Exception: # noqa: BLE001 + # Don't want to crash at import time. + _last_cpu_times = {} + +try: + _last_per_cpu_times = { + threading.current_thread().ident: cpu_times(percpu=True) + } +except Exception: # noqa: BLE001 + # Don't want to crash at import time. + _last_per_cpu_times = {} + + +def _cpu_tot_time(times): + """Given a cpu_time() ntuple calculates the total CPU time + (including idle time). + """ + tot = sum(times) + if LINUX: + # On Linux guest times are already accounted in "user" or + # "nice" times, so we subtract them from total. + # Htop does the same. References: + # https://github.com/giampaolo/psutil/pull/940 + # http://unix.stackexchange.com/questions/178045 + # https://github.com/torvalds/linux/blob/ + # 447976ef4fd09b1be88b316d1a81553f1aa7cd07/kernel/sched/ + # cputime.c#L158 + tot -= getattr(times, "guest", 0) # Linux 2.6.24+ + tot -= getattr(times, "guest_nice", 0) # Linux 3.2.0+ + return tot + + +def _cpu_busy_time(times): + """Given a cpu_time() ntuple calculates the busy CPU time. + We do so by subtracting all idle CPU times. + """ + busy = _cpu_tot_time(times) + busy -= times.idle + # Linux: "iowait" is time during which the CPU does not do anything + # (waits for IO to complete). On Linux IO wait is *not* accounted + # in "idle" time so we subtract it. Htop does the same. + # References: + # https://github.com/torvalds/linux/blob/ + # 447976ef4fd09b1be88b316d1a81553f1aa7cd07/kernel/sched/cputime.c#L244 + busy -= getattr(times, "iowait", 0) + return busy + + +def _cpu_times_deltas(t1, t2): + assert t1._fields == t2._fields, (t1, t2) + field_deltas = [] + for field in _psplatform.scputimes._fields: + field_delta = getattr(t2, field) - getattr(t1, field) + # CPU times are always supposed to increase over time + # or at least remain the same and that's because time + # cannot go backwards. + # Surprisingly sometimes this might not be the case (at + # least on Windows and Linux), see: + # https://github.com/giampaolo/psutil/issues/392 + # https://github.com/giampaolo/psutil/issues/645 + # https://github.com/giampaolo/psutil/issues/1210 + # Trim negative deltas to zero to ignore decreasing fields. + # top does the same. Reference: + # https://gitlab.com/procps-ng/procps/blob/v3.3.12/top/top.c#L5063 + field_delta = max(0, field_delta) + field_deltas.append(field_delta) + return _psplatform.scputimes(*field_deltas) + + +def cpu_percent(interval=None, percpu=False): + """Return a float representing the current system-wide CPU + utilization as a percentage. + + When *interval* is > 0.0 compares system CPU times elapsed before + and after the interval (blocking). + + When *interval* is 0.0 or None compares system CPU times elapsed + since last call or module import, returning immediately (non + blocking). That means the first time this is called it will + return a meaningless 0.0 value which you should ignore. + In this case is recommended for accuracy that this function be + called with at least 0.1 seconds between calls. + + When *percpu* is True returns a list of floats representing the + utilization as a percentage for each CPU. + First element of the list refers to first CPU, second element + to second CPU and so on. + The order of the list is consistent across calls. + + Examples: + + >>> # blocking, system-wide + >>> psutil.cpu_percent(interval=1) + 2.0 + >>> + >>> # blocking, per-cpu + >>> psutil.cpu_percent(interval=1, percpu=True) + [2.0, 1.0] + >>> + >>> # non-blocking (percentage since last call) + >>> psutil.cpu_percent(interval=None) + 2.9 + >>> + """ + tid = threading.current_thread().ident + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval})" + raise ValueError(msg) + + def calculate(t1, t2): + times_delta = _cpu_times_deltas(t1, t2) + all_delta = _cpu_tot_time(times_delta) + busy_delta = _cpu_busy_time(times_delta) + + try: + busy_perc = (busy_delta / all_delta) * 100 + except ZeroDivisionError: + return 0.0 + else: + return round(busy_perc, 1) + + # system-wide usage + if not percpu: + if blocking: + t1 = cpu_times() + time.sleep(interval) + else: + t1 = _last_cpu_times.get(tid) or cpu_times() + _last_cpu_times[tid] = cpu_times() + return calculate(t1, _last_cpu_times[tid]) + # per-cpu usage + else: + ret = [] + if blocking: + tot1 = cpu_times(percpu=True) + time.sleep(interval) + else: + tot1 = _last_per_cpu_times.get(tid) or cpu_times(percpu=True) + _last_per_cpu_times[tid] = cpu_times(percpu=True) + for t1, t2 in zip(tot1, _last_per_cpu_times[tid]): + ret.append(calculate(t1, t2)) + return ret + + +# Use a separate dict for cpu_times_percent(), so it's independent from +# cpu_percent() and they can both be used within the same program. +_last_cpu_times_2 = _last_cpu_times.copy() +_last_per_cpu_times_2 = _last_per_cpu_times.copy() + + +def cpu_times_percent(interval=None, percpu=False): + """Same as cpu_percent() but provides utilization percentages + for each specific CPU time as is returned by cpu_times(). + For instance, on Linux we'll get: + + >>> cpu_times_percent() + cpupercent(user=4.8, nice=0.0, system=4.8, idle=90.5, iowait=0.0, + irq=0.0, softirq=0.0, steal=0.0, guest=0.0, guest_nice=0.0) + >>> + + *interval* and *percpu* arguments have the same meaning as in + cpu_percent(). + """ + tid = threading.current_thread().ident + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval!r})" + raise ValueError(msg) + + def calculate(t1, t2): + nums = [] + times_delta = _cpu_times_deltas(t1, t2) + all_delta = _cpu_tot_time(times_delta) + # "scale" is the value to multiply each delta with to get percentages. + # We use "max" to avoid division by zero (if all_delta is 0, then all + # fields are 0 so percentages will be 0 too. all_delta cannot be a + # fraction because cpu times are integers) + scale = 100.0 / max(1, all_delta) + for field_delta in times_delta: + field_perc = field_delta * scale + field_perc = round(field_perc, 1) + # make sure we don't return negative values or values over 100% + field_perc = min(max(0.0, field_perc), 100.0) + nums.append(field_perc) + return _psplatform.scputimes(*nums) + + # system-wide usage + if not percpu: + if blocking: + t1 = cpu_times() + time.sleep(interval) + else: + t1 = _last_cpu_times_2.get(tid) or cpu_times() + _last_cpu_times_2[tid] = cpu_times() + return calculate(t1, _last_cpu_times_2[tid]) + # per-cpu usage + else: + ret = [] + if blocking: + tot1 = cpu_times(percpu=True) + time.sleep(interval) + else: + tot1 = _last_per_cpu_times_2.get(tid) or cpu_times(percpu=True) + _last_per_cpu_times_2[tid] = cpu_times(percpu=True) + for t1, t2 in zip(tot1, _last_per_cpu_times_2[tid]): + ret.append(calculate(t1, t2)) + return ret + + +def cpu_stats(): + """Return CPU statistics.""" + return _psplatform.cpu_stats() + + +if hasattr(_psplatform, "cpu_freq"): + + def cpu_freq(percpu=False): + """Return CPU frequency as a namedtuple including current, + min and max frequency expressed in Mhz. + + If *percpu* is True and the system supports per-cpu frequency + retrieval (Linux only) a list of frequencies is returned for + each CPU. If not a list with one element is returned. + """ + ret = _psplatform.cpu_freq() + if percpu: + return ret + else: + num_cpus = float(len(ret)) + if num_cpus == 0: + return None + elif num_cpus == 1: + return ret[0] + else: + currs, mins, maxs = 0.0, 0.0, 0.0 + set_none = False + for cpu in ret: + currs += cpu.current + # On Linux if /proc/cpuinfo is used min/max are set + # to None. + if LINUX and cpu.min is None: + set_none = True + continue + mins += cpu.min + maxs += cpu.max + + current = currs / num_cpus + + if set_none: + min_ = max_ = None + else: + min_ = mins / num_cpus + max_ = maxs / num_cpus + + return _common.scpufreq(current, min_, max_) + + __all__.append("cpu_freq") + + +if hasattr(os, "getloadavg") or hasattr(_psplatform, "getloadavg"): + # Perform this hasattr check once on import time to either use the + # platform based code or proxy straight from the os module. + if hasattr(os, "getloadavg"): + getloadavg = os.getloadavg + else: + getloadavg = _psplatform.getloadavg + + __all__.append("getloadavg") + + +# ===================================================================== +# --- system memory related functions +# ===================================================================== + + +def virtual_memory(): + """Return statistics about system memory usage as a namedtuple + including the following fields, expressed in bytes: + + - total: + total physical memory available. + + - available: + the memory that can be given instantly to processes without the + system going into swap. + This is calculated by summing different memory values depending + on the platform and it is supposed to be used to monitor actual + memory usage in a cross platform fashion. + + - percent: + the percentage usage calculated as (total - available) / total * 100 + + - used: + memory used, calculated differently depending on the platform and + designed for informational purposes only: + macOS: active + wired + BSD: active + wired + cached + Linux: total - free + + - free: + memory not being used at all (zeroed) that is readily available; + note that this doesn't reflect the actual memory available + (use 'available' instead) + + Platform-specific fields: + + - active (UNIX): + memory currently in use or very recently used, and so it is in RAM. + + - inactive (UNIX): + memory that is marked as not used. + + - buffers (BSD, Linux): + cache for things like file system metadata. + + - cached (BSD, macOS): + cache for various things. + + - wired (macOS, BSD): + memory that is marked to always stay in RAM. It is never moved to disk. + + - shared (BSD): + memory that may be simultaneously accessed by multiple processes. + + The sum of 'used' and 'available' does not necessarily equal total. + On Windows 'available' and 'free' are the same. + """ + global _TOTAL_PHYMEM + ret = _psplatform.virtual_memory() + # cached for later use in Process.memory_percent() + _TOTAL_PHYMEM = ret.total + return ret + + +def swap_memory(): + """Return system swap memory statistics as a namedtuple including + the following fields: + + - total: total swap memory in bytes + - used: used swap memory in bytes + - free: free swap memory in bytes + - percent: the percentage usage + - sin: no. of bytes the system has swapped in from disk (cumulative) + - sout: no. of bytes the system has swapped out from disk (cumulative) + + 'sin' and 'sout' on Windows are meaningless and always set to 0. + """ + return _psplatform.swap_memory() + + +# ===================================================================== +# --- disks/partitions related functions +# ===================================================================== + + +def disk_usage(path): + """Return disk usage statistics about the given *path* as a + namedtuple including total, used and free space expressed in bytes + plus the percentage usage. + """ + return _psplatform.disk_usage(path) + + +def disk_partitions(all=False): + """Return mounted partitions as a list of + (device, mountpoint, fstype, opts) namedtuple. + 'opts' field is a raw string separated by commas indicating mount + options which may vary depending on the platform. + + If *all* parameter is False return physical devices only and ignore + all others. + """ + return _psplatform.disk_partitions(all) + + +def disk_io_counters(perdisk=False, nowrap=True): + """Return system disk I/O statistics as a namedtuple including + the following fields: + + - read_count: number of reads + - write_count: number of writes + - read_bytes: number of bytes read + - write_bytes: number of bytes written + - read_time: time spent reading from disk (in ms) + - write_time: time spent writing to disk (in ms) + + Platform specific: + + - busy_time: (Linux, FreeBSD) time spent doing actual I/Os (in ms) + - read_merged_count (Linux): number of merged reads + - write_merged_count (Linux): number of merged writes + + If *perdisk* is True return the same information for every + physical disk installed on the system as a dictionary + with partition names as the keys and the namedtuple + described above as the values. + + If *nowrap* is True it detects and adjust the numbers which overflow + and wrap (restart from 0) and add "old value" to "new value" so that + the returned numbers will always be increasing or remain the same, + but never decrease. + "disk_io_counters.cache_clear()" can be used to invalidate the + cache. + + On recent Windows versions 'diskperf -y' command may need to be + executed first otherwise this function won't find any disk. + """ + kwargs = dict(perdisk=perdisk) if LINUX else {} + rawdict = _psplatform.disk_io_counters(**kwargs) + if not rawdict: + return {} if perdisk else None + if nowrap: + rawdict = _wrap_numbers(rawdict, 'psutil.disk_io_counters') + nt = getattr(_psplatform, "sdiskio", _common.sdiskio) + if perdisk: + for disk, fields in rawdict.items(): + rawdict[disk] = nt(*fields) + return rawdict + else: + return nt(*(sum(x) for x in zip(*rawdict.values()))) + + +disk_io_counters.cache_clear = functools.partial( + _wrap_numbers.cache_clear, 'psutil.disk_io_counters' +) +disk_io_counters.cache_clear.__doc__ = "Clears nowrap argument cache" + + +# ===================================================================== +# --- network related functions +# ===================================================================== + + +def net_io_counters(pernic=False, nowrap=True): + """Return network I/O statistics as a namedtuple including + the following fields: + + - bytes_sent: number of bytes sent + - bytes_recv: number of bytes received + - packets_sent: number of packets sent + - packets_recv: number of packets received + - errin: total number of errors while receiving + - errout: total number of errors while sending + - dropin: total number of incoming packets which were dropped + - dropout: total number of outgoing packets which were dropped + (always 0 on macOS and BSD) + + If *pernic* is True return the same information for every + network interface installed on the system as a dictionary + with network interface names as the keys and the namedtuple + described above as the values. + + If *nowrap* is True it detects and adjust the numbers which overflow + and wrap (restart from 0) and add "old value" to "new value" so that + the returned numbers will always be increasing or remain the same, + but never decrease. + "net_io_counters.cache_clear()" can be used to invalidate the + cache. + """ + rawdict = _psplatform.net_io_counters() + if not rawdict: + return {} if pernic else None + if nowrap: + rawdict = _wrap_numbers(rawdict, 'psutil.net_io_counters') + if pernic: + for nic, fields in rawdict.items(): + rawdict[nic] = _common.snetio(*fields) + return rawdict + else: + return _common.snetio(*[sum(x) for x in zip(*rawdict.values())]) + + +net_io_counters.cache_clear = functools.partial( + _wrap_numbers.cache_clear, 'psutil.net_io_counters' +) +net_io_counters.cache_clear.__doc__ = "Clears nowrap argument cache" + + +def net_connections(kind='inet'): + """Return system-wide socket connections as a list of + (fd, family, type, laddr, raddr, status, pid) namedtuples. + In case of limited privileges 'fd' and 'pid' may be set to -1 + and None respectively. + The *kind* parameter filters for connections that fit the + following criteria: + + +------------+----------------------------------------------------+ + | Kind Value | Connections using | + +------------+----------------------------------------------------+ + | inet | IPv4 and IPv6 | + | inet4 | IPv4 | + | inet6 | IPv6 | + | tcp | TCP | + | tcp4 | TCP over IPv4 | + | tcp6 | TCP over IPv6 | + | udp | UDP | + | udp4 | UDP over IPv4 | + | udp6 | UDP over IPv6 | + | unix | UNIX socket (both UDP and TCP protocols) | + | all | the sum of all the possible families and protocols | + +------------+----------------------------------------------------+ + + On macOS this function requires root privileges. + """ + _check_conn_kind(kind) + return _psplatform.net_connections(kind) + + +def net_if_addrs(): + """Return the addresses associated to each NIC (network interface + card) installed on the system as a dictionary whose keys are the + NIC names and value is a list of namedtuples for each address + assigned to the NIC. Each namedtuple includes 5 fields: + + - family: can be either socket.AF_INET, socket.AF_INET6 or + psutil.AF_LINK, which refers to a MAC address. + - address: is the primary address and it is always set. + - netmask: and 'broadcast' and 'ptp' may be None. + - ptp: stands for "point to point" and references the + destination address on a point to point interface + (typically a VPN). + - broadcast: and *ptp* are mutually exclusive. + + Note: you can have more than one address of the same family + associated with each interface. + """ + rawlist = _psplatform.net_if_addrs() + rawlist.sort(key=lambda x: x[1]) # sort by family + ret = collections.defaultdict(list) + for name, fam, addr, mask, broadcast, ptp in rawlist: + try: + fam = socket.AddressFamily(fam) + except ValueError: + if WINDOWS and fam == -1: + fam = _psplatform.AF_LINK + elif ( + hasattr(_psplatform, "AF_LINK") and fam == _psplatform.AF_LINK + ): + # Linux defines AF_LINK as an alias for AF_PACKET. + # We re-set the family here so that repr(family) + # will show AF_LINK rather than AF_PACKET + fam = _psplatform.AF_LINK + + if fam == _psplatform.AF_LINK: + # The underlying C function may return an incomplete MAC + # address in which case we fill it with null bytes, see: + # https://github.com/giampaolo/psutil/issues/786 + separator = ":" if POSIX else "-" + while addr.count(separator) < 5: + addr += f"{separator}00" + + nt = _common.snicaddr(fam, addr, mask, broadcast, ptp) + + # On Windows broadcast is None, so we determine it via + # ipaddress module. + if WINDOWS and fam in {socket.AF_INET, socket.AF_INET6}: + try: + broadcast = _common.broadcast_addr(nt) + except Exception as err: # noqa: BLE001 + debug(err) + else: + if broadcast is not None: + nt._replace(broadcast=broadcast) + + ret[name].append(nt) + + return dict(ret) + + +def net_if_stats(): + """Return information about each NIC (network interface card) + installed on the system as a dictionary whose keys are the + NIC names and value is a namedtuple with the following fields: + + - isup: whether the interface is up (bool) + - duplex: can be either NIC_DUPLEX_FULL, NIC_DUPLEX_HALF or + NIC_DUPLEX_UNKNOWN + - speed: the NIC speed expressed in mega bits (MB); if it can't + be determined (e.g. 'localhost') it will be set to 0. + - mtu: the maximum transmission unit expressed in bytes. + """ + return _psplatform.net_if_stats() + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +# Linux, macOS +if hasattr(_psplatform, "sensors_temperatures"): + + def sensors_temperatures(fahrenheit=False): + """Return hardware temperatures. Each entry is a namedtuple + representing a certain hardware sensor (it may be a CPU, an + hard disk or something else, depending on the OS and its + configuration). + All temperatures are expressed in celsius unless *fahrenheit* + is set to True. + """ + + def convert(n): + if n is not None: + return (float(n) * 9 / 5) + 32 if fahrenheit else n + + ret = collections.defaultdict(list) + rawdict = _psplatform.sensors_temperatures() + + for name, values in rawdict.items(): + while values: + label, current, high, critical = values.pop(0) + current = convert(current) + high = convert(high) + critical = convert(critical) + + if high and not critical: + critical = high + elif critical and not high: + high = critical + + ret[name].append( + _common.shwtemp(label, current, high, critical) + ) + + return dict(ret) + + __all__.append("sensors_temperatures") + + +# Linux +if hasattr(_psplatform, "sensors_fans"): + + def sensors_fans(): + """Return fans speed. Each entry is a namedtuple + representing a certain hardware sensor. + All speed are expressed in RPM (rounds per minute). + """ + return _psplatform.sensors_fans() + + __all__.append("sensors_fans") + + +# Linux, Windows, FreeBSD, macOS +if hasattr(_psplatform, "sensors_battery"): + + def sensors_battery(): + """Return battery information. If no battery is installed + returns None. + + - percent: battery power left as a percentage. + - secsleft: a rough approximation of how many seconds are left + before the battery runs out of power. May be + POWER_TIME_UNLIMITED or POWER_TIME_UNLIMITED. + - power_plugged: True if the AC power cable is connected. + """ + return _psplatform.sensors_battery() + + __all__.append("sensors_battery") + + +# ===================================================================== +# --- other system related functions +# ===================================================================== + + +def boot_time(): + """Return the system boot time expressed in seconds since the epoch + (seconds since January 1, 1970, at midnight UTC). The returned + value is based on the system clock, which means it may be affected + by changes such as manual adjustments or time synchronization (e.g. + NTP). + """ + return _psplatform.boot_time() + + +def users(): + """Return users currently connected on the system as a list of + namedtuples including the following fields. + + - user: the name of the user + - terminal: the tty or pseudo-tty associated with the user, if any. + - host: the host name associated with the entry, if any. + - started: the creation time as a floating point number expressed in + seconds since the epoch. + """ + return _psplatform.users() + + +# ===================================================================== +# --- Windows services +# ===================================================================== + + +if WINDOWS: + + def win_service_iter(): + """Return a generator yielding a WindowsService instance for all + Windows services installed. + """ + return _psplatform.win_service_iter() + + def win_service_get(name): + """Get a Windows service by *name*. + Raise NoSuchProcess if no service with such name exists. + """ + return _psplatform.win_service_get(name) + + +# ===================================================================== + + +def _set_debug(value): + """Enable or disable PSUTIL_DEBUG option, which prints debugging + messages to stderr. + """ + import psutil._common + + psutil._common.PSUTIL_DEBUG = bool(value) + _psplatform.cext.set_debug(bool(value)) + + +del memoize_when_activated diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..89079161b19b852bdc419c95f0fe601cd3be6c6f Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_common.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad7cecafd780f7f8412bfd0ceeb2255d5ca58e27 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psaix.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psaix.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4971c06fcf43c3977adfd24124f6eb3855fb56d Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psaix.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psbsd.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psbsd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31143e45a7d87c59c32b33d4286b8456cb5c137a Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psbsd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pslinux.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pslinux.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd68f0865078a3867287b52b4b27dfb7ab0a0465 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pslinux.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psosx.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psosx.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a33d7582805f250556b174848dc7a574e0b3477 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psosx.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psposix.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psposix.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6482609d53e54b03d17054f50d246233cada60c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_psposix.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pssunos.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pssunos.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..868d3c930f8a42926a35ecdd28e67bde4a14646e Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pssunos.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pswindows.cpython-312.pyc b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pswindows.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f45aef153c1fadd4394810f59cda953bc3733d4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/__pycache__/_pswindows.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_common.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..51c798b4c670354be1c154e8751f9bffba78a375 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_common.py @@ -0,0 +1,948 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Common objects shared by __init__.py and _ps*.py modules. + +Note: this module is imported by setup.py, so it should not import +psutil or third-party modules. +""" + +import collections +import enum +import functools +import os +import socket +import stat +import sys +import threading +import warnings +from collections import namedtuple +from socket import AF_INET +from socket import SOCK_DGRAM +from socket import SOCK_STREAM + +try: + from socket import AF_INET6 +except ImportError: + AF_INET6 = None +try: + from socket import AF_UNIX +except ImportError: + AF_UNIX = None + + +PSUTIL_DEBUG = bool(os.getenv('PSUTIL_DEBUG')) +_DEFAULT = object() + +# fmt: off +__all__ = [ + # OS constants + 'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'MACOS', 'OSX', 'POSIX', + 'SUNOS', 'WINDOWS', + # connection constants + 'CONN_CLOSE', 'CONN_CLOSE_WAIT', 'CONN_CLOSING', 'CONN_ESTABLISHED', + 'CONN_FIN_WAIT1', 'CONN_FIN_WAIT2', 'CONN_LAST_ACK', 'CONN_LISTEN', + 'CONN_NONE', 'CONN_SYN_RECV', 'CONN_SYN_SENT', 'CONN_TIME_WAIT', + # net constants + 'NIC_DUPLEX_FULL', 'NIC_DUPLEX_HALF', 'NIC_DUPLEX_UNKNOWN', # noqa: F822 + # process status constants + 'STATUS_DEAD', 'STATUS_DISK_SLEEP', 'STATUS_IDLE', 'STATUS_LOCKED', + 'STATUS_RUNNING', 'STATUS_SLEEPING', 'STATUS_STOPPED', 'STATUS_SUSPENDED', + 'STATUS_TRACING_STOP', 'STATUS_WAITING', 'STATUS_WAKE_KILL', + 'STATUS_WAKING', 'STATUS_ZOMBIE', 'STATUS_PARKED', + # other constants + 'ENCODING', 'ENCODING_ERRS', 'AF_INET6', + # named tuples + 'pconn', 'pcputimes', 'pctxsw', 'pgids', 'pio', 'pionice', 'popenfile', + 'pthread', 'puids', 'sconn', 'scpustats', 'sdiskio', 'sdiskpart', + 'sdiskusage', 'snetio', 'snicaddr', 'snicstats', 'sswap', 'suser', + # utility functions + 'conn_tmap', 'deprecated_method', 'isfile_strict', 'memoize', + 'parse_environ_block', 'path_exists_strict', 'usage_percent', + 'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum', "wrap_numbers", + 'open_text', 'open_binary', 'cat', 'bcat', + 'bytes2human', 'conn_to_ntuple', 'debug', + # shell utils + 'hilite', 'term_supports_colors', 'print_color', +] +# fmt: on + + +# =================================================================== +# --- OS constants +# =================================================================== + + +POSIX = os.name == "posix" +WINDOWS = os.name == "nt" +LINUX = sys.platform.startswith("linux") +MACOS = sys.platform.startswith("darwin") +OSX = MACOS # deprecated alias +FREEBSD = sys.platform.startswith(("freebsd", "midnightbsd")) +OPENBSD = sys.platform.startswith("openbsd") +NETBSD = sys.platform.startswith("netbsd") +BSD = FREEBSD or OPENBSD or NETBSD +SUNOS = sys.platform.startswith(("sunos", "solaris")) +AIX = sys.platform.startswith("aix") + + +# =================================================================== +# --- API constants +# =================================================================== + + +# Process.status() +STATUS_RUNNING = "running" +STATUS_SLEEPING = "sleeping" +STATUS_DISK_SLEEP = "disk-sleep" +STATUS_STOPPED = "stopped" +STATUS_TRACING_STOP = "tracing-stop" +STATUS_ZOMBIE = "zombie" +STATUS_DEAD = "dead" +STATUS_WAKE_KILL = "wake-kill" +STATUS_WAKING = "waking" +STATUS_IDLE = "idle" # Linux, macOS, FreeBSD +STATUS_LOCKED = "locked" # FreeBSD +STATUS_WAITING = "waiting" # FreeBSD +STATUS_SUSPENDED = "suspended" # NetBSD +STATUS_PARKED = "parked" # Linux + +# Process.net_connections() and psutil.net_connections() +CONN_ESTABLISHED = "ESTABLISHED" +CONN_SYN_SENT = "SYN_SENT" +CONN_SYN_RECV = "SYN_RECV" +CONN_FIN_WAIT1 = "FIN_WAIT1" +CONN_FIN_WAIT2 = "FIN_WAIT2" +CONN_TIME_WAIT = "TIME_WAIT" +CONN_CLOSE = "CLOSE" +CONN_CLOSE_WAIT = "CLOSE_WAIT" +CONN_LAST_ACK = "LAST_ACK" +CONN_LISTEN = "LISTEN" +CONN_CLOSING = "CLOSING" +CONN_NONE = "NONE" + + +# net_if_stats() +class NicDuplex(enum.IntEnum): + NIC_DUPLEX_FULL = 2 + NIC_DUPLEX_HALF = 1 + NIC_DUPLEX_UNKNOWN = 0 + + +globals().update(NicDuplex.__members__) + + +# sensors_battery() +class BatteryTime(enum.IntEnum): + POWER_TIME_UNKNOWN = -1 + POWER_TIME_UNLIMITED = -2 + + +globals().update(BatteryTime.__members__) + +# --- others + +ENCODING = sys.getfilesystemencoding() +ENCODING_ERRS = sys.getfilesystemencodeerrors() + + +# =================================================================== +# --- namedtuples +# =================================================================== + +# --- for system functions + +# fmt: off +# psutil.swap_memory() +sswap = namedtuple('sswap', ['total', 'used', 'free', 'percent', 'sin', + 'sout']) +# psutil.disk_usage() +sdiskusage = namedtuple('sdiskusage', ['total', 'used', 'free', 'percent']) +# psutil.disk_io_counters() +sdiskio = namedtuple('sdiskio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes', + 'read_time', 'write_time']) +# psutil.disk_partitions() +sdiskpart = namedtuple('sdiskpart', ['device', 'mountpoint', 'fstype', 'opts']) +# psutil.net_io_counters() +snetio = namedtuple('snetio', ['bytes_sent', 'bytes_recv', + 'packets_sent', 'packets_recv', + 'errin', 'errout', + 'dropin', 'dropout']) +# psutil.users() +suser = namedtuple('suser', ['name', 'terminal', 'host', 'started', 'pid']) +# psutil.net_connections() +sconn = namedtuple('sconn', ['fd', 'family', 'type', 'laddr', 'raddr', + 'status', 'pid']) +# psutil.net_if_addrs() +snicaddr = namedtuple('snicaddr', + ['family', 'address', 'netmask', 'broadcast', 'ptp']) +# psutil.net_if_stats() +snicstats = namedtuple('snicstats', + ['isup', 'duplex', 'speed', 'mtu', 'flags']) +# psutil.cpu_stats() +scpustats = namedtuple( + 'scpustats', ['ctx_switches', 'interrupts', 'soft_interrupts', 'syscalls']) +# psutil.cpu_freq() +scpufreq = namedtuple('scpufreq', ['current', 'min', 'max']) +# psutil.sensors_temperatures() +shwtemp = namedtuple( + 'shwtemp', ['label', 'current', 'high', 'critical']) +# psutil.sensors_battery() +sbattery = namedtuple('sbattery', ['percent', 'secsleft', 'power_plugged']) +# psutil.sensors_fans() +sfan = namedtuple('sfan', ['label', 'current']) +# fmt: on + +# --- for Process methods + +# psutil.Process.cpu_times() +pcputimes = namedtuple( + 'pcputimes', ['user', 'system', 'children_user', 'children_system'] +) +# psutil.Process.open_files() +popenfile = namedtuple('popenfile', ['path', 'fd']) +# psutil.Process.threads() +pthread = namedtuple('pthread', ['id', 'user_time', 'system_time']) +# psutil.Process.uids() +puids = namedtuple('puids', ['real', 'effective', 'saved']) +# psutil.Process.gids() +pgids = namedtuple('pgids', ['real', 'effective', 'saved']) +# psutil.Process.io_counters() +pio = namedtuple( + 'pio', ['read_count', 'write_count', 'read_bytes', 'write_bytes'] +) +# psutil.Process.ionice() +pionice = namedtuple('pionice', ['ioclass', 'value']) +# psutil.Process.ctx_switches() +pctxsw = namedtuple('pctxsw', ['voluntary', 'involuntary']) +# psutil.Process.net_connections() +pconn = namedtuple( + 'pconn', ['fd', 'family', 'type', 'laddr', 'raddr', 'status'] +) + +# psutil.net_connections() and psutil.Process.net_connections() +addr = namedtuple('addr', ['ip', 'port']) + + +# =================================================================== +# --- Process.net_connections() 'kind' parameter mapping +# =================================================================== + + +conn_tmap = { + "all": ([AF_INET, AF_INET6, AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]), + "tcp": ([AF_INET, AF_INET6], [SOCK_STREAM]), + "tcp4": ([AF_INET], [SOCK_STREAM]), + "udp": ([AF_INET, AF_INET6], [SOCK_DGRAM]), + "udp4": ([AF_INET], [SOCK_DGRAM]), + "inet": ([AF_INET, AF_INET6], [SOCK_STREAM, SOCK_DGRAM]), + "inet4": ([AF_INET], [SOCK_STREAM, SOCK_DGRAM]), + "inet6": ([AF_INET6], [SOCK_STREAM, SOCK_DGRAM]), +} + +if AF_INET6 is not None: + conn_tmap.update({ + "tcp6": ([AF_INET6], [SOCK_STREAM]), + "udp6": ([AF_INET6], [SOCK_DGRAM]), + }) + +if AF_UNIX is not None and not SUNOS: + conn_tmap.update({"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM])}) + + +# ===================================================================== +# --- Exceptions +# ===================================================================== + + +class Error(Exception): + """Base exception class. All other psutil exceptions inherit + from this one. + """ + + __module__ = 'psutil' + + def _infodict(self, attrs): + info = collections.OrderedDict() + for name in attrs: + value = getattr(self, name, None) + if value or (name == "pid" and value == 0): + info[name] = value + return info + + def __str__(self): + # invoked on `raise Error` + info = self._infodict(("pid", "ppid", "name")) + if info: + details = "({})".format( + ", ".join([f"{k}={v!r}" for k, v in info.items()]) + ) + else: + details = None + return " ".join([x for x in (getattr(self, "msg", ""), details) if x]) + + def __repr__(self): + # invoked on `repr(Error)` + info = self._infodict(("pid", "ppid", "name", "seconds", "msg")) + details = ", ".join([f"{k}={v!r}" for k, v in info.items()]) + return f"psutil.{self.__class__.__name__}({details})" + + +class NoSuchProcess(Error): + """Exception raised when a process with a certain PID doesn't + or no longer exists. + """ + + __module__ = 'psutil' + + def __init__(self, pid, name=None, msg=None): + Error.__init__(self) + self.pid = pid + self.name = name + self.msg = msg or "process no longer exists" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.msg)) + + +class ZombieProcess(NoSuchProcess): + """Exception raised when querying a zombie process. This is + raised on macOS, BSD and Solaris only, and not always: depending + on the query the OS may be able to succeed anyway. + On Linux all zombie processes are querable (hence this is never + raised). Windows doesn't have zombie processes. + """ + + __module__ = 'psutil' + + def __init__(self, pid, name=None, ppid=None, msg=None): + NoSuchProcess.__init__(self, pid, name, msg) + self.ppid = ppid + self.msg = msg or "PID still exists but it's a zombie" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.ppid, self.msg)) + + +class AccessDenied(Error): + """Exception raised when permission to perform an action is denied.""" + + __module__ = 'psutil' + + def __init__(self, pid=None, name=None, msg=None): + Error.__init__(self) + self.pid = pid + self.name = name + self.msg = msg or "" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.msg)) + + +class TimeoutExpired(Error): + """Raised on Process.wait(timeout) if timeout expires and process + is still alive. + """ + + __module__ = 'psutil' + + def __init__(self, seconds, pid=None, name=None): + Error.__init__(self) + self.seconds = seconds + self.pid = pid + self.name = name + self.msg = f"timeout after {seconds} seconds" + + def __reduce__(self): + return (self.__class__, (self.seconds, self.pid, self.name)) + + +# =================================================================== +# --- utils +# =================================================================== + + +def usage_percent(used, total, round_=None): + """Calculate percentage usage of 'used' against 'total'.""" + try: + ret = (float(used) / total) * 100 + except ZeroDivisionError: + return 0.0 + else: + if round_ is not None: + ret = round(ret, round_) + return ret + + +def memoize(fun): + """A simple memoize decorator for functions supporting (hashable) + positional arguments. + It also provides a cache_clear() function for clearing the cache: + + >>> @memoize + ... def foo() + ... return 1 + ... + >>> foo() + 1 + >>> foo.cache_clear() + >>> + + It supports: + - functions + - classes (acts as a @singleton) + - staticmethods + - classmethods + + It does NOT support: + - methods + """ + + @functools.wraps(fun) + def wrapper(*args, **kwargs): + key = (args, frozenset(sorted(kwargs.items()))) + try: + return cache[key] + except KeyError: + try: + ret = cache[key] = fun(*args, **kwargs) + except Exception as err: + raise err from None + return ret + + def cache_clear(): + """Clear cache.""" + cache.clear() + + cache = {} + wrapper.cache_clear = cache_clear + return wrapper + + +def memoize_when_activated(fun): + """A memoize decorator which is disabled by default. It can be + activated and deactivated on request. + For efficiency reasons it can be used only against class methods + accepting no arguments. + + >>> class Foo: + ... @memoize + ... def foo() + ... print(1) + ... + >>> f = Foo() + >>> # deactivated (default) + >>> foo() + 1 + >>> foo() + 1 + >>> + >>> # activated + >>> foo.cache_activate(self) + >>> foo() + 1 + >>> foo() + >>> foo() + >>> + """ + + @functools.wraps(fun) + def wrapper(self): + try: + # case 1: we previously entered oneshot() ctx + ret = self._cache[fun] + except AttributeError: + # case 2: we never entered oneshot() ctx + try: + return fun(self) + except Exception as err: + raise err from None + except KeyError: + # case 3: we entered oneshot() ctx but there's no cache + # for this entry yet + try: + ret = fun(self) + except Exception as err: + raise err from None + try: + self._cache[fun] = ret + except AttributeError: + # multi-threading race condition, see: + # https://github.com/giampaolo/psutil/issues/1948 + pass + return ret + + def cache_activate(proc): + """Activate cache. Expects a Process instance. Cache will be + stored as a "_cache" instance attribute. + """ + proc._cache = {} + + def cache_deactivate(proc): + """Deactivate and clear cache.""" + try: + del proc._cache + except AttributeError: + pass + + wrapper.cache_activate = cache_activate + wrapper.cache_deactivate = cache_deactivate + return wrapper + + +def isfile_strict(path): + """Same as os.path.isfile() but does not swallow EACCES / EPERM + exceptions, see: + http://mail.python.org/pipermail/python-dev/2012-June/120787.html. + """ + try: + st = os.stat(path) + except PermissionError: + raise + except OSError: + return False + else: + return stat.S_ISREG(st.st_mode) + + +def path_exists_strict(path): + """Same as os.path.exists() but does not swallow EACCES / EPERM + exceptions. See: + http://mail.python.org/pipermail/python-dev/2012-June/120787.html. + """ + try: + os.stat(path) + except PermissionError: + raise + except OSError: + return False + else: + return True + + +def supports_ipv6(): + """Return True if IPv6 is supported on this platform.""" + if not socket.has_ipv6 or AF_INET6 is None: + return False + try: + with socket.socket(AF_INET6, socket.SOCK_STREAM) as sock: + sock.bind(("::1", 0)) + return True + except OSError: + return False + + +def parse_environ_block(data): + """Parse a C environ block of environment variables into a dictionary.""" + # The block is usually raw data from the target process. It might contain + # trailing garbage and lines that do not look like assignments. + ret = {} + pos = 0 + + # localize global variable to speed up access. + WINDOWS_ = WINDOWS + while True: + next_pos = data.find("\0", pos) + # nul byte at the beginning or double nul byte means finish + if next_pos <= pos: + break + # there might not be an equals sign + equal_pos = data.find("=", pos, next_pos) + if equal_pos > pos: + key = data[pos:equal_pos] + value = data[equal_pos + 1 : next_pos] + # Windows expects environment variables to be uppercase only + if WINDOWS_: + key = key.upper() + ret[key] = value + pos = next_pos + 1 + + return ret + + +def sockfam_to_enum(num): + """Convert a numeric socket family value to an IntEnum member. + If it's not a known member, return the numeric value itself. + """ + try: + return socket.AddressFamily(num) + except ValueError: + return num + + +def socktype_to_enum(num): + """Convert a numeric socket type value to an IntEnum member. + If it's not a known member, return the numeric value itself. + """ + try: + return socket.SocketKind(num) + except ValueError: + return num + + +def conn_to_ntuple(fd, fam, type_, laddr, raddr, status, status_map, pid=None): + """Convert a raw connection tuple to a proper ntuple.""" + if fam in {socket.AF_INET, AF_INET6}: + if laddr: + laddr = addr(*laddr) + if raddr: + raddr = addr(*raddr) + if type_ == socket.SOCK_STREAM and fam in {AF_INET, AF_INET6}: + status = status_map.get(status, CONN_NONE) + else: + status = CONN_NONE # ignore whatever C returned to us + fam = sockfam_to_enum(fam) + type_ = socktype_to_enum(type_) + if pid is None: + return pconn(fd, fam, type_, laddr, raddr, status) + else: + return sconn(fd, fam, type_, laddr, raddr, status, pid) + + +def broadcast_addr(addr): + """Given the address ntuple returned by ``net_if_addrs()`` + calculates the broadcast address. + """ + import ipaddress + + if not addr.address or not addr.netmask: + return None + if addr.family == socket.AF_INET: + return str( + ipaddress.IPv4Network( + f"{addr.address}/{addr.netmask}", strict=False + ).broadcast_address + ) + if addr.family == socket.AF_INET6: + return str( + ipaddress.IPv6Network( + f"{addr.address}/{addr.netmask}", strict=False + ).broadcast_address + ) + + +def deprecated_method(replacement): + """A decorator which can be used to mark a method as deprecated + 'replcement' is the method name which will be called instead. + """ + + def outer(fun): + msg = ( + f"{fun.__name__}() is deprecated and will be removed; use" + f" {replacement}() instead" + ) + if fun.__doc__ is None: + fun.__doc__ = msg + + @functools.wraps(fun) + def inner(self, *args, **kwargs): + warnings.warn(msg, category=DeprecationWarning, stacklevel=2) + return getattr(self, replacement)(*args, **kwargs) + + return inner + + return outer + + +class _WrapNumbers: + """Watches numbers so that they don't overflow and wrap + (reset to zero). + """ + + def __init__(self): + self.lock = threading.Lock() + self.cache = {} + self.reminders = {} + self.reminder_keys = {} + + def _add_dict(self, input_dict, name): + assert name not in self.cache + assert name not in self.reminders + assert name not in self.reminder_keys + self.cache[name] = input_dict + self.reminders[name] = collections.defaultdict(int) + self.reminder_keys[name] = collections.defaultdict(set) + + def _remove_dead_reminders(self, input_dict, name): + """In case the number of keys changed between calls (e.g. a + disk disappears) this removes the entry from self.reminders. + """ + old_dict = self.cache[name] + gone_keys = set(old_dict.keys()) - set(input_dict.keys()) + for gone_key in gone_keys: + for remkey in self.reminder_keys[name][gone_key]: + del self.reminders[name][remkey] + del self.reminder_keys[name][gone_key] + + def run(self, input_dict, name): + """Cache dict and sum numbers which overflow and wrap. + Return an updated copy of `input_dict`. + """ + if name not in self.cache: + # This was the first call. + self._add_dict(input_dict, name) + return input_dict + + self._remove_dead_reminders(input_dict, name) + + old_dict = self.cache[name] + new_dict = {} + for key in input_dict: + input_tuple = input_dict[key] + try: + old_tuple = old_dict[key] + except KeyError: + # The input dict has a new key (e.g. a new disk or NIC) + # which didn't exist in the previous call. + new_dict[key] = input_tuple + continue + + bits = [] + for i in range(len(input_tuple)): + input_value = input_tuple[i] + old_value = old_tuple[i] + remkey = (key, i) + if input_value < old_value: + # it wrapped! + self.reminders[name][remkey] += old_value + self.reminder_keys[name][key].add(remkey) + bits.append(input_value + self.reminders[name][remkey]) + + new_dict[key] = tuple(bits) + + self.cache[name] = input_dict + return new_dict + + def cache_clear(self, name=None): + """Clear the internal cache, optionally only for function 'name'.""" + with self.lock: + if name is None: + self.cache.clear() + self.reminders.clear() + self.reminder_keys.clear() + else: + self.cache.pop(name, None) + self.reminders.pop(name, None) + self.reminder_keys.pop(name, None) + + def cache_info(self): + """Return internal cache dicts as a tuple of 3 elements.""" + with self.lock: + return (self.cache, self.reminders, self.reminder_keys) + + +def wrap_numbers(input_dict, name): + """Given an `input_dict` and a function `name`, adjust the numbers + which "wrap" (restart from zero) across different calls by adding + "old value" to "new value" and return an updated dict. + """ + with _wn.lock: + return _wn.run(input_dict, name) + + +_wn = _WrapNumbers() +wrap_numbers.cache_clear = _wn.cache_clear +wrap_numbers.cache_info = _wn.cache_info + + +# The read buffer size for open() builtin. This (also) dictates how +# much data we read(2) when iterating over file lines as in: +# >>> with open(file) as f: +# ... for line in f: +# ... ... +# Default per-line buffer size for binary files is 1K. For text files +# is 8K. We use a bigger buffer (32K) in order to have more consistent +# results when reading /proc pseudo files on Linux, see: +# https://github.com/giampaolo/psutil/issues/2050 +# https://github.com/giampaolo/psutil/issues/708 +FILE_READ_BUFFER_SIZE = 32 * 1024 + + +def open_binary(fname): + return open(fname, "rb", buffering=FILE_READ_BUFFER_SIZE) + + +def open_text(fname): + """Open a file in text mode by using the proper FS encoding and + en/decoding error handlers. + """ + # See: + # https://github.com/giampaolo/psutil/issues/675 + # https://github.com/giampaolo/psutil/pull/733 + fobj = open( # noqa: SIM115 + fname, + buffering=FILE_READ_BUFFER_SIZE, + encoding=ENCODING, + errors=ENCODING_ERRS, + ) + try: + # Dictates per-line read(2) buffer size. Defaults is 8k. See: + # https://github.com/giampaolo/psutil/issues/2050#issuecomment-1013387546 + fobj._CHUNK_SIZE = FILE_READ_BUFFER_SIZE + except AttributeError: + pass + except Exception: + fobj.close() + raise + + return fobj + + +def cat(fname, fallback=_DEFAULT, _open=open_text): + """Read entire file content and return it as a string. File is + opened in text mode. If specified, `fallback` is the value + returned in case of error, either if the file does not exist or + it can't be read(). + """ + if fallback is _DEFAULT: + with _open(fname) as f: + return f.read() + else: + try: + with _open(fname) as f: + return f.read() + except OSError: + return fallback + + +def bcat(fname, fallback=_DEFAULT): + """Same as above but opens file in binary mode.""" + return cat(fname, fallback=fallback, _open=open_binary) + + +def bytes2human(n, format="%(value).1f%(symbol)s"): + """Used by various scripts. See: https://code.activestate.com/recipes/578019-bytes-to-human-human-to-bytes-converter/?in=user-4178764. + + >>> bytes2human(10000) + '9.8K' + >>> bytes2human(100001221) + '95.4M' + """ + symbols = ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y') + prefix = {} + for i, s in enumerate(symbols[1:]): + prefix[s] = 1 << (i + 1) * 10 + for symbol in reversed(symbols[1:]): + if abs(n) >= prefix[symbol]: + value = float(n) / prefix[symbol] + return format % locals() + return format % dict(symbol=symbols[0], value=n) + + +def get_procfs_path(): + """Return updated psutil.PROCFS_PATH constant.""" + return sys.modules['psutil'].PROCFS_PATH + + +def decode(s): + return s.decode(encoding=ENCODING, errors=ENCODING_ERRS) + + +# ===================================================================== +# --- shell utils +# ===================================================================== + + +@memoize +def term_supports_colors(file=sys.stdout): # pragma: no cover + if os.name == 'nt': + return True + try: + import curses + + assert file.isatty() + curses.setupterm() + assert curses.tigetnum("colors") > 0 + except Exception: # noqa: BLE001 + return False + else: + return True + + +def hilite(s, color=None, bold=False): # pragma: no cover + """Return an highlighted version of 'string'.""" + if not term_supports_colors(): + return s + attr = [] + colors = dict( + blue='34', + brown='33', + darkgrey='30', + green='32', + grey='37', + lightblue='36', + red='91', + violet='35', + yellow='93', + ) + colors[None] = '29' + try: + color = colors[color] + except KeyError: + msg = f"invalid color {color!r}; choose amongst {list(colors.keys())}" + raise ValueError(msg) from None + attr.append(color) + if bold: + attr.append('1') + return f"\x1b[{';'.join(attr)}m{s}\x1b[0m" + + +def print_color( + s, color=None, bold=False, file=sys.stdout +): # pragma: no cover + """Print a colorized version of string.""" + if not term_supports_colors(): + print(s, file=file) + elif POSIX: + print(hilite(s, color, bold), file=file) + else: + import ctypes + + DEFAULT_COLOR = 7 + GetStdHandle = ctypes.windll.Kernel32.GetStdHandle + SetConsoleTextAttribute = ( + ctypes.windll.Kernel32.SetConsoleTextAttribute + ) + + colors = dict(green=2, red=4, brown=6, yellow=6) + colors[None] = DEFAULT_COLOR + try: + color = colors[color] + except KeyError: + msg = ( + f"invalid color {color!r}; choose between" + f" {list(colors.keys())!r}" + ) + raise ValueError(msg) from None + if bold and color <= 7: + color += 8 + + handle_id = -12 if file is sys.stderr else -11 + GetStdHandle.restype = ctypes.c_ulong + handle = GetStdHandle(handle_id) + SetConsoleTextAttribute(handle, color) + try: + print(s, file=file) + finally: + SetConsoleTextAttribute(handle, DEFAULT_COLOR) + + +def debug(msg): + """If PSUTIL_DEBUG env var is set, print a debug message to stderr.""" + if PSUTIL_DEBUG: + import inspect + + fname, lineno, _, _lines, _index = inspect.getframeinfo( + inspect.currentframe().f_back + ) + if isinstance(msg, Exception): + if isinstance(msg, OSError): + # ...because str(exc) may contain info about the file name + msg = f"ignoring {msg}" + else: + msg = f"ignoring {msg!r}" + print( # noqa: T201 + f"psutil-debug [{fname}:{lineno}]> {msg}", file=sys.stderr + ) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psaix.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psaix.py new file mode 100644 index 0000000000000000000000000000000000000000..a84abe3f53ff24c701205103032afe3f60ab4084 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psaix.py @@ -0,0 +1,563 @@ +# Copyright (c) 2009, Giampaolo Rodola' +# Copyright (c) 2017, Arnon Yaari +# All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""AIX platform implementation.""" + +import functools +import glob +import os +import re +import subprocess +import sys +from collections import namedtuple + +from . import _common +from . import _psposix +from . import _psutil_aix as cext +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_to_ntuple +from ._common import get_procfs_path +from ._common import memoize_when_activated +from ._common import usage_percent + +__extra__all__ = ["PROCFS_PATH"] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +HAS_THREADS = hasattr(cext, "proc_threads") +HAS_NET_IO_COUNTERS = hasattr(cext, "net_io_counters") +HAS_PROC_IO_COUNTERS = hasattr(cext, "proc_io_counters") + +PAGE_SIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SACTIVE: _common.STATUS_RUNNING, + cext.SSWAP: _common.STATUS_RUNNING, # TODO what status is this? + cext.SSTOP: _common.STATUS_STOPPED, +} + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RCVD: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +proc_info_map = dict( + ppid=0, + rss=1, + vms=2, + create_time=3, + nice=4, + num_threads=5, + status=6, + ttynr=7, +) + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# psutil.Process.memory_info() +pmem = namedtuple('pmem', ['rss', 'vms']) +# psutil.Process.memory_full_info() +pfullmem = pmem +# psutil.Process.cpu_times() +scputimes = namedtuple('scputimes', ['user', 'system', 'idle', 'iowait']) +# psutil.virtual_memory() +svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free']) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + total, avail, free, _pinned, inuse = cext.virtual_mem() + percent = usage_percent((total - avail), total, round_=1) + return svmem(total, avail, percent, inuse, free) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + total, free, sin, sout = cext.swap_mem() + used = total - free + percent = usage_percent(used, total, round_=1) + return _common.sswap(total, used, free, percent, sin, sout) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system-wide CPU times as a named tuple.""" + ret = cext.per_cpu_times() + return scputimes(*[sum(x) for x in zip(*ret)]) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = cext.per_cpu_times() + return [scputimes(*x) for x in ret] + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # mimic os.cpu_count() behavior + return None + + +def cpu_count_cores(): + cmd = ["lsdev", "-Cc", "processor"] + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = p.communicate() + stdout, stderr = (x.decode(sys.stdout.encoding) for x in (stdout, stderr)) + if p.returncode != 0: + msg = f"{cmd!r} command error\n{stderr}" + raise RuntimeError(msg) + processors = stdout.strip().splitlines() + return len(processors) or None + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + ctx_switches, interrupts, soft_interrupts, syscalls = cext.cpu_stats() + return _common.scpustats( + ctx_switches, interrupts, soft_interrupts, syscalls + ) + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters +disk_usage = _psposix.disk_usage + + +def disk_partitions(all=False): + """Return system disk partitions.""" + # TODO - the filtering logic should be better checked so that + # it tries to reflect 'df' as much as possible + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + # Differently from, say, Linux, we don't have a list of + # common fs types so the best we can do, AFAIK, is to + # filter by filesystem having a total size > 0. + if not disk_usage(mountpoint).total: + continue + ntuple = _common.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_if_addrs = cext.net_if_addrs + +if HAS_NET_IO_COUNTERS: + net_io_counters = cext.net_io_counters + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + """ + families, types = _common.conn_tmap[kind] + rawlist = cext.net_connections(_pid) + ret = [] + for item in rawlist: + fd, fam, type_, laddr, raddr, status, pid = item + if fam not in families: + continue + if type_ not in types: + continue + nt = conn_to_ntuple( + fd, + fam, + type_, + laddr, + raddr, + status, + TCP_STATUSES, + pid=pid if _pid == -1 else None, + ) + ret.append(nt) + return ret + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + duplex_map = {"Full": NIC_DUPLEX_FULL, "Half": NIC_DUPLEX_HALF} + names = {x[0] for x in net_if_addrs()} + ret = {} + for name in names: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + + # try to get speed and duplex + # TODO: rewrite this in C (entstat forks, so use truss -f to follow. + # looks like it is using an undocumented ioctl?) + duplex = "" + speed = 0 + p = subprocess.Popen( + ["/usr/bin/entstat", "-d", name], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if p.returncode == 0: + re_result = re.search( + r"Running: (\d+) Mbps.*?(\w+) Duplex", stdout + ) + if re_result is not None: + speed = int(re_result.group(1)) + duplex = re_result.group(2) + + output_flags = ','.join(flags) + isup = 'running' in flags + duplex = duplex_map.get(duplex, NIC_DUPLEX_UNKNOWN) + ret[name] = _common.snicstats(isup, duplex, speed, mtu, output_flags) + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + localhost = (':0.0', ':0') + for item in rawlist: + user, tty, hostname, tstamp, user_process, pid = item + # note: the underlying C function includes entries about + # system boot, run level and others. We might want + # to use them in the future. + if not user_process: + continue + if hostname in localhost: + hostname = 'localhost' + nt = _common.suser(user, tty, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + return [int(x) for x in os.listdir(get_procfs_path()) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix pid.""" + return os.path.exists(os.path.join(get_procfs_path(), str(pid), "psinfo")) + + +def wrap_exceptions(fun): + """Call callable into a try/except clause and translate ENOENT, + EACCES and EPERM in NoSuchProcess or AccessDenied exceptions. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except (FileNotFoundError, ProcessLookupError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if not pid_exists(pid): + raise NoSuchProcess(pid, name) from err + raise ZombieProcess(pid, name, ppid) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "_procfs_path", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._procfs_path = get_procfs_path() + + def oneshot_enter(self): + self._proc_basic_info.cache_activate(self) + self._proc_cred.cache_activate(self) + + def oneshot_exit(self): + self._proc_basic_info.cache_deactivate(self) + self._proc_cred.cache_deactivate(self) + + @wrap_exceptions + @memoize_when_activated + def _proc_basic_info(self): + return cext.proc_basic_info(self.pid, self._procfs_path) + + @wrap_exceptions + @memoize_when_activated + def _proc_cred(self): + return cext.proc_cred(self.pid, self._procfs_path) + + @wrap_exceptions + def name(self): + if self.pid == 0: + return "swapper" + # note: max 16 characters + return cext.proc_name(self.pid, self._procfs_path).rstrip("\x00") + + @wrap_exceptions + def exe(self): + # there is no way to get executable path in AIX other than to guess, + # and guessing is more complex than what's in the wrapping class + cmdline = self.cmdline() + if not cmdline: + return '' + exe = cmdline[0] + if os.path.sep in exe: + # relative or absolute path + if not os.path.isabs(exe): + # if cwd has changed, we're out of luck - this may be wrong! + exe = os.path.abspath(os.path.join(self.cwd(), exe)) + if ( + os.path.isabs(exe) + and os.path.isfile(exe) + and os.access(exe, os.X_OK) + ): + return exe + # not found, move to search in PATH using basename only + exe = os.path.basename(exe) + # search for exe name PATH + for path in os.environ["PATH"].split(":"): + possible_exe = os.path.abspath(os.path.join(path, exe)) + if os.path.isfile(possible_exe) and os.access( + possible_exe, os.X_OK + ): + return possible_exe + return '' + + @wrap_exceptions + def cmdline(self): + return cext.proc_args(self.pid) + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid) + + @wrap_exceptions + def create_time(self): + return self._proc_basic_info()[proc_info_map['create_time']] + + @wrap_exceptions + def num_threads(self): + return self._proc_basic_info()[proc_info_map['num_threads']] + + if HAS_THREADS: + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = _common.pthread(thread_id, utime, stime) + retlist.append(ntuple) + # The underlying C implementation retrieves all OS threads + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not retlist: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = net_connections(kind, _pid=self.pid) + # The underlying C implementation retrieves all OS connections + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not ret: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + return ret + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ppid(self): + self._ppid = self._proc_basic_info()[proc_info_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + real, effective, saved, _, _, _ = self._proc_cred() + return _common.puids(real, effective, saved) + + @wrap_exceptions + def gids(self): + _, _, _, real, effective, saved = self._proc_cred() + return _common.puids(real, effective, saved) + + @wrap_exceptions + def cpu_times(self): + t = cext.proc_cpu_times(self.pid, self._procfs_path) + return _common.pcputimes(*t) + + @wrap_exceptions + def terminal(self): + ttydev = self._proc_basic_info()[proc_info_map['ttynr']] + # convert from 64-bit dev_t to 32-bit dev_t and then map the device + ttydev = ((ttydev & 0x0000FFFF00000000) >> 16) | (ttydev & 0xFFFF) + # try to match rdev of /dev/pts/* files ttydev + for dev in glob.glob("/dev/**/*"): + if os.stat(dev).st_rdev == ttydev: + return dev + return None + + @wrap_exceptions + def cwd(self): + procfs_path = self._procfs_path + try: + result = os.readlink(f"{procfs_path}/{self.pid}/cwd") + return result.rstrip('/') + except FileNotFoundError: + os.stat(f"{procfs_path}/{self.pid}") # raise NSP or AD + return "" + + @wrap_exceptions + def memory_info(self): + ret = self._proc_basic_info() + rss = ret[proc_info_map['rss']] * 1024 + vms = ret[proc_info_map['vms']] * 1024 + return pmem(rss, vms) + + memory_full_info = memory_info + + @wrap_exceptions + def status(self): + code = self._proc_basic_info()[proc_info_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + def open_files(self): + # TODO rewrite without using procfiles (stat /proc/pid/fd/* and then + # find matching name of the inode) + p = subprocess.Popen( + ["/usr/bin/procfiles", "-n", str(self.pid)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if "no such process" in stderr.lower(): + raise NoSuchProcess(self.pid, self._name) + procfiles = re.findall(r"(\d+): S_IFREG.*name:(.*)\n", stdout) + retlist = [] + for fd, path in procfiles: + path = path.strip() + if path.startswith("//"): + path = path[1:] + if path.lower() == "cannot be retrieved": + continue + retlist.append(_common.popenfile(path, int(fd))) + return retlist + + @wrap_exceptions + def num_fds(self): + if self.pid == 0: # no /proc/0/fd + return 0 + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def num_ctx_switches(self): + return _common.pctxsw(*cext.proc_num_ctx_switches(self.pid)) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout, self._name) + + if HAS_PROC_IO_COUNTERS: + + @wrap_exceptions + def io_counters(self): + try: + rc, wc, rb, wb = cext.proc_io_counters(self.pid) + except OSError as err: + # if process is terminated, proc_io_counters returns OSError + # instead of NSP + if not pid_exists(self.pid): + raise NoSuchProcess(self.pid, self._name) from err + raise + return _common.pio(rc, wc, rb, wb) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psbsd.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psbsd.py new file mode 100644 index 0000000000000000000000000000000000000000..ea58c76cad8af7151560148c8ea993ce0f701f37 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psbsd.py @@ -0,0 +1,939 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""FreeBSD, OpenBSD and NetBSD platforms implementation.""" + +import contextlib +import errno +import functools +import os +from collections import defaultdict +from collections import namedtuple +from xml.etree import ElementTree # noqa: ICN001 + +from . import _common +from . import _psposix +from . import _psutil_bsd as cext +from ._common import FREEBSD +from ._common import NETBSD +from ._common import OPENBSD +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import memoize +from ._common import memoize_when_activated +from ._common import usage_percent + +__extra__all__ = [] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +if FREEBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SRUN: _common.STATUS_RUNNING, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SWAIT: _common.STATUS_WAITING, + cext.SLOCK: _common.STATUS_LOCKED, + } +elif OPENBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + # According to /usr/include/sys/proc.h SZOMB is unused. + # test_zombie_process() shows that SDEAD is the right + # equivalent. Also it appears there's no equivalent of + # psutil.STATUS_DEAD. SDEAD really means STATUS_ZOMBIE. + # cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SDEAD: _common.STATUS_ZOMBIE, + cext.SZOMB: _common.STATUS_ZOMBIE, + # From http://www.eecs.harvard.edu/~margo/cs161/videos/proc.h.txt + # OpenBSD has SRUN and SONPROC: SRUN indicates that a process + # is runnable but *not* yet running, i.e. is on a run queue. + # SONPROC indicates that the process is actually executing on + # a CPU, i.e. it is no longer on a run queue. + # As such we'll map SRUN to STATUS_WAKING and SONPROC to + # STATUS_RUNNING + cext.SRUN: _common.STATUS_WAKING, + cext.SONPROC: _common.STATUS_RUNNING, + } +elif NETBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SRUN: _common.STATUS_WAKING, + cext.SONPROC: _common.STATUS_RUNNING, + } + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +PAGESIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +HAS_PROC_NUM_THREADS = hasattr(cext, "proc_num_threads") + +kinfo_proc_map = dict( + ppid=0, + status=1, + real_uid=2, + effective_uid=3, + saved_uid=4, + real_gid=5, + effective_gid=6, + saved_gid=7, + ttynr=8, + create_time=9, + ctx_switches_vol=10, + ctx_switches_unvol=11, + read_io_count=12, + write_io_count=13, + user_time=14, + sys_time=15, + ch_user_time=16, + ch_sys_time=17, + rss=18, + vms=19, + memtext=20, + memdata=21, + memstack=22, + cpunum=23, + name=24, +) + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# fmt: off +# psutil.virtual_memory() +svmem = namedtuple( + 'svmem', ['total', 'available', 'percent', 'used', 'free', + 'active', 'inactive', 'buffers', 'cached', 'shared', 'wired']) +# psutil.cpu_times() +scputimes = namedtuple( + 'scputimes', ['user', 'nice', 'system', 'idle', 'irq']) +# psutil.Process.memory_info() +pmem = namedtuple('pmem', ['rss', 'vms', 'text', 'data', 'stack']) +# psutil.Process.memory_full_info() +pfullmem = pmem +# psutil.Process.cpu_times() +pcputimes = namedtuple('pcputimes', + ['user', 'system', 'children_user', 'children_system']) +# psutil.Process.memory_maps(grouped=True) +pmmap_grouped = namedtuple( + 'pmmap_grouped', 'path rss, private, ref_count, shadow_count') +# psutil.Process.memory_maps(grouped=False) +pmmap_ext = namedtuple( + 'pmmap_ext', 'addr, perms path rss, private, ref_count, shadow_count') +# psutil.disk_io_counters() +if FREEBSD: + sdiskio = namedtuple('sdiskio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes', + 'read_time', 'write_time', + 'busy_time']) +else: + sdiskio = namedtuple('sdiskio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes']) +# fmt: on + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + mem = cext.virtual_mem() + if NETBSD: + total, free, active, inactive, wired, cached = mem + # On NetBSD buffers and shared mem is determined via /proc. + # The C ext set them to 0. + with open('/proc/meminfo', 'rb') as f: + for line in f: + if line.startswith(b'Buffers:'): + buffers = int(line.split()[1]) * 1024 + elif line.startswith(b'MemShared:'): + shared = int(line.split()[1]) * 1024 + # Before avail was calculated as (inactive + cached + free), + # same as zabbix, but it turned out it could exceed total (see + # #2233), so zabbix seems to be wrong. Htop calculates it + # differently, and the used value seem more realistic, so let's + # match htop. + # https://github.com/htop-dev/htop/blob/e7f447b/netbsd/NetBSDProcessList.c#L162 + # https://github.com/zabbix/zabbix/blob/af5e0f8/src/libs/zbxsysinfo/netbsd/memory.c#L135 + used = active + wired + avail = total - used + else: + total, free, active, inactive, wired, cached, buffers, shared = mem + # matches freebsd-memory CLI: + # * https://people.freebsd.org/~rse/dist/freebsd-memory + # * https://www.cyberciti.biz/files/scripts/freebsd-memory.pl.txt + # matches zabbix: + # * https://github.com/zabbix/zabbix/blob/af5e0f8/src/libs/zbxsysinfo/freebsd/memory.c#L143 + avail = inactive + cached + free + used = active + wired + cached + + percent = usage_percent((total - avail), total, round_=1) + return svmem( + total, + avail, + percent, + used, + free, + active, + inactive, + buffers, + cached, + shared, + wired, + ) + + +def swap_memory(): + """System swap memory as (total, used, free, sin, sout) namedtuple.""" + total, used, free, sin, sout = cext.swap_mem() + percent = usage_percent(used, total, round_=1) + return _common.sswap(total, used, free, percent, sin, sout) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system per-CPU times as a namedtuple.""" + user, nice, system, idle, irq = cext.cpu_times() + return scputimes(user, nice, system, idle, irq) + + +def per_cpu_times(): + """Return system CPU times as a namedtuple.""" + ret = [] + for cpu_t in cext.per_cpu_times(): + user, nice, system, idle, irq = cpu_t + item = scputimes(user, nice, system, idle, irq) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +if OPENBSD or NETBSD: + + def cpu_count_cores(): + # OpenBSD and NetBSD do not implement this. + return 1 if cpu_count_logical() == 1 else None + +else: + + def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + # From the C module we'll get an XML string similar to this: + # http://manpages.ubuntu.com/manpages/precise/man4/smp.4freebsd.html + # We may get None in case "sysctl kern.sched.topology_spec" + # is not supported on this BSD version, in which case we'll mimic + # os.cpu_count() and return None. + ret = None + s = cext.cpu_topology() + if s is not None: + # get rid of padding chars appended at the end of the string + index = s.rfind("") + if index != -1: + s = s[: index + 9] + root = ElementTree.fromstring(s) + try: + ret = len(root.findall('group/children/group/cpu')) or None + finally: + # needed otherwise it will memleak + root.clear() + if not ret: + # If logical CPUs == 1 it's obvious we' have only 1 core. + if cpu_count_logical() == 1: + return 1 + return ret + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + if FREEBSD: + # Note: the C ext is returning some metrics we are not exposing: + # traps. + ctxsw, intrs, soft_intrs, syscalls, _traps = cext.cpu_stats() + elif NETBSD: + # XXX + # Note about intrs: the C extension returns 0. intrs + # can be determined via /proc/stat; it has the same value as + # soft_intrs thought so the kernel is faking it (?). + # + # Note about syscalls: the C extension always sets it to 0 (?). + # + # Note: the C ext is returning some metrics we are not exposing: + # traps, faults and forks. + ctxsw, intrs, soft_intrs, syscalls, _traps, _faults, _forks = ( + cext.cpu_stats() + ) + with open('/proc/stat', 'rb') as f: + for line in f: + if line.startswith(b'intr'): + intrs = int(line.split()[1]) + elif OPENBSD: + # Note: the C ext is returning some metrics we are not exposing: + # traps, faults and forks. + ctxsw, intrs, soft_intrs, syscalls, _traps, _faults, _forks = ( + cext.cpu_stats() + ) + return _common.scpustats(ctxsw, intrs, soft_intrs, syscalls) + + +if FREEBSD: + + def cpu_freq(): + """Return frequency metrics for CPUs. As of Dec 2018 only + CPU 0 appears to be supported by FreeBSD and all other cores + match the frequency of CPU 0. + """ + ret = [] + num_cpus = cpu_count_logical() + for cpu in range(num_cpus): + try: + current, available_freq = cext.cpu_freq(cpu) + except NotImplementedError: + continue + if available_freq: + try: + min_freq = int(available_freq.split(" ")[-1].split("/")[0]) + except (IndexError, ValueError): + min_freq = None + try: + max_freq = int(available_freq.split(" ")[0].split("/")[0]) + except (IndexError, ValueError): + max_freq = None + ret.append(_common.scpufreq(current, min_freq, max_freq)) + return ret + +elif OPENBSD: + + def cpu_freq(): + curr = float(cext.cpu_freq()) + return [_common.scpufreq(curr, 0.0, 0.0)] + + +# ===================================================================== +# --- disks +# ===================================================================== + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples. + 'all' argument is ignored, see: + https://github.com/giampaolo/psutil/issues/906. + """ + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + ntuple = _common.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +disk_usage = _psposix.disk_usage +disk_io_counters = cext.disk_io_counters + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + else: + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = _common.snicstats( + isup, duplex, speed, mtu, output_flags + ) + return ret + + +def net_connections(kind): + """System-wide network connections.""" + families, types = conn_tmap[kind] + ret = set() + if OPENBSD: + rawlist = cext.net_connections(-1, families, types) + elif NETBSD: + rawlist = cext.net_connections(-1, kind) + else: # FreeBSD + rawlist = cext.net_connections(families, types) + + for item in rawlist: + fd, fam, type, laddr, raddr, status, pid = item + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES, pid + ) + ret.add(nt) + return list(ret) + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +if FREEBSD: + + def sensors_battery(): + """Return battery info.""" + try: + percent, minsleft, power_plugged = cext.sensors_battery() + except NotImplementedError: + # See: https://github.com/giampaolo/psutil/issues/1074 + return None + power_plugged = power_plugged == 1 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif minsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = minsleft * 60 + return _common.sbattery(percent, secsleft, power_plugged) + + def sensors_temperatures(): + """Return CPU cores temperatures if available, else an empty dict.""" + ret = defaultdict(list) + num_cpus = cpu_count_logical() + for cpu in range(num_cpus): + try: + current, high = cext.sensors_cpu_temperature(cpu) + if high <= 0: + high = None + name = f"Core {cpu}" + ret["coretemp"].append( + _common.shwtemp(name, current, high, high) + ) + except NotImplementedError: + pass + + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +if NETBSD: + + try: + INIT_BOOT_TIME = boot_time() + except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + INIT_BOOT_TIME = 0 + + def adjust_proc_create_time(ctime): + """Account for system clock updates.""" + if INIT_BOOT_TIME == 0: + return ctime + + diff = INIT_BOOT_TIME - boot_time() + if diff == 0 or abs(diff) < 1: + return ctime + + debug("system clock was updated; adjusting process create_time()") + if diff < 0: + return ctime - diff + return ctime + diff + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + if tty == '~': + continue # reboot or shutdown + nt = _common.suser(user, tty or None, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +@memoize +def _pid_0_exists(): + try: + Process(0).name() + except NoSuchProcess: + return False + except AccessDenied: + return True + else: + return True + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + ret = cext.pids() + if OPENBSD and (0 not in ret) and _pid_0_exists(): + # On OpenBSD the kernel does not return PID 0 (neither does + # ps) but it's actually querable (Process(0) will succeed). + ret.insert(0, 0) + return ret + + +if NETBSD: + + def pid_exists(pid): + exists = _psposix.pid_exists(pid) + if not exists: + # We do this because _psposix.pid_exists() lies in case of + # zombie processes. + return pid in pids() + else: + return True + +elif OPENBSD: + + def pid_exists(pid): + exists = _psposix.pid_exists(pid) + if not exists: + return False + else: + # OpenBSD seems to be the only BSD platform where + # _psposix.pid_exists() returns True for thread IDs (tids), + # so we can't use it. + return pid in pids() + +else: # FreeBSD + pid_exists = _psposix.pid_exists + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError exceptions into + NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except ProcessLookupError as err: + if cext.proc_is_zombie(pid): + raise ZombieProcess(pid, name, ppid) from err + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except cext.ZombieProcessError as err: + raise ZombieProcess(pid, name, ppid) from err + except OSError as err: + if pid == 0 and 0 in pids(): + raise AccessDenied(pid, name) from err + raise err from None + + return wrapper + + +@contextlib.contextmanager +def wrap_exceptions_procfs(inst): + """Same as above, for routines relying on reading /proc fs.""" + pid, name, ppid = inst.pid, inst._name, inst._ppid + try: + yield + except (ProcessLookupError, FileNotFoundError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if cext.proc_is_zombie(inst.pid): + raise ZombieProcess(pid, name, ppid) from err + else: + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + def _assert_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + cext.proc_name(self.pid) + + @wrap_exceptions + @memoize_when_activated + def oneshot(self): + """Retrieves multiple process info in one shot as a raw tuple.""" + ret = cext.proc_oneshot_info(self.pid) + assert len(ret) == len(kinfo_proc_map) + return ret + + def oneshot_enter(self): + self.oneshot.cache_activate(self) + + def oneshot_exit(self): + self.oneshot.cache_deactivate(self) + + @wrap_exceptions + def name(self): + name = self.oneshot()[kinfo_proc_map['name']] + return name if name is not None else cext.proc_name(self.pid) + + @wrap_exceptions + def exe(self): + if FREEBSD: + if self.pid == 0: + return '' # else NSP + return cext.proc_exe(self.pid) + elif NETBSD: + if self.pid == 0: + # /proc/0 dir exists but /proc/0/exe doesn't + return "" + with wrap_exceptions_procfs(self): + return os.readlink(f"/proc/{self.pid}/exe") + else: + # OpenBSD: exe cannot be determined; references: + # https://chromium.googlesource.com/chromium/src/base/+/ + # master/base_paths_posix.cc + # We try our best guess by using which against the first + # cmdline arg (may return None). + import shutil + + cmdline = self.cmdline() + if cmdline: + return shutil.which(cmdline[0]) or "" + else: + return "" + + @wrap_exceptions + def cmdline(self): + if OPENBSD and self.pid == 0: + return [] # ...else it crashes + elif NETBSD: + # XXX - most of the times the underlying sysctl() call on + # NetBSD and OpenBSD returns a truncated string. Also + # /proc/pid/cmdline behaves the same so it looks like this + # is a kernel bug. + try: + return cext.proc_cmdline(self.pid) + except OSError as err: + if err.errno == errno.EINVAL: + pid, name, ppid = self.pid, self._name, self._ppid + if cext.proc_is_zombie(self.pid): + raise ZombieProcess(pid, name, ppid) from err + if not pid_exists(self.pid): + raise NoSuchProcess(pid, name, ppid) from err + # XXX: this happens with unicode tests. It means the C + # routine is unable to decode invalid unicode chars. + debug(f"ignoring {err!r} and returning an empty list") + return [] + else: + raise + else: + return cext.proc_cmdline(self.pid) + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid) + + @wrap_exceptions + def terminal(self): + tty_nr = self.oneshot()[kinfo_proc_map['ttynr']] + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + @wrap_exceptions + def ppid(self): + self._ppid = self.oneshot()[kinfo_proc_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + rawtuple = self.oneshot() + return _common.puids( + rawtuple[kinfo_proc_map['real_uid']], + rawtuple[kinfo_proc_map['effective_uid']], + rawtuple[kinfo_proc_map['saved_uid']], + ) + + @wrap_exceptions + def gids(self): + rawtuple = self.oneshot() + return _common.pgids( + rawtuple[kinfo_proc_map['real_gid']], + rawtuple[kinfo_proc_map['effective_gid']], + rawtuple[kinfo_proc_map['saved_gid']], + ) + + @wrap_exceptions + def cpu_times(self): + rawtuple = self.oneshot() + return _common.pcputimes( + rawtuple[kinfo_proc_map['user_time']], + rawtuple[kinfo_proc_map['sys_time']], + rawtuple[kinfo_proc_map['ch_user_time']], + rawtuple[kinfo_proc_map['ch_sys_time']], + ) + + if FREEBSD: + + @wrap_exceptions + def cpu_num(self): + return self.oneshot()[kinfo_proc_map['cpunum']] + + @wrap_exceptions + def memory_info(self): + rawtuple = self.oneshot() + return pmem( + rawtuple[kinfo_proc_map['rss']], + rawtuple[kinfo_proc_map['vms']], + rawtuple[kinfo_proc_map['memtext']], + rawtuple[kinfo_proc_map['memdata']], + rawtuple[kinfo_proc_map['memstack']], + ) + + memory_full_info = memory_info + + @wrap_exceptions + def create_time(self, monotonic=False): + ctime = self.oneshot()[kinfo_proc_map['create_time']] + if NETBSD and not monotonic: + # NetBSD: ctime subject to system clock updates. + ctime = adjust_proc_create_time(ctime) + return ctime + + @wrap_exceptions + def num_threads(self): + if HAS_PROC_NUM_THREADS: + # FreeBSD / NetBSD + return cext.proc_num_threads(self.pid) + else: + return len(self.threads()) + + @wrap_exceptions + def num_ctx_switches(self): + rawtuple = self.oneshot() + return _common.pctxsw( + rawtuple[kinfo_proc_map['ctx_switches_vol']], + rawtuple[kinfo_proc_map['ctx_switches_unvol']], + ) + + @wrap_exceptions + def threads(self): + # Note: on OpenSBD this (/dev/mem) requires root access. + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = _common.pthread(thread_id, utime, stime) + retlist.append(ntuple) + if OPENBSD: + self._assert_alive() + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + families, types = conn_tmap[kind] + ret = [] + + if NETBSD: + rawlist = cext.net_connections(self.pid, kind) + elif OPENBSD: + rawlist = cext.net_connections(self.pid, families, types) + else: + rawlist = cext.proc_net_connections(self.pid, families, types) + + for item in rawlist: + fd, fam, type, laddr, raddr, status = item[:6] + if FREEBSD: + if (fam not in families) or (type not in types): + continue + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES + ) + ret.append(nt) + + self._assert_alive() + return ret + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout, self._name) + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def status(self): + code = self.oneshot()[kinfo_proc_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def io_counters(self): + rawtuple = self.oneshot() + return _common.pio( + rawtuple[kinfo_proc_map['read_io_count']], + rawtuple[kinfo_proc_map['write_io_count']], + -1, + -1, + ) + + @wrap_exceptions + def cwd(self): + """Return process current working directory.""" + # sometimes we get an empty string, in which case we turn + # it into None + if OPENBSD and self.pid == 0: + return "" # ...else it would raise EINVAL + return cext.proc_cwd(self.pid) + + nt_mmap_grouped = namedtuple( + 'mmap', 'path rss, private, ref_count, shadow_count' + ) + nt_mmap_ext = namedtuple( + 'mmap', 'addr, perms path rss, private, ref_count, shadow_count' + ) + + @wrap_exceptions + def open_files(self): + """Return files opened by process as a list of namedtuples.""" + rawlist = cext.proc_open_files(self.pid) + return [_common.popenfile(path, fd) for path, fd in rawlist] + + @wrap_exceptions + def num_fds(self): + """Return the number of file descriptors opened by this process.""" + ret = cext.proc_num_fds(self.pid) + if NETBSD: + self._assert_alive() + return ret + + # --- FreeBSD only APIs + + if FREEBSD: + + @wrap_exceptions + def cpu_affinity_get(self): + return cext.proc_cpu_affinity_get(self.pid) + + @wrap_exceptions + def cpu_affinity_set(self, cpus): + # Pre-emptively check if CPUs are valid because the C + # function has a weird behavior in case of invalid CPUs, + # see: https://github.com/giampaolo/psutil/issues/586 + allcpus = set(range(len(per_cpu_times()))) + for cpu in cpus: + if cpu not in allcpus: + msg = f"invalid CPU {cpu!r} (choose between {allcpus})" + raise ValueError(msg) + try: + cext.proc_cpu_affinity_set(self.pid, cpus) + except OSError as err: + # 'man cpuset_setaffinity' about EDEADLK: + # <> + if err.errno in {errno.EINVAL, errno.EDEADLK}: + for cpu in cpus: + if cpu not in allcpus: + msg = ( + f"invalid CPU {cpu!r} (choose between" + f" {allcpus})" + ) + raise ValueError(msg) from err + raise + + @wrap_exceptions + def memory_maps(self): + return cext.proc_memory_maps(self.pid) + + @wrap_exceptions + def rlimit(self, resource, limits=None): + if limits is None: + return cext.proc_getrlimit(self.pid, resource) + else: + if len(limits) != 2: + msg = ( + "second argument must be a (soft, hard) tuple, got" + f" {limits!r}" + ) + raise ValueError(msg) + soft, hard = limits + return cext.proc_setrlimit(self.pid, resource, soft, hard) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pslinux.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pslinux.py new file mode 100644 index 0000000000000000000000000000000000000000..821c4eb7178bcf98a8a071c4f5f781101ee325a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pslinux.py @@ -0,0 +1,2306 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Linux platform implementation.""" + + +import base64 +import collections +import enum +import errno +import functools +import glob +import os +import re +import resource +import socket +import struct +import sys +import warnings +from collections import defaultdict +from collections import namedtuple + +from . import _common +from . import _psposix +from . import _psutil_linux as cext +from ._common import ENCODING +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import bcat +from ._common import cat +from ._common import debug +from ._common import decode +from ._common import get_procfs_path +from ._common import isfile_strict +from ._common import memoize +from ._common import memoize_when_activated +from ._common import open_binary +from ._common import open_text +from ._common import parse_environ_block +from ._common import path_exists_strict +from ._common import supports_ipv6 +from ._common import usage_percent + +# fmt: off +__extra__all__ = [ + 'PROCFS_PATH', + # io prio constants + "IOPRIO_CLASS_NONE", "IOPRIO_CLASS_RT", "IOPRIO_CLASS_BE", + "IOPRIO_CLASS_IDLE", + # connection status constants + "CONN_ESTABLISHED", "CONN_SYN_SENT", "CONN_SYN_RECV", "CONN_FIN_WAIT1", + "CONN_FIN_WAIT2", "CONN_TIME_WAIT", "CONN_CLOSE", "CONN_CLOSE_WAIT", + "CONN_LAST_ACK", "CONN_LISTEN", "CONN_CLOSING", +] +# fmt: on + + +# ===================================================================== +# --- globals +# ===================================================================== + + +POWER_SUPPLY_PATH = "/sys/class/power_supply" +HAS_PROC_SMAPS = os.path.exists(f"/proc/{os.getpid()}/smaps") +HAS_PROC_SMAPS_ROLLUP = os.path.exists(f"/proc/{os.getpid()}/smaps_rollup") +HAS_PROC_IO_PRIORITY = hasattr(cext, "proc_ioprio_get") +HAS_CPU_AFFINITY = hasattr(cext, "proc_cpu_affinity_get") + +# Number of clock ticks per second +CLOCK_TICKS = os.sysconf("SC_CLK_TCK") +PAGESIZE = cext.getpagesize() +LITTLE_ENDIAN = sys.byteorder == 'little' +UNSET = object() + +# "man iostat" states that sectors are equivalent with blocks and have +# a size of 512 bytes. Despite this value can be queried at runtime +# via /sys/block/{DISK}/queue/hw_sector_size and results may vary +# between 1k, 2k, or 4k... 512 appears to be a magic constant used +# throughout Linux source code: +# * https://stackoverflow.com/a/38136179/376587 +# * https://lists.gt.net/linux/kernel/2241060 +# * https://github.com/giampaolo/psutil/issues/1305 +# * https://github.com/torvalds/linux/blob/ +# 4f671fe2f9523a1ea206f63fe60a7c7b3a56d5c7/include/linux/bio.h#L99 +# * https://lkml.org/lkml/2015/8/17/234 +DISK_SECTOR_SIZE = 512 + +AddressFamily = enum.IntEnum( + 'AddressFamily', {'AF_LINK': int(socket.AF_PACKET)} +) +AF_LINK = AddressFamily.AF_LINK + + +# ioprio_* constants http://linux.die.net/man/2/ioprio_get +class IOPriority(enum.IntEnum): + IOPRIO_CLASS_NONE = 0 + IOPRIO_CLASS_RT = 1 + IOPRIO_CLASS_BE = 2 + IOPRIO_CLASS_IDLE = 3 + + +globals().update(IOPriority.__members__) + +# See: +# https://github.com/torvalds/linux/blame/master/fs/proc/array.c +# ...and (TASK_* constants): +# https://github.com/torvalds/linux/blob/master/include/linux/sched.h +PROC_STATUSES = { + "R": _common.STATUS_RUNNING, + "S": _common.STATUS_SLEEPING, + "D": _common.STATUS_DISK_SLEEP, + "T": _common.STATUS_STOPPED, + "t": _common.STATUS_TRACING_STOP, + "Z": _common.STATUS_ZOMBIE, + "X": _common.STATUS_DEAD, + "x": _common.STATUS_DEAD, + "K": _common.STATUS_WAKE_KILL, + "W": _common.STATUS_WAKING, + "I": _common.STATUS_IDLE, + "P": _common.STATUS_PARKED, +} + +# https://github.com/torvalds/linux/blob/master/include/net/tcp_states.h +TCP_STATUSES = { + "01": _common.CONN_ESTABLISHED, + "02": _common.CONN_SYN_SENT, + "03": _common.CONN_SYN_RECV, + "04": _common.CONN_FIN_WAIT1, + "05": _common.CONN_FIN_WAIT2, + "06": _common.CONN_TIME_WAIT, + "07": _common.CONN_CLOSE, + "08": _common.CONN_CLOSE_WAIT, + "09": _common.CONN_LAST_ACK, + "0A": _common.CONN_LISTEN, + "0B": _common.CONN_CLOSING, +} + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# fmt: off +# psutil.virtual_memory() +svmem = namedtuple( + 'svmem', ['total', 'available', 'percent', 'used', 'free', + 'active', 'inactive', 'buffers', 'cached', 'shared', 'slab']) +# psutil.disk_io_counters() +sdiskio = namedtuple( + 'sdiskio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes', + 'read_time', 'write_time', + 'read_merged_count', 'write_merged_count', + 'busy_time']) +# psutil.Process().open_files() +popenfile = namedtuple( + 'popenfile', ['path', 'fd', 'position', 'mode', 'flags']) +# psutil.Process().memory_info() +pmem = namedtuple('pmem', 'rss vms shared text lib data dirty') +# psutil.Process().memory_full_info() +pfullmem = namedtuple('pfullmem', pmem._fields + ('uss', 'pss', 'swap')) +# psutil.Process().memory_maps(grouped=True) +pmmap_grouped = namedtuple( + 'pmmap_grouped', + ['path', 'rss', 'size', 'pss', 'shared_clean', 'shared_dirty', + 'private_clean', 'private_dirty', 'referenced', 'anonymous', 'swap']) +# psutil.Process().memory_maps(grouped=False) +pmmap_ext = namedtuple( + 'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields)) +# psutil.Process.io_counters() +pio = namedtuple('pio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes', + 'read_chars', 'write_chars']) +# psutil.Process.cpu_times() +pcputimes = namedtuple('pcputimes', + ['user', 'system', 'children_user', 'children_system', + 'iowait']) +# fmt: on + + +# ===================================================================== +# --- utils +# ===================================================================== + + +def readlink(path): + """Wrapper around os.readlink().""" + assert isinstance(path, str), path + path = os.readlink(path) + # readlink() might return paths containing null bytes ('\x00') + # resulting in "TypeError: must be encoded string without NULL + # bytes, not str" errors when the string is passed to other + # fs-related functions (os.*, open(), ...). + # Apparently everything after '\x00' is garbage (we can have + # ' (deleted)', 'new' and possibly others), see: + # https://github.com/giampaolo/psutil/issues/717 + path = path.split('\x00')[0] + # Certain paths have ' (deleted)' appended. Usually this is + # bogus as the file actually exists. Even if it doesn't we + # don't care. + if path.endswith(' (deleted)') and not path_exists_strict(path): + path = path[:-10] + return path + + +def file_flags_to_mode(flags): + """Convert file's open() flags into a readable string. + Used by Process.open_files(). + """ + modes_map = {os.O_RDONLY: 'r', os.O_WRONLY: 'w', os.O_RDWR: 'w+'} + mode = modes_map[flags & (os.O_RDONLY | os.O_WRONLY | os.O_RDWR)] + if flags & os.O_APPEND: + mode = mode.replace('w', 'a', 1) + mode = mode.replace('w+', 'r+') + # possible values: r, w, a, r+, a+ + return mode + + +def is_storage_device(name): + """Return True if the given name refers to a root device (e.g. + "sda", "nvme0n1") as opposed to a logical partition (e.g. "sda1", + "nvme0n1p1"). If name is a virtual device (e.g. "loop1", "ram") + return True. + """ + # Re-adapted from iostat source code, see: + # https://github.com/sysstat/sysstat/blob/ + # 97912938cd476645b267280069e83b1c8dc0e1c7/common.c#L208 + # Some devices may have a slash in their name (e.g. cciss/c0d0...). + name = name.replace('/', '!') + including_virtual = True + if including_virtual: + path = f"/sys/block/{name}" + else: + path = f"/sys/block/{name}/device" + return os.access(path, os.F_OK) + + +@memoize +def set_scputimes_ntuple(procfs_path): + """Set a namedtuple of variable fields depending on the CPU times + available on this Linux kernel version which may be: + (user, nice, system, idle, iowait, irq, softirq, [steal, [guest, + [guest_nice]]]) + Used by cpu_times() function. + """ + global scputimes + with open_binary(f"{procfs_path}/stat") as f: + values = f.readline().split()[1:] + fields = ['user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq'] + vlen = len(values) + if vlen >= 8: + # Linux >= 2.6.11 + fields.append('steal') + if vlen >= 9: + # Linux >= 2.6.24 + fields.append('guest') + if vlen >= 10: + # Linux >= 3.2.0 + fields.append('guest_nice') + scputimes = namedtuple('scputimes', fields) + + +try: + set_scputimes_ntuple("/proc") +except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + scputimes = namedtuple('scputimes', 'user system idle')(0.0, 0.0, 0.0) + + +# ===================================================================== +# --- system memory +# ===================================================================== + + +def calculate_avail_vmem(mems): + """Fallback for kernels < 3.14 where /proc/meminfo does not provide + "MemAvailable", see: + https://blog.famzah.net/2014/09/24/. + + This code reimplements the algorithm outlined here: + https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/ + commit/?id=34e431b0ae398fc54ea69ff85ec700722c9da773 + + We use this function also when "MemAvailable" returns 0 (possibly a + kernel bug, see: https://github.com/giampaolo/psutil/issues/1915). + In that case this routine matches "free" CLI tool result ("available" + column). + + XXX: on recent kernels this calculation may differ by ~1.5% compared + to "MemAvailable:", as it's calculated slightly differently. + It is still way more realistic than doing (free + cached) though. + See: + * https://gitlab.com/procps-ng/procps/issues/42 + * https://github.com/famzah/linux-memavailable-procfs/issues/2 + """ + # Note about "fallback" value. According to: + # https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/ + # commit/?id=34e431b0ae398fc54ea69ff85ec700722c9da773 + # ...long ago "available" memory was calculated as (free + cached), + # We use fallback when one of these is missing from /proc/meminfo: + # "Active(file)": introduced in 2.6.28 / Dec 2008 + # "Inactive(file)": introduced in 2.6.28 / Dec 2008 + # "SReclaimable": introduced in 2.6.19 / Nov 2006 + # /proc/zoneinfo: introduced in 2.6.13 / Aug 2005 + free = mems[b'MemFree:'] + fallback = free + mems.get(b"Cached:", 0) + try: + lru_active_file = mems[b'Active(file):'] + lru_inactive_file = mems[b'Inactive(file):'] + slab_reclaimable = mems[b'SReclaimable:'] + except KeyError as err: + debug( + f"{err.args[0]} is missing from /proc/meminfo; using an" + " approximation for calculating available memory" + ) + return fallback + try: + f = open_binary(f"{get_procfs_path()}/zoneinfo") + except OSError: + return fallback # kernel 2.6.13 + + watermark_low = 0 + with f: + for line in f: + line = line.strip() + if line.startswith(b'low'): + watermark_low += int(line.split()[1]) + watermark_low *= PAGESIZE + + avail = free - watermark_low + pagecache = lru_active_file + lru_inactive_file + pagecache -= min(pagecache / 2, watermark_low) + avail += pagecache + avail += slab_reclaimable - min(slab_reclaimable / 2.0, watermark_low) + return int(avail) + + +def virtual_memory(): + """Report virtual memory stats. + This implementation mimics procps-ng-3.3.12, aka "free" CLI tool: + https://gitlab.com/procps-ng/procps/blob/ + 24fd2605c51fccc375ab0287cec33aa767f06718/proc/sysinfo.c#L778-791 + The returned values are supposed to match both "free" and "vmstat -s" + CLI tools. + """ + missing_fields = [] + mems = {} + with open_binary(f"{get_procfs_path()}/meminfo") as f: + for line in f: + fields = line.split() + mems[fields[0]] = int(fields[1]) * 1024 + + # /proc doc states that the available fields in /proc/meminfo vary + # by architecture and compile options, but these 3 values are also + # returned by sysinfo(2); as such we assume they are always there. + total = mems[b'MemTotal:'] + free = mems[b'MemFree:'] + try: + buffers = mems[b'Buffers:'] + except KeyError: + # https://github.com/giampaolo/psutil/issues/1010 + buffers = 0 + missing_fields.append('buffers') + try: + cached = mems[b"Cached:"] + except KeyError: + cached = 0 + missing_fields.append('cached') + else: + # "free" cmdline utility sums reclaimable to cached. + # Older versions of procps used to add slab memory instead. + # This got changed in: + # https://gitlab.com/procps-ng/procps/commit/ + # 05d751c4f076a2f0118b914c5e51cfbb4762ad8e + cached += mems.get(b"SReclaimable:", 0) # since kernel 2.6.19 + + try: + shared = mems[b'Shmem:'] # since kernel 2.6.32 + except KeyError: + try: + shared = mems[b'MemShared:'] # kernels 2.4 + except KeyError: + shared = 0 + missing_fields.append('shared') + + try: + active = mems[b"Active:"] + except KeyError: + active = 0 + missing_fields.append('active') + + try: + inactive = mems[b"Inactive:"] + except KeyError: + try: + inactive = ( + mems[b"Inact_dirty:"] + + mems[b"Inact_clean:"] + + mems[b"Inact_laundry:"] + ) + except KeyError: + inactive = 0 + missing_fields.append('inactive') + + try: + slab = mems[b"Slab:"] + except KeyError: + slab = 0 + + # - starting from 4.4.0 we match free's "available" column. + # Before 4.4.0 we calculated it as (free + buffers + cached) + # which matched htop. + # - free and htop available memory differs as per: + # http://askubuntu.com/a/369589 + # http://unix.stackexchange.com/a/65852/168884 + # - MemAvailable has been introduced in kernel 3.14 + try: + avail = mems[b'MemAvailable:'] + except KeyError: + avail = calculate_avail_vmem(mems) + else: + if avail == 0: + # Yes, it can happen (probably a kernel bug): + # https://github.com/giampaolo/psutil/issues/1915 + # In this case "free" CLI tool makes an estimate. We do the same, + # and it matches "free" CLI tool. + avail = calculate_avail_vmem(mems) + + if avail < 0: + avail = 0 + missing_fields.append('available') + elif avail > total: + # If avail is greater than total or our calculation overflows, + # that's symptomatic of running within a LCX container where such + # values will be dramatically distorted over those of the host. + # https://gitlab.com/procps-ng/procps/blob/ + # 24fd2605c51fccc375ab0287cec33aa767f06718/proc/sysinfo.c#L764 + avail = free + + used = total - avail + + percent = usage_percent((total - avail), total, round_=1) + + # Warn about missing metrics which are set to 0. + if missing_fields: + msg = "{} memory stats couldn't be determined and {} set to 0".format( + ", ".join(missing_fields), + "was" if len(missing_fields) == 1 else "were", + ) + warnings.warn(msg, RuntimeWarning, stacklevel=2) + + return svmem( + total, + avail, + percent, + used, + free, + active, + inactive, + buffers, + cached, + shared, + slab, + ) + + +def swap_memory(): + """Return swap memory metrics.""" + mems = {} + with open_binary(f"{get_procfs_path()}/meminfo") as f: + for line in f: + fields = line.split() + mems[fields[0]] = int(fields[1]) * 1024 + # We prefer /proc/meminfo over sysinfo() syscall so that + # psutil.PROCFS_PATH can be used in order to allow retrieval + # for linux containers, see: + # https://github.com/giampaolo/psutil/issues/1015 + try: + total = mems[b'SwapTotal:'] + free = mems[b'SwapFree:'] + except KeyError: + _, _, _, _, total, free, unit_multiplier = cext.linux_sysinfo() + total *= unit_multiplier + free *= unit_multiplier + + used = total - free + percent = usage_percent(used, total, round_=1) + # get pgin/pgouts + try: + f = open_binary(f"{get_procfs_path()}/vmstat") + except OSError as err: + # see https://github.com/giampaolo/psutil/issues/722 + msg = ( + "'sin' and 'sout' swap memory stats couldn't " + f"be determined and were set to 0 ({err})" + ) + warnings.warn(msg, RuntimeWarning, stacklevel=2) + sin = sout = 0 + else: + with f: + sin = sout = None + for line in f: + # values are expressed in 4 kilo bytes, we want + # bytes instead + if line.startswith(b'pswpin'): + sin = int(line.split(b' ')[1]) * 4 * 1024 + elif line.startswith(b'pswpout'): + sout = int(line.split(b' ')[1]) * 4 * 1024 + if sin is not None and sout is not None: + break + else: + # we might get here when dealing with exotic Linux + # flavors, see: + # https://github.com/giampaolo/psutil/issues/313 + msg = "'sin' and 'sout' swap memory stats couldn't " + msg += "be determined and were set to 0" + warnings.warn(msg, RuntimeWarning, stacklevel=2) + sin = sout = 0 + return _common.sswap(total, used, free, percent, sin, sout) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return a named tuple representing the following system-wide + CPU times: + (user, nice, system, idle, iowait, irq, softirq [steal, [guest, + [guest_nice]]]) + Last 3 fields may not be available on all Linux kernel versions. + """ + procfs_path = get_procfs_path() + set_scputimes_ntuple(procfs_path) + with open_binary(f"{procfs_path}/stat") as f: + values = f.readline().split() + fields = values[1 : len(scputimes._fields) + 1] + fields = [float(x) / CLOCK_TICKS for x in fields] + return scputimes(*fields) + + +def per_cpu_times(): + """Return a list of namedtuple representing the CPU times + for every CPU available on the system. + """ + procfs_path = get_procfs_path() + set_scputimes_ntuple(procfs_path) + cpus = [] + with open_binary(f"{procfs_path}/stat") as f: + # get rid of the first line which refers to system wide CPU stats + f.readline() + for line in f: + if line.startswith(b'cpu'): + values = line.split() + fields = values[1 : len(scputimes._fields) + 1] + fields = [float(x) / CLOCK_TICKS for x in fields] + entry = scputimes(*fields) + cpus.append(entry) + return cpus + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # as a second fallback we try to parse /proc/cpuinfo + num = 0 + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + for line in f: + if line.lower().startswith(b'processor'): + num += 1 + + # unknown format (e.g. amrel/sparc architectures), see: + # https://github.com/giampaolo/psutil/issues/200 + # try to parse /proc/stat as a last resort + if num == 0: + search = re.compile(r'cpu\d') + with open_text(f"{get_procfs_path()}/stat") as f: + for line in f: + line = line.split(' ')[0] + if search.match(line): + num += 1 + + if num == 0: + # mimic os.cpu_count() + return None + return num + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + # Method #1 + ls = set() + # These 2 files are the same but */core_cpus_list is newer while + # */thread_siblings_list is deprecated and may disappear in the future. + # https://www.kernel.org/doc/Documentation/admin-guide/cputopology.rst + # https://github.com/giampaolo/psutil/pull/1727#issuecomment-707624964 + # https://lkml.org/lkml/2019/2/26/41 + p1 = "/sys/devices/system/cpu/cpu[0-9]*/topology/core_cpus_list" + p2 = "/sys/devices/system/cpu/cpu[0-9]*/topology/thread_siblings_list" + for path in glob.glob(p1) or glob.glob(p2): + with open_binary(path) as f: + ls.add(f.read().strip()) + result = len(ls) + if result != 0: + return result + + # Method #2 + mapping = {} + current_info = {} + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + for line in f: + line = line.strip().lower() + if not line: + # new section + try: + mapping[current_info[b'physical id']] = current_info[ + b'cpu cores' + ] + except KeyError: + pass + current_info = {} + elif line.startswith((b'physical id', b'cpu cores')): + # ongoing section + key, value = line.split(b'\t:', 1) + current_info[key] = int(value) + + result = sum(mapping.values()) + return result or None # mimic os.cpu_count() + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + with open_binary(f"{get_procfs_path()}/stat") as f: + ctx_switches = None + interrupts = None + soft_interrupts = None + for line in f: + if line.startswith(b'ctxt'): + ctx_switches = int(line.split()[1]) + elif line.startswith(b'intr'): + interrupts = int(line.split()[1]) + elif line.startswith(b'softirq'): + soft_interrupts = int(line.split()[1]) + if ( + ctx_switches is not None + and soft_interrupts is not None + and interrupts is not None + ): + break + syscalls = 0 + return _common.scpustats( + ctx_switches, interrupts, soft_interrupts, syscalls + ) + + +def _cpu_get_cpuinfo_freq(): + """Return current CPU frequency from cpuinfo if available.""" + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + return [ + float(line.split(b':', 1)[1]) + for line in f + if line.lower().startswith(b'cpu mhz') + ] + + +if os.path.exists("/sys/devices/system/cpu/cpufreq/policy0") or os.path.exists( + "/sys/devices/system/cpu/cpu0/cpufreq" +): + + def cpu_freq(): + """Return frequency metrics for all CPUs. + Contrarily to other OSes, Linux updates these values in + real-time. + """ + cpuinfo_freqs = _cpu_get_cpuinfo_freq() + paths = glob.glob( + "/sys/devices/system/cpu/cpufreq/policy[0-9]*" + ) or glob.glob("/sys/devices/system/cpu/cpu[0-9]*/cpufreq") + paths.sort(key=lambda x: int(re.search(r"[0-9]+", x).group())) + ret = [] + pjoin = os.path.join + for i, path in enumerate(paths): + if len(paths) == len(cpuinfo_freqs): + # take cached value from cpuinfo if available, see: + # https://github.com/giampaolo/psutil/issues/1851 + curr = cpuinfo_freqs[i] * 1000 + else: + curr = bcat(pjoin(path, "scaling_cur_freq"), fallback=None) + if curr is None: + # Likely an old RedHat, see: + # https://github.com/giampaolo/psutil/issues/1071 + curr = bcat(pjoin(path, "cpuinfo_cur_freq"), fallback=None) + if curr is None: + online_path = f"/sys/devices/system/cpu/cpu{i}/online" + # if cpu core is offline, set to all zeroes + if cat(online_path, fallback=None) == "0\n": + ret.append(_common.scpufreq(0.0, 0.0, 0.0)) + continue + msg = "can't find current frequency file" + raise NotImplementedError(msg) + curr = int(curr) / 1000 + max_ = int(bcat(pjoin(path, "scaling_max_freq"))) / 1000 + min_ = int(bcat(pjoin(path, "scaling_min_freq"))) / 1000 + ret.append(_common.scpufreq(curr, min_, max_)) + return ret + +else: + + def cpu_freq(): + """Alternate implementation using /proc/cpuinfo. + min and max frequencies are not available and are set to None. + """ + return [_common.scpufreq(x, 0.0, 0.0) for x in _cpu_get_cpuinfo_freq()] + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_if_addrs = cext.net_if_addrs + + +class _Ipv6UnsupportedError(Exception): + pass + + +class NetConnections: + """A wrapper on top of /proc/net/* files, retrieving per-process + and system-wide open connections (TCP, UDP, UNIX) similarly to + "netstat -an". + + Note: in case of UNIX sockets we're only able to determine the + local endpoint/path, not the one it's connected to. + According to [1] it would be possible but not easily. + + [1] http://serverfault.com/a/417946 + """ + + def __init__(self): + # The string represents the basename of the corresponding + # /proc/net/{proto_name} file. + tcp4 = ("tcp", socket.AF_INET, socket.SOCK_STREAM) + tcp6 = ("tcp6", socket.AF_INET6, socket.SOCK_STREAM) + udp4 = ("udp", socket.AF_INET, socket.SOCK_DGRAM) + udp6 = ("udp6", socket.AF_INET6, socket.SOCK_DGRAM) + unix = ("unix", socket.AF_UNIX, None) + self.tmap = { + "all": (tcp4, tcp6, udp4, udp6, unix), + "tcp": (tcp4, tcp6), + "tcp4": (tcp4,), + "tcp6": (tcp6,), + "udp": (udp4, udp6), + "udp4": (udp4,), + "udp6": (udp6,), + "unix": (unix,), + "inet": (tcp4, tcp6, udp4, udp6), + "inet4": (tcp4, udp4), + "inet6": (tcp6, udp6), + } + self._procfs_path = None + + def get_proc_inodes(self, pid): + inodes = defaultdict(list) + for fd in os.listdir(f"{self._procfs_path}/{pid}/fd"): + try: + inode = readlink(f"{self._procfs_path}/{pid}/fd/{fd}") + except (FileNotFoundError, ProcessLookupError): + # ENOENT == file which is gone in the meantime; + # os.stat(f"/proc/{self.pid}") will be done later + # to force NSP (if it's the case) + continue + except OSError as err: + if err.errno == errno.EINVAL: + # not a link + continue + if err.errno == errno.ENAMETOOLONG: + # file name too long + debug(err) + continue + raise + else: + if inode.startswith('socket:['): + # the process is using a socket + inode = inode[8:][:-1] + inodes[inode].append((pid, int(fd))) + return inodes + + def get_all_inodes(self): + inodes = {} + for pid in pids(): + try: + inodes.update(self.get_proc_inodes(pid)) + except (FileNotFoundError, ProcessLookupError, PermissionError): + # os.listdir() is gonna raise a lot of access denied + # exceptions in case of unprivileged user; that's fine + # as we'll just end up returning a connection with PID + # and fd set to None anyway. + # Both netstat -an and lsof does the same so it's + # unlikely we can do any better. + # ENOENT just means a PID disappeared on us. + continue + return inodes + + @staticmethod + def decode_address(addr, family): + """Accept an "ip:port" address as displayed in /proc/net/* + and convert it into a human readable form, like: + + "0500000A:0016" -> ("10.0.0.5", 22) + "0000000000000000FFFF00000100007F:9E49" -> ("::ffff:127.0.0.1", 40521) + + The IP address portion is a little or big endian four-byte + hexadecimal number; that is, the least significant byte is listed + first, so we need to reverse the order of the bytes to convert it + to an IP address. + The port is represented as a two-byte hexadecimal number. + + Reference: + http://linuxdevcenter.com/pub/a/linux/2000/11/16/LinuxAdmin.html + """ + ip, port = addr.split(':') + port = int(port, 16) + # this usually refers to a local socket in listen mode with + # no end-points connected + if not port: + return () + ip = ip.encode('ascii') + if family == socket.AF_INET: + # see: https://github.com/giampaolo/psutil/issues/201 + if LITTLE_ENDIAN: + ip = socket.inet_ntop(family, base64.b16decode(ip)[::-1]) + else: + ip = socket.inet_ntop(family, base64.b16decode(ip)) + else: # IPv6 + ip = base64.b16decode(ip) + try: + # see: https://github.com/giampaolo/psutil/issues/201 + if LITTLE_ENDIAN: + ip = socket.inet_ntop( + socket.AF_INET6, + struct.pack('>4I', *struct.unpack('<4I', ip)), + ) + else: + ip = socket.inet_ntop( + socket.AF_INET6, + struct.pack('<4I', *struct.unpack('<4I', ip)), + ) + except ValueError: + # see: https://github.com/giampaolo/psutil/issues/623 + if not supports_ipv6(): + raise _Ipv6UnsupportedError from None + raise + return _common.addr(ip, port) + + @staticmethod + def process_inet(file, family, type_, inodes, filter_pid=None): + """Parse /proc/net/tcp* and /proc/net/udp* files.""" + if file.endswith('6') and not os.path.exists(file): + # IPv6 not supported + return + with open_text(file) as f: + f.readline() # skip the first line + for lineno, line in enumerate(f, 1): + try: + _, laddr, raddr, status, _, _, _, _, _, inode = ( + line.split()[:10] + ) + except ValueError: + msg = ( + f"error while parsing {file}; malformed line" + f" {lineno} {line!r}" + ) + raise RuntimeError(msg) from None + if inode in inodes: + # # We assume inet sockets are unique, so we error + # # out if there are multiple references to the + # # same inode. We won't do this for UNIX sockets. + # if len(inodes[inode]) > 1 and family != socket.AF_UNIX: + # raise ValueError("ambiguous inode with multiple " + # "PIDs references") + pid, fd = inodes[inode][0] + else: + pid, fd = None, -1 + if filter_pid is not None and filter_pid != pid: + continue + else: + if type_ == socket.SOCK_STREAM: + status = TCP_STATUSES[status] + else: + status = _common.CONN_NONE + try: + laddr = NetConnections.decode_address(laddr, family) + raddr = NetConnections.decode_address(raddr, family) + except _Ipv6UnsupportedError: + continue + yield (fd, family, type_, laddr, raddr, status, pid) + + @staticmethod + def process_unix(file, family, inodes, filter_pid=None): + """Parse /proc/net/unix files.""" + with open_text(file) as f: + f.readline() # skip the first line + for line in f: + tokens = line.split() + try: + _, _, _, _, type_, _, inode = tokens[0:7] + except ValueError: + if ' ' not in line: + # see: https://github.com/giampaolo/psutil/issues/766 + continue + msg = ( + f"error while parsing {file}; malformed line {line!r}" + ) + raise RuntimeError(msg) # noqa: B904 + if inode in inodes: # noqa: SIM108 + # With UNIX sockets we can have a single inode + # referencing many file descriptors. + pairs = inodes[inode] + else: + pairs = [(None, -1)] + for pid, fd in pairs: + if filter_pid is not None and filter_pid != pid: + continue + else: + path = tokens[-1] if len(tokens) == 8 else '' + type_ = _common.socktype_to_enum(int(type_)) + # XXX: determining the remote endpoint of a + # UNIX socket on Linux is not possible, see: + # https://serverfault.com/questions/252723/ + raddr = "" + status = _common.CONN_NONE + yield (fd, family, type_, path, raddr, status, pid) + + def retrieve(self, kind, pid=None): + self._procfs_path = get_procfs_path() + if pid is not None: + inodes = self.get_proc_inodes(pid) + if not inodes: + # no connections for this process + return [] + else: + inodes = self.get_all_inodes() + ret = set() + for proto_name, family, type_ in self.tmap[kind]: + path = f"{self._procfs_path}/net/{proto_name}" + if family in {socket.AF_INET, socket.AF_INET6}: + ls = self.process_inet( + path, family, type_, inodes, filter_pid=pid + ) + else: + ls = self.process_unix(path, family, inodes, filter_pid=pid) + for fd, family, type_, laddr, raddr, status, bound_pid in ls: + if pid: + conn = _common.pconn( + fd, family, type_, laddr, raddr, status + ) + else: + conn = _common.sconn( + fd, family, type_, laddr, raddr, status, bound_pid + ) + ret.add(conn) + return list(ret) + + +_net_connections = NetConnections() + + +def net_connections(kind='inet'): + """Return system-wide open connections.""" + return _net_connections.retrieve(kind) + + +def net_io_counters(): + """Return network I/O statistics for every network interface + installed on the system as a dict of raw tuples. + """ + with open_text(f"{get_procfs_path()}/net/dev") as f: + lines = f.readlines() + retdict = {} + for line in lines[2:]: + colon = line.rfind(':') + assert colon > 0, repr(line) + name = line[:colon].strip() + fields = line[colon + 1 :].strip().split() + + ( + # in + bytes_recv, + packets_recv, + errin, + dropin, + _fifoin, # unused + _framein, # unused + _compressedin, # unused + _multicastin, # unused + # out + bytes_sent, + packets_sent, + errout, + dropout, + _fifoout, # unused + _collisionsout, # unused + _carrierout, # unused + _compressedout, # unused + ) = map(int, fields) + + retdict[name] = ( + bytes_sent, + bytes_recv, + packets_sent, + packets_recv, + errin, + errout, + dropin, + dropout, + ) + return retdict + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + duplex_map = { + cext.DUPLEX_FULL: NIC_DUPLEX_FULL, + cext.DUPLEX_HALF: NIC_DUPLEX_HALF, + cext.DUPLEX_UNKNOWN: NIC_DUPLEX_UNKNOWN, + } + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + debug(err) + else: + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = _common.snicstats( + isup, duplex_map[duplex], speed, mtu, output_flags + ) + return ret + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_usage = _psposix.disk_usage + + +def disk_io_counters(perdisk=False): + """Return disk I/O statistics for every disk installed on the + system as a dict of raw tuples. + """ + + def read_procfs(): + # OK, this is a bit confusing. The format of /proc/diskstats can + # have 3 variations. + # On Linux 2.4 each line has always 15 fields, e.g.: + # "3 0 8 hda 8 8 8 8 8 8 8 8 8 8 8" + # On Linux 2.6+ each line *usually* has 14 fields, and the disk + # name is in another position, like this: + # "3 0 hda 8 8 8 8 8 8 8 8 8 8 8" + # ...unless (Linux 2.6) the line refers to a partition instead + # of a disk, in which case the line has less fields (7): + # "3 1 hda1 8 8 8 8" + # 4.18+ has 4 fields added: + # "3 0 hda 8 8 8 8 8 8 8 8 8 8 8 0 0 0 0" + # 5.5 has 2 more fields. + # See: + # https://www.kernel.org/doc/Documentation/iostats.txt + # https://www.kernel.org/doc/Documentation/ABI/testing/procfs-diskstats + with open_text(f"{get_procfs_path()}/diskstats") as f: + lines = f.readlines() + for line in lines: + fields = line.split() + flen = len(fields) + # fmt: off + if flen == 15: + # Linux 2.4 + name = fields[3] + reads = int(fields[2]) + (reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time, _) = map(int, fields[4:14]) + elif flen == 14 or flen >= 18: + # Linux 2.6+, line referring to a disk + name = fields[2] + (reads, reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time, _) = map(int, fields[3:14]) + elif flen == 7: + # Linux 2.6+, line referring to a partition + name = fields[2] + reads, rbytes, writes, wbytes = map(int, fields[3:]) + rtime = wtime = reads_merged = writes_merged = busy_time = 0 + else: + msg = f"not sure how to interpret line {line!r}" + raise ValueError(msg) + yield (name, reads, writes, rbytes, wbytes, rtime, wtime, + reads_merged, writes_merged, busy_time) + # fmt: on + + def read_sysfs(): + for block in os.listdir('/sys/block'): + for root, _, files in os.walk(os.path.join('/sys/block', block)): + if 'stat' not in files: + continue + with open_text(os.path.join(root, 'stat')) as f: + fields = f.read().strip().split() + name = os.path.basename(root) + # fmt: off + (reads, reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time) = map(int, fields[:10]) + yield (name, reads, writes, rbytes, wbytes, rtime, + wtime, reads_merged, writes_merged, busy_time) + # fmt: on + + if os.path.exists(f"{get_procfs_path()}/diskstats"): + gen = read_procfs() + elif os.path.exists('/sys/block'): + gen = read_sysfs() + else: + msg = ( + f"{get_procfs_path()}/diskstats nor /sys/block are available on" + " this system" + ) + raise NotImplementedError(msg) + + retdict = {} + for entry in gen: + # fmt: off + (name, reads, writes, rbytes, wbytes, rtime, wtime, reads_merged, + writes_merged, busy_time) = entry + if not perdisk and not is_storage_device(name): + # perdisk=False means we want to calculate totals so we skip + # partitions (e.g. 'sda1', 'nvme0n1p1') and only include + # base disk devices (e.g. 'sda', 'nvme0n1'). Base disks + # include a total of all their partitions + some extra size + # of their own: + # $ cat /proc/diskstats + # 259 0 sda 10485760 ... + # 259 1 sda1 5186039 ... + # 259 1 sda2 5082039 ... + # See: + # https://github.com/giampaolo/psutil/pull/1313 + continue + + rbytes *= DISK_SECTOR_SIZE + wbytes *= DISK_SECTOR_SIZE + retdict[name] = (reads, writes, rbytes, wbytes, rtime, wtime, + reads_merged, writes_merged, busy_time) + # fmt: on + + return retdict + + +class RootFsDeviceFinder: + """disk_partitions() may return partitions with device == "/dev/root" + or "rootfs". This container class uses different strategies to try to + obtain the real device path. Resources: + https://bootlin.com/blog/find-root-device/ + https://www.systutorials.com/how-to-find-the-disk-where-root-is-on-in-bash-on-linux/. + """ + + __slots__ = ['major', 'minor'] + + def __init__(self): + dev = os.stat("/").st_dev + self.major = os.major(dev) + self.minor = os.minor(dev) + + def ask_proc_partitions(self): + with open_text(f"{get_procfs_path()}/partitions") as f: + for line in f.readlines()[2:]: + fields = line.split() + if len(fields) < 4: # just for extra safety + continue + major = int(fields[0]) if fields[0].isdigit() else None + minor = int(fields[1]) if fields[1].isdigit() else None + name = fields[3] + if major == self.major and minor == self.minor: + if name: # just for extra safety + return f"/dev/{name}" + + def ask_sys_dev_block(self): + path = f"/sys/dev/block/{self.major}:{self.minor}/uevent" + with open_text(path) as f: + for line in f: + if line.startswith("DEVNAME="): + name = line.strip().rpartition("DEVNAME=")[2] + if name: # just for extra safety + return f"/dev/{name}" + + def ask_sys_class_block(self): + needle = f"{self.major}:{self.minor}" + files = glob.iglob("/sys/class/block/*/dev") + for file in files: + try: + f = open_text(file) + except FileNotFoundError: # race condition + continue + else: + with f: + data = f.read().strip() + if data == needle: + name = os.path.basename(os.path.dirname(file)) + return f"/dev/{name}" + + def find(self): + path = None + if path is None: + try: + path = self.ask_proc_partitions() + except OSError as err: + debug(err) + if path is None: + try: + path = self.ask_sys_dev_block() + except OSError as err: + debug(err) + if path is None: + try: + path = self.ask_sys_class_block() + except OSError as err: + debug(err) + # We use exists() because the "/dev/*" part of the path is hard + # coded, so we want to be sure. + if path is not None and os.path.exists(path): + return path + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples.""" + fstypes = set() + procfs_path = get_procfs_path() + if not all: + with open_text(f"{procfs_path}/filesystems") as f: + for line in f: + line = line.strip() + if not line.startswith("nodev"): + fstypes.add(line.strip()) + else: + # ignore all lines starting with "nodev" except "nodev zfs" + fstype = line.split("\t")[1] + if fstype == "zfs": + fstypes.add("zfs") + + # See: https://github.com/giampaolo/psutil/issues/1307 + if procfs_path == "/proc" and os.path.isfile('/etc/mtab'): + mounts_path = os.path.realpath("/etc/mtab") + else: + mounts_path = os.path.realpath(f"{procfs_path}/self/mounts") + + retlist = [] + partitions = cext.disk_partitions(mounts_path) + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if device in {"/dev/root", "rootfs"}: + device = RootFsDeviceFinder().find() or device + if not all: + if not device or fstype not in fstypes: + continue + ntuple = _common.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + + return retlist + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_temperatures(): + """Return hardware (CPU and others) temperatures as a dict + including hardware name, label, current, max and critical + temperatures. + + Implementation notes: + - /sys/class/hwmon looks like the most recent interface to + retrieve this info, and this implementation relies on it + only (old distros will probably use something else) + - lm-sensors on Ubuntu 16.04 relies on /sys/class/hwmon + - /sys/class/thermal/thermal_zone* is another one but it's more + difficult to parse + """ + ret = collections.defaultdict(list) + basenames = glob.glob('/sys/class/hwmon/hwmon*/temp*_*') + # CentOS has an intermediate /device directory: + # https://github.com/giampaolo/psutil/issues/971 + # https://github.com/nicolargo/glances/issues/1060 + basenames.extend(glob.glob('/sys/class/hwmon/hwmon*/device/temp*_*')) + basenames = sorted({x.split('_')[0] for x in basenames}) + + # Only add the coretemp hwmon entries if they're not already in + # /sys/class/hwmon/ + # https://github.com/giampaolo/psutil/issues/1708 + # https://github.com/giampaolo/psutil/pull/1648 + basenames2 = glob.glob( + '/sys/devices/platform/coretemp.*/hwmon/hwmon*/temp*_*' + ) + repl = re.compile(r"/sys/devices/platform/coretemp.*/hwmon/") + for name in basenames2: + altname = repl.sub('/sys/class/hwmon/', name) + if altname not in basenames: + basenames.append(name) + + for base in basenames: + try: + path = base + '_input' + current = float(bcat(path)) / 1000.0 + path = os.path.join(os.path.dirname(base), 'name') + unit_name = cat(path).strip() + except (OSError, ValueError): + # A lot of things can go wrong here, so let's just skip the + # whole entry. Sure thing is Linux's /sys/class/hwmon really + # is a stinky broken mess. + # https://github.com/giampaolo/psutil/issues/1009 + # https://github.com/giampaolo/psutil/issues/1101 + # https://github.com/giampaolo/psutil/issues/1129 + # https://github.com/giampaolo/psutil/issues/1245 + # https://github.com/giampaolo/psutil/issues/1323 + continue + + high = bcat(base + '_max', fallback=None) + critical = bcat(base + '_crit', fallback=None) + label = cat(base + '_label', fallback='').strip() + + if high is not None: + try: + high = float(high) / 1000.0 + except ValueError: + high = None + if critical is not None: + try: + critical = float(critical) / 1000.0 + except ValueError: + critical = None + + ret[unit_name].append((label, current, high, critical)) + + # Indication that no sensors were detected in /sys/class/hwmon/ + if not basenames: + basenames = glob.glob('/sys/class/thermal/thermal_zone*') + basenames = sorted(set(basenames)) + + for base in basenames: + try: + path = os.path.join(base, 'temp') + current = float(bcat(path)) / 1000.0 + path = os.path.join(base, 'type') + unit_name = cat(path).strip() + except (OSError, ValueError) as err: + debug(err) + continue + + trip_paths = glob.glob(base + '/trip_point*') + trip_points = { + '_'.join(os.path.basename(p).split('_')[0:3]) + for p in trip_paths + } + critical = None + high = None + for trip_point in trip_points: + path = os.path.join(base, trip_point + "_type") + trip_type = cat(path, fallback='').strip() + if trip_type == 'critical': + critical = bcat( + os.path.join(base, trip_point + "_temp"), fallback=None + ) + elif trip_type == 'high': + high = bcat( + os.path.join(base, trip_point + "_temp"), fallback=None + ) + + if high is not None: + try: + high = float(high) / 1000.0 + except ValueError: + high = None + if critical is not None: + try: + critical = float(critical) / 1000.0 + except ValueError: + critical = None + + ret[unit_name].append(('', current, high, critical)) + + return dict(ret) + + +def sensors_fans(): + """Return hardware fans info (for CPU and other peripherals) as a + dict including hardware label and current speed. + + Implementation notes: + - /sys/class/hwmon looks like the most recent interface to + retrieve this info, and this implementation relies on it + only (old distros will probably use something else) + - lm-sensors on Ubuntu 16.04 relies on /sys/class/hwmon + """ + ret = collections.defaultdict(list) + basenames = glob.glob('/sys/class/hwmon/hwmon*/fan*_*') + if not basenames: + # CentOS has an intermediate /device directory: + # https://github.com/giampaolo/psutil/issues/971 + basenames = glob.glob('/sys/class/hwmon/hwmon*/device/fan*_*') + + basenames = sorted({x.split("_")[0] for x in basenames}) + for base in basenames: + try: + current = int(bcat(base + '_input')) + except OSError as err: + debug(err) + continue + unit_name = cat(os.path.join(os.path.dirname(base), 'name')).strip() + label = cat(base + '_label', fallback='').strip() + ret[unit_name].append(_common.sfan(label, current)) + + return dict(ret) + + +def sensors_battery(): + """Return battery information. + Implementation note: it appears /sys/class/power_supply/BAT0/ + directory structure may vary and provide files with the same + meaning but under different names, see: + https://github.com/giampaolo/psutil/issues/966. + """ + null = object() + + def multi_bcat(*paths): + """Attempt to read the content of multiple files which may + not exist. If none of them exist return None. + """ + for path in paths: + ret = bcat(path, fallback=null) + if ret != null: + try: + return int(ret) + except ValueError: + return ret.strip() + return None + + bats = [ + x + for x in os.listdir(POWER_SUPPLY_PATH) + if x.startswith('BAT') or 'battery' in x.lower() + ] + if not bats: + return None + # Get the first available battery. Usually this is "BAT0", except + # some rare exceptions: + # https://github.com/giampaolo/psutil/issues/1238 + root = os.path.join(POWER_SUPPLY_PATH, min(bats)) + + # Base metrics. + energy_now = multi_bcat(root + "/energy_now", root + "/charge_now") + power_now = multi_bcat(root + "/power_now", root + "/current_now") + energy_full = multi_bcat(root + "/energy_full", root + "/charge_full") + time_to_empty = multi_bcat(root + "/time_to_empty_now") + + # Percent. If we have energy_full the percentage will be more + # accurate compared to reading /capacity file (float vs. int). + if energy_full is not None and energy_now is not None: + try: + percent = 100.0 * energy_now / energy_full + except ZeroDivisionError: + percent = 0.0 + else: + percent = int(cat(root + "/capacity", fallback=-1)) + if percent == -1: + return None + + # Is AC power cable plugged in? + # Note: AC0 is not always available and sometimes (e.g. CentOS7) + # it's called "AC". + power_plugged = None + online = multi_bcat( + os.path.join(POWER_SUPPLY_PATH, "AC0/online"), + os.path.join(POWER_SUPPLY_PATH, "AC/online"), + ) + if online is not None: + power_plugged = online == 1 + else: + status = cat(root + "/status", fallback="").strip().lower() + if status == "discharging": + power_plugged = False + elif status in {"charging", "full"}: + power_plugged = True + + # Seconds left. + # Note to self: we may also calculate the charging ETA as per: + # https://github.com/thialfihar/dotfiles/blob/ + # 013937745fd9050c30146290e8f963d65c0179e6/bin/battery.py#L55 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif energy_now is not None and power_now is not None: + try: + secsleft = int(energy_now / abs(power_now) * 3600) + except ZeroDivisionError: + secsleft = _common.POWER_TIME_UNKNOWN + elif time_to_empty is not None: + secsleft = int(time_to_empty * 60) + if secsleft < 0: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = _common.POWER_TIME_UNKNOWN + + return _common.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + nt = _common.suser(user, tty or None, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +def boot_time(): + """Return the system boot time expressed in seconds since the epoch.""" + path = f"{get_procfs_path()}/stat" + with open_binary(path) as f: + for line in f: + if line.startswith(b'btime'): + return float(line.strip().split()[1]) + msg = f"line 'btime' not found in {path}" + raise RuntimeError(msg) + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + path = get_procfs_path().encode(ENCODING) + return [int(x) for x in os.listdir(path) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix PID. Linux TIDs are not + supported (always return False). + """ + if not _psposix.pid_exists(pid): + return False + else: + # Linux's apparently does not distinguish between PIDs and TIDs + # (thread IDs). + # listdir("/proc") won't show any TID (only PIDs) but + # os.stat("/proc/{tid}") will succeed if {tid} exists. + # os.kill() can also be passed a TID. This is quite confusing. + # In here we want to enforce this distinction and support PIDs + # only, see: + # https://github.com/giampaolo/psutil/issues/687 + try: + # Note: already checked that this is faster than using a + # regular expr. Also (a lot) faster than doing + # 'return pid in pids()' + path = f"{get_procfs_path()}/{pid}/status" + with open_binary(path) as f: + for line in f: + if line.startswith(b"Tgid:"): + tgid = int(line.split()[1]) + # If tgid and pid are the same then we're + # dealing with a process PID. + return tgid == pid + msg = f"'Tgid' line not found in {path}" + raise ValueError(msg) + except (OSError, ValueError): + return pid in pids() + + +def ppid_map(): + """Obtain a {pid: ppid, ...} dict for all running processes in + one shot. Used to speed up Process.children(). + """ + ret = {} + procfs_path = get_procfs_path() + for pid in pids(): + try: + with open_binary(f"{procfs_path}/{pid}/stat") as f: + data = f.read() + except (FileNotFoundError, ProcessLookupError): + pass + except PermissionError as err: + raise AccessDenied(pid) from err + else: + rpar = data.rfind(b')') + dset = data[rpar + 2 :].split() + ppid = int(dset[1]) + ret[pid] = ppid + return ret + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError and OSError exceptions + into NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, name = self.pid, self._name + try: + return fun(self, *args, **kwargs) + except PermissionError as err: + raise AccessDenied(pid, name) from err + except ProcessLookupError as err: + self._raise_if_zombie() + raise NoSuchProcess(pid, name) from err + except FileNotFoundError as err: + self._raise_if_zombie() + # /proc/PID directory may still exist, but the files within + # it may not, indicating the process is gone, see: + # https://github.com/giampaolo/psutil/issues/2418 + if not os.path.exists(f"{self._procfs_path}/{pid}/stat"): + raise NoSuchProcess(pid, name) from err + raise + + return wrapper + + +class Process: + """Linux process implementation.""" + + __slots__ = [ + "_cache", + "_ctime", + "_name", + "_ppid", + "_procfs_path", + "pid", + ] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._ctime = None + self._procfs_path = get_procfs_path() + + def _is_zombie(self): + # Note: most of the times Linux is able to return info about the + # process even if it's a zombie, and /proc/{pid} will exist. + # There are some exceptions though, like exe(), cmdline() and + # memory_maps(). In these cases /proc/{pid}/{file} exists but + # it's empty. Instead of returning a "null" value we'll raise an + # exception. + try: + data = bcat(f"{self._procfs_path}/{self.pid}/stat") + except OSError: + return False + else: + rpar = data.rfind(b')') + status = data[rpar + 2 : rpar + 3] + return status == b"Z" + + def _raise_if_zombie(self): + if self._is_zombie(): + raise ZombieProcess(self.pid, self._name, self._ppid) + + def _raise_if_not_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + os.stat(f"{self._procfs_path}/{self.pid}") + + def _readlink(self, path, fallback=UNSET): + # * https://github.com/giampaolo/psutil/issues/503 + # os.readlink('/proc/pid/exe') may raise ESRCH (ProcessLookupError) + # instead of ENOENT (FileNotFoundError) when it races. + # * ENOENT may occur also if the path actually exists if PID is + # a low PID (~0-20 range). + # * https://github.com/giampaolo/psutil/issues/2514 + try: + return readlink(path) + except (FileNotFoundError, ProcessLookupError): + if os.path.lexists(f"{self._procfs_path}/{self.pid}"): + self._raise_if_zombie() + if fallback is not UNSET: + return fallback + raise + + @wrap_exceptions + @memoize_when_activated + def _parse_stat_file(self): + """Parse /proc/{pid}/stat file and return a dict with various + process info. + Using "man proc" as a reference: where "man proc" refers to + position N always subtract 3 (e.g ppid position 4 in + 'man proc' == position 1 in here). + The return value is cached in case oneshot() ctx manager is + in use. + """ + data = bcat(f"{self._procfs_path}/{self.pid}/stat") + # Process name is between parentheses. It can contain spaces and + # other parentheses. This is taken into account by looking for + # the first occurrence of "(" and the last occurrence of ")". + rpar = data.rfind(b')') + name = data[data.find(b'(') + 1 : rpar] + fields = data[rpar + 2 :].split() + + ret = {} + ret['name'] = name + ret['status'] = fields[0] + ret['ppid'] = fields[1] + ret['ttynr'] = fields[4] + ret['utime'] = fields[11] + ret['stime'] = fields[12] + ret['children_utime'] = fields[13] + ret['children_stime'] = fields[14] + ret['create_time'] = fields[19] + ret['cpu_num'] = fields[36] + try: + ret['blkio_ticks'] = fields[39] # aka 'delayacct_blkio_ticks' + except IndexError: + # https://github.com/giampaolo/psutil/issues/2455 + debug("can't get blkio_ticks, set iowait to 0") + ret['blkio_ticks'] = 0 + + return ret + + @wrap_exceptions + @memoize_when_activated + def _read_status_file(self): + """Read /proc/{pid}/stat file and return its content. + The return value is cached in case oneshot() ctx manager is + in use. + """ + with open_binary(f"{self._procfs_path}/{self.pid}/status") as f: + return f.read() + + @wrap_exceptions + @memoize_when_activated + def _read_smaps_file(self): + with open_binary(f"{self._procfs_path}/{self.pid}/smaps") as f: + return f.read().strip() + + def oneshot_enter(self): + self._parse_stat_file.cache_activate(self) + self._read_status_file.cache_activate(self) + self._read_smaps_file.cache_activate(self) + + def oneshot_exit(self): + self._parse_stat_file.cache_deactivate(self) + self._read_status_file.cache_deactivate(self) + self._read_smaps_file.cache_deactivate(self) + + @wrap_exceptions + def name(self): + # XXX - gets changed later and probably needs refactoring + return decode(self._parse_stat_file()['name']) + + @wrap_exceptions + def exe(self): + return self._readlink( + f"{self._procfs_path}/{self.pid}/exe", fallback="" + ) + + @wrap_exceptions + def cmdline(self): + with open_text(f"{self._procfs_path}/{self.pid}/cmdline") as f: + data = f.read() + if not data: + # may happen in case of zombie process + self._raise_if_zombie() + return [] + # 'man proc' states that args are separated by null bytes '\0' + # and last char is supposed to be a null byte. Nevertheless + # some processes may change their cmdline after being started + # (via setproctitle() or similar), they are usually not + # compliant with this rule and use spaces instead. Google + # Chrome process is an example. See: + # https://github.com/giampaolo/psutil/issues/1179 + sep = '\x00' if data.endswith('\x00') else ' ' + if data.endswith(sep): + data = data[:-1] + cmdline = data.split(sep) + # Sometimes last char is a null byte '\0' but the args are + # separated by spaces, see: https://github.com/giampaolo/psutil/ + # issues/1179#issuecomment-552984549 + if sep == '\x00' and len(cmdline) == 1 and ' ' in data: + cmdline = data.split(' ') + return cmdline + + @wrap_exceptions + def environ(self): + with open_text(f"{self._procfs_path}/{self.pid}/environ") as f: + data = f.read() + return parse_environ_block(data) + + @wrap_exceptions + def terminal(self): + tty_nr = int(self._parse_stat_file()['ttynr']) + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + # May not be available on old kernels. + if os.path.exists(f"/proc/{os.getpid()}/io"): + + @wrap_exceptions + def io_counters(self): + fname = f"{self._procfs_path}/{self.pid}/io" + fields = {} + with open_binary(fname) as f: + for line in f: + # https://github.com/giampaolo/psutil/issues/1004 + line = line.strip() + if line: + try: + name, value = line.split(b': ') + except ValueError: + # https://github.com/giampaolo/psutil/issues/1004 + continue + else: + fields[name] = int(value) + if not fields: + msg = f"{fname} file was empty" + raise RuntimeError(msg) + try: + return pio( + fields[b'syscr'], # read syscalls + fields[b'syscw'], # write syscalls + fields[b'read_bytes'], # read bytes + fields[b'write_bytes'], # write bytes + fields[b'rchar'], # read chars + fields[b'wchar'], # write chars + ) + except KeyError as err: + msg = ( + f"{err.args[0]!r} field was not found in {fname}; found" + f" fields are {fields!r}" + ) + raise ValueError(msg) from None + + @wrap_exceptions + def cpu_times(self): + values = self._parse_stat_file() + utime = float(values['utime']) / CLOCK_TICKS + stime = float(values['stime']) / CLOCK_TICKS + children_utime = float(values['children_utime']) / CLOCK_TICKS + children_stime = float(values['children_stime']) / CLOCK_TICKS + iowait = float(values['blkio_ticks']) / CLOCK_TICKS + return pcputimes(utime, stime, children_utime, children_stime, iowait) + + @wrap_exceptions + def cpu_num(self): + """What CPU the process is on.""" + return int(self._parse_stat_file()['cpu_num']) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout, self._name) + + @wrap_exceptions + def create_time(self, monotonic=False): + # The 'starttime' field in /proc/[pid]/stat is expressed in + # jiffies (clock ticks per second), a relative value which + # represents the number of clock ticks that have passed since + # the system booted until the process was created. It never + # changes and is unaffected by system clock updates. + if self._ctime is None: + self._ctime = ( + float(self._parse_stat_file()['create_time']) / CLOCK_TICKS + ) + if monotonic: + return self._ctime + # Add the boot time, returning time expressed in seconds since + # the epoch. This is subject to system clock updates. + return self._ctime + boot_time() + + @wrap_exceptions + def memory_info(self): + # ============================================================ + # | FIELD | DESCRIPTION | AKA | TOP | + # ============================================================ + # | rss | resident set size | | RES | + # | vms | total program size | size | VIRT | + # | shared | shared pages (from shared mappings) | | SHR | + # | text | text ('code') | trs | CODE | + # | lib | library (unused in Linux 2.6) | lrs | | + # | data | data + stack | drs | DATA | + # | dirty | dirty pages (unused in Linux 2.6) | dt | | + # ============================================================ + with open_binary(f"{self._procfs_path}/{self.pid}/statm") as f: + vms, rss, shared, text, lib, data, dirty = ( + int(x) * PAGESIZE for x in f.readline().split()[:7] + ) + return pmem(rss, vms, shared, text, lib, data, dirty) + + if HAS_PROC_SMAPS_ROLLUP or HAS_PROC_SMAPS: + + def _parse_smaps_rollup(self): + # /proc/pid/smaps_rollup was added to Linux in 2017. Faster + # than /proc/pid/smaps. It reports higher PSS than */smaps + # (from 1k up to 200k higher; tested against all processes). + # IMPORTANT: /proc/pid/smaps_rollup is weird, because it + # raises ESRCH / ENOENT for many PIDs, even if they're alive + # (also as root). In that case we'll use /proc/pid/smaps as + # fallback, which is slower but has a +50% success rate + # compared to /proc/pid/smaps_rollup. + uss = pss = swap = 0 + with open_binary( + f"{self._procfs_path}/{self.pid}/smaps_rollup" + ) as f: + for line in f: + if line.startswith(b"Private_"): + # Private_Clean, Private_Dirty, Private_Hugetlb + uss += int(line.split()[1]) * 1024 + elif line.startswith(b"Pss:"): + pss = int(line.split()[1]) * 1024 + elif line.startswith(b"Swap:"): + swap = int(line.split()[1]) * 1024 + return (uss, pss, swap) + + @wrap_exceptions + def _parse_smaps( + self, + # Gets Private_Clean, Private_Dirty, Private_Hugetlb. + _private_re=re.compile(br"\nPrivate.*:\s+(\d+)"), + _pss_re=re.compile(br"\nPss\:\s+(\d+)"), + _swap_re=re.compile(br"\nSwap\:\s+(\d+)"), + ): + # /proc/pid/smaps does not exist on kernels < 2.6.14 or if + # CONFIG_MMU kernel configuration option is not enabled. + + # Note: using 3 regexes is faster than reading the file + # line by line. + # + # You might be tempted to calculate USS by subtracting + # the "shared" value from the "resident" value in + # /proc//statm. But at least on Linux, statm's "shared" + # value actually counts pages backed by files, which has + # little to do with whether the pages are actually shared. + # /proc/self/smaps on the other hand appears to give us the + # correct information. + smaps_data = self._read_smaps_file() + # Note: smaps file can be empty for certain processes. + # The code below will not crash though and will result to 0. + uss = sum(map(int, _private_re.findall(smaps_data))) * 1024 + pss = sum(map(int, _pss_re.findall(smaps_data))) * 1024 + swap = sum(map(int, _swap_re.findall(smaps_data))) * 1024 + return (uss, pss, swap) + + @wrap_exceptions + def memory_full_info(self): + if HAS_PROC_SMAPS_ROLLUP: # faster + try: + uss, pss, swap = self._parse_smaps_rollup() + except (ProcessLookupError, FileNotFoundError): + uss, pss, swap = self._parse_smaps() + else: + uss, pss, swap = self._parse_smaps() + basic_mem = self.memory_info() + return pfullmem(*basic_mem + (uss, pss, swap)) + + else: + memory_full_info = memory_info + + if HAS_PROC_SMAPS: + + @wrap_exceptions + def memory_maps(self): + """Return process's mapped memory regions as a list of named + tuples. Fields are explained in 'man proc'; here is an updated + (Apr 2012) version: https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/Documentation/filesystems/proc.txt?id=b76437579d1344b612cf1851ae610c636cec7db0. + + /proc/{PID}/smaps does not exist on kernels < 2.6.14 or if + CONFIG_MMU kernel configuration option is not enabled. + """ + + def get_blocks(lines, current_block): + data = {} + for line in lines: + fields = line.split(None, 5) + if not fields[0].endswith(b':'): + # new block section + yield (current_block.pop(), data) + current_block.append(line) + else: + try: + data[fields[0]] = int(fields[1]) * 1024 + except (ValueError, IndexError): + if fields[0].startswith(b'VmFlags:'): + # see issue #369 + continue + msg = f"don't know how to interpret line {line!r}" + raise ValueError(msg) from None + yield (current_block.pop(), data) + + data = self._read_smaps_file() + # Note: smaps file can be empty for certain processes or for + # zombies. + if not data: + self._raise_if_zombie() + return [] + lines = data.split(b'\n') + ls = [] + first_line = lines.pop(0) + current_block = [first_line] + for header, data in get_blocks(lines, current_block): + hfields = header.split(None, 5) + try: + addr, perms, _offset, _dev, _inode, path = hfields + except ValueError: + addr, perms, _offset, _dev, _inode, path = hfields + [''] + if not path: + path = '[anon]' + else: + path = decode(path) + path = path.strip() + if path.endswith(' (deleted)') and not path_exists_strict( + path + ): + path = path[:-10] + item = ( + decode(addr), + decode(perms), + path, + data.get(b'Rss:', 0), + data.get(b'Size:', 0), + data.get(b'Pss:', 0), + data.get(b'Shared_Clean:', 0), + data.get(b'Shared_Dirty:', 0), + data.get(b'Private_Clean:', 0), + data.get(b'Private_Dirty:', 0), + data.get(b'Referenced:', 0), + data.get(b'Anonymous:', 0), + data.get(b'Swap:', 0), + ) + ls.append(item) + return ls + + @wrap_exceptions + def cwd(self): + return self._readlink( + f"{self._procfs_path}/{self.pid}/cwd", fallback="" + ) + + @wrap_exceptions + def num_ctx_switches( + self, _ctxsw_re=re.compile(br'ctxt_switches:\t(\d+)') + ): + data = self._read_status_file() + ctxsw = _ctxsw_re.findall(data) + if not ctxsw: + msg = ( + "'voluntary_ctxt_switches' and" + " 'nonvoluntary_ctxt_switches'lines were not found in" + f" {self._procfs_path}/{self.pid}/status; the kernel is" + " probably older than 2.6.23" + ) + raise NotImplementedError(msg) + return _common.pctxsw(int(ctxsw[0]), int(ctxsw[1])) + + @wrap_exceptions + def num_threads(self, _num_threads_re=re.compile(br'Threads:\t(\d+)')): + # Using a re is faster than iterating over file line by line. + data = self._read_status_file() + return int(_num_threads_re.findall(data)[0]) + + @wrap_exceptions + def threads(self): + thread_ids = os.listdir(f"{self._procfs_path}/{self.pid}/task") + thread_ids.sort() + retlist = [] + hit_enoent = False + for thread_id in thread_ids: + fname = f"{self._procfs_path}/{self.pid}/task/{thread_id}/stat" + try: + with open_binary(fname) as f: + st = f.read().strip() + except (FileNotFoundError, ProcessLookupError): + # no such file or directory or no such process; + # it means thread disappeared on us + hit_enoent = True + continue + # ignore the first two values ("pid (exe)") + st = st[st.find(b')') + 2 :] + values = st.split(b' ') + utime = float(values[11]) / CLOCK_TICKS + stime = float(values[12]) / CLOCK_TICKS + ntuple = _common.pthread(int(thread_id), utime, stime) + retlist.append(ntuple) + if hit_enoent: + self._raise_if_not_alive() + return retlist + + @wrap_exceptions + def nice_get(self): + # with open_text(f"{self._procfs_path}/{self.pid}/stat") as f: + # data = f.read() + # return int(data.split()[18]) + + # Use C implementation + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + # starting from CentOS 6. + if HAS_CPU_AFFINITY: + + @wrap_exceptions + def cpu_affinity_get(self): + return cext.proc_cpu_affinity_get(self.pid) + + def _get_eligible_cpus( + self, _re=re.compile(br"Cpus_allowed_list:\t(\d+)-(\d+)") + ): + # See: https://github.com/giampaolo/psutil/issues/956 + data = self._read_status_file() + match = _re.findall(data) + if match: + return list(range(int(match[0][0]), int(match[0][1]) + 1)) + else: + return list(range(len(per_cpu_times()))) + + @wrap_exceptions + def cpu_affinity_set(self, cpus): + try: + cext.proc_cpu_affinity_set(self.pid, cpus) + except (OSError, ValueError) as err: + if isinstance(err, ValueError) or err.errno == errno.EINVAL: + eligible_cpus = self._get_eligible_cpus() + all_cpus = tuple(range(len(per_cpu_times()))) + for cpu in cpus: + if cpu not in all_cpus: + msg = ( + f"invalid CPU {cpu!r}; choose between" + f" {eligible_cpus!r}" + ) + raise ValueError(msg) from None + if cpu not in eligible_cpus: + msg = ( + f"CPU number {cpu} is not eligible; choose" + f" between {eligible_cpus}" + ) + raise ValueError(msg) from err + raise + + # only starting from kernel 2.6.13 + if HAS_PROC_IO_PRIORITY: + + @wrap_exceptions + def ionice_get(self): + ioclass, value = cext.proc_ioprio_get(self.pid) + ioclass = IOPriority(ioclass) + return _common.pionice(ioclass, value) + + @wrap_exceptions + def ionice_set(self, ioclass, value): + if value is None: + value = 0 + if value and ioclass in { + IOPriority.IOPRIO_CLASS_IDLE, + IOPriority.IOPRIO_CLASS_NONE, + }: + msg = f"{ioclass!r} ioclass accepts no value" + raise ValueError(msg) + if value < 0 or value > 7: + msg = "value not in 0-7 range" + raise ValueError(msg) + return cext.proc_ioprio_set(self.pid, ioclass, value) + + if hasattr(resource, "prlimit"): + + @wrap_exceptions + def rlimit(self, resource_, limits=None): + # If pid is 0 prlimit() applies to the calling process and + # we don't want that. We should never get here though as + # PID 0 is not supported on Linux. + if self.pid == 0: + msg = "can't use prlimit() against PID 0 process" + raise ValueError(msg) + try: + if limits is None: + # get + return resource.prlimit(self.pid, resource_) + else: + # set + if len(limits) != 2: + msg = ( + "second argument must be a (soft, hard) " + f"tuple, got {limits!r}" + ) + raise ValueError(msg) + resource.prlimit(self.pid, resource_, limits) + except OSError as err: + if err.errno == errno.ENOSYS: + # I saw this happening on Travis: + # https://travis-ci.org/giampaolo/psutil/jobs/51368273 + self._raise_if_zombie() + raise + + @wrap_exceptions + def status(self): + letter = self._parse_stat_file()['status'] + letter = letter.decode() + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(letter, '?') + + @wrap_exceptions + def open_files(self): + retlist = [] + files = os.listdir(f"{self._procfs_path}/{self.pid}/fd") + hit_enoent = False + for fd in files: + file = f"{self._procfs_path}/{self.pid}/fd/{fd}" + try: + path = readlink(file) + except (FileNotFoundError, ProcessLookupError): + # ENOENT == file which is gone in the meantime + hit_enoent = True + continue + except OSError as err: + if err.errno == errno.EINVAL: + # not a link + continue + if err.errno == errno.ENAMETOOLONG: + # file name too long + debug(err) + continue + raise + else: + # If path is not an absolute there's no way to tell + # whether it's a regular file or not, so we skip it. + # A regular file is always supposed to be have an + # absolute path though. + if path.startswith('/') and isfile_strict(path): + # Get file position and flags. + file = f"{self._procfs_path}/{self.pid}/fdinfo/{fd}" + try: + with open_binary(file) as f: + pos = int(f.readline().split()[1]) + flags = int(f.readline().split()[1], 8) + except (FileNotFoundError, ProcessLookupError): + # fd gone in the meantime; process may + # still be alive + hit_enoent = True + else: + mode = file_flags_to_mode(flags) + ntuple = popenfile( + path, int(fd), int(pos), mode, flags + ) + retlist.append(ntuple) + if hit_enoent: + self._raise_if_not_alive() + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = _net_connections.retrieve(kind, self.pid) + self._raise_if_not_alive() + return ret + + @wrap_exceptions + def num_fds(self): + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def ppid(self): + return int(self._parse_stat_file()['ppid']) + + @wrap_exceptions + def uids(self, _uids_re=re.compile(br'Uid:\t(\d+)\t(\d+)\t(\d+)')): + data = self._read_status_file() + real, effective, saved = _uids_re.findall(data)[0] + return _common.puids(int(real), int(effective), int(saved)) + + @wrap_exceptions + def gids(self, _gids_re=re.compile(br'Gid:\t(\d+)\t(\d+)\t(\d+)')): + data = self._read_status_file() + real, effective, saved = _gids_re.findall(data)[0] + return _common.pgids(int(real), int(effective), int(saved)) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psosx.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psosx.py new file mode 100644 index 0000000000000000000000000000000000000000..e1030b654b0558b2328602424f9c24c8cb9bf5ad --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psosx.py @@ -0,0 +1,565 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""macOS platform implementation.""" + +import errno +import functools +import os +from collections import namedtuple + +from . import _common +from . import _psposix +from . import _psutil_osx as cext +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import isfile_strict +from ._common import memoize_when_activated +from ._common import parse_environ_block +from ._common import usage_percent + +__extra__all__ = [] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +PAGESIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SRUN: _common.STATUS_RUNNING, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, +} + +kinfo_proc_map = dict( + ppid=0, + ruid=1, + euid=2, + suid=3, + rgid=4, + egid=5, + sgid=6, + ttynr=7, + ctime=8, + status=9, + name=10, +) + +pidtaskinfo_map = dict( + cpuutime=0, + cpustime=1, + rss=2, + vms=3, + pfaults=4, + pageins=5, + numthreads=6, + volctxsw=7, +) + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# fmt: off +# psutil.cpu_times() +scputimes = namedtuple('scputimes', ['user', 'nice', 'system', 'idle']) +# psutil.virtual_memory() +svmem = namedtuple( + 'svmem', ['total', 'available', 'percent', 'used', 'free', + 'active', 'inactive', 'wired']) +# psutil.Process.memory_info() +pmem = namedtuple('pmem', ['rss', 'vms', 'pfaults', 'pageins']) +# psutil.Process.memory_full_info() +pfullmem = namedtuple('pfullmem', pmem._fields + ('uss', )) +# fmt: on + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """System virtual memory as a namedtuple.""" + total, active, inactive, wired, free, speculative = cext.virtual_mem() + # This is how Zabbix calculate avail and used mem: + # https://github.com/zabbix/zabbix/blob/master/src/libs/zbxsysinfo/osx/memory.c + # Also see: https://github.com/giampaolo/psutil/issues/1277 + avail = inactive + free + used = active + wired + # This is NOT how Zabbix calculates free mem but it matches "free" + # cmdline utility. + free -= speculative + percent = usage_percent((total - avail), total, round_=1) + return svmem(total, avail, percent, used, free, active, inactive, wired) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + total, used, free, sin, sout = cext.swap_mem() + percent = usage_percent(used, total, round_=1) + return _common.sswap(total, used, free, percent, sin, sout) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system CPU times as a namedtuple.""" + user, nice, system, idle = cext.cpu_times() + return scputimes(user, nice, system, idle) + + +def per_cpu_times(): + """Return system CPU times as a named tuple.""" + ret = [] + for cpu_t in cext.per_cpu_times(): + user, nice, system, idle = cpu_t + item = scputimes(user, nice, system, idle) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + ctx_switches, interrupts, soft_interrupts, syscalls, _traps = ( + cext.cpu_stats() + ) + return _common.scpustats( + ctx_switches, interrupts, soft_interrupts, syscalls + ) + + +if cext.has_cpu_freq(): # not always available on ARM64 + + def cpu_freq(): + """Return CPU frequency. + On macOS per-cpu frequency is not supported. + Also, the returned frequency never changes, see: + https://arstechnica.com/civis/viewtopic.php?f=19&t=465002. + """ + curr, min_, max_ = cext.cpu_freq() + return [_common.scpufreq(curr, min_, max_)] + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_usage = _psposix.disk_usage +disk_io_counters = cext.disk_io_counters + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples.""" + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + if not os.path.isabs(device) or not os.path.exists(device): + continue + ntuple = _common.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_battery(): + """Return battery information.""" + try: + percent, minsleft, power_plugged = cext.sensors_battery() + except NotImplementedError: + # no power source - return None according to interface + return None + power_plugged = power_plugged == 1 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif minsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = minsleft * 60 + return _common.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_connections(kind='inet'): + """System-wide network connections.""" + # Note: on macOS this will fail with AccessDenied unless + # the process is owned by root. + ret = [] + for pid in pids(): + try: + cons = Process(pid).net_connections(kind) + except NoSuchProcess: + continue + else: + if cons: + for c in cons: + c = list(c) + [pid] + ret.append(_common.sconn(*c)) + return ret + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + else: + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = _common.snicstats( + isup, duplex, speed, mtu, output_flags + ) + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +try: + INIT_BOOT_TIME = boot_time() +except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + INIT_BOOT_TIME = 0 + + +def adjust_proc_create_time(ctime): + """Account for system clock updates.""" + if INIT_BOOT_TIME == 0: + return ctime + + diff = INIT_BOOT_TIME - boot_time() + if diff == 0 or abs(diff) < 1: + return ctime + + debug("system clock was updated; adjusting process create_time()") + if diff < 0: + return ctime - diff + return ctime + diff + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + if tty == '~': + continue # reboot or shutdown + if not tstamp: + continue + nt = _common.suser(user, tty or None, hostname or None, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + ls = cext.pids() + if 0 not in ls: + # On certain macOS versions pids() C doesn't return PID 0 but + # "ps" does and the process is querable via sysctl(): + # https://travis-ci.org/giampaolo/psutil/jobs/309619941 + try: + Process(0).create_time() + ls.insert(0, 0) + except NoSuchProcess: + pass + except AccessDenied: + ls.insert(0, 0) + return ls + + +pid_exists = _psposix.pid_exists + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError exceptions into + NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except ProcessLookupError as err: + if cext.proc_is_zombie(pid): + raise ZombieProcess(pid, name, ppid) from err + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except cext.ZombieProcessError as err: + raise ZombieProcess(pid, name, ppid) from err + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + @wrap_exceptions + @memoize_when_activated + def _get_kinfo_proc(self): + # Note: should work with all PIDs without permission issues. + ret = cext.proc_kinfo_oneshot(self.pid) + assert len(ret) == len(kinfo_proc_map) + return ret + + @wrap_exceptions + @memoize_when_activated + def _get_pidtaskinfo(self): + # Note: should work for PIDs owned by user only. + ret = cext.proc_pidtaskinfo_oneshot(self.pid) + assert len(ret) == len(pidtaskinfo_map) + return ret + + def oneshot_enter(self): + self._get_kinfo_proc.cache_activate(self) + self._get_pidtaskinfo.cache_activate(self) + + def oneshot_exit(self): + self._get_kinfo_proc.cache_deactivate(self) + self._get_pidtaskinfo.cache_deactivate(self) + + @wrap_exceptions + def name(self): + name = self._get_kinfo_proc()[kinfo_proc_map['name']] + return name if name is not None else cext.proc_name(self.pid) + + @wrap_exceptions + def exe(self): + return cext.proc_exe(self.pid) + + @wrap_exceptions + def cmdline(self): + return cext.proc_cmdline(self.pid) + + @wrap_exceptions + def environ(self): + return parse_environ_block(cext.proc_environ(self.pid)) + + @wrap_exceptions + def ppid(self): + self._ppid = self._get_kinfo_proc()[kinfo_proc_map['ppid']] + return self._ppid + + @wrap_exceptions + def cwd(self): + return cext.proc_cwd(self.pid) + + @wrap_exceptions + def uids(self): + rawtuple = self._get_kinfo_proc() + return _common.puids( + rawtuple[kinfo_proc_map['ruid']], + rawtuple[kinfo_proc_map['euid']], + rawtuple[kinfo_proc_map['suid']], + ) + + @wrap_exceptions + def gids(self): + rawtuple = self._get_kinfo_proc() + return _common.puids( + rawtuple[kinfo_proc_map['rgid']], + rawtuple[kinfo_proc_map['egid']], + rawtuple[kinfo_proc_map['sgid']], + ) + + @wrap_exceptions + def terminal(self): + tty_nr = self._get_kinfo_proc()[kinfo_proc_map['ttynr']] + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + @wrap_exceptions + def memory_info(self): + rawtuple = self._get_pidtaskinfo() + return pmem( + rawtuple[pidtaskinfo_map['rss']], + rawtuple[pidtaskinfo_map['vms']], + rawtuple[pidtaskinfo_map['pfaults']], + rawtuple[pidtaskinfo_map['pageins']], + ) + + @wrap_exceptions + def memory_full_info(self): + basic_mem = self.memory_info() + uss = cext.proc_memory_uss(self.pid) + return pfullmem(*basic_mem + (uss,)) + + @wrap_exceptions + def cpu_times(self): + rawtuple = self._get_pidtaskinfo() + return _common.pcputimes( + rawtuple[pidtaskinfo_map['cpuutime']], + rawtuple[pidtaskinfo_map['cpustime']], + # children user / system times are not retrievable (set to 0) + 0.0, + 0.0, + ) + + @wrap_exceptions + def create_time(self, monotonic=False): + ctime = self._get_kinfo_proc()[kinfo_proc_map['ctime']] + if not monotonic: + ctime = adjust_proc_create_time(ctime) + return ctime + + @wrap_exceptions + def num_ctx_switches(self): + # Unvoluntary value seems not to be available; + # getrusage() numbers seems to confirm this theory. + # We set it to 0. + vol = self._get_pidtaskinfo()[pidtaskinfo_map['volctxsw']] + return _common.pctxsw(vol, 0) + + @wrap_exceptions + def num_threads(self): + return self._get_pidtaskinfo()[pidtaskinfo_map['numthreads']] + + @wrap_exceptions + def open_files(self): + if self.pid == 0: + return [] + files = [] + rawlist = cext.proc_open_files(self.pid) + for path, fd in rawlist: + if isfile_strict(path): + ntuple = _common.popenfile(path, fd) + files.append(ntuple) + return files + + @wrap_exceptions + def net_connections(self, kind='inet'): + families, types = conn_tmap[kind] + rawlist = cext.proc_net_connections(self.pid, families, types) + ret = [] + for item in rawlist: + fd, fam, type, laddr, raddr, status = item + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES + ) + ret.append(nt) + return ret + + @wrap_exceptions + def num_fds(self): + if self.pid == 0: + return 0 + return cext.proc_num_fds(self.pid) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout, self._name) + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def status(self): + code = self._get_kinfo_proc()[kinfo_proc_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = _common.pthread(thread_id, utime, stime) + retlist.append(ntuple) + return retlist diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psposix.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psposix.py new file mode 100644 index 0000000000000000000000000000000000000000..83f1acdbe31d4ac0e4fb237bc865226602b1131f --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_psposix.py @@ -0,0 +1,206 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Routines common to all posix systems.""" + +import enum +import glob +import os +import signal +import time + +from ._common import MACOS +from ._common import TimeoutExpired +from ._common import memoize +from ._common import sdiskusage +from ._common import usage_percent + +if MACOS: + from . import _psutil_osx + + +__all__ = ['pid_exists', 'wait_pid', 'disk_usage', 'get_terminal_map'] + + +def pid_exists(pid): + """Check whether pid exists in the current process table.""" + if pid == 0: + # According to "man 2 kill" PID 0 has a special meaning: + # it refers to <> so we don't want to go any further. + # If we get here it means this UNIX platform *does* have + # a process with id 0. + return True + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + # EPERM clearly means there's a process to deny access to + return True + # According to "man 2 kill" possible error values are + # (EINVAL, EPERM, ESRCH) + else: + return True + + +Negsignal = enum.IntEnum( + 'Negsignal', {x.name: -x.value for x in signal.Signals} +) + + +def negsig_to_enum(num): + """Convert a negative signal value to an enum.""" + try: + return Negsignal(num) + except ValueError: + return num + + +def wait_pid( + pid, + timeout=None, + proc_name=None, + _waitpid=os.waitpid, + _timer=getattr(time, 'monotonic', time.time), # noqa: B008 + _min=min, + _sleep=time.sleep, + _pid_exists=pid_exists, +): + """Wait for a process PID to terminate. + + If the process terminated normally by calling exit(3) or _exit(2), + or by returning from main(), the return value is the positive integer + passed to *exit(). + + If it was terminated by a signal it returns the negated value of the + signal which caused the termination (e.g. -SIGTERM). + + If PID is not a children of os.getpid() (current process) just + wait until the process disappears and return None. + + If PID does not exist at all return None immediately. + + If *timeout* != None and process is still alive raise TimeoutExpired. + timeout=0 is also possible (either return immediately or raise). + """ + if pid <= 0: + # see "man waitpid" + msg = "can't wait for PID 0" + raise ValueError(msg) + interval = 0.0001 + flags = 0 + if timeout is not None: + flags |= os.WNOHANG + stop_at = _timer() + timeout + + def sleep(interval): + # Sleep for some time and return a new increased interval. + if timeout is not None: + if _timer() >= stop_at: + raise TimeoutExpired(timeout, pid=pid, name=proc_name) + _sleep(interval) + return _min(interval * 2, 0.04) + + # See: https://linux.die.net/man/2/waitpid + while True: + try: + retpid, status = os.waitpid(pid, flags) + except InterruptedError: + interval = sleep(interval) + except ChildProcessError: + # This has two meanings: + # - PID is not a child of os.getpid() in which case + # we keep polling until it's gone + # - PID never existed in the first place + # In both cases we'll eventually return None as we + # can't determine its exit status code. + while _pid_exists(pid): + interval = sleep(interval) + return None + else: + if retpid == 0: + # WNOHANG flag was used and PID is still running. + interval = sleep(interval) + continue + + if os.WIFEXITED(status): + # Process terminated normally by calling exit(3) or _exit(2), + # or by returning from main(). The return value is the + # positive integer passed to *exit(). + return os.WEXITSTATUS(status) + elif os.WIFSIGNALED(status): + # Process exited due to a signal. Return the negative value + # of that signal. + return negsig_to_enum(-os.WTERMSIG(status)) + # elif os.WIFSTOPPED(status): + # # Process was stopped via SIGSTOP or is being traced, and + # # waitpid() was called with WUNTRACED flag. PID is still + # # alive. From now on waitpid() will keep returning (0, 0) + # # until the process state doesn't change. + # # It may make sense to catch/enable this since stopped PIDs + # # ignore SIGTERM. + # interval = sleep(interval) + # continue + # elif os.WIFCONTINUED(status): + # # Process was resumed via SIGCONT and waitpid() was called + # # with WCONTINUED flag. + # interval = sleep(interval) + # continue + else: + # Should never happen. + msg = f"unknown process exit status {status!r}" + raise ValueError(msg) + + +def disk_usage(path): + """Return disk usage associated with path. + Note: UNIX usually reserves 5% disk space which is not accessible + by user. In this function "total" and "used" values reflect the + total and used disk space whereas "free" and "percent" represent + the "free" and "used percent" user disk space. + """ + st = os.statvfs(path) + # Total space which is only available to root (unless changed + # at system level). + total = st.f_blocks * st.f_frsize + # Remaining free space usable by root. + avail_to_root = st.f_bfree * st.f_frsize + # Remaining free space usable by user. + avail_to_user = st.f_bavail * st.f_frsize + # Total space being used in general. + used = total - avail_to_root + if MACOS: + # see: https://github.com/giampaolo/psutil/pull/2152 + used = _psutil_osx.disk_usage_used(path, used) + # Total space which is available to user (same as 'total' but + # for the user). + total_user = used + avail_to_user + # User usage percent compared to the total amount of space + # the user can use. This number would be higher if compared + # to root's because the user has less space (usually -5%). + usage_percent_user = usage_percent(used, total_user, round_=1) + + # NB: the percentage is -5% than what shown by df due to + # reserved blocks that we are currently not considering: + # https://github.com/giampaolo/psutil/issues/829#issuecomment-223750462 + return sdiskusage( + total=total, used=used, free=avail_to_user, percent=usage_percent_user + ) + + +@memoize +def get_terminal_map(): + """Get a map of device-id -> path as a dict. + Used by Process.terminal(). + """ + ret = {} + ls = glob.glob('/dev/tty*') + glob.glob('/dev/pts/*') + for name in ls: + assert name not in ret, name + try: + ret[os.stat(name).st_rdev] = name + except FileNotFoundError: + pass + return ret diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pssunos.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pssunos.py new file mode 100644 index 0000000000000000000000000000000000000000..e3268932f6423dc73ece883a97345f1defdd6635 --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pssunos.py @@ -0,0 +1,732 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Sun OS Solaris platform implementation.""" + +import errno +import functools +import os +import socket +import subprocess +import sys +from collections import namedtuple +from socket import AF_INET + +from . import _common +from . import _psposix +from . import _psutil_sunos as cext +from ._common import AF_INET6 +from ._common import ENCODING +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import debug +from ._common import get_procfs_path +from ._common import isfile_strict +from ._common import memoize_when_activated +from ._common import sockfam_to_enum +from ._common import socktype_to_enum +from ._common import usage_percent + +__extra__all__ = ["CONN_IDLE", "CONN_BOUND", "PROCFS_PATH"] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +PAGE_SIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK +IS_64_BIT = sys.maxsize > 2**32 + +CONN_IDLE = "IDLE" +CONN_BOUND = "BOUND" + +PROC_STATUSES = { + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SRUN: _common.STATUS_RUNNING, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SIDL: _common.STATUS_IDLE, + cext.SONPROC: _common.STATUS_RUNNING, # same as run + cext.SWAIT: _common.STATUS_WAITING, +} + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RCVD: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, + cext.TCPS_IDLE: CONN_IDLE, # sunos specific + cext.TCPS_BOUND: CONN_BOUND, # sunos specific +} + +proc_info_map = dict( + ppid=0, + rss=1, + vms=2, + create_time=3, + nice=4, + num_threads=5, + status=6, + ttynr=7, + uid=8, + euid=9, + gid=10, + egid=11, +) + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# psutil.cpu_times() +scputimes = namedtuple('scputimes', ['user', 'system', 'idle', 'iowait']) +# psutil.cpu_times(percpu=True) +pcputimes = namedtuple( + 'pcputimes', ['user', 'system', 'children_user', 'children_system'] +) +# psutil.virtual_memory() +svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free']) +# psutil.Process.memory_info() +pmem = namedtuple('pmem', ['rss', 'vms']) +pfullmem = pmem +# psutil.Process.memory_maps(grouped=True) +pmmap_grouped = namedtuple( + 'pmmap_grouped', ['path', 'rss', 'anonymous', 'locked'] +) +# psutil.Process.memory_maps(grouped=False) +pmmap_ext = namedtuple( + 'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields) +) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """Report virtual memory metrics.""" + # we could have done this with kstat, but IMHO this is good enough + total = os.sysconf('SC_PHYS_PAGES') * PAGE_SIZE + # note: there's no difference on Solaris + free = avail = os.sysconf('SC_AVPHYS_PAGES') * PAGE_SIZE + used = total - free + percent = usage_percent(used, total, round_=1) + return svmem(total, avail, percent, used, free) + + +def swap_memory(): + """Report swap memory metrics.""" + sin, sout = cext.swap_mem() + # XXX + # we are supposed to get total/free by doing so: + # http://cvs.opensolaris.org/source/xref/onnv/onnv-gate/ + # usr/src/cmd/swap/swap.c + # ...nevertheless I can't manage to obtain the same numbers as 'swap' + # cmdline utility, so let's parse its output (sigh!) + p = subprocess.Popen( + [ + '/usr/bin/env', + f"PATH=/usr/sbin:/sbin:{os.environ['PATH']}", + 'swap', + '-l', + ], + stdout=subprocess.PIPE, + ) + stdout, _ = p.communicate() + stdout = stdout.decode(sys.stdout.encoding) + if p.returncode != 0: + msg = f"'swap -l' failed (retcode={p.returncode})" + raise RuntimeError(msg) + + lines = stdout.strip().split('\n')[1:] + if not lines: + msg = 'no swap device(s) configured' + raise RuntimeError(msg) + total = free = 0 + for line in lines: + line = line.split() + t, f = line[3:5] + total += int(int(t) * 512) + free += int(int(f) * 512) + used = total - free + percent = usage_percent(used, total, round_=1) + return _common.sswap( + total, used, free, percent, sin * PAGE_SIZE, sout * PAGE_SIZE + ) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system-wide CPU times as a named tuple.""" + ret = cext.per_cpu_times() + return scputimes(*[sum(x) for x in zip(*ret)]) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = cext.per_cpu_times() + return [scputimes(*x) for x in ret] + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # mimic os.cpu_count() behavior + return None + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + ctx_switches, interrupts, syscalls, _traps = cext.cpu_stats() + soft_interrupts = 0 + return _common.scpustats( + ctx_switches, interrupts, soft_interrupts, syscalls + ) + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters +disk_usage = _psposix.disk_usage + + +def disk_partitions(all=False): + """Return system disk partitions.""" + # TODO - the filtering logic should be better checked so that + # it tries to reflect 'df' as much as possible + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + # Differently from, say, Linux, we don't have a list of + # common fs types so the best we can do, AFAIK, is to + # filter by filesystem having a total size > 0. + try: + if not disk_usage(mountpoint).total: + continue + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1674 + debug(f"skipping {mountpoint!r}: {err}") + continue + ntuple = _common.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + Only INET sockets are returned (UNIX are not). + """ + families, types = _common.conn_tmap[kind] + rawlist = cext.net_connections(_pid) + ret = set() + for item in rawlist: + fd, fam, type_, laddr, raddr, status, pid = item + if fam not in families: + continue + if type_ not in types: + continue + # TODO: refactor and use _common.conn_to_ntuple. + if fam in {AF_INET, AF_INET6}: + if laddr: + laddr = _common.addr(*laddr) + if raddr: + raddr = _common.addr(*raddr) + status = TCP_STATUSES[status] + fam = sockfam_to_enum(fam) + type_ = socktype_to_enum(type_) + if _pid == -1: + nt = _common.sconn(fd, fam, type_, laddr, raddr, status, pid) + else: + nt = _common.pconn(fd, fam, type_, laddr, raddr, status) + ret.add(nt) + return list(ret) + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + ret = cext.net_if_stats() + for name, items in ret.items(): + isup, duplex, speed, mtu = items + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + ret[name] = _common.snicstats(isup, duplex, speed, mtu, '') + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + localhost = (':0.0', ':0') + for item in rawlist: + user, tty, hostname, tstamp, user_process, pid = item + # note: the underlying C function includes entries about + # system boot, run level and others. We might want + # to use them in the future. + if not user_process: + continue + if hostname in localhost: + hostname = 'localhost' + nt = _common.suser(user, tty, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + path = get_procfs_path().encode(ENCODING) + return [int(x) for x in os.listdir(path) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix pid.""" + return _psposix.pid_exists(pid) + + +def wrap_exceptions(fun): + """Call callable into a try/except clause and translate ENOENT, + EACCES and EPERM in NoSuchProcess or AccessDenied exceptions. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except (FileNotFoundError, ProcessLookupError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if not pid_exists(pid): + raise NoSuchProcess(pid, name) from err + raise ZombieProcess(pid, name, ppid) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except OSError as err: + if pid == 0: + if 0 in pids(): + raise AccessDenied(pid, name) from err + raise + raise + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "_procfs_path", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._procfs_path = get_procfs_path() + + def _assert_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + os.stat(f"{self._procfs_path}/{self.pid}") + + def oneshot_enter(self): + self._proc_name_and_args.cache_activate(self) + self._proc_basic_info.cache_activate(self) + self._proc_cred.cache_activate(self) + + def oneshot_exit(self): + self._proc_name_and_args.cache_deactivate(self) + self._proc_basic_info.cache_deactivate(self) + self._proc_cred.cache_deactivate(self) + + @wrap_exceptions + @memoize_when_activated + def _proc_name_and_args(self): + return cext.proc_name_and_args(self.pid, self._procfs_path) + + @wrap_exceptions + @memoize_when_activated + def _proc_basic_info(self): + if self.pid == 0 and not os.path.exists( + f"{self._procfs_path}/{self.pid}/psinfo" + ): + raise AccessDenied(self.pid) + ret = cext.proc_basic_info(self.pid, self._procfs_path) + assert len(ret) == len(proc_info_map) + return ret + + @wrap_exceptions + @memoize_when_activated + def _proc_cred(self): + return cext.proc_cred(self.pid, self._procfs_path) + + @wrap_exceptions + def name(self): + # note: max len == 15 + return self._proc_name_and_args()[0] + + @wrap_exceptions + def exe(self): + try: + return os.readlink(f"{self._procfs_path}/{self.pid}/path/a.out") + except OSError: + pass # continue and guess the exe name from the cmdline + # Will be guessed later from cmdline but we want to explicitly + # invoke cmdline here in order to get an AccessDenied + # exception if the user has not enough privileges. + self.cmdline() + return "" + + @wrap_exceptions + def cmdline(self): + return self._proc_name_and_args()[1] + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid, self._procfs_path) + + @wrap_exceptions + def create_time(self): + return self._proc_basic_info()[proc_info_map['create_time']] + + @wrap_exceptions + def num_threads(self): + return self._proc_basic_info()[proc_info_map['num_threads']] + + @wrap_exceptions + def nice_get(self): + # Note #1: getpriority(3) doesn't work for realtime processes. + # Psinfo is what ps uses, see: + # https://github.com/giampaolo/psutil/issues/1194 + return self._proc_basic_info()[proc_info_map['nice']] + + @wrap_exceptions + def nice_set(self, value): + if self.pid in {2, 3}: + # Special case PIDs: internally setpriority(3) return ESRCH + # (no such process), no matter what. + # The process actually exists though, as it has a name, + # creation time, etc. + raise AccessDenied(self.pid, self._name) + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ppid(self): + self._ppid = self._proc_basic_info()[proc_info_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + try: + real, effective, saved, _, _, _ = self._proc_cred() + except AccessDenied: + real = self._proc_basic_info()[proc_info_map['uid']] + effective = self._proc_basic_info()[proc_info_map['euid']] + saved = None + return _common.puids(real, effective, saved) + + @wrap_exceptions + def gids(self): + try: + _, _, _, real, effective, saved = self._proc_cred() + except AccessDenied: + real = self._proc_basic_info()[proc_info_map['gid']] + effective = self._proc_basic_info()[proc_info_map['egid']] + saved = None + return _common.puids(real, effective, saved) + + @wrap_exceptions + def cpu_times(self): + try: + times = cext.proc_cpu_times(self.pid, self._procfs_path) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + times = (0.0, 0.0, 0.0, 0.0) + else: + raise + return _common.pcputimes(*times) + + @wrap_exceptions + def cpu_num(self): + return cext.proc_cpu_num(self.pid, self._procfs_path) + + @wrap_exceptions + def terminal(self): + procfs_path = self._procfs_path + hit_enoent = False + tty = wrap_exceptions(self._proc_basic_info()[proc_info_map['ttynr']]) + if tty != cext.PRNODEV: + for x in (0, 1, 2, 255): + try: + return os.readlink(f"{procfs_path}/{self.pid}/path/{x}") + except FileNotFoundError: + hit_enoent = True + continue + if hit_enoent: + self._assert_alive() + + @wrap_exceptions + def cwd(self): + # /proc/PID/path/cwd may not be resolved by readlink() even if + # it exists (ls shows it). If that's the case and the process + # is still alive return None (we can return None also on BSD). + # Reference: https://groups.google.com/g/comp.unix.solaris/c/tcqvhTNFCAs + procfs_path = self._procfs_path + try: + return os.readlink(f"{procfs_path}/{self.pid}/path/cwd") + except FileNotFoundError: + os.stat(f"{procfs_path}/{self.pid}") # raise NSP or AD + return "" + + @wrap_exceptions + def memory_info(self): + ret = self._proc_basic_info() + rss = ret[proc_info_map['rss']] * 1024 + vms = ret[proc_info_map['vms']] * 1024 + return pmem(rss, vms) + + memory_full_info = memory_info + + @wrap_exceptions + def status(self): + code = self._proc_basic_info()[proc_info_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def threads(self): + procfs_path = self._procfs_path + ret = [] + tids = os.listdir(f"{procfs_path}/{self.pid}/lwp") + hit_enoent = False + for tid in tids: + tid = int(tid) + try: + utime, stime = cext.query_process_thread( + self.pid, tid, procfs_path + ) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + continue + # ENOENT == thread gone in meantime + if err.errno == errno.ENOENT: + hit_enoent = True + continue + raise + else: + nt = _common.pthread(tid, utime, stime) + ret.append(nt) + if hit_enoent: + self._assert_alive() + return ret + + @wrap_exceptions + def open_files(self): + retlist = [] + hit_enoent = False + procfs_path = self._procfs_path + pathdir = f"{procfs_path}/{self.pid}/path" + for fd in os.listdir(f"{procfs_path}/{self.pid}/fd"): + path = os.path.join(pathdir, fd) + if os.path.islink(path): + try: + file = os.readlink(path) + except FileNotFoundError: + hit_enoent = True + continue + else: + if isfile_strict(file): + retlist.append(_common.popenfile(file, int(fd))) + if hit_enoent: + self._assert_alive() + return retlist + + def _get_unix_sockets(self, pid): + """Get UNIX sockets used by process by parsing 'pfiles' output.""" + # TODO: rewrite this in C (...but the damn netstat source code + # does not include this part! Argh!!) + cmd = ["pfiles", str(pid)] + p = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if p.returncode != 0: + if 'permission denied' in stderr.lower(): + raise AccessDenied(self.pid, self._name) + if 'no such process' in stderr.lower(): + raise NoSuchProcess(self.pid, self._name) + msg = f"{cmd!r} command error\n{stderr}" + raise RuntimeError(msg) + + lines = stdout.split('\n')[2:] + for i, line in enumerate(lines): + line = line.lstrip() + if line.startswith('sockname: AF_UNIX'): + path = line.split(' ', 2)[2] + type = lines[i - 2].strip() + if type == 'SOCK_STREAM': + type = socket.SOCK_STREAM + elif type == 'SOCK_DGRAM': + type = socket.SOCK_DGRAM + else: + type = -1 + yield (-1, socket.AF_UNIX, type, path, "", _common.CONN_NONE) + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = net_connections(kind, _pid=self.pid) + # The underlying C implementation retrieves all OS connections + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not ret: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + + # UNIX sockets + if kind in {'all', 'unix'}: + ret.extend([ + _common.pconn(*conn) + for conn in self._get_unix_sockets(self.pid) + ]) + return ret + + nt_mmap_grouped = namedtuple('mmap', 'path rss anon locked') + nt_mmap_ext = namedtuple('mmap', 'addr perms path rss anon locked') + + @wrap_exceptions + def memory_maps(self): + def toaddr(start, end): + return "{}-{}".format( + hex(start)[2:].strip('L'), hex(end)[2:].strip('L') + ) + + procfs_path = self._procfs_path + retlist = [] + try: + rawlist = cext.proc_memory_maps(self.pid, procfs_path) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + return [] + else: + raise + hit_enoent = False + for item in rawlist: + addr, addrsize, perm, name, rss, anon, locked = item + addr = toaddr(addr, addrsize) + if not name.startswith('['): + try: + name = os.readlink(f"{procfs_path}/{self.pid}/path/{name}") + except OSError as err: + if err.errno == errno.ENOENT: + # sometimes the link may not be resolved by + # readlink() even if it exists (ls shows it). + # If that's the case we just return the + # unresolved link path. + # This seems an inconsistency with /proc similar + # to: http://goo.gl/55XgO + name = f"{procfs_path}/{self.pid}/path/{name}" + hit_enoent = True + else: + raise + retlist.append((addr, perm, name, rss, anon, locked)) + if hit_enoent: + self._assert_alive() + return retlist + + @wrap_exceptions + def num_fds(self): + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def num_ctx_switches(self): + return _common.pctxsw( + *cext.proc_num_ctx_switches(self.pid, self._procfs_path) + ) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout, self._name) diff --git a/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pswindows.py b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pswindows.py new file mode 100644 index 0000000000000000000000000000000000000000..7d343f26bbb41c205cfb48cbcaf05531d1fef1ef --- /dev/null +++ b/lib/python3.12/site-packages/ray/thirdparty_files/psutil/_pswindows.py @@ -0,0 +1,1122 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Windows platform implementation.""" + +import contextlib +import enum +import functools +import os +import signal +import sys +import threading +import time +from collections import namedtuple + +from . import _common +from ._common import ENCODING +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import TimeoutExpired +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import isfile_strict +from ._common import memoize +from ._common import memoize_when_activated +from ._common import parse_environ_block +from ._common import usage_percent +from ._psutil_windows import ABOVE_NORMAL_PRIORITY_CLASS +from ._psutil_windows import BELOW_NORMAL_PRIORITY_CLASS +from ._psutil_windows import HIGH_PRIORITY_CLASS +from ._psutil_windows import IDLE_PRIORITY_CLASS +from ._psutil_windows import NORMAL_PRIORITY_CLASS +from ._psutil_windows import REALTIME_PRIORITY_CLASS + +try: + from . import _psutil_windows as cext +except ImportError as err: + if ( + str(err).lower().startswith("dll load failed") + and sys.getwindowsversion()[0] < 6 + ): + # We may get here if: + # 1) we are on an old Windows version + # 2) psutil was installed via pip + wheel + # See: https://github.com/giampaolo/psutil/issues/811 + msg = "this Windows version is too old (< Windows Vista); " + msg += "psutil 3.4.2 is the latest version which supports Windows " + msg += "2000, XP and 2003 server" + raise RuntimeError(msg) from err + else: + raise + + +# process priority constants, import from __init__.py: +# http://msdn.microsoft.com/en-us/library/ms686219(v=vs.85).aspx +# fmt: off +__extra__all__ = [ + "win_service_iter", "win_service_get", + # Process priority + "ABOVE_NORMAL_PRIORITY_CLASS", "BELOW_NORMAL_PRIORITY_CLASS", + "HIGH_PRIORITY_CLASS", "IDLE_PRIORITY_CLASS", "NORMAL_PRIORITY_CLASS", + "REALTIME_PRIORITY_CLASS", + # IO priority + "IOPRIO_VERYLOW", "IOPRIO_LOW", "IOPRIO_NORMAL", "IOPRIO_HIGH", + # others + "CONN_DELETE_TCB", "AF_LINK", +] +# fmt: on + + +# ===================================================================== +# --- globals +# ===================================================================== + +CONN_DELETE_TCB = "DELETE_TCB" +ERROR_PARTIAL_COPY = 299 +PYPY = '__pypy__' in sys.builtin_module_names + +AddressFamily = enum.IntEnum('AddressFamily', {'AF_LINK': -1}) +AF_LINK = AddressFamily.AF_LINK + +TCP_STATUSES = { + cext.MIB_TCP_STATE_ESTAB: _common.CONN_ESTABLISHED, + cext.MIB_TCP_STATE_SYN_SENT: _common.CONN_SYN_SENT, + cext.MIB_TCP_STATE_SYN_RCVD: _common.CONN_SYN_RECV, + cext.MIB_TCP_STATE_FIN_WAIT1: _common.CONN_FIN_WAIT1, + cext.MIB_TCP_STATE_FIN_WAIT2: _common.CONN_FIN_WAIT2, + cext.MIB_TCP_STATE_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.MIB_TCP_STATE_CLOSED: _common.CONN_CLOSE, + cext.MIB_TCP_STATE_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.MIB_TCP_STATE_LAST_ACK: _common.CONN_LAST_ACK, + cext.MIB_TCP_STATE_LISTEN: _common.CONN_LISTEN, + cext.MIB_TCP_STATE_CLOSING: _common.CONN_CLOSING, + cext.MIB_TCP_STATE_DELETE_TCB: CONN_DELETE_TCB, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + + +class Priority(enum.IntEnum): + ABOVE_NORMAL_PRIORITY_CLASS = ABOVE_NORMAL_PRIORITY_CLASS + BELOW_NORMAL_PRIORITY_CLASS = BELOW_NORMAL_PRIORITY_CLASS + HIGH_PRIORITY_CLASS = HIGH_PRIORITY_CLASS + IDLE_PRIORITY_CLASS = IDLE_PRIORITY_CLASS + NORMAL_PRIORITY_CLASS = NORMAL_PRIORITY_CLASS + REALTIME_PRIORITY_CLASS = REALTIME_PRIORITY_CLASS + + +globals().update(Priority.__members__) + + +class IOPriority(enum.IntEnum): + IOPRIO_VERYLOW = 0 + IOPRIO_LOW = 1 + IOPRIO_NORMAL = 2 + IOPRIO_HIGH = 3 + + +globals().update(IOPriority.__members__) + +pinfo_map = dict( + num_handles=0, + ctx_switches=1, + user_time=2, + kernel_time=3, + create_time=4, + num_threads=5, + io_rcount=6, + io_wcount=7, + io_rbytes=8, + io_wbytes=9, + io_count_others=10, + io_bytes_others=11, + num_page_faults=12, + peak_wset=13, + wset=14, + peak_paged_pool=15, + paged_pool=16, + peak_non_paged_pool=17, + non_paged_pool=18, + pagefile=19, + peak_pagefile=20, + mem_private=21, +) + + +# ===================================================================== +# --- named tuples +# ===================================================================== + + +# fmt: off +# psutil.cpu_times() +scputimes = namedtuple('scputimes', + ['user', 'system', 'idle', 'interrupt', 'dpc']) +# psutil.virtual_memory() +svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free']) +# psutil.Process.memory_info() +pmem = namedtuple( + 'pmem', ['rss', 'vms', + 'num_page_faults', 'peak_wset', 'wset', 'peak_paged_pool', + 'paged_pool', 'peak_nonpaged_pool', 'nonpaged_pool', + 'pagefile', 'peak_pagefile', 'private']) +# psutil.Process.memory_full_info() +pfullmem = namedtuple('pfullmem', pmem._fields + ('uss', )) +# psutil.Process.memory_maps(grouped=True) +pmmap_grouped = namedtuple('pmmap_grouped', ['path', 'rss']) +# psutil.Process.memory_maps(grouped=False) +pmmap_ext = namedtuple( + 'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields)) +# psutil.Process.io_counters() +pio = namedtuple('pio', ['read_count', 'write_count', + 'read_bytes', 'write_bytes', + 'other_count', 'other_bytes']) +# fmt: on + + +# ===================================================================== +# --- utils +# ===================================================================== + + +@functools.lru_cache(maxsize=512) +def convert_dos_path(s): + r"""Convert paths using native DOS format like: + "\Device\HarddiskVolume1\Windows\systemew\file.txt" or + "\??\C:\Windows\systemew\file.txt" + into: + "C:\Windows\systemew\file.txt". + """ + if s.startswith('\\\\'): + return s + rawdrive = '\\'.join(s.split('\\')[:3]) + if rawdrive in {"\\??\\UNC", "\\Device\\Mup"}: + rawdrive = '\\'.join(s.split('\\')[:5]) + driveletter = '\\\\' + '\\'.join(s.split('\\')[3:5]) + elif rawdrive.startswith('\\??\\'): + driveletter = s.split('\\')[2] + else: + driveletter = cext.QueryDosDevice(rawdrive) + remainder = s[len(rawdrive) :] + return os.path.join(driveletter, remainder) + + +@memoize +def getpagesize(): + return cext.getpagesize() + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """System virtual memory as a namedtuple.""" + mem = cext.virtual_mem() + totphys, availphys, _totsys, _availsys = mem + total = totphys + avail = availphys + free = availphys + used = total - avail + percent = usage_percent((total - avail), total, round_=1) + return svmem(total, avail, percent, used, free) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + mem = cext.virtual_mem() + + total_phys = mem[0] + total_system = mem[2] + + # system memory (commit total/limit) is the sum of physical and swap + # thus physical memory values need to be subtracted to get swap values + total = total_system - total_phys + # commit total is incremented immediately (decrementing free_system) + # while the corresponding free physical value is not decremented until + # pages are accessed, so we can't use free system memory for swap. + # instead, we calculate page file usage based on performance counter + if total > 0: + percentswap = cext.swap_percent() + used = int(0.01 * percentswap * total) + else: + percentswap = 0.0 + used = 0 + + free = total - used + percent = round(percentswap, 1) + return _common.sswap(total, used, free, percent, 0, 0) + + +# ===================================================================== +# --- disk +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters + + +def disk_usage(path): + """Return disk usage associated with path.""" + if isinstance(path, bytes): + # XXX: do we want to use "strict"? Probably yes, in order + # to fail immediately. After all we are accepting input here... + path = path.decode(ENCODING, errors="strict") + total, used, free = cext.disk_usage(path) + percent = usage_percent(used, total, round_=1) + return _common.sdiskusage(total, used, free, percent) + + +def disk_partitions(all): + """Return disk partitions.""" + rawlist = cext.disk_partitions(all) + return [_common.sdiskpart(*x) for x in rawlist] + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system CPU times as a named tuple.""" + user, system, idle = cext.cpu_times() + # Internally, GetSystemTimes() is used, and it doesn't return + # interrupt and dpc times. cext.per_cpu_times() does, so we + # rely on it to get those only. + percpu_summed = scputimes(*[sum(n) for n in zip(*cext.per_cpu_times())]) + return scputimes( + user, system, idle, percpu_summed.interrupt, percpu_summed.dpc + ) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = [] + for user, system, idle, interrupt, dpc in cext.per_cpu_times(): + item = scputimes(user, system, idle, interrupt, dpc) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + """Return CPU statistics.""" + ctx_switches, interrupts, _dpcs, syscalls = cext.cpu_stats() + soft_interrupts = 0 + return _common.scpustats( + ctx_switches, interrupts, soft_interrupts, syscalls + ) + + +def cpu_freq(): + """Return CPU frequency. + On Windows per-cpu frequency is not supported. + """ + curr, max_ = cext.cpu_freq() + min_ = 0.0 + return [_common.scpufreq(float(curr), min_, float(max_))] + + +_loadavg_initialized = False +_lock = threading.Lock() + + +def _getloadavg_impl(): + # Drop to 2 decimal points which is what Linux does + raw_loads = cext.getloadavg() + return tuple(round(load, 2) for load in raw_loads) + + +def getloadavg(): + """Return the number of processes in the system run queue averaged + over the last 1, 5, and 15 minutes respectively as a tuple. + """ + global _loadavg_initialized + + if _loadavg_initialized: + return _getloadavg_impl() + + with _lock: + if not _loadavg_initialized: + cext.init_loadavg_counter() + _loadavg_initialized = True + + return _getloadavg_impl() + + +# ===================================================================== +# --- network +# ===================================================================== + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + """ + families, types = conn_tmap[kind] + rawlist = cext.net_connections(_pid, families, types) + ret = set() + for item in rawlist: + fd, fam, type, laddr, raddr, status, pid = item + nt = conn_to_ntuple( + fd, + fam, + type, + laddr, + raddr, + status, + TCP_STATUSES, + pid=pid if _pid == -1 else None, + ) + ret.add(nt) + return list(ret) + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + ret = {} + rawdict = cext.net_if_stats() + for name, items in rawdict.items(): + isup, duplex, speed, mtu = items + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + ret[name] = _common.snicstats(isup, duplex, speed, mtu, '') + return ret + + +def net_io_counters(): + """Return network I/O statistics for every network interface + installed on the system as a dict of raw tuples. + """ + return cext.net_io_counters() + + +def net_if_addrs(): + """Return the addresses associated to each NIC.""" + return cext.net_if_addrs() + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_battery(): + """Return battery information.""" + # For constants meaning see: + # https://msdn.microsoft.com/en-us/library/windows/desktop/ + # aa373232(v=vs.85).aspx + acline_status, flags, percent, secsleft = cext.sensors_battery() + power_plugged = acline_status == 1 + no_battery = bool(flags & 128) + charging = bool(flags & 8) + + if no_battery: + return None + if power_plugged or charging: + secsleft = _common.POWER_TIME_UNLIMITED + elif secsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + + return _common.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +_last_btime = 0 + + +def boot_time(): + """The system boot time expressed in seconds since the epoch. This + also includes the time spent during hybernate / suspend. + """ + # This dirty hack is to adjust the precision of the returned + # value which may have a 1 second fluctuation, see: + # https://github.com/giampaolo/psutil/issues/1007 + global _last_btime + ret = time.time() - cext.uptime() + if abs(ret - _last_btime) <= 1: + return _last_btime + else: + _last_btime = ret + return ret + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, hostname, tstamp = item + nt = _common.suser(user, None, hostname, tstamp, None) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- Windows services +# ===================================================================== + + +def win_service_iter(): + """Yields a list of WindowsService instances.""" + for name, display_name in cext.winservice_enumerate(): + yield WindowsService(name, display_name) + + +def win_service_get(name): + """Open a Windows service and return it as a WindowsService instance.""" + service = WindowsService(name, None) + service._display_name = service._query_config()['display_name'] + return service + + +class WindowsService: # noqa: PLW1641 + """Represents an installed Windows service.""" + + def __init__(self, name, display_name): + self._name = name + self._display_name = display_name + + def __str__(self): + details = f"(name={self._name!r}, display_name={self._display_name!r})" + return f"{self.__class__.__name__}{details}" + + def __repr__(self): + return f"<{self.__str__()} at {id(self)}>" + + def __eq__(self, other): + # Test for equality with another WindosService object based + # on name. + if not isinstance(other, WindowsService): + return NotImplemented + return self._name == other._name + + def __ne__(self, other): + return not self == other + + def _query_config(self): + with self._wrap_exceptions(): + display_name, binpath, username, start_type = ( + cext.winservice_query_config(self._name) + ) + # XXX - update _self.display_name? + return dict( + display_name=display_name, + binpath=binpath, + username=username, + start_type=start_type, + ) + + def _query_status(self): + with self._wrap_exceptions(): + status, pid = cext.winservice_query_status(self._name) + if pid == 0: + pid = None + return dict(status=status, pid=pid) + + @contextlib.contextmanager + def _wrap_exceptions(self): + """Ctx manager which translates bare OSError and WindowsError + exceptions into NoSuchProcess and AccessDenied. + """ + try: + yield + except OSError as err: + name = self._name + if is_permission_err(err): + msg = ( + f"service {name!r} is not querable (not enough privileges)" + ) + raise AccessDenied(pid=None, name=name, msg=msg) from err + elif err.winerror in { + cext.ERROR_INVALID_NAME, + cext.ERROR_SERVICE_DOES_NOT_EXIST, + }: + msg = f"service {name!r} does not exist" + raise NoSuchProcess(pid=None, name=name, msg=msg) from err + else: + raise + + # config query + + def name(self): + """The service name. This string is how a service is referenced + and can be passed to win_service_get() to get a new + WindowsService instance. + """ + return self._name + + def display_name(self): + """The service display name. The value is cached when this class + is instantiated. + """ + return self._display_name + + def binpath(self): + """The fully qualified path to the service binary/exe file as + a string, including command line arguments. + """ + return self._query_config()['binpath'] + + def username(self): + """The name of the user that owns this service.""" + return self._query_config()['username'] + + def start_type(self): + """A string which can either be "automatic", "manual" or + "disabled". + """ + return self._query_config()['start_type'] + + # status query + + def pid(self): + """The process PID, if any, else None. This can be passed + to Process class to control the service's process. + """ + return self._query_status()['pid'] + + def status(self): + """Service status as a string.""" + return self._query_status()['status'] + + def description(self): + """Service long description.""" + return cext.winservice_query_descr(self.name()) + + # utils + + def as_dict(self): + """Utility method retrieving all the information above as a + dictionary. + """ + d = self._query_config() + d.update(self._query_status()) + d['name'] = self.name() + d['display_name'] = self.display_name() + d['description'] = self.description() + return d + + # actions + # XXX: the necessary C bindings for start() and stop() are + # implemented but for now I prefer not to expose them. + # I may change my mind in the future. Reasons: + # - they require Administrator privileges + # - can't implement a timeout for stop() (unless by using a thread, + # which sucks) + # - would require adding ServiceAlreadyStarted and + # ServiceAlreadyStopped exceptions, adding two new APIs. + # - we might also want to have modify(), which would basically mean + # rewriting win32serviceutil.ChangeServiceConfig, which involves a + # lot of stuff (and API constants which would pollute the API), see: + # http://pyxr.sourceforge.net/PyXR/c/python24/lib/site-packages/ + # win32/lib/win32serviceutil.py.html#0175 + # - psutil is typically about "read only" monitoring stuff; + # win_service_* APIs should only be used to retrieve a service and + # check whether it's running + + # def start(self, timeout=None): + # with self._wrap_exceptions(): + # cext.winservice_start(self.name()) + # if timeout: + # giveup_at = time.time() + timeout + # while True: + # if self.status() == "running": + # return + # else: + # if time.time() > giveup_at: + # raise TimeoutExpired(timeout) + # else: + # time.sleep(.1) + + # def stop(self): + # # Note: timeout is not implemented because it's just not + # # possible, see: + # # http://stackoverflow.com/questions/11973228/ + # with self._wrap_exceptions(): + # return cext.winservice_stop(self.name()) + + +# ===================================================================== +# --- processes +# ===================================================================== + + +pids = cext.pids +pid_exists = cext.pid_exists +ppid_map = cext.ppid_map # used internally by Process.children() + + +def is_permission_err(exc): + """Return True if this is a permission error.""" + assert isinstance(exc, OSError), exc + return isinstance(exc, PermissionError) or exc.winerror in { + cext.ERROR_ACCESS_DENIED, + cext.ERROR_PRIVILEGE_NOT_HELD, + } + + +def convert_oserror(exc, pid=None, name=None): + """Convert OSError into NoSuchProcess or AccessDenied.""" + assert isinstance(exc, OSError), exc + if is_permission_err(exc): + return AccessDenied(pid=pid, name=name) + if isinstance(exc, ProcessLookupError): + return NoSuchProcess(pid=pid, name=name) + raise exc + + +def wrap_exceptions(fun): + """Decorator which converts OSError into NoSuchProcess or AccessDenied.""" + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + try: + return fun(self, *args, **kwargs) + except OSError as err: + raise convert_oserror(err, pid=self.pid, name=self._name) from err + + return wrapper + + +def retry_error_partial_copy(fun): + """Workaround for https://github.com/giampaolo/psutil/issues/875. + See: https://stackoverflow.com/questions/4457745#4457745. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + delay = 0.0001 + times = 33 + for _ in range(times): # retries for roughly 1 second + try: + return fun(self, *args, **kwargs) + except OSError as _: + err = _ + if err.winerror == ERROR_PARTIAL_COPY: + time.sleep(delay) + delay = min(delay * 2, 0.04) + continue + raise + msg = ( + f"{fun} retried {times} times, converted to AccessDenied as it's " + f"still returning {err}" + ) + raise AccessDenied(pid=self.pid, name=self._name, msg=msg) + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + # --- oneshot() stuff + + def oneshot_enter(self): + self._proc_info.cache_activate(self) + self.exe.cache_activate(self) + + def oneshot_exit(self): + self._proc_info.cache_deactivate(self) + self.exe.cache_deactivate(self) + + @memoize_when_activated + def _proc_info(self): + """Return multiple information about this process as a + raw tuple. + """ + ret = cext.proc_info(self.pid) + assert len(ret) == len(pinfo_map) + return ret + + def name(self): + """Return process name, which on Windows is always the final + part of the executable. + """ + # This is how PIDs 0 and 4 are always represented in taskmgr + # and process-hacker. + if self.pid == 0: + return "System Idle Process" + if self.pid == 4: + return "System" + return os.path.basename(self.exe()) + + @wrap_exceptions + @memoize_when_activated + def exe(self): + if PYPY: + try: + exe = cext.proc_exe(self.pid) + except OSError as err: + # 24 = ERROR_TOO_MANY_OPEN_FILES. Not sure why this happens + # (perhaps PyPy's JIT delaying garbage collection of files?). + if err.errno == 24: + debug(f"{err!r} translated into AccessDenied") + raise AccessDenied(self.pid, self._name) from err + raise + else: + exe = cext.proc_exe(self.pid) + if exe.startswith('\\'): + return convert_dos_path(exe) + return exe # May be "Registry", "MemCompression", ... + + @wrap_exceptions + @retry_error_partial_copy + def cmdline(self): + if cext.WINVER >= cext.WINDOWS_8_1: + # PEB method detects cmdline changes but requires more + # privileges: https://github.com/giampaolo/psutil/pull/1398 + try: + return cext.proc_cmdline(self.pid, use_peb=True) + except OSError as err: + if is_permission_err(err): + return cext.proc_cmdline(self.pid, use_peb=False) + else: + raise + else: + return cext.proc_cmdline(self.pid, use_peb=True) + + @wrap_exceptions + @retry_error_partial_copy + def environ(self): + s = cext.proc_environ(self.pid) + return parse_environ_block(s) + + def ppid(self): + try: + return ppid_map()[self.pid] + except KeyError: + raise NoSuchProcess(self.pid, self._name) from None + + def _get_raw_meminfo(self): + try: + return cext.proc_memory_info(self.pid) + except OSError as err: + if is_permission_err(err): + # TODO: the C ext can probably be refactored in order + # to get this from cext.proc_info() + debug("attempting memory_info() fallback (slower)") + info = self._proc_info() + return ( + info[pinfo_map['num_page_faults']], + info[pinfo_map['peak_wset']], + info[pinfo_map['wset']], + info[pinfo_map['peak_paged_pool']], + info[pinfo_map['paged_pool']], + info[pinfo_map['peak_non_paged_pool']], + info[pinfo_map['non_paged_pool']], + info[pinfo_map['pagefile']], + info[pinfo_map['peak_pagefile']], + info[pinfo_map['mem_private']], + ) + raise + + @wrap_exceptions + def memory_info(self): + # on Windows RSS == WorkingSetSize and VSM == PagefileUsage. + # Underlying C function returns fields of PROCESS_MEMORY_COUNTERS + # struct. + t = self._get_raw_meminfo() + rss = t[2] # wset + vms = t[7] # pagefile + return pmem(*(rss, vms) + t) + + @wrap_exceptions + def memory_full_info(self): + basic_mem = self.memory_info() + uss = cext.proc_memory_uss(self.pid) + uss *= getpagesize() + return pfullmem(*basic_mem + (uss,)) + + def memory_maps(self): + try: + raw = cext.proc_memory_maps(self.pid) + except OSError as err: + # XXX - can't use wrap_exceptions decorator as we're + # returning a generator; probably needs refactoring. + raise convert_oserror(err, self.pid, self._name) from err + else: + for addr, perm, path, rss in raw: + path = convert_dos_path(path) + addr = hex(addr) + yield (addr, perm, path, rss) + + @wrap_exceptions + def kill(self): + return cext.proc_kill(self.pid) + + @wrap_exceptions + def send_signal(self, sig): + if sig == signal.SIGTERM: + cext.proc_kill(self.pid) + elif sig in {signal.CTRL_C_EVENT, signal.CTRL_BREAK_EVENT}: + os.kill(self.pid, sig) + else: + msg = ( + "only SIGTERM, CTRL_C_EVENT and CTRL_BREAK_EVENT signals " + "are supported on Windows" + ) + raise ValueError(msg) + + @wrap_exceptions + def wait(self, timeout=None): + if timeout is None: + cext_timeout = cext.INFINITE + else: + # WaitForSingleObject() expects time in milliseconds. + cext_timeout = int(timeout * 1000) + + timer = getattr(time, 'monotonic', time.time) + stop_at = timer() + timeout if timeout is not None else None + + try: + # Exit code is supposed to come from GetExitCodeProcess(). + # May also be None if OpenProcess() failed with + # ERROR_INVALID_PARAMETER, meaning PID is already gone. + exit_code = cext.proc_wait(self.pid, cext_timeout) + except cext.TimeoutExpired as err: + # WaitForSingleObject() returned WAIT_TIMEOUT. Just raise. + raise TimeoutExpired(timeout, self.pid, self._name) from err + except cext.TimeoutAbandoned: + # WaitForSingleObject() returned WAIT_ABANDONED, see: + # https://github.com/giampaolo/psutil/issues/1224 + # We'll just rely on the internal polling and return None + # when the PID disappears. Subprocess module does the same + # (return None): + # https://github.com/python/cpython/blob/ + # be50a7b627d0aa37e08fa8e2d5568891f19903ce/ + # Lib/subprocess.py#L1193-L1194 + exit_code = None + + # At this point WaitForSingleObject() returned WAIT_OBJECT_0, + # meaning the process is gone. Stupidly there are cases where + # its PID may still stick around so we do a further internal + # polling. + delay = 0.0001 + while True: + if not pid_exists(self.pid): + return exit_code + if stop_at and timer() >= stop_at: + raise TimeoutExpired(timeout, pid=self.pid, name=self._name) + time.sleep(delay) + delay = min(delay * 2, 0.04) # incremental delay + + @wrap_exceptions + def username(self): + if self.pid in {0, 4}: + return 'NT AUTHORITY\\SYSTEM' + domain, user = cext.proc_username(self.pid) + return f"{domain}\\{user}" + + @wrap_exceptions + def create_time(self, fast_only=False): + # Note: proc_times() not put under oneshot() 'cause create_time() + # is already cached by the main Process class. + try: + _user, _system, created = cext.proc_times(self.pid) + return created + except OSError as err: + if is_permission_err(err): + if fast_only: + raise + debug("attempting create_time() fallback (slower)") + return self._proc_info()[pinfo_map['create_time']] + raise + + @wrap_exceptions + def num_threads(self): + return self._proc_info()[pinfo_map['num_threads']] + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = _common.pthread(thread_id, utime, stime) + retlist.append(ntuple) + return retlist + + @wrap_exceptions + def cpu_times(self): + try: + user, system, _created = cext.proc_times(self.pid) + except OSError as err: + if not is_permission_err(err): + raise + debug("attempting cpu_times() fallback (slower)") + info = self._proc_info() + user = info[pinfo_map['user_time']] + system = info[pinfo_map['kernel_time']] + # Children user/system times are not retrievable (set to 0). + return _common.pcputimes(user, system, 0.0, 0.0) + + @wrap_exceptions + def suspend(self): + cext.proc_suspend_or_resume(self.pid, True) + + @wrap_exceptions + def resume(self): + cext.proc_suspend_or_resume(self.pid, False) + + @wrap_exceptions + @retry_error_partial_copy + def cwd(self): + if self.pid in {0, 4}: + raise AccessDenied(self.pid, self._name) + # return a normalized pathname since the native C function appends + # "\\" at the and of the path + path = cext.proc_cwd(self.pid) + return os.path.normpath(path) + + @wrap_exceptions + def open_files(self): + if self.pid in {0, 4}: + return [] + ret = set() + # Filenames come in in native format like: + # "\Device\HarddiskVolume1\Windows\systemew\file.txt" + # Convert the first part in the corresponding drive letter + # (e.g. "C:\") by using Windows's QueryDosDevice() + raw_file_names = cext.proc_open_files(self.pid) + for file in raw_file_names: + file = convert_dos_path(file) + if isfile_strict(file): + ntuple = _common.popenfile(file, -1) + ret.add(ntuple) + return list(ret) + + @wrap_exceptions + def net_connections(self, kind='inet'): + return net_connections(kind, _pid=self.pid) + + @wrap_exceptions + def nice_get(self): + value = cext.proc_priority_get(self.pid) + value = Priority(value) + return value + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ionice_get(self): + ret = cext.proc_io_priority_get(self.pid) + ret = IOPriority(ret) + return ret + + @wrap_exceptions + def ionice_set(self, ioclass, value): + if value: + msg = "value argument not accepted on Windows" + raise TypeError(msg) + if ioclass not in { + IOPriority.IOPRIO_VERYLOW, + IOPriority.IOPRIO_LOW, + IOPriority.IOPRIO_NORMAL, + IOPriority.IOPRIO_HIGH, + }: + msg = f"{ioclass} is not a valid priority" + raise ValueError(msg) + cext.proc_io_priority_set(self.pid, ioclass) + + @wrap_exceptions + def io_counters(self): + try: + ret = cext.proc_io_counters(self.pid) + except OSError as err: + if not is_permission_err(err): + raise + debug("attempting io_counters() fallback (slower)") + info = self._proc_info() + ret = ( + info[pinfo_map['io_rcount']], + info[pinfo_map['io_wcount']], + info[pinfo_map['io_rbytes']], + info[pinfo_map['io_wbytes']], + info[pinfo_map['io_count_others']], + info[pinfo_map['io_bytes_others']], + ) + return pio(*ret) + + @wrap_exceptions + def status(self): + suspended = cext.proc_is_suspended(self.pid) + if suspended: + return _common.STATUS_STOPPED + else: + return _common.STATUS_RUNNING + + @wrap_exceptions + def cpu_affinity_get(self): + def from_bitmask(x): + return [i for i in range(64) if (1 << i) & x] + + bitmask = cext.proc_cpu_affinity_get(self.pid) + return from_bitmask(bitmask) + + @wrap_exceptions + def cpu_affinity_set(self, value): + def to_bitmask(ls): + if not ls: + msg = f"invalid argument {ls!r}" + raise ValueError(msg) + out = 0 + for b in ls: + out |= 2**b + return out + + # SetProcessAffinityMask() states that ERROR_INVALID_PARAMETER + # is returned for an invalid CPU but this seems not to be true, + # therefore we check CPUs validy beforehand. + allcpus = list(range(len(per_cpu_times()))) + for cpu in value: + if cpu not in allcpus: + if not isinstance(cpu, int): + msg = f"invalid CPU {cpu!r}; an integer is required" + raise TypeError(msg) + msg = f"invalid CPU {cpu!r}" + raise ValueError(msg) + + bitmask = to_bitmask(value) + cext.proc_cpu_affinity_set(self.pid, bitmask) + + @wrap_exceptions + def num_handles(self): + try: + return cext.proc_num_handles(self.pid) + except OSError as err: + if is_permission_err(err): + debug("attempting num_handles() fallback (slower)") + return self._proc_info()[pinfo_map['num_handles']] + raise + + @wrap_exceptions + def num_ctx_switches(self): + ctx_switches = self._proc_info()[pinfo_map['ctx_switches']] + # only voluntary ctx switches are supported + return _common.pctxsw(ctx_switches, 0) diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e03cfad883fa229b62332f262840d73c1412515e Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/callback.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/callback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e939a47db1a0a878dcfb1600ecb34153404cc7eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/callback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01f28cd6dba4f17fa88f5ee8958231e224e0fdd7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/context.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6fd8b0c24a001b4c754c07e7d81757633de4ec6a Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/error.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/error.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e19e7a1c456e76e743816a0fe4658a7a69e13ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/error.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/progress_reporter.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/progress_reporter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0b3f4b978ba01da56075dfb6b13f6c744a82f88 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/progress_reporter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/registry.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/registry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e31d89045534534a28400217cf0d78e10cefb71b Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/registry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/resources.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/resources.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0562ae9c190897f5dffa3ee17edb7940be5e367e Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/resources.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/result.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/result.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2db1ab8fb13fb084cc458fed72f96d7aec5a64b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/result.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/result_grid.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/result_grid.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d61e495bc6e7f75fc87c862f8de92bcffeda268 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/result_grid.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/syncer.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/syncer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a3280ca03f560b17dcdb80892c3c3eaad001b95 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/syncer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/tune.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/tune.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fbc433add3b851feb88c2b2fc857ba3f2acc604f Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/tune.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/tune_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/tune_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..156ed40734c59c674e9a449ae015ac19cba7e182 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/tune_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/__pycache__/tuner.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/__pycache__/tuner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f89e7f18391ffd1d3f84ffbaf9ee1c8cbc76954c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/__pycache__/tuner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/cli/__init__.py b/lib/python3.12/site-packages/ray/tune/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/tune/cli/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76333734f40cc28432694946d6ecae3adb9f03c0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/cli/__pycache__/commands.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/commands.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9da6410af99b6caefb0c9f76676391f0a3ee9381 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/commands.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/cli/__pycache__/scripts.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/scripts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3ae903d378cb46d5d30675282693bbdfb99fc71 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/cli/__pycache__/scripts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/cli/commands.py b/lib/python3.12/site-packages/ray/tune/cli/commands.py new file mode 100644 index 0000000000000000000000000000000000000000..09070124124eaeb9a9e959dfc66e3c79d5fa7734 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/cli/commands.py @@ -0,0 +1,306 @@ +import logging +import operator +import os +import shutil +import subprocess +from datetime import datetime +from pathlib import Path +from typing import List, Optional + +import click +import pandas as pd +from pandas.api.types import is_numeric_dtype, is_string_dtype + +from ray._private.thirdparty.tabulate.tabulate import tabulate +from ray.air.constants import EXPR_RESULT_FILE +from ray.tune import TuneError +from ray.tune.analysis import ExperimentAnalysis +from ray.tune.result import ( + CONFIG_PREFIX, + DEFAULT_EXPERIMENT_INFO_KEYS, + DEFAULT_RESULT_KEYS, +) + +logger = logging.getLogger(__name__) + +EDITOR = os.getenv("EDITOR", "vim") + +TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S (%A)" + +DEFAULT_CLI_KEYS = DEFAULT_EXPERIMENT_INFO_KEYS + DEFAULT_RESULT_KEYS + +DEFAULT_PROJECT_INFO_KEYS = ( + "name", + "total_trials", + "last_updated", +) + +TERM_WIDTH, TERM_HEIGHT = shutil.get_terminal_size(fallback=(100, 100)) + +OPERATORS = { + "<": operator.lt, + "<=": operator.le, + "==": operator.eq, + "!=": operator.ne, + ">=": operator.ge, + ">": operator.gt, +} + + +def _check_tabulate(): + """Checks whether tabulate is installed.""" + if tabulate is None: + raise ImportError("Tabulate not installed. Please run `pip install tabulate`.") + + +def print_format_output(dataframe): + """Prints output of given dataframe to fit into terminal. + + Returns: + table: Final outputted dataframe. + dropped_cols: Columns dropped due to terminal size. + empty_cols: Empty columns (dropped on default). + """ + print_df = pd.DataFrame() + dropped_cols = [] + empty_cols = [] + # column display priority is based on the info_keys passed in + for i, col in enumerate(dataframe): + if dataframe[col].isnull().all(): + # Don't add col to print_df if is fully empty + empty_cols += [col] + continue + + print_df[col] = dataframe[col] + test_table = tabulate(print_df, headers="keys", tablefmt="psql") + if str(test_table).index("\n") > TERM_WIDTH: + # Drop all columns beyond terminal width + print_df.drop(col, axis=1, inplace=True) + dropped_cols += list(dataframe.columns)[i:] + break + + table = tabulate(print_df, headers="keys", tablefmt="psql", showindex="never") + + print(table) + if dropped_cols: + click.secho("Dropped columns: {}".format(dropped_cols), fg="yellow") + click.secho("Please increase your terminal size to view remaining columns.") + if empty_cols: + click.secho("Empty columns: {}".format(empty_cols), fg="yellow") + + return table, dropped_cols, empty_cols + + +def list_trials( + experiment_path: str, + sort: Optional[List[str]] = None, + output: Optional[str] = None, + filter_op: Optional[str] = None, + info_keys: Optional[List[str]] = None, + limit: int = None, + desc: bool = False, +): + """Lists trials in the directory subtree starting at the given path. + + Args: + experiment_path: Directory where trials are located. + Like Experiment.local_dir/Experiment.name/experiment*.json. + sort: Keys to sort by. + output: Name of file where output is saved. + filter_op: Filter operation in the format + " ". + info_keys: Keys that are displayed. + limit: Number of rows to display. + desc: Sort ascending vs. descending. + """ + _check_tabulate() + + try: + checkpoints_df = ExperimentAnalysis(experiment_path).dataframe() # last result + except TuneError as e: + raise click.ClickException("No trial data found!") from e + + config_prefix = CONFIG_PREFIX + "/" + + def key_filter(k): + return k in DEFAULT_CLI_KEYS or k.startswith(config_prefix) + + col_keys = [k for k in checkpoints_df.columns if key_filter(k)] + + if info_keys: + for k in info_keys: + if k not in checkpoints_df.columns: + raise click.ClickException( + "Provided key invalid: {}. " + "Available keys: {}.".format(k, checkpoints_df.columns) + ) + col_keys = [k for k in checkpoints_df.columns if k in info_keys] + + if not col_keys: + raise click.ClickException("No columns to output.") + + checkpoints_df = checkpoints_df[col_keys] + if "last_update_time" in checkpoints_df: + with pd.option_context("mode.use_inf_as_null", True): + datetime_series = checkpoints_df["last_update_time"].dropna() + + datetime_series = datetime_series.apply( + lambda t: datetime.fromtimestamp(t).strftime(TIMESTAMP_FORMAT) + ) + checkpoints_df["last_update_time"] = datetime_series + + if "logdir" in checkpoints_df: + # logdir often too long to view in table, so drop experiment_path + checkpoints_df["logdir"] = checkpoints_df["logdir"].str.replace( + experiment_path, "" + ) + + if filter_op: + col, op, val = filter_op.split(" ") + col_type = checkpoints_df[col].dtype + if is_numeric_dtype(col_type): + val = float(val) + elif is_string_dtype(col_type): + val = str(val) + # TODO(Andrew): add support for datetime and boolean + else: + raise click.ClickException( + "Unsupported dtype for {}: {}".format(val, col_type) + ) + op = OPERATORS[op] + filtered_index = op(checkpoints_df[col], val) + checkpoints_df = checkpoints_df[filtered_index] + + if sort: + for key in sort: + if key not in checkpoints_df: + raise click.ClickException( + "{} not in: {}".format(key, list(checkpoints_df)) + ) + ascending = not desc + checkpoints_df = checkpoints_df.sort_values(by=sort, ascending=ascending) + + if limit: + checkpoints_df = checkpoints_df[:limit] + + print_format_output(checkpoints_df) + + if output: + file_extension = os.path.splitext(output)[1].lower() + if file_extension in (".p", ".pkl", ".pickle"): + checkpoints_df.to_pickle(output) + elif file_extension == ".csv": + checkpoints_df.to_csv(output, index=False) + else: + raise click.ClickException("Unsupported filetype: {}".format(output)) + click.secho("Output saved at {}".format(output), fg="green") + + +def list_experiments( + project_path: str, + sort: Optional[List[str]] = None, + output: str = None, + filter_op: str = None, + info_keys: Optional[List[str]] = None, + limit: int = None, + desc: bool = False, +): + """Lists experiments in the directory subtree. + + Args: + project_path: Directory where experiments are located. + Corresponds to Experiment.local_dir. + sort: Keys to sort by. + output: Name of file where output is saved. + filter_op: Filter operation in the format + " ". + info_keys: Keys that are displayed. + limit: Number of rows to display. + desc: Sort ascending vs. descending. + """ + _check_tabulate() + base, experiment_folders, _ = next(os.walk(project_path)) + + experiment_data_collection = [] + + for experiment_dir in experiment_folders: + num_trials = sum( + EXPR_RESULT_FILE in files + for _, _, files in os.walk(os.path.join(base, experiment_dir)) + ) + + experiment_data = {"name": experiment_dir, "total_trials": num_trials} + experiment_data_collection.append(experiment_data) + + if not experiment_data_collection: + raise click.ClickException("No experiments found!") + + info_df = pd.DataFrame(experiment_data_collection) + if not info_keys: + info_keys = DEFAULT_PROJECT_INFO_KEYS + col_keys = [k for k in list(info_keys) if k in info_df] + if not col_keys: + raise click.ClickException( + "None of keys {} in experiment data!".format(info_keys) + ) + info_df = info_df[col_keys] + + if filter_op: + col, op, val = filter_op.split(" ") + col_type = info_df[col].dtype + if is_numeric_dtype(col_type): + val = float(val) + elif is_string_dtype(col_type): + val = str(val) + # TODO(Andrew): add support for datetime and boolean + else: + raise click.ClickException( + "Unsupported dtype for {}: {}".format(val, col_type) + ) + op = OPERATORS[op] + filtered_index = op(info_df[col], val) + info_df = info_df[filtered_index] + + if sort: + for key in sort: + if key not in info_df: + raise click.ClickException("{} not in: {}".format(key, list(info_df))) + ascending = not desc + info_df = info_df.sort_values(by=sort, ascending=ascending) + + if limit: + info_df = info_df[:limit] + + print_format_output(info_df) + + if output: + file_extension = os.path.splitext(output)[1].lower() + if file_extension in (".p", ".pkl", ".pickle"): + info_df.to_pickle(output) + elif file_extension == ".csv": + info_df.to_csv(output, index=False) + else: + raise click.ClickException("Unsupported filetype: {}".format(output)) + click.secho("Output saved at {}".format(output), fg="green") + + +def add_note(path: str, filename: str = "note.txt"): + """Opens a txt file at the given path where user can add and save notes. + + Args: + path: Directory where note will be saved. + filename: Name of note. Defaults to "note.txt" + """ + path = Path(path).expanduser() + assert path.is_dir(), "{} is not a valid directory.".format(path) + + filepath = path / filename + + try: + subprocess.call([EDITOR, filepath.as_posix()]) + except Exception as exc: + click.secho("Editing note failed: {}".format(str(exc)), fg="red") + if filepath.exists(): + print("Note updated at:", filepath.as_posix()) + else: + print("Note created at:", filepath.as_posix()) diff --git a/lib/python3.12/site-packages/ray/tune/cli/scripts.py b/lib/python3.12/site-packages/ray/tune/cli/scripts.py new file mode 100644 index 0000000000000000000000000000000000000000..5401d091c3ba622e78532985e60a18307a2484d9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/cli/scripts.py @@ -0,0 +1,101 @@ +import click + +import ray.tune.cli.commands as commands + + +@click.group() +def cli(): + pass + + +@cli.command() +@click.argument("experiment_path", required=True, type=str) +@click.option("--sort", default=None, type=str, help="Select which column to sort on.") +@click.option( + "--output", + "-o", + default=None, + type=str, + help="Select file to output information to.", +) +@click.option( + "--filter", + "filter_op", + default=None, + type=str, + help="Select filter in the format ' '.", +) +@click.option( + "--columns", default=None, type=str, help="Select columns to be displayed." +) +@click.option( + "--limit", default=None, type=int, help="Select number of rows to display." +) +@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.") +def list_trials(experiment_path, sort, output, filter_op, columns, limit, desc): + """Lists trials in the directory subtree starting at the given path.""" + if sort: + sort = sort.split(",") + if columns: + columns = columns.split(",") + commands.list_trials(experiment_path, sort, output, filter_op, columns, limit, desc) + + +@cli.command() +@click.argument("project_path", required=True, type=str) +@click.option("--sort", default=None, type=str, help="Select which column to sort on.") +@click.option( + "--output", + "-o", + default=None, + type=str, + help="Select file to output information to.", +) +@click.option( + "--filter", + "filter_op", + default=None, + type=str, + help="Select filter in the format ' '.", +) +@click.option( + "--columns", default=None, type=str, help="Select columns to be displayed." +) +@click.option( + "--limit", default=None, type=int, help="Select number of rows to display." +) +@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.") +def list_experiments(project_path, sort, output, filter_op, columns, limit, desc): + """Lists experiments in the directory subtree.""" + if sort: + sort = sort.split(",") + if columns: + columns = columns.split(",") + commands.list_experiments( + project_path, sort, output, filter_op, columns, limit, desc + ) + + +@cli.command() +@click.argument("path", required=True, type=str) +@click.option( + "--filename", default="note.txt", type=str, help="Specify filename for note." +) +def add_note(path, filename): + """Adds user notes as a text file at the given path.""" + commands.add_note(path, filename) + + +cli.add_command(list_trials, name="ls") +cli.add_command(list_trials, name="list-trials") +cli.add_command(list_experiments, name="lsx") +cli.add_command(list_experiments, name="list-experiments") +cli.add_command(add_note, name="add-note") + + +def main(): + return cli() + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa75b22b4f2d16d27d7f8f44d5a119451f43cf6e Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/async_hyperband_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/async_hyperband_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee8f0d3406508a12d88f831805c9da01379cdb15 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/async_hyperband_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/ax_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/ax_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0db222c3e977c78bc1e71f0c9579bd9d354c155b Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/ax_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bayesopt_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bayesopt_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..303c41359379b5bd725f9e83d4cdd4f9fc629a70 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bayesopt_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bohb_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bohb_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80d2c497c47d1c87289ea5773addac1d67da43a1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/bohb_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/cifar10_pytorch.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/cifar10_pytorch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcc138f9f00dbd006969c9aba8d21721cd75682c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/cifar10_pytorch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_checkpointing_with_callback.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_checkpointing_with_callback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af3dab1bc0420b961238e88e9241403717be249c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_checkpointing_with_callback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_func_checkpointing.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_func_checkpointing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79720fdab00a862ac32406fffa403d34c269f23a Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/custom_func_checkpointing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23db060084c833daba74a397109ac0b693c11ad1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_function_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_function_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..809c7a143e55ada8e0fba12c6c080a8cfaa2a141 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperband_function_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperopt_conditional_search_space_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperopt_conditional_search_space_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..175ccc898440da1da6198ab902076c8f78f8fdb7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/hyperopt_conditional_search_space_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/lightgbm_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/lightgbm_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a3f625560c49561bcea0b31386a94bb98e83dd5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/lightgbm_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/logging_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/logging_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4daba2a7dfc43395850a8b20fa53fb37c5f79b81 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/logging_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..932fe312bb573f40d34d1337b1f2394c9a65df7f Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_ptl.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_ptl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98d243be9dc61621cda9f0349b8618a24ea656bc Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mlflow_ptl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_ptl_mini.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_ptl_mini.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0576cced8f2bfc9927781ad543bed0ca34ea45ed Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_ptl_mini.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7c55c413d92828e1d4d2069e2e5c03538e1dec2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch_trainable.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch_trainable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e00c9fe6cf1588136b3535aec7b31a15c370cd10 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/mnist_pytorch_trainable.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/nevergrad_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/nevergrad_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..abf0c31e3e69e9810bb34f1a9986e51b0d7616e4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/nevergrad_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_define_by_run_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_define_by_run_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c48da9a01385099eafd5e5f849a210902ed54b74 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_define_by_run_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c7d1cd71a99d9d037151bdcf5cc24f56b3170da Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_multiobjective_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_multiobjective_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df7b91d9085ef7ee183497f36acb3e5855ce3f8c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/optuna_multiobjective_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee4724f2e62ea258ec69657a277bdc3a4e2b52eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_ppo_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_ppo_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf6c092e10981790564027967716f3b951c84bdf Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pb2_ppo_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..250fcd4a10ee16d70094af521cd73a00f3ba0ca8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_function_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_function_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c07e130b7ff01136ab26f9d60209885b6e34b06d Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_convnet_function_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffa93ecdb0d0c629d7242344188c1808e8020aa6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_function.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9677dec9b73487ff3119c22307e89db5e58917f Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_memnn_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_memnn_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a82f927db4daf99e70215ecc36b1884d90ef1539 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_memnn_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_ppo_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_ppo_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..846b1238ab25d1aacb4e1407f29e0feb38a04aaf Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_ppo_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_tune_cifar10_with_keras.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_tune_cifar10_with_keras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb5c32c16f9e3d3156b65598b88d1383ceaf80de Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/pbt_tune_cifar10_with_keras.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tf_mnist_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tf_mnist_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..02e6d7ae10154ad0bb91d92beaee1cda88183d8c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tf_mnist_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_basic_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_basic_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd5afd1cb0fd3e6c639537750a64963f837fc774 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_basic_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_mnist_keras.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_mnist_keras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb0e2391cddebab52d347b88e2b2e65291420e57 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/tune_mnist_keras.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad1bc1dd00283036fd028e3c10fa3835e9d0e440 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_dynamic_resources_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_dynamic_resources_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d0dacf0ff5dac33b37271fd8012d7e245b97036 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_dynamic_resources_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..647a5b873da228422fe36ff3007fbfbd73551645 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/__pycache__/xgboost_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/async_hyperband_example.py b/lib/python3.12/site-packages/ray/tune/examples/async_hyperband_example.py new file mode 100644 index 0000000000000000000000000000000000000000..bd372e2062d8e2623f35c8a92aff78ef7c990204 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/async_hyperband_example.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python + +import argparse +import time +from typing import Any, Dict + +from ray import tune +from ray.tune.schedulers import AsyncHyperBandScheduler + + +def evaluation_fn(step, width, height) -> float: + # simulate model evaluation + time.sleep(0.1) + return (0.1 + width * step / 100) ** (-1) + height * 0.1 + + +def easy_objective(config: Dict[str, Any]) -> None: + # Config contains the hyperparameters to tune + width, height = config["width"], config["height"] + + for step in range(config["steps"]): + # Iterative training function - can be an arbitrary training procedure + intermediate_score = evaluation_fn(step, width, height) + # Feed the score back back to Tune. + tune.report({"iterations": step, "mean_loss": intermediate_score}) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="AsyncHyperBand optimization example") + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + args, _ = parser.parse_known_args() + + # AsyncHyperBand enables aggressive early stopping of poorly performing trials + scheduler = AsyncHyperBandScheduler( + grace_period=5, # Minimum training iterations before stopping + max_t=100, # Maximum training iterations + ) + + tuner = tune.Tuner( + tune.with_resources(easy_objective, {"cpu": 1, "gpu": 0}), + run_config=tune.RunConfig( + name="asynchyperband_test", + stop={"training_iteration": 1 if args.smoke_test else 9999}, + verbose=1, + ), + tune_config=tune.TuneConfig( + metric="mean_loss", + mode="min", + scheduler=scheduler, + num_samples=20, # Number of trials to run + ), + param_space={ + "steps": 100, + "width": tune.uniform(10, 100), + "height": tune.uniform(0, 100), + }, + ) + + # Run the hyperparameter optimization + results = tuner.fit() + print(f"Best hyperparameters found: {results.get_best_result().config}") diff --git a/lib/python3.12/site-packages/ray/tune/examples/lightgbm_example.py b/lib/python3.12/site-packages/ray/tune/examples/lightgbm_example.py new file mode 100644 index 0000000000000000000000000000000000000000..94bdb29aba25500dea7678287ea9b983b73b81ed --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/lightgbm_example.py @@ -0,0 +1,103 @@ +import lightgbm as lgb +import sklearn.datasets +import sklearn.metrics +from sklearn.model_selection import train_test_split + +from ray import tune +from ray.tune.integration.lightgbm import TuneReportCheckpointCallback +from ray.tune.schedulers import ASHAScheduler + + +def train_breast_cancer(config: dict): + # This is a simple training function to be passed into Tune + + # Load dataset + data, target = sklearn.datasets.load_breast_cancer(return_X_y=True) + + # Split into train and test set + train_x, test_x, train_y, test_y = train_test_split(data, target, test_size=0.25) + + # Build input Datasets for LightGBM + train_set = lgb.Dataset(train_x, label=train_y) + test_set = lgb.Dataset(test_x, label=test_y) + + # Train the classifier, using the Tune callback + lgb.train( + config, + train_set, + valid_sets=[test_set], + valid_names=["eval"], + callbacks=[ + TuneReportCheckpointCallback( + { + "binary_error": "eval-binary_error", + "binary_logloss": "eval-binary_logloss", + } + ) + ], + ) + + +def train_breast_cancer_cv(config: dict): + # This is a simple training function to be passed into Tune, using + # lightgbm's cross validation functionality + + # Load dataset + data, target = sklearn.datasets.load_breast_cancer(return_X_y=True) + + train_set = lgb.Dataset(data, label=target) + + # Run CV, using the Tune callback + lgb.cv( + config, + train_set, + stratified=True, + # Checkpointing is not supported for CV + # LightGBM aggregates metrics over folds automatically + # with the cv_agg key. Both mean and standard deviation + # are provided. + callbacks=[ + TuneReportCheckpointCallback( + { + "binary_error": "valid-binary_error-mean", + "binary_logloss": "valid-binary_logloss-mean", + "binary_error_stdv": "valid-binary_error-stdv", + "binary_logloss_stdv": "valid-binary_logloss-stdv", + }, + frequency=0, + ) + ], + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--use-cv", action="store_true", help="Use `lgb.cv` instead of `lgb.train`." + ) + args, _ = parser.parse_known_args() + + config = { + "objective": "binary", + "metric": ["binary_error", "binary_logloss"], + "verbose": -1, + "boosting_type": tune.grid_search(["gbdt", "dart"]), + "num_leaves": tune.randint(10, 1000), + "learning_rate": tune.loguniform(1e-8, 1e-1), + } + + tuner = tune.Tuner( + train_breast_cancer if not args.use_cv else train_breast_cancer_cv, + tune_config=tune.TuneConfig( + metric="binary_error", + mode="min", + num_samples=2, + scheduler=ASHAScheduler(), + ), + param_space=config, + ) + results = tuner.fit() + + print("Best hyperparameters found were: ", results.get_best_result().config) diff --git a/lib/python3.12/site-packages/ray/tune/examples/optuna_example.py b/lib/python3.12/site-packages/ray/tune/examples/optuna_example.py new file mode 100644 index 0000000000000000000000000000000000000000..5c4bb7c809fa56b4358231f4e6332a674151b29d --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/optuna_example.py @@ -0,0 +1,73 @@ +"""This example demonstrates the usage of Optuna with Ray Tune. + +It also checks that it is usable with a separate scheduler. + +Requires the Optuna library to be installed (`pip install optuna`). + +For an example of using an Optuna define-by-run function, see +:doc:`/tune/examples/optuna_define_by_run_example`. +""" + +import time + +import ray +from ray import tune +from ray.tune.schedulers import AsyncHyperBandScheduler +from ray.tune.search import ConcurrencyLimiter +from ray.tune.search.optuna import OptunaSearch + + +def evaluation_fn(step, width, height): + return (0.1 + width * step / 100) ** (-1) + height * 0.1 + + +def easy_objective(config): + # Hyperparameters + width, height = config["width"], config["height"] + + for step in range(config["steps"]): + # Iterative training function - can be any arbitrary training procedure + intermediate_score = evaluation_fn(step, width, height) + # Feed the score back back to Tune. + tune.report({"iterations": step, "mean_loss": intermediate_score}) + time.sleep(0.1) + + +def run_optuna_tune(smoke_test=False): + algo = OptunaSearch() + algo = ConcurrencyLimiter(algo, max_concurrent=4) + scheduler = AsyncHyperBandScheduler() + tuner = tune.Tuner( + easy_objective, + tune_config=tune.TuneConfig( + metric="mean_loss", + mode="min", + search_alg=algo, + scheduler=scheduler, + num_samples=10 if smoke_test else 100, + ), + param_space={ + "steps": 100, + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100), + # This is an ignored parameter. + "activation": tune.choice(["relu", "tanh"]), + }, + ) + results = tuner.fit() + + print("Best hyperparameters found were: ", results.get_best_result().config) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + args, _ = parser.parse_known_args() + + ray.init(configure_logging=False) + + run_optuna_tune(smoke_test=args.smoke_test) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_convnet_example.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_convnet_example.py new file mode 100644 index 0000000000000000000000000000000000000000..9365ad73ebcfdc60084c15c15b2ce2854a21d0cb --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_convnet_example.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python + +# ruff: noqa +# fmt: off + +# __tutorial_imports_begin__ +import argparse +import os + +import numpy as np +import torch +import torch.optim as optim +from torchvision import datasets + +import ray +from ray import tune +from ray.tune.examples.mnist_pytorch import ( + ConvNet, + get_data_loaders, + test_func, + train_func, +) +from ray.tune.schedulers import PopulationBasedTraining +from ray.tune.utils import validate_save_restore + +# __tutorial_imports_end__ + + +# __trainable_begin__ +class PytorchTrainable(tune.Trainable): + """Train a Pytorch ConvNet with Trainable and PopulationBasedTraining + scheduler. The example reuse some of the functions in mnist_pytorch, + and is a good demo for how to add the tuning function without + changing the original training code. + """ + + def setup(self, config): + self.train_loader, self.test_loader = get_data_loaders() + self.model = ConvNet() + self.optimizer = optim.SGD( + self.model.parameters(), + lr=config.get("lr", 0.01), + momentum=config.get("momentum", 0.9)) + + def step(self): + train_func(self.model, self.optimizer, self.train_loader) + acc = test_func(self.model, self.test_loader) + return {"mean_accuracy": acc} + + def save_checkpoint(self, checkpoint_dir): + checkpoint_path = os.path.join(checkpoint_dir, "model.pth") + torch.save(self.model.state_dict(), checkpoint_path) + + def load_checkpoint(self, checkpoint_dir): + checkpoint_path = os.path.join(checkpoint_dir, "model.pth") + self.model.load_state_dict(torch.load(checkpoint_path)) + + def reset_config(self, new_config): + for param_group in self.optimizer.param_groups: + if "lr" in new_config: + param_group["lr"] = new_config["lr"] + if "momentum" in new_config: + param_group["momentum"] = new_config["momentum"] + + self.config = new_config + return True +# __trainable_end__ + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing") + args, _ = parser.parse_known_args() + + ray.init(num_cpus=2) + datasets.MNIST("~/data", train=True, download=True) + + # check if PytorchTrainble will save/restore correctly before execution + validate_save_restore(PytorchTrainable) + + # __pbt_begin__ + scheduler = PopulationBasedTraining( + time_attr="training_iteration", + perturbation_interval=5, + hyperparam_mutations={ + # distribution for resampling + "lr": lambda: np.random.uniform(0.0001, 1), + # allow perturbations within this set of categorical values + "momentum": [0.8, 0.9, 0.99], + }) + # __pbt_end__ + + # __tune_begin__ + class CustomStopper(tune.Stopper): + def __init__(self): + self.should_stop = False + + def __call__(self, trial_id, result): + max_iter = 5 if args.smoke_test else 100 + if not self.should_stop and result["mean_accuracy"] > 0.96: + self.should_stop = True + return self.should_stop or result["training_iteration"] >= max_iter + + def stop_all(self): + return self.should_stop + + stopper = CustomStopper() + + tuner = tune.Tuner( + PytorchTrainable, + run_config=tune.RunConfig( + name="pbt_test", + stop=stopper, + verbose=1, + checkpoint_config=tune.CheckpointConfig( + checkpoint_score_attribute="mean_accuracy", + checkpoint_frequency=5, + num_to_keep=4, + ), + ), + tune_config=tune.TuneConfig( + scheduler=scheduler, + metric="mean_accuracy", + mode="max", + num_samples=4, + reuse_actors=True, + ), + param_space={ + "lr": tune.uniform(0.001, 1), + "momentum": tune.uniform(0.001, 1), + }, + ) + results = tuner.fit() + # __tune_end__ + + best_result = results.get_best_result() + best_checkpoint = best_result.checkpoint diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__init__.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f94401b5487466bf250308673d6d76cadbeb0486 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..148a30750f4eb991f1887f3f88e13255da066385 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_func.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_func.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7646acf0f1492da9beb4dbd5382b27d5453adeb4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_func.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_trainable.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_trainable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47dd5920d2804053be3cb1e3a7debc34d798fd5e Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/__pycache__/pbt_dcgan_mnist_trainable.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/common.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/common.py new file mode 100644 index 0000000000000000000000000000000000000000..c8ee8ad8ff621d5c63db7060909ea18421ca044b --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/common.py @@ -0,0 +1,285 @@ +import os + +import matplotlib.animation as animation +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn as nn +import torch.nn.parallel +import torch.utils.data +import torchvision.datasets as dset +import torchvision.transforms as transforms +import torchvision.utils as vutils +from scipy.stats import entropy +from torch.autograd import Variable +from torch.nn import functional as F + +import ray + +# Training parameters +workers = 2 +batch_size = 64 +image_size = 32 + +# Number of channels in the training images. For color images this is 3 +nc = 1 + +# Size of z latent vector (i.e. size of generator input) +nz = 100 + +# Size of feature maps in generator +ngf = 32 + +# Size of feature maps in discriminator +ndf = 32 + +# Beta1 hyperparam for Adam optimizers +beta1 = 0.5 + +# iterations of actual training in each Trainable _train +train_iterations_per_step = 5 + +MODEL_PATH = os.path.expanduser("~/.ray/models/mnist_cnn.pt") + + +def get_data_loader(data_dir="~/data"): + dataset = dset.MNIST( + root=data_dir, + download=True, + transform=transforms.Compose( + [ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize((0.13066,), (0.30131,)), + ] + ), + ) + + # Create the dataloader + dataloader = torch.utils.data.DataLoader( + dataset, batch_size=batch_size, shuffle=True, num_workers=workers + ) + + return dataloader + + +# __GANmodel_begin__ +# custom weights initialization called on netG and netD +def weights_init(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + nn.init.normal_(m.weight.data, 0.0, 0.02) + elif classname.find("BatchNorm") != -1: + nn.init.normal_(m.weight.data, 1.0, 0.02) + nn.init.constant_(m.bias.data, 0) + + +# Generator Code +class Generator(nn.Module): + def __init__(self): + super(Generator, self).__init__() + self.main = nn.Sequential( + # input is Z, going into a convolution + nn.ConvTranspose2d(nz, ngf * 4, 4, 1, 0, bias=False), + nn.BatchNorm2d(ngf * 4), + nn.ReLU(True), + nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False), + nn.BatchNorm2d(ngf * 2), + nn.ReLU(True), + nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False), + nn.BatchNorm2d(ngf), + nn.ReLU(True), + nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False), + nn.Tanh(), + ) + + def forward(self, input): + return self.main(input) + + +class Discriminator(nn.Module): + def __init__(self): + super(Discriminator, self).__init__() + self.main = nn.Sequential( + nn.Conv2d(nc, ndf, 4, 2, 1, bias=False), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False), + nn.BatchNorm2d(ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False), + nn.BatchNorm2d(ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False), + nn.Sigmoid(), + ) + + def forward(self, input): + return self.main(input) + + +# __GANmodel_end__ + + +# __INCEPTION_SCORE_begin__ +class Net(nn.Module): + """ + LeNet for MNist classification, used for inception_score + """ + + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) + + +def inception_score(imgs, mnist_model_ref, batch_size=32, splits=1): + N = len(imgs) + dtype = torch.FloatTensor + dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size) + cm = ray.get(mnist_model_ref) # Get the mnist model from Ray object store. + up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype) + + def get_pred(x): + x = up(x) + x = cm(x) + return F.softmax(x).data.cpu().numpy() + + preds = np.zeros((N, 10)) + for i, batch in enumerate(dataloader, 0): + batch = batch.type(dtype) + batchv = Variable(batch) + batch_size_i = batch.size()[0] + preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batchv) + + # Now compute the mean kl-div + split_scores = [] + for k in range(splits): + part = preds[k * (N // splits) : (k + 1) * (N // splits), :] + py = np.mean(part, axis=0) + scores = [] + for i in range(part.shape[0]): + pyx = part[i, :] + scores.append(entropy(pyx, py)) + split_scores.append(np.exp(np.mean(scores))) + + return np.mean(split_scores), np.std(split_scores) + + +# __INCEPTION_SCORE_end__ + + +def train_func( + netD, + netG, + optimG, + optimD, + criterion, + dataloader, + iteration, + device, + mnist_model_ref, +): + real_label = 1 + fake_label = 0 + + for i, data in enumerate(dataloader, 0): + if i >= train_iterations_per_step: + break + + netD.zero_grad() + real_cpu = data[0].to(device) + b_size = real_cpu.size(0) + label = torch.full((b_size,), real_label, dtype=torch.float, device=device) + output = netD(real_cpu).view(-1) + errD_real = criterion(output, label) + errD_real.backward() + D_x = output.mean().item() + + noise = torch.randn(b_size, nz, 1, 1, device=device) + fake = netG(noise) + label.fill_(fake_label) + output = netD(fake.detach()).view(-1) + errD_fake = criterion(output, label) + errD_fake.backward() + D_G_z1 = output.mean().item() + errD = errD_real + errD_fake + optimD.step() + + netG.zero_grad() + label.fill_(real_label) + output = netD(fake).view(-1) + errG = criterion(output, label) + errG.backward() + D_G_z2 = output.mean().item() + optimG.step() + + is_score, is_std = inception_score(fake, mnist_model_ref) + + # Output training stats + if iteration % 10 == 0: + print( + "[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z))" + ": %.4f / %.4f \tInception score: %.4f" + % ( + iteration, + len(dataloader), + errD.item(), + errG.item(), + D_x, + D_G_z1, + D_G_z2, + is_score, + ) + ) + + return errG.item(), errD.item(), is_score + + +def plot_images(dataloader): + # Plot some training images + real_batch = next(iter(dataloader)) + plt.figure(figsize=(8, 8)) + plt.axis("off") + plt.title("Original Images") + plt.imshow( + np.transpose( + vutils.make_grid(real_batch[0][:64], padding=2, normalize=True).cpu(), + (1, 2, 0), + ) + ) + + plt.show() + + +def demo_gan(checkpoint_paths): + img_list = [] + fixed_noise = torch.randn(64, nz, 1, 1) + for path in checkpoint_paths: + checkpoint_dict = torch.load(os.path.join(path, "checkpoint.pt")) + + loadedG = Generator() + loadedG.load_state_dict(checkpoint_dict["netGmodel"]) + with torch.no_grad(): + fake = loadedG(fixed_noise).detach().cpu() + img_list.append(vutils.make_grid(fake, padding=2, normalize=True)) + + fig = plt.figure(figsize=(8, 8)) + plt.axis("off") + ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list] + ani = animation.ArtistAnimation( + fig, ims, interval=1000, repeat_delay=1000, blit=True + ) + ani.save("./generated.gif", writer="imagemagick", dpi=72) + plt.show() diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_func.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_func.py new file mode 100644 index 0000000000000000000000000000000000000000..57d28cb41f5e7fdb35a2c6a86a0cd5c5ec8dd673 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_func.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python +""" +Example of training DCGAN on MNIST using PBT with Tune's function API. +""" +import argparse +import os +import tempfile + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.parallel +import torch.optim as optim +import torch.utils.data +from filelock import FileLock + +import ray +from ray import tune +from ray.tune import Checkpoint +from ray.tune.examples.pbt_dcgan_mnist.common import ( + MODEL_PATH, + Discriminator, + Generator, + Net, + beta1, + demo_gan, + get_data_loader, + plot_images, + train_func, + weights_init, +) +from ray.tune.schedulers import PopulationBasedTraining + + +# __Train_begin__ +def dcgan_train(config): + use_cuda = config.get("use_gpu") and torch.cuda.is_available() + device = torch.device("cuda" if use_cuda else "cpu") + netD = Discriminator().to(device) + netD.apply(weights_init) + netG = Generator().to(device) + netG.apply(weights_init) + criterion = nn.BCELoss() + optimizerD = optim.Adam( + netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999) + ) + optimizerG = optim.Adam( + netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999) + ) + with FileLock(os.path.expanduser("~/ray_results/.data.lock")): + dataloader = get_data_loader() + + step = 1 + checkpoint = tune.get_checkpoint() + if checkpoint: + with checkpoint.as_directory() as checkpoint_dir: + checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt")) + netD.load_state_dict(checkpoint_dict["netDmodel"]) + netG.load_state_dict(checkpoint_dict["netGmodel"]) + optimizerD.load_state_dict(checkpoint_dict["optimD"]) + optimizerG.load_state_dict(checkpoint_dict["optimG"]) + # Note: Make sure to increment the loaded step by 1 to get the + # current step. + last_step = checkpoint_dict["step"] + step = last_step + 1 + + # NOTE: It's important to set the optimizer learning rates + # again, since we want to explore the parameters passed in by PBT. + # Without this, we would continue using the exact same + # configuration as the trial whose checkpoint we are exploiting. + if "netD_lr" in config: + for param_group in optimizerD.param_groups: + param_group["lr"] = config["netD_lr"] + if "netG_lr" in config: + for param_group in optimizerG.param_groups: + param_group["lr"] = config["netG_lr"] + + while True: + lossG, lossD, is_score = train_func( + netD, + netG, + optimizerG, + optimizerD, + criterion, + dataloader, + step, + device, + config["mnist_model_ref"], + ) + metrics = {"lossg": lossG, "lossd": lossD, "is_score": is_score} + + if step % config["checkpoint_interval"] == 0: + with tempfile.TemporaryDirectory() as tmpdir: + torch.save( + { + "netDmodel": netD.state_dict(), + "netGmodel": netG.state_dict(), + "optimD": optimizerD.state_dict(), + "optimG": optimizerG.state_dict(), + "step": step, + }, + os.path.join(tmpdir, "checkpoint.pt"), + ) + tune.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir)) + else: + tune.report(metrics) + + step += 1 + + +# __Train_end__ + + +def download_mnist_cnn(): + import urllib.request + + # Download a pre-trained MNIST model for inception score calculation. + # This is a tiny model (<100kb). + if not os.path.exists(MODEL_PATH): + print("downloading model") + os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True) + urllib.request.urlretrieve( + "https://github.com/ray-project/ray/raw/master/python/ray/tune/" + "examples/pbt_dcgan_mnist/mnist_cnn.pt", + MODEL_PATH, + ) + return MODEL_PATH + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + parser.add_argument( + "--data-dir", type=str, default="~/data/", help="Set the path of the dataset." + ) + args, _ = parser.parse_known_args() + ray.init() + + download_mnist_cnn() + + dataloader = get_data_loader(args.data_dir) + if not args.smoke_test: + plot_images(dataloader) + + # __tune_begin__ + + # load the pretrained mnist classification model for inception_score + mnist_cnn = Net() + mnist_cnn.load_state_dict(torch.load(MODEL_PATH)) + mnist_cnn.eval() + # Put the model in Ray object store. + mnist_model_ref = ray.put(mnist_cnn) + + scheduler = PopulationBasedTraining( + perturbation_interval=5, + hyperparam_mutations={ + # distribution for resampling + "netG_lr": lambda: np.random.uniform(1e-2, 1e-5), + "netD_lr": lambda: np.random.uniform(1e-2, 1e-5), + }, + ) + + tune_iter = 5 if args.smoke_test else 300 + tuner = tune.Tuner( + dcgan_train, + run_config=tune.RunConfig( + name="pbt_dcgan_mnist", + stop={"training_iteration": tune_iter}, + verbose=1, + ), + tune_config=tune.TuneConfig( + metric="is_score", + mode="max", + num_samples=8, + scheduler=scheduler, + ), + param_space={ + "netG_lr": tune.choice([0.0001, 0.0002, 0.0005]), + "netD_lr": tune.choice([0.0001, 0.0002, 0.0005]), + "mnist_model_ref": mnist_model_ref, + }, + ) + results = tuner.fit() + # __tune_end__ + + # demo of the trained Generators + if not args.smoke_test: + checkpoint_paths = [result.checkpoint.to_directory() for result in results] + demo_gan(checkpoint_paths) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_trainable.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_trainable.py new file mode 100644 index 0000000000000000000000000000000000000000..51909033f658f83ad5c08610e730efdd5b413520 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_dcgan_mnist/pbt_dcgan_mnist_trainable.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python +""" +Example of training DCGAN on MNIST using PBT with Tune's Trainable Class +API. +""" +import argparse +import os +import random + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.parallel +import torch.optim as optim +import torch.utils.data +from filelock import FileLock + +import ray +from ray import tune +from ray.tune.examples.pbt_dcgan_mnist.common import ( + MODEL_PATH, + Discriminator, + Generator, + Net, + beta1, + demo_gan, + get_data_loader, + plot_images, + train_func, + weights_init, +) +from ray.tune.schedulers import PopulationBasedTraining + + +# __Trainable_begin__ +class PytorchTrainable(tune.Trainable): + def setup(self, config): + use_cuda = config.get("use_gpu") and torch.cuda.is_available() + self.device = torch.device("cuda" if use_cuda else "cpu") + self.netD = Discriminator().to(self.device) + self.netD.apply(weights_init) + self.netG = Generator().to(self.device) + self.netG.apply(weights_init) + self.criterion = nn.BCELoss() + self.optimizerD = optim.Adam( + self.netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999) + ) + self.optimizerG = optim.Adam( + self.netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999) + ) + with FileLock(os.path.expanduser("~/.data.lock")): + self.dataloader = get_data_loader(config.get("data_dir", "~/data")) + self.mnist_model_ref = config["mnist_model_ref"] + + def step(self): + lossG, lossD, is_score = train_func( + self.netD, + self.netG, + self.optimizerG, + self.optimizerD, + self.criterion, + self.dataloader, + self._iteration, + self.device, + self.mnist_model_ref, + ) + return {"lossg": lossG, "lossd": lossD, "is_score": is_score} + + def save_checkpoint(self, checkpoint_dir): + path = os.path.join(checkpoint_dir, "checkpoint.pt") + torch.save( + { + "netDmodel": self.netD.state_dict(), + "netGmodel": self.netG.state_dict(), + "optimD": self.optimizerD.state_dict(), + "optimG": self.optimizerG.state_dict(), + }, + path, + ) + + return checkpoint_dir + + def load_checkpoint(self, checkpoint_dir): + path = os.path.join(checkpoint_dir, "checkpoint.pt") + checkpoint = torch.load(path) + self.netD.load_state_dict(checkpoint["netDmodel"]) + self.netG.load_state_dict(checkpoint["netGmodel"]) + self.optimizerD.load_state_dict(checkpoint["optimD"]) + self.optimizerG.load_state_dict(checkpoint["optimG"]) + + def reset_config(self, new_config): + if "netD_lr" in new_config: + for param_group in self.optimizerD.param_groups: + param_group["lr"] = new_config["netD_lr"] + if "netG_lr" in new_config: + for param_group in self.optimizerG.param_groups: + param_group["lr"] = new_config["netG_lr"] + + self.config = new_config + return True + + +# __Trainable_end__ + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + parser.add_argument( + "--data-dir", type=str, default="~/data/", help="Set the path of the dataset." + ) + args, _ = parser.parse_known_args() + ray.init() + + import urllib.request + + # Download a pre-trained MNIST model for inception score calculation. + # This is a tiny model (<100kb). + if not os.path.exists(MODEL_PATH): + print("downloading model") + os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True) + urllib.request.urlretrieve( + "https://github.com/ray-project/ray/raw/master/python/ray/tune/" + "examples/pbt_dcgan_mnist/mnist_cnn.pt", + MODEL_PATH, + ) + + dataloader = get_data_loader() + if not args.smoke_test: + plot_images(dataloader) + + # load the pretrained mnist classification model for inception_score + mnist_cnn = Net() + mnist_cnn.load_state_dict(torch.load(MODEL_PATH)) + mnist_cnn.eval() + mnist_model_ref = ray.put(mnist_cnn) + + # __tune_begin__ + scheduler = PopulationBasedTraining( + time_attr="training_iteration", + perturbation_interval=5, + hyperparam_mutations={ + # distribution for resampling + "netG_lr": lambda: np.random.uniform(1e-2, 1e-5), + "netD_lr": lambda: np.random.uniform(1e-2, 1e-5), + }, + ) + + tune_iter = 10 if args.smoke_test else 300 + tuner = tune.Tuner( + PytorchTrainable, + run_config=tune.RunConfig( + name="pbt_dcgan_mnist", + stop={"training_iteration": tune_iter}, + verbose=1, + checkpoint_config=tune.CheckpointConfig(checkpoint_at_end=True), + ), + tune_config=tune.TuneConfig( + metric="is_score", + mode="max", + num_samples=8, + scheduler=scheduler, + reuse_actors=True, + ), + param_space={ + "netG_lr": tune.sample_from( + lambda spec: random.choice([0.0001, 0.0002, 0.0005]) + ), + "netD_lr": tune.sample_from( + lambda spec: random.choice([0.0001, 0.0002, 0.0005]) + ), + "mnist_model_ref": mnist_model_ref, + "data_dir": args.data_dir, + }, + ) + results = tuner.fit() + + # export_formats=[ExportFormat.MODEL] + # __tune_end__ + + # demo of the trained Generators + if not args.smoke_test: + checkpoint_paths = [result.checkpoint.to_directory() for result in results] + demo_gan(checkpoint_paths) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_ppo_example.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_ppo_example.py new file mode 100644 index 0000000000000000000000000000000000000000..5e3f4216f47a3f01023b9cc456c7f78b542fd524 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_ppo_example.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +"""Example of using PBT with RLlib. + +Note that this requires a cluster with at least 8 GPUs in order for all trials +to run concurrently, otherwise PBT will round-robin train the trials which +is less efficient (or you can set {"gpu": 0} to use CPUs for SGD instead). + +Note that Tune in general does not need 8 GPUs, and this is just a more +computationally demanding example. +""" + +import random + +from ray import tune +from ray.rllib.algorithms.ppo import PPO +from ray.tune.schedulers import PopulationBasedTraining + +if __name__ == "__main__": + # Postprocess the perturbed config to ensure it's still valid + def explore(config): + # ensure we collect enough timesteps to do sgd + if config["train_batch_size"] < config["sgd_minibatch_size"] * 2: + config["train_batch_size"] = config["sgd_minibatch_size"] * 2 + # ensure we run at least one sgd iter + if config["num_sgd_iter"] < 1: + config["num_sgd_iter"] = 1 + return config + + pbt = PopulationBasedTraining( + time_attr="time_total_s", + perturbation_interval=120, + resample_probability=0.25, + # Specifies the mutations of these hyperparams + hyperparam_mutations={ + "lambda": lambda: random.uniform(0.9, 1.0), + "clip_param": lambda: random.uniform(0.01, 0.5), + "lr": [1e-3, 5e-4, 1e-4, 5e-5, 1e-5], + "num_sgd_iter": lambda: random.randint(1, 30), + "sgd_minibatch_size": lambda: random.randint(128, 16384), + "train_batch_size": lambda: random.randint(2000, 160000), + }, + custom_explore_fn=explore, + ) + + tuner = tune.Tuner( + PPO, + run_config=tune.RunConfig( + name="pbt_humanoid_test", + ), + tune_config=tune.TuneConfig( + scheduler=pbt, + num_samples=8, + metric="episode_reward_mean", + mode="max", + reuse_actors=True, + ), + param_space={ + "env": "Humanoid-v1", + "kl_coeff": 1.0, + "num_workers": 8, + "num_gpus": 1, + "model": {"free_log_std": True}, + # These params are tuned from a fixed starting value. + "lambda": 0.95, + "clip_param": 0.2, + "lr": 1e-4, + # These params start off randomly drawn from a set. + "num_sgd_iter": tune.choice([10, 20, 30]), + "sgd_minibatch_size": tune.choice([128, 512, 2048]), + "train_batch_size": tune.choice([10000, 20000, 40000]), + }, + ) + results = tuner.fit() + + print("best hyperparameters: ", results.get_best_result().config) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__init__.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e97e1a9d5b73417f1baf1f36eefaee05377f9ffd Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/pbt_transformers.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/pbt_transformers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c78bf1131d1fea75ffa7e4add8dcf158778cf12b Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/pbt_transformers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76d7f7ef77a4033bbdfecab71264bfff89b5a1ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/pbt_transformers.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/pbt_transformers.py new file mode 100644 index 0000000000000000000000000000000000000000..2688c057cd6e2ae3974eee0ac7fabde1f31275ae --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/pbt_transformers.py @@ -0,0 +1,165 @@ +""" +This example is uses the official +huggingface transformers `hyperparameter_search` API. +""" + +import os + +from transformers import ( + AutoConfig, + AutoModelForSequenceClassification, + AutoTokenizer, + GlueDataset, + GlueDataTrainingArguments, + Trainer, + TrainingArguments, + glue_tasks_num_labels, +) + +from ray import tune +from ray.tune import CheckpointConfig, CLIReporter +from ray.tune.examples.pbt_transformers.utils import ( + build_compute_metrics_fn, + download_data, +) +from ray.tune.schedulers import PopulationBasedTraining + + +def tune_transformer(num_samples=8, gpus_per_trial=0, smoke_test=False): + data_dir_name = "./data" if not smoke_test else "./test_data" + data_dir = os.path.abspath(os.path.join(os.getcwd(), data_dir_name)) + if not os.path.exists(data_dir): + os.mkdir(data_dir, 0o755) + + # Change these as needed. + model_name = ( + "bert-base-uncased" if not smoke_test else "sshleifer/tiny-distilroberta-base" + ) + task_name = "rte" + + task_data_dir = os.path.join(data_dir, task_name.upper()) + + num_labels = glue_tasks_num_labels[task_name] + + config = AutoConfig.from_pretrained( + model_name, num_labels=num_labels, finetuning_task=task_name + ) + + # Download and cache tokenizer, model, and features + print("Downloading and caching Tokenizer") + tokenizer = AutoTokenizer.from_pretrained(model_name) + + # Triggers tokenizer download to cache + print("Downloading and caching pre-trained model") + AutoModelForSequenceClassification.from_pretrained( + model_name, + config=config, + ) + + def get_model(): + return AutoModelForSequenceClassification.from_pretrained( + model_name, + config=config, + ) + + # Download data. + download_data(task_name, data_dir) + + data_args = GlueDataTrainingArguments(task_name=task_name, data_dir=task_data_dir) + + train_dataset = GlueDataset( + data_args, tokenizer=tokenizer, mode="train", cache_dir=task_data_dir + ) + eval_dataset = GlueDataset( + data_args, tokenizer=tokenizer, mode="dev", cache_dir=task_data_dir + ) + + training_args = TrainingArguments( + output_dir=".", + learning_rate=1e-5, # config + do_train=True, + do_eval=True, + no_cuda=gpus_per_trial <= 0, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, + num_train_epochs=2, # config + max_steps=-1, + per_device_train_batch_size=16, # config + per_device_eval_batch_size=16, # config + warmup_steps=0, + weight_decay=0.1, # config + logging_dir="./logs", + skip_memory_metrics=True, + report_to="none", + ) + + trainer = Trainer( + model_init=get_model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + compute_metrics=build_compute_metrics_fn(task_name), + ) + + tune_config = { + "per_device_train_batch_size": 32, + "per_device_eval_batch_size": 32, + "num_train_epochs": tune.choice([2, 3, 4, 5]), + "max_steps": 1 if smoke_test else -1, # Used for smoke test. + } + + scheduler = PopulationBasedTraining( + time_attr="training_iteration", + metric="eval_acc", + mode="max", + perturbation_interval=1, + hyperparam_mutations={ + "weight_decay": tune.uniform(0.0, 0.3), + "learning_rate": tune.uniform(1e-5, 5e-5), + "per_device_train_batch_size": [16, 32, 64], + }, + ) + + reporter = CLIReporter( + parameter_columns={ + "weight_decay": "w_decay", + "learning_rate": "lr", + "per_device_train_batch_size": "train_bs/gpu", + "num_train_epochs": "num_epochs", + }, + metric_columns=["eval_acc", "eval_loss", "epoch", "training_iteration"], + ) + + trainer.hyperparameter_search( + hp_space=lambda _: tune_config, + backend="ray", + n_trials=num_samples, + resources_per_trial={"cpu": 1, "gpu": gpus_per_trial}, + scheduler=scheduler, + checkpoint_config=CheckpointConfig( + num_to_keep=1, + checkpoint_score_attribute="training_iteration", + ), + stop={"training_iteration": 1} if smoke_test else None, + progress_reporter=reporter, + local_dir="~/ray_results/", + name="tune_transformer_pbt", + log_to_file=True, + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + args, _ = parser.parse_known_args() + + if args.smoke_test: + tune_transformer(num_samples=1, gpus_per_trial=0, smoke_test=True) + else: + # You can change the number of GPUs here: + tune_transformer(num_samples=8, gpus_per_trial=1) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/utils.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ff304664b98b2fe7835d18a5f7dc33a4a6563fd8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_transformers/utils.py @@ -0,0 +1,46 @@ +"""Utilities to load and cache data.""" + +import os +from typing import Callable, Dict + +import numpy as np +from transformers import EvalPrediction, glue_compute_metrics, glue_output_modes + + +def build_compute_metrics_fn(task_name: str) -> Callable[[EvalPrediction], Dict]: + """Function from transformers/examples/text-classification/run_glue.py""" + output_mode = glue_output_modes[task_name] + + def compute_metrics_fn(p: EvalPrediction): + if output_mode == "classification": + preds = np.argmax(p.predictions, axis=1) + elif output_mode == "regression": + preds = np.squeeze(p.predictions) + metrics = glue_compute_metrics(task_name, preds, p.label_ids) + return metrics + + return compute_metrics_fn + + +def download_data(task_name, data_dir="./data"): + # Download RTE training data + print("Downloading dataset.") + import urllib + import zipfile + + if task_name == "rte": + url = "https://dl.fbaipublicfiles.com/glue/data/RTE.zip" + else: + raise ValueError("Unknown task: {}".format(task_name)) + data_file = os.path.join(data_dir, "{}.zip".format(task_name)) + if not os.path.exists(data_file): + urllib.request.urlretrieve(url, data_file) + with zipfile.ZipFile(data_file) as zip_ref: + zip_ref.extractall(data_dir) + print("Downloaded data for task {} to {}".format(task_name, data_dir)) + else: + print( + "Data already exists. Using downloaded data for task {} from {}".format( + task_name, data_dir + ) + ) diff --git a/lib/python3.12/site-packages/ray/tune/examples/pbt_tune_cifar10_with_keras.py b/lib/python3.12/site-packages/ray/tune/examples/pbt_tune_cifar10_with_keras.py new file mode 100644 index 0000000000000000000000000000000000000000..06638036b311afc32a6565240223450682f7e7c9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/pbt_tune_cifar10_with_keras.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +"""Train keras CNN on the CIFAR10 small images dataset. + +The model comes from: https://zhuanlan.zhihu.com/p/29214791, +and it gets to about 87% validation accuracy in 100 epochs. + +Note that the script requires a machine with 4 GPUs. You +can set {"gpu": 0} to use CPUs for training, although +it is less efficient. +""" + +from __future__ import print_function + +import argparse + +import numpy as np +import tensorflow as tf +from tensorflow.keras.datasets import cifar10 +from tensorflow.keras.layers import ( + Convolution2D, + Dense, + Dropout, + Flatten, + Input, + MaxPooling2D, +) +from tensorflow.keras.models import Model, load_model +from tensorflow.keras.preprocessing.image import ImageDataGenerator + +from ray import tune +from ray.tune import Trainable +from ray.tune.schedulers import PopulationBasedTraining + +num_classes = 10 +NUM_SAMPLES = 128 + + +class Cifar10Model(Trainable): + def _read_data(self): + # The data, split between train and test sets: + (x_train, y_train), (x_test, y_test) = cifar10.load_data() + + # Convert class vectors to binary class matrices. + y_train = tf.keras.utils.to_categorical(y_train, num_classes) + y_test = tf.keras.utils.to_categorical(y_test, num_classes) + + x_train = x_train.astype("float32") + x_train /= 255 + x_test = x_test.astype("float32") + x_test /= 255 + + return (x_train, y_train), (x_test, y_test) + + def _build_model(self, input_shape): + x = Input(shape=(32, 32, 3)) + y = x + y = Convolution2D( + filters=64, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = Convolution2D( + filters=64, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y) + + y = Convolution2D( + filters=128, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = Convolution2D( + filters=128, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y) + + y = Convolution2D( + filters=256, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = Convolution2D( + filters=256, + kernel_size=3, + strides=1, + padding="same", + activation="relu", + kernel_initializer="he_normal", + )(y) + y = MaxPooling2D(pool_size=2, strides=2, padding="same")(y) + + y = Flatten()(y) + y = Dropout(self.config.get("dropout", 0.5))(y) + y = Dense(units=10, activation="softmax", kernel_initializer="he_normal")(y) + + model = Model(inputs=x, outputs=y, name="model1") + return model + + def setup(self, config): + self.train_data, self.test_data = self._read_data() + x_train = self.train_data[0] + model = self._build_model(x_train.shape[1:]) + + opt = tf.keras.optimizers.Adadelta( + lr=self.config.get("lr", 1e-4), weight_decay=self.config.get("decay", 1e-4) + ) + model.compile( + loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"] + ) + self.model = model + + def step(self): + x_train, y_train = self.train_data + x_train, y_train = x_train[:NUM_SAMPLES], y_train[:NUM_SAMPLES] + x_test, y_test = self.test_data + x_test, y_test = x_test[:NUM_SAMPLES], y_test[:NUM_SAMPLES] + + aug_gen = ImageDataGenerator( + # set input mean to 0 over the dataset + featurewise_center=False, + # set each sample mean to 0 + samplewise_center=False, + # divide inputs by dataset std + featurewise_std_normalization=False, + # divide each input by its std + samplewise_std_normalization=False, + # apply ZCA whitening + zca_whitening=False, + # randomly rotate images in the range (degrees, 0 to 180) + rotation_range=0, + # randomly shift images horizontally (fraction of total width) + width_shift_range=0.1, + # randomly shift images vertically (fraction of total height) + height_shift_range=0.1, + # randomly flip images + horizontal_flip=True, + # randomly flip images + vertical_flip=False, + ) + + aug_gen.fit(x_train) + batch_size = self.config.get("batch_size", 64) + gen = aug_gen.flow(x_train, y_train, batch_size=batch_size) + self.model.fit_generator( + generator=gen, epochs=self.config.get("epochs", 1), validation_data=None + ) + + # loss, accuracy + _, accuracy = self.model.evaluate(x_test, y_test, verbose=0) + return {"mean_accuracy": accuracy} + + def save_checkpoint(self, checkpoint_dir): + file_path = checkpoint_dir + "/model" + self.model.save(file_path) + + def load_checkpoint(self, checkpoint_dir): + # See https://stackoverflow.com/a/42763323 + del self.model + file_path = checkpoint_dir + "/model" + self.model = load_model(file_path) + + def cleanup(self): + # If need, save your model when exit. + # saved_path = self.model.save(self.logdir) + # print("save model at: ", saved_path) + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + args, _ = parser.parse_known_args() + + space = { + "epochs": 1, + "batch_size": 64, + "lr": tune.grid_search([10**-4, 10**-5]), + "decay": tune.sample_from(lambda spec: spec.config.lr / 100.0), + "dropout": tune.grid_search([0.25, 0.5]), + } + if args.smoke_test: + space["lr"] = 10**-4 + space["dropout"] = 0.5 + + perturbation_interval = 10 + pbt = PopulationBasedTraining( + time_attr="training_iteration", + perturbation_interval=perturbation_interval, + hyperparam_mutations={ + "dropout": lambda _: np.random.uniform(0, 1), + }, + ) + + tuner = tune.Tuner( + tune.with_resources( + Cifar10Model, + resources={"cpu": 1, "gpu": 1}, + ), + run_config=tune.RunConfig( + name="pbt_cifar10", + stop={ + "mean_accuracy": 0.80, + "training_iteration": 30, + }, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=perturbation_interval, + checkpoint_score_attribute="mean_accuracy", + num_to_keep=2, + ), + ), + tune_config=tune.TuneConfig( + scheduler=pbt, + num_samples=4, + metric="mean_accuracy", + mode="max", + reuse_actors=True, + ), + param_space=space, + ) + results = tuner.fit() + print("Best hyperparameters found were: ", results.get_best_result().config) diff --git a/lib/python3.12/site-packages/ray/tune/examples/tune_basic_example.py b/lib/python3.12/site-packages/ray/tune/examples/tune_basic_example.py new file mode 100644 index 0000000000000000000000000000000000000000..7834010c59f39f532b43ae946d6d75f21395a582 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/examples/tune_basic_example.py @@ -0,0 +1,57 @@ +"""This example demonstrates basic Ray Tune random search and grid search.""" +import time + +import ray +from ray import tune + + +def evaluation_fn(step, width, height): + time.sleep(0.1) + return (0.1 + width * step / 100) ** (-1) + height * 0.1 + + +def easy_objective(config): + # Hyperparameters + width, height = config["width"], config["height"] + + for step in range(config["steps"]): + # Iterative training function - can be any arbitrary training procedure + intermediate_score = evaluation_fn(step, width, height) + # Feed the score back back to Tune. + tune.report({"iterations": step, "mean_loss": intermediate_score}) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing" + ) + args, _ = parser.parse_known_args() + + ray.init(configure_logging=False) + + # This will do a grid search over the `activation` parameter. This means + # that each of the two values (`relu` and `tanh`) will be sampled once + # for each sample (`num_samples`). We end up with 2 * 50 = 100 samples. + # The `width` and `height` parameters are sampled randomly. + # `steps` is a constant parameter. + + tuner = tune.Tuner( + easy_objective, + tune_config=tune.TuneConfig( + metric="mean_loss", + mode="min", + num_samples=5 if args.smoke_test else 50, + ), + param_space={ + "steps": 5 if args.smoke_test else 100, + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100), + "activation": tune.grid_search(["relu", "tanh"]), + }, + ) + results = tuner.fit() + + print("Best hyperparameters found were: ", results.get_best_result().config) diff --git a/lib/python3.12/site-packages/ray/tune/execution/__init__.py b/lib/python3.12/site-packages/ray/tune/execution/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b7ba5992659cd1df61b1b42721aea96a4137fb7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/class_cache.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/class_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f9784133024d363b31390f8b8e61417ecf760eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/class_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/cluster_info.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/cluster_info.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f4a86161a2f5b07a8b4da273d07b3c2aa7db752 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/cluster_info.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/experiment_state.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/experiment_state.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06b17caa32ace8488e96a8fe2b6d38077f9159f8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/experiment_state.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/insufficient_resources_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/insufficient_resources_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b2205eedd0b6f009b8df9ddf6d2906677007fd5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/insufficient_resources_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/placement_groups.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/placement_groups.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f4e1bfd4ca2de68754c4307c4041467e9b6b91b Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/placement_groups.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/__pycache__/tune_controller.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/tune_controller.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85b668679a5d166e923da8556a4cfb4beaba04ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/execution/__pycache__/tune_controller.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/execution/class_cache.py b/lib/python3.12/site-packages/ray/tune/execution/class_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..94c4b5148a4db18ce506515469cc5272595f2af3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/class_cache.py @@ -0,0 +1,68 @@ +import os + +import ray +from ray.air.constants import COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV +from ray.train.constants import ( + ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, + RAY_CHDIR_TO_TRIAL_DIR, +) +from ray.train.v2._internal.constants import ( + ENV_VARS_TO_PROPAGATE as TRAIN_ENV_VARS_TO_PROPAGATE, +) + +DEFAULT_ENV_VARS = { + # https://github.com/ray-project/ray/issues/28197 + "PL_DISABLE_FORK": "1" +} +ENV_VARS_TO_PROPAGATE = ( + { + COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV, + RAY_CHDIR_TO_TRIAL_DIR, + ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SECURITY_TOKEN", + "AWS_SESSION_TOKEN", + } + # Propagate the Ray Train environment variables from the driver process + # to the trainable process so that Tune + Train v2 can be used together. + | TRAIN_ENV_VARS_TO_PROPAGATE +) + + +class _ActorClassCache: + """Caches actor classes. + + ray.remote is a registration call. It sends the serialized object to the + key value store (redis), and will be fetched at an arbitrary worker + later. Registration does not use any Ray scheduling resources. + + Later, class.remote() actually creates the remote actor. The + actor will be instantiated on some arbitrary machine, + according to the underlying Ray scheduler. + + Without this cache, you would register the same serialized object + over and over again. Naturally, since redis doesn’t spill to disk, + this can easily nuke the redis instance (and basically blow up Ray). + This cache instead allows us to register once and only once. + + Note that we assume there can be multiple trainables in the + system at once. + """ + + def __init__(self): + self._cache = {} + + def get(self, trainable_cls): + """Gets the wrapped trainable_cls, otherwise calls ray.remote.""" + env_vars = DEFAULT_ENV_VARS.copy() + + for env_var_to_propagate in ENV_VARS_TO_PROPAGATE: + if env_var_to_propagate in os.environ: + env_vars[env_var_to_propagate] = os.environ[env_var_to_propagate] + + runtime_env = {"env_vars": env_vars} + if trainable_cls not in self._cache: + remote_cls = ray.remote(runtime_env=runtime_env)(trainable_cls) + self._cache[trainable_cls] = remote_cls + return self._cache[trainable_cls] diff --git a/lib/python3.12/site-packages/ray/tune/execution/cluster_info.py b/lib/python3.12/site-packages/ray/tune/execution/cluster_info.py new file mode 100644 index 0000000000000000000000000000000000000000..23bb9d03dbd04a7ee8ff99ee07a1f455ef95ac9f --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/cluster_info.py @@ -0,0 +1,12 @@ +from functools import lru_cache +from pathlib import Path + + +@lru_cache() +def _is_ray_cluster(): + """Checks if the bootstrap config file exists. + + This will always exist if using an autoscaling cluster/started + with the ray cluster launcher. + """ + return Path("~/ray_bootstrap_config.yaml").expanduser().exists() diff --git a/lib/python3.12/site-packages/ray/tune/execution/experiment_state.py b/lib/python3.12/site-packages/ray/tune/execution/experiment_state.py new file mode 100644 index 0000000000000000000000000000000000000000..8afda2a3a2751b6bb6f44f4e59a6bc24c89b81d3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/experiment_state.py @@ -0,0 +1,287 @@ +import fnmatch +import logging +import os +import time +from collections import Counter +from pathlib import Path +from typing import Callable, Dict, Optional, Union + +import pyarrow.fs + +from ray.train._internal.storage import ( + StorageContext, + _download_from_fs_path, + _list_at_fs_path, + get_fs_and_path, +) +from ray.tune.experiment.trial import Trial +from ray.tune.impl.out_of_band_serialize_dataset import out_of_band_serialize_dataset + +logger = logging.getLogger(__name__) + + +_SLOW_SYNC_WARNING = ( + "This could be due to a large number of trials, " + "large logfiles from lots of reported metrics, or throttling from the " + "remote storage if uploading too frequently.\n" + "You may want to consider switching the `RunConfig(storage_filesystem)`" + " to a more performant storage backend such as s3fs for a " + "S3 storage path.\n" + "You can suppress this error by setting the environment variable " + "TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a higher " + "value than the current threshold ({threshold})." +) + + +def _find_newest_experiment_checkpoint( + experiment_path: str, fs: Optional[pyarrow.fs.FileSystem] = None +) -> Optional[str]: + """Returns file name of most recently created experiment checkpoint. + + Args: + experiment_path: Local or remote path to the experiment directory + containing at least one experiment checkpoint file. + + Returns: + str: The local or remote path to the latest experiment checkpoint file + based on timestamp. None if no experiment checkpoints were found. + """ + from ray.tune.execution.tune_controller import TuneController + + fs, experiment_fs_path = get_fs_and_path(experiment_path, storage_filesystem=fs) + filenames = _list_at_fs_path(fs=fs, fs_path=experiment_fs_path) + pattern = TuneController.CKPT_FILE_TMPL.format("*") + matching = fnmatch.filter(filenames, pattern) + if not matching: + return None + filename = max(matching) + return Path(experiment_fs_path, filename).as_posix() + + +class _ExperimentCheckpointManager: + """Helper class for managing experiment-level checkpoints. + + This class implements the ``checkpoint()`` method used to checkpoint + experiment state. When called, this will serialize and write to disk + the state of the trial runner, trial executor, and search algorithm, to + a specified checkpoint file. + + The checkpoint period is automatically adjusted to + ``max(10, time_per_checkpoint * 19)``. This means that at most 5% of the + time (1/20) will be used for writing checkpoints, while 95% of the time + (19/20) will be used to handle the rest of the training loop. + """ + + def __init__( + self, + *, + storage: Optional[StorageContext], + checkpoint_period: Union[int, float, str], + sync_every_n_trial_checkpoints: Optional[int] = None, + ): + self._storage = storage + + self._last_save_time = float("-inf") + self._last_sync_time = None + + # Dynamic checkpointing period + self._auto_checkpoint_enabled = checkpoint_period == "auto" + if self._auto_checkpoint_enabled: + self._checkpoint_period = 10.0 # Initial value + else: + self._checkpoint_period = float(checkpoint_period) + + # TODO(justinvyu): This is a non-performant workaround to force sync + # every num_to_keep checkpoints in order to maintain consistency + # between the experiment state's view of the latest checkpoint, + # and the actual latest checkpoint that was uploaded. + self._sync_every_n_trial_checkpoints = sync_every_n_trial_checkpoints + self._trial_num_checkpoints_since_last_sync: Dict[Trial, int] = Counter() + self._should_force_sync_up: bool = False + + self._excessive_sync_threshold = float( + os.environ.get( + "TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "5" + ) + ) + self._slow_sync_threshold = float( + os.environ.get( + "TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "30" + ) + ) + + @property + def auto_checkpoint_enabled(self): + return self._auto_checkpoint_enabled + + def _update_auto_checkpoint_time(self, time_taken: float): + if self._auto_checkpoint_enabled: + # Multiplying this time by 19 means we spend ~5% of the time + # writing global checkpoints and 95% of the time processing trials + self._checkpoint_period = max(10.0, time_taken * 19) + logger.debug( + f"Experiment state snapshotting took " + f"{time_taken:.2f} seconds. " + f"Adjusting snapshotting period to " + f"{self._checkpoint_period:.2f} seconds." + ) + + def sync_up_experiment_state( + self, + save_fn: Callable[[], None], + force: bool = False, + wait: bool = False, + ): + """Saves execution state to the experiment directory on the storage path. + This includes an experiment checkpoint file that contains trial statuses + and the searcher state. + + Overwrites the current session checkpoint, which starts when self + is instantiated. Throttle depends on self._checkpoint_period. + + Args: + save_fn: Function to call to actually save data to the driver + staging path. The files in the driver staging path will be + uploaded to the storage path. + force: Forces an experiment checkpoint and launches a sync to storage. + This happens regardless of checkpoint_period + wait: Waits for the sync up to complete before returning. + """ + driver_staging_path = self._storage.experiment_driver_staging_path + + force = force or self._should_force_sync_up + + now = time.monotonic() + if now - self._last_save_time < self._checkpoint_period and not force: + return + + # Checkpoint + checkpoint_time_start = time.monotonic() + + # NOTE: This context manager is for Datasets captured in a trial config. + # This is the case when *tuning over datasets*. + # If the datasets have already been full executed, then serializing + # block refs means that this checkpoint is not usable in a new Ray cluster. + # This context will serialize the dataset execution plan instead, if available. + with out_of_band_serialize_dataset(): + save_fn() + + def wait_for_sync(): + try: + self._storage.syncer.wait() + except Exception: + logger.error( + "Saving experiment state to storage at " + f"'{self._storage.experiment_fs_path}' failed with exception: ", + exc_info=True, + ) + + if force: + start_time = time.monotonic() + wait_for_sync() + wait_time = time.monotonic() - start_time + if wait_time > self._slow_sync_threshold: + logger.warning( + "Saving the experiment state (which holds a global view " + "of trial statuses and is used to restore the experiment) " + f"took ~{wait_time:.2f} seconds, which may be a performance " + "bottleneck.\n" + f"{_SLOW_SYNC_WARNING.format(threshold=self._slow_sync_threshold)}" + ) + + time_since_last_sync = ( + time.monotonic() - self._last_sync_time + if self._last_sync_time is not None + else None + ) + launched_sync = self._storage.syncer.sync_up( + driver_staging_path, self._storage.experiment_fs_path + ) + if launched_sync: + if ( + time_since_last_sync is not None + and time_since_last_sync < self._excessive_sync_threshold + and self._should_force_sync_up + ): + logger.warning( + "Experiment state snapshotting has been triggered multiple " + f"times in the last {self._excessive_sync_threshold} seconds " + "and may become a bottleneck. " + "A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, " + "and a trial has checkpointed >= `num_to_keep` times " + "since the last snapshot.\n" + "You may want to consider increasing the " + "`CheckpointConfig(num_to_keep)` or decreasing the frequency of " + "saving checkpoints.\n" + "You can suppress this warning by setting the environment variable " + "TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S " + "to a smaller value than the current threshold " + f"({self._excessive_sync_threshold}). " + "Set it to 0 to completely suppress this warning." + ) + + self._last_sync_time = time.monotonic() + + # We just synced, so reset the force flag + self._trial_num_checkpoints_since_last_sync.clear() + self._should_force_sync_up = False + else: + if ( + time_since_last_sync is not None + and time_since_last_sync > self._slow_sync_threshold + ): + logger.warning( + "Saving the experiment state (which holds a global view " + "of trial statuses and is used to restore the experiment) " + f"has already taken {time_since_last_sync:.2f} seconds, " + "which may cause consistency issues upon restoration if your " + "driver script ungracefully exits.\n" + f"{_SLOW_SYNC_WARNING.format(threshold=self._slow_sync_threshold)}" + ) + + if wait: + wait_for_sync() + + checkpoint_time_taken = time.monotonic() - checkpoint_time_start + + # Adjust dynamic checkpointing + self._update_auto_checkpoint_time(time_taken=checkpoint_time_taken) + + # Finish + self._last_save_time = time.monotonic() + + def sync_down_experiment_state(self) -> None: + fs = self._storage.storage_filesystem + filepaths = _list_at_fs_path(fs=fs, fs_path=self._storage.experiment_fs_path) + # TODO(ekl) we should refactor our restore code to read the necessary data + # directly from the storage context. As a temporary hack, restore all the + # serialized files from the root dir where other modules expect them to be. + matches = [ + path + for path in filepaths + if path.endswith(".json") or path.endswith(".pkl") + ] + for relpath in matches: + fs_path = Path(self._storage.experiment_fs_path, relpath).as_posix() + local_path = Path( + self._storage.experiment_driver_staging_path, relpath + ).as_posix() + _download_from_fs_path(fs=fs, fs_path=fs_path, local_path=local_path) + logger.debug( + f"Copied {matches} from:\n(fs, path) = " + f"({self._storage.storage_filesystem.type_name}, " + f"{self._storage.experiment_fs_path})\n" + f"-> {self._storage.experiment_driver_staging_path}" + ) + + def on_trial_checkpoint(self, trial: Trial): + if not self._sync_every_n_trial_checkpoints: + return + + self._trial_num_checkpoints_since_last_sync[trial] += 1 + + if ( + self._trial_num_checkpoints_since_last_sync[trial] + >= self._sync_every_n_trial_checkpoints + ): + self._should_force_sync_up = True diff --git a/lib/python3.12/site-packages/ray/tune/execution/insufficient_resources_manager.py b/lib/python3.12/site-packages/ray/tune/execution/insufficient_resources_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..0f788d2df7df64c3f34c12fef8b39d68ec758e66 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/insufficient_resources_manager.py @@ -0,0 +1,167 @@ +import logging +import os +import time +from functools import lru_cache +from typing import Dict, Optional, Tuple + +import ray +from ray.tune.execution.cluster_info import _is_ray_cluster +from ray.tune.experiment import Trial + +logger = logging.getLogger(__name__) + + +# Ideally we want to use @cache; but it's only available for python 3.9. +# Caching is only helpful/correct for no autoscaler case. +@lru_cache() +def _get_cluster_resources_no_autoscaler() -> Dict: + return ray.cluster_resources() + + +def _get_trial_cpu_and_gpu(trial: Trial) -> Tuple[int, int]: + cpu = trial.placement_group_factory.required_resources.get("CPU", 0) + gpu = trial.placement_group_factory.required_resources.get("GPU", 0) + return cpu, gpu + + +def _can_fulfill_no_autoscaler(trial: Trial) -> bool: + """Calculates if there is enough resources for a PENDING trial. + + For no autoscaler case. + """ + assert trial.status == Trial.PENDING + asked_cpus, asked_gpus = _get_trial_cpu_and_gpu(trial) + + return asked_cpus <= _get_cluster_resources_no_autoscaler().get( + "CPU", 0 + ) and asked_gpus <= _get_cluster_resources_no_autoscaler().get("GPU", 0) + + +@lru_cache() +def _get_insufficient_resources_warning_threshold() -> float: + if _is_ray_cluster(): + return float( + os.environ.get( + "TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S_AUTOSCALER", "60" + ) + ) + else: + # Set the default to 10s so that we don't prematurely determine that + # a cluster cannot fulfill the resources requirements. + # TODO(xwjiang): Change it back once #18608 is resolved. + return float(os.environ.get("TUNE_WARN_INSUFFICENT_RESOURCE_THRESHOLD_S", "60")) + + +MSG_TRAIN_START = ( + "Training has not started in the last {wait_time:.0f} seconds. " + "This could be due to the cluster not having enough resources available. " +) +MSG_TRAIN_INSUFFICIENT = ( + "You asked for {asked_cpus} CPUs and {asked_gpus} GPUs, but the cluster only " + "has {cluster_cpus} CPUs and {cluster_gpus} GPUs available. " +) +MSG_TRAIN_END = ( + "Stop the training and adjust the required resources (e.g. via the " + "`ScalingConfig` or `resources_per_trial`, or `num_workers` for rllib), " + "or add more resources to your cluster." +) + +MSG_TUNE_START = ( + "No trial is running and no new trial has been started within " + "the last {wait_time:.0f} seconds. " + "This could be due to the cluster not having enough resources available. " +) +MSG_TUNE_INSUFFICIENT = ( + "You asked for {asked_cpus} CPUs and {asked_gpus} GPUs per trial, " + "but the cluster only has {cluster_cpus} CPUs and {cluster_gpus} GPUs available. " +) +MSG_TUNE_END = ( + "Stop the tuning and adjust the required resources (e.g. via the " + "`ScalingConfig` or `resources_per_trial`, or `num_workers` for rllib), " + "or add more resources to your cluster." +) + + +# TODO(xwjiang): Consider having a help page with more detailed instructions. +@lru_cache() +def _get_insufficient_resources_warning_msg( + for_train: bool = False, trial: Optional[Trial] = None +) -> str: + msg = "Ignore this message if the cluster is autoscaling. " + + if for_train: + start = MSG_TRAIN_START + insufficient = MSG_TRAIN_INSUFFICIENT + end = MSG_TRAIN_END + else: + start = MSG_TUNE_START + insufficient = MSG_TUNE_INSUFFICIENT + end = MSG_TUNE_END + + msg += start.format(wait_time=_get_insufficient_resources_warning_threshold()) + + if trial: + asked_cpus, asked_gpus = _get_trial_cpu_and_gpu(trial) + cluster_resources = _get_cluster_resources_no_autoscaler() + + msg += insufficient.format( + asked_cpus=asked_cpus, + asked_gpus=asked_gpus, + cluster_cpus=cluster_resources.get("CPU", 0), + cluster_gpus=cluster_resources.get("GPU", 0), + ) + + msg += end + + return msg + + +class _InsufficientResourcesManager: + """Insufficient resources manager. + + Makes best effort, conservative guesses about if Tune loop is stuck due to + infeasible resources. If so, outputs usability messages for users to + act upon. + """ + + def __init__(self, for_train: bool = False): + # The information tracked across the life time of Tune loop. + self._no_running_trials_since = -1 + self._last_trial_num = -1 + self._for_train = for_train + + def on_no_available_trials(self, all_trials): + """Tracks information across the life of Tune loop and makes guesses + about if Tune loop is stuck due to infeasible resources. + If so, outputs certain warning messages. + The logic should be conservative, non-intrusive and informative. + For example, rate limiting is applied so that the message is not + spammy. + """ + # This is approximately saying we are not making progress. + if len(all_trials) == self._last_trial_num: + if self._no_running_trials_since == -1: + self._no_running_trials_since = time.monotonic() + elif ( + time.monotonic() - self._no_running_trials_since + > _get_insufficient_resources_warning_threshold() + ): + can_fulfill_any = any( + trial.status == Trial.PENDING and _can_fulfill_no_autoscaler(trial) + for trial in all_trials + ) + + if can_fulfill_any: + # If one trial can be fulfilled, it will be fulfilled eventually + self._no_running_trials_since = -1 + return + + # Otherwise, can fulfill none + msg = _get_insufficient_resources_warning_msg( + for_train=self._for_train, trial=all_trials[0] + ) + logger.warning(msg) + self._no_running_trials_since = time.monotonic() + else: + self._no_running_trials_since = -1 + self._last_trial_num = len(all_trials) diff --git a/lib/python3.12/site-packages/ray/tune/execution/placement_groups.py b/lib/python3.12/site-packages/ray/tune/execution/placement_groups.py new file mode 100644 index 0000000000000000000000000000000000000000..a2eaf548fd35af7ffa9df7a78b224987e5b19e15 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/placement_groups.py @@ -0,0 +1,131 @@ +import warnings +from typing import Dict, Optional + +from ray.air.execution.resources.request import ResourceRequest +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.placement_group import placement_group + + +@PublicAPI(stability="beta") +class PlacementGroupFactory(ResourceRequest): + """Wrapper class that creates placement groups for trials. + + This function should be used to define resource requests for Ray Tune + trials. It holds the parameters to create + :ref:`placement groups `. + At a minimum, this will hold at least one bundle specifying the + resource requirements for each trial: + + .. code-block:: python + + from ray import tune + + tuner = tune.Tuner( + tune.with_resources( + train, + resources=tune.PlacementGroupFactory([ + {"CPU": 1, "GPU": 0.5, "custom_resource": 2} + ]) + ) + ) + tuner.fit() + + If the trial itself schedules further remote workers, the resource + requirements should be specified in additional bundles. You can also + pass the placement strategy for these bundles, e.g. to enforce + co-located placement: + + .. code-block:: python + + from ray import tune + + tuner = tune.Tuner( + tune.with_resources( + train, + resources=tune.PlacementGroupFactory([ + {"CPU": 1, "GPU": 0.5, "custom_resource": 2}, + {"CPU": 2}, + {"CPU": 2}, + ], strategy="PACK") + ) + ) + tuner.fit() + + The example above will reserve 1 CPU, 0.5 GPUs and 2 custom_resources + for the trainable itself, and reserve another 2 bundles of 2 CPUs each. + The trial will only start when all these resources are available. This + could be used e.g. if you had one learner running in the main trainable + that schedules two remote workers that need access to 2 CPUs each. + + If the trainable itself doesn't require resources. + You can specify it as: + + .. code-block:: python + + from ray import tune + + tuner = tune.Tuner( + tune.with_resources( + train, + resources=tune.PlacementGroupFactory([ + {}, + {"CPU": 2}, + {"CPU": 2}, + ], strategy="PACK") + ) + ) + tuner.fit() + + Args: + bundles: A list of bundles which + represent the resources requirements. + strategy: The strategy to create the placement group. + + - "PACK": Packs Bundles into as few nodes as possible. + - "SPREAD": Places Bundles across distinct nodes as even as possible. + - "STRICT_PACK": Packs Bundles into one node. The group is + not allowed to span multiple nodes. + - "STRICT_SPREAD": Packs Bundles across distinct nodes. + *args: Passed to the call of ``placement_group()`` + **kwargs: Passed to the call of ``placement_group()`` + + """ + + def __call__(self, *args, **kwargs): + warnings.warn( + "Calling PlacementGroupFactory objects is deprecated. Use " + "`to_placement_group()` instead.", + DeprecationWarning, + ) + kwargs.update(self._bound.kwargs) + # Call with bounded *args and **kwargs + return placement_group(*self._bound.args, **kwargs) + + +@DeveloperAPI +def resource_dict_to_pg_factory(spec: Optional[Dict[str, float]] = None): + """Translates resource dict into PlacementGroupFactory.""" + spec = spec or {"cpu": 1} + + spec = spec.copy() + + cpus = spec.pop("cpu", spec.pop("CPU", 0.0)) + gpus = spec.pop("gpu", spec.pop("GPU", 0.0)) + memory = spec.pop("memory", 0.0) + + # If there is a custom_resources key, use as base for bundle + bundle = dict(spec.pop("custom_resources", {})) + + # Otherwise, consider all other keys as custom resources + if not bundle: + bundle = spec + + bundle.update( + { + "CPU": cpus, + "GPU": gpus, + "memory": memory, + } + ) + + return PlacementGroupFactory([bundle]) diff --git a/lib/python3.12/site-packages/ray/tune/execution/tune_controller.py b/lib/python3.12/site-packages/ray/tune/execution/tune_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..2e6a9438261a84440b3d50556d6d5536dca7306a --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/execution/tune_controller.py @@ -0,0 +1,2184 @@ +import copy +import json +import logging +import os +import time +import traceback +import warnings +from collections import defaultdict, deque +from datetime import datetime +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union + +import ray +from ray.air import ResourceRequest +from ray.air.constants import TIME_THIS_ITER_S +from ray.air.execution import PlacementGroupResourceManager, ResourceManager +from ray.air.execution._internal import RayActorManager, TrackedActor +from ray.exceptions import RayActorError, RayTaskError +from ray.train._internal.session import _FutureTrainingResult, _TrainingResult +from ray.train._internal.storage import StorageContext +from ray.tune import CheckpointConfig +from ray.tune.callback import Callback, CallbackList +from ray.tune.error import TuneError, _AbortTrialExecution, _TuneStopTrialError +from ray.tune.execution.class_cache import _ActorClassCache +from ray.tune.execution.experiment_state import ( + _ExperimentCheckpointManager, + _find_newest_experiment_checkpoint, +) +from ray.tune.execution.insufficient_resources_manager import ( + _InsufficientResourcesManager, +) +from ray.tune.execution.placement_groups import PlacementGroupFactory +from ray.tune.experiment import Experiment, Trial +from ray.tune.experiment.trial import ( + _change_working_directory, + _get_trainable_kwargs, + _Location, + _noop_logger_creator, + _TrialInfo, +) +from ray.tune.result import ( + DEBUG_METRICS, + DEFAULT_METRIC, + DONE, + RESULT_DUPLICATE, + SHOULD_CHECKPOINT, + STDERR_FILE, + STDOUT_FILE, + TRIAL_INFO, +) +from ray.tune.schedulers import FIFOScheduler, TrialScheduler +from ray.tune.search import BasicVariantGenerator, SearchAlgorithm +from ray.tune.stopper import NoopStopper, Stopper +from ray.tune.tune_config import ResumeConfig +from ray.tune.utils import flatten_dict, warn_if_slow +from ray.tune.utils.log import Verbosity, _dedup_logs, has_verbosity +from ray.tune.utils.object_cache import _ObjectCache +from ray.tune.utils.resource_updater import _ResourceUpdater +from ray.tune.utils.serialization import TuneFunctionDecoder, TuneFunctionEncoder +from ray.util.annotations import DeveloperAPI +from ray.util.debug import log_once + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class TuneController: + CKPT_FILE_TMPL = "experiment_state-{}.json" + RAISE = "RAISE" + + def __init__( + self, + *, + search_alg: Optional[SearchAlgorithm] = None, + placeholder_resolvers: Optional[Dict[Tuple, Any]] = None, + scheduler: Optional[TrialScheduler] = None, + stopper: Optional[Stopper] = None, + resume_config: Optional[ResumeConfig] = None, + fail_fast: bool = False, + checkpoint_period: Union[str, int] = None, + callbacks: Optional[List[Callback]] = None, + metric: Optional[str] = None, + trial_checkpoint_config: Optional[CheckpointConfig] = None, + storage: Optional[StorageContext] = None, + reuse_actors: bool = False, + resource_manager_factory: Optional[Callable[[], ResourceManager]] = None, + _trainer_api: bool = False, + ): + if resource_manager_factory: + resource_manager = resource_manager_factory() + else: + resource_manager = PlacementGroupResourceManager() + + self._actor_manager = RayActorManager(resource_manager=resource_manager) + + self._class_cache = _ActorClassCache() + + # Resource status + self._resource_updater = _ResourceUpdater(None) + + # Actor <-> Trial mappings + self._actor_to_trial: Dict[TrackedActor, Trial] = {} + self._trial_to_actor: Dict[Trial, TrackedActor] = {} + + # Resources <-> Trial + self._resources_to_pending_trials: Dict[ + ResourceRequest, Set[Trial] + ] = defaultdict(set) + + # Keep track of actor states + self._pending_trials: Set[Trial] = set() + self._pending_trials_list: List[Trial] = [] + + self._running_trials: Set[Trial] = set() + + self._paused_trials: Set[Trial] = set() + + self._stopped_trials: Set[Trial] = set() + self._failed_trials: Set[Trial] = set() + + self._resetting_trials: Set[Trial] = set() + self._staged_trials: Set[Trial] = set() + + # Removed actors + self._started_actors: Set[TrackedActor] = set() + + # Map of tracked actors -> timestamp + # The timestamp is when we requested the stop. + # We track these actors here to force a + # cleanup after some time (as they might be hanging). + # Todo: This timeout logic should be moved into the actor manager. + # This map is populated whenever we request an actor stop: + # - Regular STOP decision + # - Removing an actor because its trial REUSEs a different trial's actor + # - Removing a cached actor because it's not needed anymore + # Actors are only tracked in this map if they actually started (not if they + # were only requested but never started). + # Actors are removed from this map: + # - When the STOP resolved and the actor actually stopped + # - When they are forcefully cleaned up after the timeout. + self._stopping_actors: Dict[TrackedActor, float] = {} + self._earliest_stopping_actor: float = float("inf") + self._actor_cleanup_timeout: int = int( + os.environ.get("TUNE_FORCE_TRIAL_CLEANUP_S", "600") + ) + self._actor_force_cleanup_timeout: int = 10 + + # Reuse actors + self._reuse_actors = reuse_actors + self._actor_cache = _ObjectCache(may_keep_one=True) + + # Trial metadata for experiment checkpoints + self._trials_to_cache: Set[Trial] = set() + self._trial_metadata: Dict[str, str] = {} + + # TRAINING + self._buffer_length = int(os.getenv("TUNE_RESULT_BUFFER_LENGTH", 1)) + self._buffer_min_time_s = float(os.getenv("TUNE_RESULT_BUFFER_MIN_TIME_S", 0.0)) + self._buffer_max_time_s = float( + os.getenv("TUNE_RESULT_BUFFER_MAX_TIME_S", 100.0) + ) + + # Legacy TrialRunner init + self._search_alg = search_alg or BasicVariantGenerator() + self._placeholder_resolvers = placeholder_resolvers + self._scheduler_alg = scheduler or FIFOScheduler() + self._callbacks = CallbackList(callbacks or []) + self._insufficient_resources_manager = _InsufficientResourcesManager( + for_train=_trainer_api + ) + self._pending_trial_queue_times = {} + + self._max_pending_trials = _get_max_pending_trials(self._search_alg) + + self._storage = storage + self._metric = metric + + self._total_time = 0 + self._iteration = 0 + self._has_errored = False + self._fail_fast = fail_fast + if isinstance(self._fail_fast, str): + self._fail_fast = self._fail_fast.upper() + if self._fail_fast == self.RAISE: + warnings.warn( + "fail_fast='raise' detected. Be careful when using this " + "mode as resources (such as Ray processes, " + "file descriptors, and temporary files) may not be " + "cleaned up properly. To use " + "a safer mode, use fail_fast=True." + ) + else: + raise ValueError( + "fail_fast must be one of {bool, RAISE}. " f"Got {self._fail_fast}." + ) + + self._print_trial_errors = bool( + int(os.environ.get("TUNE_PRINT_ALL_TRIAL_ERRORS", "1")) + ) + + self._trials: List[Trial] = [] + self._live_trials: Set[Trial] = set() # Set of non-terminated trials + self._cached_trial_decisions = {} + self._queued_trial_decisions = {} + + self._stop_queue = [] + self._should_stop_experiment = False # used by TuneServer + + self._stopper = stopper or NoopStopper() + + self._start_time = time.time() + + self._session_str = datetime.fromtimestamp(self._start_time).strftime( + "%Y-%m-%d_%H-%M-%S" + ) + + if checkpoint_period is None: + checkpoint_period = os.getenv("TUNE_GLOBAL_CHECKPOINT_S", "auto") + + self._checkpoint_period = checkpoint_period + self._trial_checkpoint_config = trial_checkpoint_config or CheckpointConfig() + self._checkpoint_manager = self._create_checkpoint_manager() + + self._resumed = False + + if resume_config is not None: + # Use the metadata file to restore TuneController state + try: + self.resume(resume_config=resume_config) + self._resumed = True + except Exception as e: + if has_verbosity(Verbosity.V3_TRIAL_DETAILS): + logger.error(str(e)) + logger.exception("Failed to restore the run state.") + if self._fail_fast: + raise + logger.info("Restarting experiment.") + else: + logger.debug("Starting a new experiment.") + + def _wrapped(self): + """Return wrapped tune controller to be passed to scheduler/searchers.""" + return TrialRunnerWrapper( + self, + trial_executor=_FakeRayTrialExecutor(self), + runner_whitelist_attr={ + "search_alg", + "get_trials", + "get_live_trials", + "_set_trial_status", + "pause_trial", + "stop_trial", + "_schedule_trial_save", + }, + executor_whitelist_attr={ + "has_resources_for_trial", + "pause_trial", + "save", + "_resource_updater", + }, + ) + + @property + def resumed(self): + return self._resumed + + @property + def search_alg(self): + return self._search_alg + + @property + def scheduler_alg(self): + return self._scheduler_alg + + def setup_experiments( + self, experiments: List[Experiment], total_num_samples: int + ) -> None: + """Obtains any necessary information from experiments. + + Mainly used to setup callbacks. + + Args: + experiments: List of Experiments + to use. + total_num_samples: Total number of samples + factoring in grid search samplers. + """ + experiment = experiments[0] + spec = experiment.public_spec if experiment else {} + spec["total_num_samples"] = total_num_samples + self._callbacks.setup(**spec) + + def end_experiment_callbacks(self) -> None: + """Calls ``on_experiment_end`` method in callbacks.""" + self._callbacks.on_experiment_end(trials=self._trials) + + @property + def experiment_state_file_name(self) -> str: + return self.CKPT_FILE_TMPL.format(self._session_str) + + @property + def experiment_state_path(self) -> str: + """Returns the local experiment checkpoint path.""" + return Path( + self._storage.experiment_driver_staging_path, + self.experiment_state_file_name, + ).as_posix() + + @property + def experiment_path(self) -> str: + return self._storage.experiment_fs_path + + def _create_checkpoint_manager(self): + return _ExperimentCheckpointManager( + storage=self._storage, + checkpoint_period=self._checkpoint_period, + sync_every_n_trial_checkpoints=self._trial_checkpoint_config.num_to_keep, + ) + + def save_to_dir(self): + """Save TuneController state to the local staging experiment directory. + + This includes: + - trial states + - TuneController internal state (all the serializable attributes) + - the searcher state + - the callback states + """ + # Get state from trial executor and runner + runner_state = { + # Trials + "trial_data": list(self._get_trial_checkpoints().values()), + # Experiment data + "runner_data": self.__getstate__(), + # Metadata + "stats": {"start_time": self._start_time}, + } + + driver_staging_path = self._storage.experiment_driver_staging_path + os.makedirs(driver_staging_path, exist_ok=True) + with open( + Path(driver_staging_path, self.experiment_state_file_name), + "w", + ) as f: + json.dump(runner_state, f, cls=TuneFunctionEncoder) + + self._search_alg.save_to_dir(driver_staging_path, session_str=self._session_str) + self._callbacks.save_to_dir(driver_staging_path, session_str=self._session_str) + + def checkpoint(self, force: bool = False, wait: bool = False): + self._checkpoint_manager.sync_up_experiment_state( + save_fn=self.save_to_dir, force=force, wait=wait + ) + + def _requeue_restored_trials( + self, trials: List[Trial], resume_config: ResumeConfig + ): + # Set trial statuses according to the resume configuration + for trial in sorted( + trials, key=lambda t: t.run_metadata.last_result_time, reverse=True + ): + if trial.status == Trial.ERROR: + resume_type = resume_config.errored + elif trial.status == Trial.TERMINATED: + resume_type = resume_config.finished + else: # Unfinished (PENDING, RUNNING, PAUSED) + resume_type = resume_config.unfinished + + trial_to_add = None + if resume_type == ResumeConfig.ResumeType.RESUME: + # Keep trial ID on resume + trial_to_add = trial + trial_to_add.run_metadata.error_filename = None + trial_to_add.run_metadata.pickled_error_filename = None + trial_to_add.set_status(Trial.PENDING) + elif resume_type == ResumeConfig.ResumeType.RESTART: + trial_to_add = trial.reset() + trial_to_add.restore_path = None + elif resume_type == ResumeConfig.ResumeType.SKIP: + trial_to_add = trial + if trial_to_add.status != Trial.ERROR: + # Set the status to terminated to skip it. + # Keep errored trial status as ERROR. + trial_to_add.set_status(Trial.TERMINATED) + else: + raise ValueError(f"Unknown resume type: {resume_type}") + assert trial_to_add is not None + + self.add_trial(trial_to_add) + + def _restore_trials(self, experiment_state: Dict) -> List[Trial]: + trials = [] + for trial_json_state, trial_runtime_metadata in experiment_state["trial_data"]: + trial = Trial.from_json_state(trial_json_state) + trial.restore_run_metadata(trial_runtime_metadata) + + # The following properties may be updated on restoration + # Ex: moved local/cloud experiment directory + + # Propagate updated storage ctx properties to the trial's restored copy. + new_storage = copy.copy(trial.storage) + new_storage.storage_filesystem = self._storage.storage_filesystem + new_storage.storage_fs_path = self._storage.storage_fs_path + new_storage.experiment_dir_name = self._storage.experiment_dir_name + + # ATTN: `trial.set_storage` is used intentionally, since it + # also updates the absolute paths and filesystem of tracked checkpoints. + trial.set_storage(new_storage) + + # Avoid creating logdir in client mode for returned trial results, + # since the dir might not be creatable locally. + # TODO(ekl) this is kind of a hack. + if not ray.util.client.ray.is_connected(): + trial.init_local_path() # Create logdir if it does not exist + + trials.append(trial) + + # NOTE: The restored run should reuse the same driver staging directory. + self._storage._timestamp = trials[0].storage._timestamp + + return trials + + def resume(self, resume_config: ResumeConfig): + """Resumes all checkpointed trials from previous run. + + Requires user to manually re-register their objects. Also stops + all ongoing trials. + """ + # 1. Restore TuneController state + # Find newest state file + newest_state_path = _find_newest_experiment_checkpoint( + self._storage.experiment_fs_path, fs=self._storage.storage_filesystem + ) + + if newest_state_path is None: + raise ValueError( + f"Tried to resume experiment from directory " + f"'{self._storage.experiment_fs_path}', but no " + f"experiment state file of the form '{TuneController.CKPT_FILE_TMPL}' " + "was found. This is expected if you are launching a new experiment." + ) + + logger.info( + "Restoring the run from the latest experiment state file: " + f"{Path(newest_state_path).name}" + ) + with self._storage.storage_filesystem.open_input_stream(newest_state_path) as f: + experiment_state = json.loads(f.readall(), cls=TuneFunctionDecoder) + + self.__setstate__(experiment_state["runner_data"]) + + # 2. Get the trial states that the run left off at. + trials = self._restore_trials(experiment_state) + + # 3. Restore search algorithm and callback state + # Download the search algorithm and callback state to the driver staging dir. + self._checkpoint_manager.sync_down_experiment_state() + + driver_staging_dir = self._storage.experiment_driver_staging_path + if self._search_alg.has_checkpoint(driver_staging_dir): + self._search_alg.restore_from_dir(driver_staging_dir) + + if self._callbacks.can_restore(driver_staging_dir): + self._callbacks.restore_from_dir(driver_staging_dir) + + # 4. Re-queue trials as needed, depending on their status. + self._requeue_restored_trials(trials, resume_config) + + def update_max_pending_trials(self, max_pending_trials: Optional[int] = None): + self._max_pending_trials = max_pending_trials or _get_max_pending_trials( + self._search_alg + ) + + def update_pending_trial_resources( + self, resources: Union[dict, PlacementGroupFactory] + ): + """Update trial resources when resuming from checkpoint. + + Only updating the pending ones. + """ + assert resources + if isinstance(resources, dict) and "gpu" not in resources: + resources["gpu"] = 0 + for trial in self._trials: + if trial.status == Trial.PENDING: + trial.update_resources(resources=resources) + + def is_finished(self): + """Returns whether all trials have finished running.""" + # The checks here are partly redundant but optimized for quick + # evaluation. Specifically, if there are live trials, we check + # these live trials first. Only if none of the live trials is + # live anymore do we loop over all trials for a final check. + trials_done = ( + len(self._live_trials) == 0 + or all(trial.is_finished() for trial in self._live_trials) + ) and all(trial.is_finished() for trial in self._trials) + return trials_done and self._search_alg.is_finished() + + def get_trial(self, tid): + trial = [t for t in self._trials if t.trial_id == tid] + return trial[0] if trial else None + + def get_trials(self): + """Returns the list of trials managed by this TrialRunner. + + Note that the caller usually should not mutate trial state directly. + """ + return self._trials + + def get_live_trials(self): + """Returns the set of trials that are not in Trial.TERMINATED state.""" + return self._live_trials + + def add_trial(self, trial: Trial): + """Adds a new trial to this TrialRunner. + + Trials may be added at any time. + + Args: + trial: Trial to queue. + """ + # If the config map has had all the references replaced with placeholders, + # resolve them before adding the trial. + if self._placeholder_resolvers: + trial.resolve_config_placeholders(self._placeholder_resolvers) + + # With trial.config resolved, create placement group factory if needed. + trial.create_placement_group_factory() + + self._trials.append(trial) + if trial.status != Trial.TERMINATED: + self._live_trials.add(trial) + with warn_if_slow("scheduler.on_trial_add"): + self._scheduler_alg.on_trial_add(self._wrapped(), trial) + self._mark_trial_to_checkpoint(trial) + + logger.debug(f"Adding trial {trial} with status {trial.status}") + + status_str_map = { + Trial.PENDING: self._pending_trials, + Trial.RUNNING: self._running_trials, + Trial.PAUSED: self._paused_trials, + Trial.TERMINATED: self._stopped_trials, + Trial.ERROR: self._failed_trials, + } + + status_str_map[trial.status].add(trial) + + if trial.status == Trial.PENDING: + self._pending_trials_list.append(trial) + self._resources_to_pending_trials[trial.placement_group_factory].add(trial) + + def _update_trial_queue(self, blocking: bool = False, timeout: int = 600) -> bool: + """Adds next trials to queue if possible. + + Note that the timeout is currently unexposed to the user. + + Args: + blocking: Blocks until either a trial is available + or is_finished (timeout or search algorithm finishes). + timeout: Seconds before blocking times out. + + Returns: + Boolean indicating if a new trial was created or not. + """ + trial = self._search_alg.next_trial() + if blocking and not trial: + start = time.time() + # Checking `is_finished` instead of _search_alg.is_finished + # is fine because blocking only occurs if all trials are + # finished and search_algorithm is not yet finished + while ( + not trial and not self.is_finished() and time.time() - start < timeout + ): + logger.debug("Blocking for next trial...") + trial = self._search_alg.next_trial() + time.sleep(1) + + if trial: + self.add_trial(trial) + return True + + return False + + def _used_resources_string(self) -> str: + allocated_resources = self._actor_manager.get_live_actors_resources() + + return self._resource_updater.debug_string(allocated_resources) + + def on_step_begin(self): + self._resource_updater.update_avail_resources() + + def on_step_end(self): + self._cleanup_cached_actors(force_all=False) + self._cleanup_stopping_actors(force_all=False) + + def _cleanup_cached_actors(self, force_all: bool = False): + if ( + self._search_alg.is_finished() + and not self._staged_trials + and self._actor_cache.total_max_objects == 0 + ): + # If there are no more trials coming in, no trials are pending execution, + # and we don't explicitly want to cache objects, we can evict the full + # cache. + force_all = True + + for tracked_actor in self._actor_cache.flush_cached_objects( + force_all=force_all + ): + logger.debug(f"Cleaning up cached actor: {tracked_actor}") + # Unset termination callbacks as no trial is associated + tracked_actor.set_on_stop(None) + tracked_actor.set_on_error(None) + self._remove_actor(tracked_actor=tracked_actor) + + def _cleanup_stopping_actors(self, force_all: bool = False): + now = time.monotonic() + + if ( + not force_all + and now - self._earliest_stopping_actor <= self._actor_cleanup_timeout + ): + # If the earliest actor to timeout has not reached the timeout, return + return + + # This is a bit costly, so we want to avoid running it too often + times = deque( + sorted( + [ + (timestamp, tracked_actor) + for tracked_actor, timestamp in self._stopping_actors.items() + ], + key=lambda item: item[0], + ) + ) + + while times and ( + force_all or time.monotonic() - times[0][0] > self._actor_cleanup_timeout + ): + if ( + time.monotonic() - times[0][0] < self._actor_force_cleanup_timeout + ) and self._actor_manager.is_actor_started(tracked_actor=times[0][1]): + # Even if force_all=True, we give the actors time to clean up + self._actor_manager.next(timeout=1) + continue + + _, tracked_actor = times.popleft() + + if tracked_actor not in self._stopping_actors: + # Actor stopping has been handled by the block above + continue + + if self._actor_manager.is_actor_started(tracked_actor=tracked_actor): + logger.debug(f"Forcefully killing actor: {tracked_actor}") + self._actor_manager.remove_actor(tracked_actor=tracked_actor, kill=True) + self._stopping_actors.pop(tracked_actor) + + if times: + self._earliest_stopping_actor = times[0][0] + else: + self._earliest_stopping_actor = float("inf") + + def step(self): + if self.is_finished(): + raise TuneError("Called step when all trials finished?") + + with warn_if_slow("on_step_begin"): + self.on_step_begin() + + with warn_if_slow("callbacks.on_step_begin"): + self._callbacks.on_step_begin( + iteration=self._iteration, trials=self._trials + ) + + # Ask searcher for more trials + self._maybe_update_trial_queue() + + # Start actors for added trials + self._maybe_add_actors() + + # Handle one event + if not self._actor_manager.next(timeout=0.1): + # If there are no actors running, warn about potentially + # insufficient resources + if not self._actor_manager.num_live_actors: + self._insufficient_resources_manager.on_no_available_trials( + self.get_trials() + ) + + # Maybe stop whole experiment + self._stop_experiment_if_needed() + + # Maybe save experiment state + try: + self.checkpoint() + except Exception as e: + logger.warning(f"Trial controller checkpointing failed: {str(e)}") + raise e + + self._iteration += 1 + + with warn_if_slow("on_step_end"): + self.on_step_end() + with warn_if_slow("callbacks.on_step_end"): + self._callbacks.on_step_end(iteration=self._iteration, trials=self._trials) + + def _set_trial_status(self, trial: Trial, status: str): + """Set trial to a specific status. + + This will keep track of trials with specific statuses in sets. + + For PENDING and PAUSED trials we also keep a list of trials to be able + to retain FIFO ordering. See ``_maybe_add_actors`` for details. + + Lastly we also keep a mapping from resources to pending/paused trials + to be able to efficiently start trials for cached actors. + """ + current_status = trial.status + + if current_status == status: + logger.debug(f"Trial {trial} already has status {status}. Skipping update.") + return + + status_str_map = { + Trial.PENDING: self._pending_trials, + Trial.RUNNING: self._running_trials, + Trial.PAUSED: self._paused_trials, + Trial.TERMINATED: self._stopped_trials, + Trial.ERROR: self._failed_trials, + } + + logger.debug( + f"Setting status for trial {trial} from {current_status} to {status}" + ) + + assert trial in status_str_map[current_status], (trial, current_status) + assert trial not in status_str_map[status], (trial, status) + + status_str_map[current_status].remove(trial) + status_str_map[status].add(trial) + + # We keep a log for pending trials for FIFO scheduling. + # We do not need to remove from this list as we will just discard + # items that are in this list but not in the respective set. + if status == Trial.PENDING: + self._pending_trials_list.append(trial) + self._resources_to_pending_trials[trial.placement_group_factory].add(trial) + else: + self._resources_to_pending_trials[trial.placement_group_factory].discard( + trial + ) + + trial.set_status(status) + + def _get_trial_checkpoints(self) -> Dict[str, str]: + for trial in self._trials_to_cache: + self._trial_metadata[trial.trial_id] = trial.get_json_state() + self._trials_to_cache.clear() + return self._trial_metadata + + def _mark_trial_to_checkpoint(self, trial: Trial): + self._trials_to_cache.add(trial) + + ### + # UPDATE TRIALS + def _maybe_update_trial_queue(self): + """Ask the searcher for more trials.""" + if self._search_alg.is_finished(): + return + + dont_wait_for_trial = ( + self._pending_trials or self._running_trials or self._paused_trials + ) + + while len(self._pending_trials) < self._max_pending_trials: + if not self._update_trial_queue(blocking=not dont_wait_for_trial): + break + dont_wait_for_trial = True + + def _cleanup_trials(self): + logger.debug("CLEANING UP all trials") + + for tracked_actor in list(self._actor_to_trial): + trial = self._actor_to_trial[tracked_actor] + logger.debug( + f"Scheduling trial stop at end of experiment (trial {trial}): " + f"{tracked_actor}" + ) + self._schedule_trial_stop(trial) + + # Clean up cached actors now + self._cleanup_cached_actors(force_all=True) + + start = time.monotonic() + while time.monotonic() - start < 5 and self._actor_manager.num_total_actors: + if _dedup_logs("actor_manager_cleanup", str(start)): + logger.debug( + "Waiting for actor manager to clean up final state [dedup]" + ) + self._actor_manager.next(timeout=1) + + logger.debug("Force cleanup of remaining actors") + self._cleanup_stopping_actors(force_all=True) + + self._actor_manager.cleanup() + + def _remove_actor(self, tracked_actor: TrackedActor): + stop_future = self._actor_manager.schedule_actor_task( + tracked_actor, "stop", _return_future=True + ) + now = time.monotonic() + + if self._actor_manager.remove_actor( + tracked_actor, kill=False, stop_future=stop_future + ): + # If the actor was previously alive, track + self._stopping_actors[tracked_actor] = now + self._earliest_stopping_actor = min(self._earliest_stopping_actor, now) + + ### + # ADD ACTORS + def _maybe_add_actors(self) -> None: + """Add actors for pending and paused trials. + + For actors that have not been staged, yet, we request an actor. + + For actors that have been staged, already, we try to reuse a cached actor. + + First, we handle the trial that the scheduler chooses to run. + + Then, we handle all trials that are pending. + + Lastly, we see if we have cached actors that we can assign to a pending or + paused trial. This can be the case when a trial has not been staged, yet, + for instance because the number of staging trials was too large. + """ + + ### + # 1: Start trial that the scheduler wants to run + with warn_if_slow("choose_trial_to_run"): + trial_to_run = self._scheduler_alg.choose_trial_to_run(self._wrapped()) + + if trial_to_run: + if _dedup_logs("trial_to_run_chosen", trial_to_run.trial_id): + logger.debug( + f"Chose trial to run from scheduler: {trial_to_run} [dedup]" + ) + if ( + trial_to_run not in self._staged_trials + and trial_to_run not in self._trial_to_actor + ): + logger.debug(f"Staging trial to run: {trial_to_run}") + self._set_trial_status(trial_to_run, Trial.PENDING) + self._staged_trials.add(trial_to_run) + self._actor_cache.increase_max(trial_to_run.placement_group_factory) + # schedule_trial_actor also potentially uses cached actors + self._schedule_trial_actor(trial_to_run) + else: + # Otherwise, only try to use the cached actor + if _dedup_logs("trial_to_run_reuse", trial_to_run.trial_id): + logger.debug( + f"Trying to re-use actor for trial to run: {trial_to_run} " + f"[dedup]" + ) + self._maybe_reuse_cached_actor(trial_to_run) + + ### + # 2: Start trials that are PENDING + def _maybe_add_actors(candidates: List[Trial]): + new_candidates = [] + + while candidates: + if self._actor_manager.num_pending_actors >= self._max_pending_trials: + break + + trial = candidates.pop(0) + + # If the trial is part of the list, but not of the set, + # we just ignore it. Removing it from the list on status + # change is too expensive. + if trial not in self._pending_trials: + continue + + if trial in self._trial_to_actor: + new_candidates.append(trial) + continue + + if trial in self._staged_trials: + self._maybe_reuse_cached_actor(trial) + continue + + logger.debug(f"Scheduling actor for enqueued trial: {trial}") + self._staged_trials.add(trial) + self._actor_cache.increase_max(trial.placement_group_factory) + self._schedule_trial_actor(trial) + + return new_candidates + candidates + + self._pending_trials_list = _maybe_add_actors(self._pending_trials_list) + + ### + # 3: Start any trial that can be started with a cached actor + if self._actor_cache.num_cached_objects: + for resource in self._resources_to_pending_trials: + if not self._resources_to_pending_trials[resource]: + continue + + if not self._actor_cache.has_cached_object(resource): + continue + + start_trial = self._resources_to_pending_trials[resource].pop() + logger.debug( + f"Trying to re-use actor for enqueued trial: {start_trial}" + ) + if not self._maybe_reuse_cached_actor(start_trial): + self._resources_to_pending_trials[resource].add(start_trial) + else: + if start_trial not in self._staged_trials: + self._staged_trials.add(start_trial) + self._actor_cache.increase_max( + start_trial.placement_group_factory + ) + + def _maybe_reuse_cached_actor(self, trial: Trial) -> bool: + """Maybe reuse a cached actor for a trial. + + If an actor has been scheduled for the trial already, + this will remove the original actor. + """ + if trial in self._resetting_trials: + return True + + resource_request = trial.placement_group_factory + + if not self._actor_cache.has_cached_object(resource_request): + return False + + cached_actor = self._actor_cache.pop_cached_object(resource_request) + logger.debug(f"Reusing ACTOR for trial {trial}: {cached_actor}") + + if trial in self._trial_to_actor: + original_actor = self._trial_to_actor.pop(trial) + self._actor_to_trial.pop(original_actor) + + logger.debug(f"Removing ORIGINAL ACTOR for trial {trial}: {original_actor}") + self._remove_actor(tracked_actor=original_actor) + + self._trial_to_actor[trial] = cached_actor + self._actor_to_trial[cached_actor] = trial + + # Todo: get rid of Trial.runner + ray_actor = self._actor_manager._live_actors_to_ray_actors_resources[ + cached_actor + ][0] + trial.set_ray_actor(ray_actor) + + self._schedule_trial_reset(trial, trial.config, trial.experiment_tag) + + return True + + def _schedule_trial_actor(self, trial: Trial): + """Schedule an actor for a trial. + + If a cached actor is available, use it. Otherwise, request a + new actor. + """ + logger.debug(f"Trying to schedule new ACTOR for trial {trial}") + + assert trial.status == Trial.PENDING + + trial.init_local_path() + # We checkpoint metadata here to try mitigating logdir duplication + self._mark_trial_to_checkpoint(trial) + + if self._maybe_reuse_cached_actor(trial): + return + + # Safeguard + if trial in self._trial_to_actor: + raise RuntimeError( + f"Tried to request a new actor for trial {trial}, but an old " + f"actor still exists. This can lead to leaked resources. The old " + f"actor should be removed first. " + f"This is an internal problem in Ray Tune. If you encounter this " + f"error, please raise an issue on " + f"https://github.com/ray-project/ray/issues" + ) + + trainable_cls = trial.get_trainable_cls() + if not trainable_cls: + exception = _AbortTrialExecution( + f"Invalid trainable: {trial.trainable_name}. If you passed " + f"a string, make sure the trainable was registered before." + ) + trial.handle_error(exception) + self._schedule_trial_stop(trial, exception=exception) + return + + _actor_cls = self._class_cache.get(trainable_cls) + + trial.set_location(_Location()) + trainable_kwargs = _get_trainable_kwargs(trial=trial) + + with _change_working_directory(trial): + tracked_actor = self._actor_manager.add_actor( + cls=_actor_cls, + resource_request=trial.placement_group_factory, + kwargs=trainable_kwargs, + on_start=self._actor_started, + on_stop=self._actor_stopped, + on_error=self._actor_failed, + ) + self._trial_to_actor[trial] = tracked_actor + self._actor_to_trial[tracked_actor] = trial + + logger.debug( + f"Scheduled new ACTOR for trial {trial}: {tracked_actor}. " + f"Resources: {trial.placement_group_factory}" + ) + + def _unstage_trial_with_resources(self, trial: Trial): + """Unstage trial, or one with the same resources as ``trial``.""" + # Case 1: The trial we started was staged. Just remove it + if trial in self._staged_trials: + self._staged_trials.remove(trial) + self._actor_cache.decrease_max(trial.placement_group_factory) + return + + # Case 2: We staged a trial "A" with the same resources, but our trial "B" + # was selected by the scheduler to run. The resource manager does not care + # about "trials", it just cares about resources being available. Thus we + # look for a staged trial with the same resource requirements and remove it + + resource_request = trial.placement_group_factory + # Remove staged trial with same resource requirements + candidate_trial = None + for staged_trial in self._staged_trials: + staged_resources = staged_trial.placement_group_factory + if staged_resources == resource_request: + candidate_trial = staged_trial + break + + if candidate_trial: + self._staged_trials.remove(candidate_trial) + self._actor_cache.decrease_max(candidate_trial.placement_group_factory) + return + + raise RuntimeError( + "Started a trial with resources requested by a different trial, but " + "this trial was lost. This is an error in Ray Tune's execution " + "logic. Please raise a GitHub issue at " + "https://github.com/ray-project/ray/issues" + ) + + def _maybe_cache_trial_actor(self, trial: Trial) -> bool: + """Cache trial actor for reuse, if needed. + + We will only cache as many actors as are needed to fulfill any pending + resource requests for actors with the same resource requirements. + E.g. if we have 6 running trials and 4 additional staged actors, we will only + cache up to 4 of the running trial actors when they finish. + + One exception is the case when we have no cached actors, yet. In that case, + we will always cache the actor in this method. + + Later, in `_cleanup_cached_actors`, we will check again if we need this cached + actor. That method will keep the actor if we don't have any staged trials, + because we don't know at that point if the next trial might require the same + resources. But because there is no staged trial, it is safe to keep the actor + around, as it won't occupy resources needed by another trial until it's staged. + """ + if not self._reuse_actors: + return False + + if self._search_alg.is_finished() and not self._staged_trials: + logger.debug( + f"Not caching actor of trial {trial} as the search is over " + f"and no more trials are staged." + ) + return False + + tracked_actor = self._trial_to_actor[trial] + + if ( + not self._actor_manager.is_actor_started(tracked_actor) + or self._actor_manager.is_actor_failed(tracked_actor) + or tracked_actor not in self._started_actors + ): + logger.debug( + f"Not caching actor of trial {trial} as it has not been started, yet: " + f"{tracked_actor}" + ) + return False + + if not self._actor_cache.cache_object( + trial.placement_group_factory, tracked_actor + ): + logger.debug( + f"Could not cache actor of trial {trial} for " + "reuse, as there are no pending trials " + "requiring its resources." + ) + return False + + logger.debug(f"Caching actor of trial {trial} for re-use: {tracked_actor}") + + tracked_actor = self._trial_to_actor.pop(trial) + self._actor_to_trial.pop(tracked_actor) + + trial.set_ray_actor(None) + + return True + + def _actor_started(self, tracked_actor: TrackedActor, log: str = "STARTED"): + self._started_actors.add(tracked_actor) + + trial = self._actor_to_trial[tracked_actor] + + logger.debug(f"Actor {log} for trial {trial}: {tracked_actor}") + + self._unstage_trial_with_resources(trial) + + ray_actor = self._actor_manager._live_actors_to_ray_actors_resources[ + tracked_actor + ][0] + trial.set_ray_actor(ray_actor) + + self._callbacks.on_trial_start( + iteration=self._iteration, trials=self._trials, trial=trial + ) + + self._set_trial_status(trial, Trial.RUNNING) + + self._mark_trial_to_checkpoint(trial) + + if not self._schedule_trial_restore(trial): + self._schedule_trial_train(trial) + + def _actor_stopped(self, tracked_actor: TrackedActor): + if tracked_actor in self._actor_to_trial: + trial = self._actor_to_trial.pop(tracked_actor) + logger.debug(f"Actor STOPPED for trial {trial}: {tracked_actor}") + self._trial_to_actor.pop(trial) + trial.set_ray_actor(None) + + logger.debug(f"Actor STOPPED: {tracked_actor}") + + self._stopping_actors.pop(tracked_actor, None) + self._started_actors.discard(tracked_actor) + + def _actor_failed(self, tracked_actor: TrackedActor, exception: Exception): + trial = self._actor_to_trial[tracked_actor] + + logger.debug( + f"Actor FAILED for trial {trial}: {tracked_actor}. " + f"Exception: {exception}" + ) + + if trial in (self._pending_trials | self._paused_trials): + # First, set to running (needed downstream in _process_trial_failure) + self._set_trial_status(trial, Trial.RUNNING) + + logger.debug( + f"Trial {trial} failed in its creation task. Unstaging " + f"to allow it to be re-scheduled." + ) + + self._unstage_trial_with_resources(trial) + self._trial_task_failure(trial, exception=exception) + + self._actor_manager.clear_actor_task_futures(tracked_actor) + + # Clean up actor + tracked_actor.set_on_stop(None) + tracked_actor.set_on_error(None) + self._actor_manager.remove_actor(tracked_actor, kill=False) + + # Trigger actor stopped callback + self._actor_stopped(tracked_actor) + + def _schedule_trial_task( + self, + trial: Trial, + method_name: str, + args: Optional[Tuple] = None, + kwargs: Optional[Dict] = None, + on_result: Optional[Callable[[Trial, Any], None]] = None, + on_error: Optional[Callable[[Trial, Exception], None]] = None, + _return_future: bool = False, + ) -> Optional[ray.ObjectRef]: + """Schedule an actor task future for a trial. + + This is a wrapper around ``ActorManager.schedule_actor_task``. This method + retrieves the tracked actor for a trial to kick off the task. + + It also wraps around the callbacks, retrieving the trial object given the + tracked actor. + """ + + tracked_actor = self._trial_to_actor[trial] + + _on_result = None + _on_error = None + + args = args or tuple() + kwargs = kwargs or {} + + if on_result: + + def _on_result(tracked_actor: TrackedActor, *args, **kwargs): + assert trial == self._actor_to_trial[tracked_actor] + logger.debug( + f"Future {method_name.upper()} RESOLVED for trial {trial}: " + f"{args}, {kwargs}" + ) + try: + on_result(trial, *args, **kwargs) + except Exception as e: + logger.debug( + f"Error handling {method_name.upper()} result " + f"for trial {trial}: {e}" + ) + if e is TuneError or self._fail_fast == self.RAISE: + raise e + else: + raise TuneError(traceback.format_exc()) + + if on_error: + + def _on_error(tracked_actor: TrackedActor, exception: Exception): + # If the actor failed, it has already been cleaned up. + if tracked_actor not in self._actor_to_trial: + assert isinstance(exception, RayActorError), type(exception) + else: + assert trial == self._actor_to_trial[tracked_actor] + + logger.debug( + f"Future {method_name.upper()} FAILED for trial {trial}: " + f"{exception}" + ) + try: + on_error(trial, exception) + except Exception as e: + logger.debug( + f"Error handling {method_name.upper()} failure " + f"for trial {trial}: {e}" + ) + if e is TuneError or self._fail_fast == self.RAISE: + raise e + else: + raise TuneError(traceback.format_exc()) + + logger.debug(f"Future {method_name.upper()} SCHEDULED for trial {trial}") + + with _change_working_directory(trial): + future = self._actor_manager.schedule_actor_task( + tracked_actor=tracked_actor, + method_name=method_name, + args=args, + kwargs=kwargs, + on_result=_on_result, + on_error=_on_error, + _return_future=_return_future, + ) + if _return_future: + return future + + def _queue_decision(self, trial, decision): + # Get old decision, setting it to the current decision if it isn't set + old_decision = self._queued_trial_decisions.setdefault(trial.trial_id, decision) + + # Stopping always takes precedence. If we decided to stop, just quit + if old_decision is TrialScheduler.STOP: + return + + # The old decision wasn't STOP. We update the decision only if it is + # STOP or PAUSE. The action will only be CONTINUE if it was set by + # the first received result and was never updated after that. + if decision is TrialScheduler.STOP or decision is TrialScheduler.PAUSE: + self._queued_trial_decisions[trial.trial_id] = decision + + def _execute_action(self, trial: Trial, decision: str, after_save: bool = False): + """Executes action based on decision. + + Args: + trial: Trial to act on. + decision: Scheduling decision to undertake. + """ + if decision == TrialScheduler.CONTINUE: + self._schedule_trial_train(trial) + elif decision == TrialScheduler.PAUSE: + self.pause_trial(trial, should_checkpoint=not after_save) + elif decision == TrialScheduler.STOP: + self.stop_trial(trial) + elif decision == TrialScheduler.NOOP: + pass + else: + raise ValueError("Invalid decision: {}".format(decision)) + + def _maybe_execute_queued_decision(self, trial: Trial, after_save: bool = False): + # `self._queued_trial_decisions` now contains a final decision + # based on all results + final_decision = self._queued_trial_decisions.pop(trial.trial_id, None) + if final_decision: + logger.debug( + f"Executing final queued decision for {trial}: {final_decision}" + ) + self._execute_action(trial, final_decision, after_save=after_save) + + def _stop_experiment_if_needed(self): + """Stops all trials.""" + fail_fast = self._fail_fast and self._has_errored + if self._stopper.stop_all() or fail_fast or self._should_stop_experiment: + self._search_alg.set_finished() + [ + self._schedule_trial_stop(t) + for t in self._trials + if t.status not in {Trial.ERROR, Trial.TERMINATED} + ] + + ### + # Failure + def _trial_task_failure(self, trial: Trial, exception: Exception): + if self._fail_fast == self.RAISE: + raise exception + else: + if self._print_trial_errors: + logger.error(f"Trial task failed for trial {trial}", exc_info=exception) + self._process_trial_failure(trial, exception=exception) + + def _process_trial_failure( + self, + trial: Trial, + exception: Union[TuneError, RayTaskError, RayActorError], + ): + """Handle trial failure. + + Attempt trial recovery if possible, clean up state otherwise. + + Args: + trial: Failed trial. + exception: Exception prior to invoking this method. + """ + self._has_errored = True + trial.handle_error(exception) + if trial.status == Trial.RUNNING and trial.should_recover(): + self._try_recover(trial, exc=exception) + self._callbacks.on_trial_recover( + iteration=self._iteration, trials=self._trials, trial=trial + ) + elif trial.status in {Trial.RUNNING, Trial.PENDING}: + self._scheduler_alg.on_trial_error(self, trial) + self._search_alg.on_trial_complete(trial.trial_id, error=True) + self._schedule_trial_stop(trial, exception=exception) + self._callbacks.on_trial_error( + iteration=self._iteration, trials=self._trials, trial=trial + ) + + def _schedule_trial_stop(self, trial: Trial, exception: Optional[Exception] = None): + if trial.status == Trial.ERROR: + logger.debug(f"Not requesting trial STOP as it is ERROR already: {trial}") + return + + logger.debug(f"Requesting to STOP actor for trial {trial}") + + if trial.is_saving: + logger.debug( + f"Trial {trial} is currently saving/pausing. Scheduling STOP after " + f"save resolved." + ) + self._cached_trial_decisions[trial.trial_id] = TrialScheduler.STOP + + trial.temporary_state.saving_to = None + trial.temporary_state.restoring_from = None + + self._set_trial_status(trial, Trial.ERROR if exception else Trial.TERMINATED) + trial.set_location(_Location()) + + if trial not in self._trial_to_actor: + logger.debug(f"Will not STOP trial actor as it is not live: {trial}") + return + + tracked_actor = self._trial_to_actor[trial] + + self._actor_manager.clear_actor_task_futures(tracked_actor=tracked_actor) + + self._mark_trial_to_checkpoint(trial) + + if not exception and self._maybe_cache_trial_actor(trial): + # Trial runner has been cached + return + + logger.debug(f"Terminating actor for trial {trial}: {tracked_actor}") + + tracked_actor = self._trial_to_actor.pop(trial) + self._actor_to_trial.pop(tracked_actor) + + trial.set_ray_actor(None) + + self._remove_actor(tracked_actor=tracked_actor) + + def stop_trial(self, trial): + """The canonical implementation of stopping a trial. + + Trials may be in any external status when this function is called. + If trial is in state PENDING or PAUSED, calls `on_trial_remove` for + scheduler and `on_trial_complete()` for search_alg. + If trial is in state RUNNING, calls `on_trial_complete` for scheduler + and search_alg if RUNNING. Caller to ensure that there is no + outstanding future to be handled for the trial. If there is, the future + would be discarded. + """ + try: + if trial.status in [Trial.ERROR, Trial.TERMINATED]: + return + elif trial.status in [Trial.PENDING, Trial.PAUSED]: + self._scheduler_alg.on_trial_remove(self, trial) + self._search_alg.on_trial_complete(trial.trial_id) + elif trial.status is Trial.RUNNING: + # By this time trial.last_result should have been + # updated already. + self._scheduler_alg.on_trial_complete( + self, trial, flatten_dict(trial.last_result) + ) + self._search_alg.on_trial_complete( + trial.trial_id, result=flatten_dict(trial.last_result) + ) + self._callbacks.on_trial_complete( + iteration=self._iteration, trials=self._trials, trial=trial + ) + self._schedule_graceful_trial_stop(trial) + self._live_trials.discard(trial) + except Exception as e: + logger.exception("Trial %s: Error stopping trial.", trial) + if self._fail_fast == self.RAISE: + raise + if isinstance(e, TuneError): + self._process_trial_failure(trial, exception=e) + else: + self._process_trial_failure( + trial, _TuneStopTrialError(traceback.format_exc()) + ) + + def _schedule_graceful_trial_stop(self, trial: Trial): + self._schedule_trial_export(trial) + if trial.status != "ERROR": + self._schedule_trial_stop(trial) + + def _schedule_trial_pause(self, trial: Trial, should_checkpoint: bool = True): + if trial not in self._trial_to_actor: + logger.debug( + f"Trial PAUSE requested for trial {trial} but trial is already " + f"stopping. Ignoring." + ) + return + + if should_checkpoint: + self._cached_trial_decisions[trial.trial_id] = TrialScheduler.PAUSE + self._schedule_trial_save(trial=trial) + else: + self._schedule_trial_stop(trial) + self._set_trial_status(trial, Trial.PAUSED) + + ### + # TRAIN + + def _schedule_trial_train(self, trial: Trial): + args = () + method_name = "train" + + buffer_length, buffer_time_s = self._maybe_buffer_training(trial) + + if buffer_length > 1: + method_name = "train_buffered" + args = (buffer_length, buffer_time_s) + + logger.debug(f"Scheduling future {method_name.upper()} for trial {trial}") + + self._schedule_trial_task( + trial=trial, + method_name=method_name, + args=args, + on_result=self._on_training_result, + on_error=self._trial_task_failure, + ) + + def _maybe_buffer_training(self, trial: Trial) -> Tuple[int, float]: + buffer_time_s = max( + self._buffer_min_time_s, + min(self._buffer_max_time_s, self._actor_manager.num_actor_tasks // 10), + ) + buffer_length = self._buffer_length + + if buffer_length > 1 and trial.checkpoint_at_end: + # If a trial checkpoint can be triggered externally, + # it is not safe to buffer results. + if log_once("trial_executor_buffer_checkpoint"): + logger.warning( + "Disabling buffered training as you passed " + "`checkpoint_at_end` to `tune.CheckpointConfig()`." + ) + return 1, buffer_time_s + + if buffer_length > 1 and trial.checkpoint_freq > 0: + return min(buffer_length, trial.checkpoint_freq), buffer_time_s + + return buffer_length, buffer_time_s + + ### + # RESULT + + def _on_training_result(self, trial, result): + if not isinstance(result, list): + result = [result] + with warn_if_slow("process_trial_result"): + self._process_trial_results(trial, result) + self._maybe_execute_queued_decision(trial, after_save=False) + + def _process_trial_results(self, trial, results): + logger.debug(f"Processing trial results for trial {trial}: {results}") + with warn_if_slow( + "process_trial_results", + message="Processing trial results took {duration:.3f} s, " + "which may be a performance bottleneck. Please consider " + "reporting results less frequently to Ray Tune.", + ): + for i, result in enumerate(results): + with warn_if_slow("process_trial_result"): + decision = self._process_trial_result(trial, result) + if decision is None: + # If we didn't get a decision, this means a + # non-training future (e.g. a save) was scheduled. + # We do not allow processing more results then. + if i < len(results) - 1: + if log_once("tune_controller_buffer_checkpoint"): + logger.warning( + f"Trial {trial} has a non-training future " + f"scheduled but {len(results) - i} results " + f"left to process. This means that a " + f"checkpoint was requested, but buffered " + f"training was continued before it was " + f"saved. Consider using non-buffered " + f"training by setting the env variable " + f"`TUNE_RESULT_BUFFER_LENGTH=1`." + ) + elif decision == TrialScheduler.STOP: + # If the decision is to stop the trial, + # ignore all results that came after that. + break + + def _process_trial_result(self, trial: Trial, result: dict[str, Any]): + result.update(trial_id=trial.trial_id) + is_duplicate = RESULT_DUPLICATE in result + force_checkpoint = False + + # TrialScheduler and SearchAlgorithm still receive a + # notification because there may be special handling for + # the `on_trial_complete` hook. + if is_duplicate: + logger.debug("Trial finished without logging 'done'.") + result = trial.last_result + result.update(done=True) + + self._total_time += result.get(TIME_THIS_ITER_S, 0) + + flat_result = flatten_dict(result) + self._validate_result_metrics(flat_result) + + if self._stopper(trial.trial_id, result) or trial.should_stop(flat_result): + decision = TrialScheduler.STOP + else: + with warn_if_slow("scheduler.on_trial_result"): + decision = self._scheduler_alg.on_trial_result( + self._wrapped(), trial, flat_result + ) + if decision == TrialScheduler.STOP: + result.update(done=True) + else: + # Only updating search alg if the trial is not to be stopped. + with warn_if_slow("search_alg.on_trial_result"): + self._search_alg.on_trial_result(trial.trial_id, flat_result) + + # If this is not a duplicate result, the callbacks should + # be informed about the result. + if not is_duplicate: + with warn_if_slow("callbacks.on_trial_result"): + self._callbacks.on_trial_result( + iteration=self._iteration, + trials=self._trials, + trial=trial, + # NOTE: Allow user callbacks to modify the Trial result in place. + result=result, + ) + force_checkpoint = result.get(SHOULD_CHECKPOINT, False) + trial.update_last_result(result) + # Include in next experiment checkpoint + self._mark_trial_to_checkpoint(trial) + + # Checkpoints to disk. This should be checked even if + # the scheduler decision is STOP or PAUSE. Note that + # PAUSE only checkpoints to memory and does not update + # the global checkpoint state. + if decision != TrialScheduler.PAUSE: + # TODO(justinvyu): This is a temporary hack to fix pausing trials. + # We already schedule a save task in `pause_trial`, so no need + # to do it again here. + self._checkpoint_trial_if_needed(trial, force=force_checkpoint) + + if trial.is_saving: + logger.debug(f"Caching trial decision for trial {trial}: {decision}") + # Cache decision to execute on after the save is processed. + # This prevents changing the trial's state or kicking off + # another training step prematurely. + if not self._cached_trial_decisions.get(trial.trial_id) or decision in { + TrialScheduler.PAUSE, + TrialScheduler.STOP, + }: + # If already set, only overwrite if it's a PAUSE or STOP. This is + # to avoid that CONTINUE decisions from a training step that resolve + # late overwrite PAUSE/STOP decision. + self._cached_trial_decisions[trial.trial_id] = decision + return None + else: + self._queue_decision(trial, decision) + return decision + + def _validate_result_metrics(self, result): + """ + Check if any of the required metrics was not reported + in the last result. If the only items are ``done`` or any of + DEBUG_METRICS, this means that no result was ever received and + the trial just returned. This is also okay and will not raise + an error. + + This will ignore checking for the DEFAULT_METRIC. + """ + if int(os.environ.get("TUNE_DISABLE_STRICT_METRIC_CHECKING", 0)) != 1 and ( + len({k for k in result if k not in list(DEBUG_METRICS) + [DONE]}) > 1 + ): + base_metric = self._metric if self._metric != DEFAULT_METRIC else None + scheduler_metric = ( + self._scheduler_alg.metric + if self._scheduler_alg.metric != DEFAULT_METRIC + else None + ) + search_metrics = ( + self._search_alg.metric + if self._search_alg.metric != DEFAULT_METRIC + else None + ) + + if isinstance(search_metrics, str): + search_metrics = [search_metrics] + + if base_metric and base_metric not in result: + report_metric = base_metric + location = "tune.TuneConfig()" + elif scheduler_metric and scheduler_metric not in result: + report_metric = scheduler_metric + location = type(self._scheduler_alg).__name__ + elif search_metrics and any( + search_metric not in result for search_metric in search_metrics + ): + report_metric = list( + filter( + lambda search_metric: search_metric not in result, + search_metrics, + ) + ) + if len(report_metric) == 1: + report_metric = report_metric[0] + location = type(self._search_alg).__name__ + else: + report_metric = None + location = None + + if report_metric: + raise ValueError( + "Trial returned a result which did not include the " + "specified metric(s) `{}` that `{}` expects. " + "Make sure your calls to `tune.report()` include the " + "metric, or set the " + "TUNE_DISABLE_STRICT_METRIC_CHECKING " + "environment variable to 1. Result: {}".format( + report_metric, location, result + ) + ) + + ### + # SAVE + def _schedule_trial_save( + self, + trial: Trial, + result: Optional[Dict] = None, + ) -> Optional[_FutureTrainingResult]: + if trial not in self._trial_to_actor: + logger.debug( + f"Trial SAVE requested for trial {trial} but trial is already " + f"stopping. Ignoring." + ) + return None + + result = result or trial.last_result + + future = self._schedule_trial_task( + trial=trial, + method_name="save", + on_result=self._on_saving_result, + on_error=self._trial_task_failure, + _return_future=True, + ) + # TODO(justinvyu): `trial.saving_to` (and trial.is_saving) is needed + # in order to prevent a done=True result from executing a STOP decision + # (which clears all futures) before the save gets processed. + # Keep this in for now while `train` and `save` are 2 separate steps. + trial.temporary_state.saving_to = _FutureTrainingResult(future) + + # `trial.saving_to` holds a future training result -- this is only used + # in the case of PBT to block until the checkpoint is ready. + # In all other situations, the checkpoint future is processed by the + # actor event manager when it is ready. + return trial.temporary_state.saving_to + + def _on_saving_result(self, trial, checkpoint_value: _TrainingResult): + with warn_if_slow("process_trial_save"): + self._process_trial_save(trial, checkpoint_value) + + with warn_if_slow("callbacks.on_trial_save"): + self._callbacks.on_trial_save( + iteration=self._iteration, trials=self._trials, trial=trial + ) + + self._maybe_execute_queued_decision(trial, after_save=True) + + def _process_trial_save(self, trial: Trial, checkpoint_value: _TrainingResult): + """Processes a trial save. + + Acts on the decision cached during the last `_process_trial` call. + + Args: + trial: Trial being saved. + """ + logger.debug("Trial %s: Processing trial save.", trial) + + try: + if not checkpoint_value.checkpoint: + logger.debug(f"Got empty checkpoint for trial {trial}") + else: + try: + self._callbacks.on_checkpoint( + iteration=self._iteration, + trials=self._trials, + trial=trial, + checkpoint=checkpoint_value.checkpoint, + ) + except Exception: + logger.warning( + "Error encountered during processing of callbacks. " + "Ray Train/Tune recently changed the checkpoint interface " + "that is passed to callbacks. If you implemented your own " + "callback with an `on_checkpoint` handler, please review " + "the checkpoint interface and adjust your code " + "accordingly." + ) + raise + + trial.on_checkpoint(checkpoint_value) + + self._checkpoint_manager.on_trial_checkpoint(trial) + + self._mark_trial_to_checkpoint(trial) + except Exception: + logger.exception( + "Trial %s: Error handling checkpoint %s", trial, checkpoint_value + ) + + trial.temporary_state.saving_to = None + decision = self._cached_trial_decisions.pop(trial.trial_id, None) + if decision and checkpoint_value: + self._queue_decision(trial, decision) + + def _checkpoint_trial_if_needed(self, trial, force=False): + """Checkpoints trial based off trial.last_result.""" + if trial.should_checkpoint() or force: + # Save trial runtime if possible. + if trial.temporary_state.ray_actor: + self._schedule_trial_save(trial) + + ### + # RESTORE + def _schedule_trial_restore(self, trial: Trial) -> bool: + checkpoint_result = trial.latest_checkpoint_result + + if not checkpoint_result: + logger.debug(f"Not restoring trial {trial}: No checkpoint found.") + return False + + # TODO(justinvyu): Is this really needed? + trial.temporary_state.restoring_from = checkpoint_result + + method_name = "restore" + args = (checkpoint_result,) + self._schedule_trial_task( + trial=trial, + method_name=method_name, + args=args, + kwargs={}, + on_result=self._on_restoring_result, + on_error=self._trial_task_failure, + ) + return True + + def _on_restoring_result(self, trial: Trial, result: Any): + self._process_trial_restore(trial) + + def _process_trial_restore(self, trial: Trial): + """Processes a trial restore. + + Args: + trial: Trial being restored. + """ + logger.debug("Trial %s: Processing trial restore.", trial) + trial.on_restore() + logger.debug("Trial %s: Restore processed successfully", trial) + self._set_trial_status(trial, Trial.RUNNING) + self._schedule_trial_train(trial) + self._live_trials.add(trial) + + def _try_recover( + self, trial: Trial, exc: Union[TuneError, RayTaskError, RayActorError] + ): + """Tries to recover trial. + + Notifies SearchAlgorithm and Scheduler if failure to recover. + + Args: + trial: Trial to recover. + exc: Exception prior to invoking this method. + """ + self._cached_trial_decisions.pop(trial.trial_id, None) + # Resetting this, in case that the trial is in saving status when it crashes. + if trial.is_saving: + trial.temporary_state.saving_to = None + self._schedule_trial_stop(trial, exception=exc) + + logger.debug("Trial %s: Notifying Scheduler and requeueing.", trial) + self._requeue_trial(trial) + + def _requeue_trial(self, trial): + """Notification to TrialScheduler and requeue trial. + + This does not notify the SearchAlgorithm because the function + evaluation is still in progress. + + """ + self._scheduler_alg.on_trial_error(self, trial) + self._set_trial_status(trial, status=Trial.PENDING) + + # TODO(rliaw): Right now, this pushes the trial to the end of queue + # because restoration can be expensive. However, this is not + # ideal since it just hides the issue - a better fix would + # be to use an actor table to detect the IP of the Trainable + # and rsync the files there. + # See https://github.com/ray-project/ray/issues/5168 + self._trials.pop(self._trials.index(trial)) + self._trials.append(trial) + self._live_trials.add(trial) + + with warn_if_slow("scheduler.on_trial_add"): + self._scheduler_alg.on_trial_add(self._wrapped(), trial) + + ### + # EXPORT + def _schedule_trial_export(self, trial: Trial): + if not trial.export_formats or len(trial.export_formats) <= 0: + return + + # Todo: We are waiting here synchronously until the task resolved. + # Instead, we should schedule the trial stop after the export resolved. + # This requires changes in TrialRunner, which we can remove once the + # legacy execution path has been removed. + future = self._schedule_trial_task( + trial=trial, + method_name="export_model", + args=(trial.export_formats,), + on_result=None, + on_error=self._trial_task_failure, + _return_future=True, + ) + self._actor_manager._actor_task_events.resolve_future(future) + + ### + # RESET + def _schedule_trial_reset( + self, + trial: Trial, + new_config: Dict, + new_experiment_tag: str, + ): + trial.set_experiment_tag(new_experiment_tag) + trial.set_config(new_config) + + # Pass magic variables + extra_config = copy.deepcopy(new_config) + extra_config[TRIAL_INFO] = _TrialInfo(trial) + + stdout_file, stderr_file = trial.log_to_file + extra_config[STDOUT_FILE] = stdout_file + extra_config[STDERR_FILE] = stderr_file + + logger_creator = partial( + _noop_logger_creator, logdir=trial.storage.trial_working_directory + ) + + self._resetting_trials.add(trial) + self._schedule_trial_task( + trial=trial, + method_name="reset", + args=(extra_config,), + kwargs={ + "logger_creator": logger_creator, + "storage": trial.storage, + }, + on_result=self._on_trial_reset, + on_error=self._trial_task_failure, + ) + + def _on_trial_reset(self, trial: Trial, success: bool): + self._resetting_trials.remove(trial) + + if not success: + info = ( + "Trainable runner reuse requires reset_config() to be " + "implemented and return True." + ) + + logger.error(f"Could not re-use actor for trial {trial}: {info}") + + exception = _AbortTrialExecution(info) + + trial.handle_error(exception) + self._schedule_trial_stop(trial, exception=exception) + return + + tracked_actor = self._trial_to_actor[trial] + + self._actor_started(tracked_actor, log="REUSED") + + def request_stop_trial(self, trial): + self._stop_queue.append(trial) + + def request_stop_experiment(self): + self._should_stop_experiment = True + + def _process_stop_requests(self): + while self._stop_queue: + t = self._stop_queue.pop() + self.stop_trial(t) + + def pause_trial(self, trial: Trial, should_checkpoint: bool = True): + """Pause a trial and reset the necessary state variables for resuming later. + + Args: + trial: Trial to pause. + should_checkpoint: Whether or not an in-memory checkpoint should be created + for this paused trial. Defaults to True. + """ + # NOTE: The cached trial decision is not needed since we will overrule this + # decision with PAUSE. + self._cached_trial_decisions.pop(trial.trial_id, None) + self._schedule_trial_pause(trial, should_checkpoint=should_checkpoint) + + def cleanup(self): + """Cleanup trials and callbacks.""" + self._cleanup_trials() + self.end_experiment_callbacks() + + def __getstate__(self): + """Gets state for trial. + + Note that this is not used as a pickling override as + does not have all fields. + """ + state = self.__dict__.copy() + for k in [ + "_trials", + "_live_trials", + "_stop_queue", + "_search_alg", + "_placeholder_resolvers", + "_scheduler_alg", + "_pending_trial_queue_times", + "_callbacks", + "_checkpoint_manager", + "_storage", + "_insufficient_resources_manager", + "_actor_manager", + "_class_cache", + "_resource_updater", + "_trials_to_cache", + "_trial_metadata", + "_actor_to_trial", + "_trial_to_actor", + "_resources_to_pending_trials", + "_pending_trials", + "_pending_trials_list", + "_running_trials", + "_paused_trials", + "_stopped_trials", + "_failed_trials", + "_resetting_trials", + "_started_actors", + "_stopping_actors", + "_staged_trials", + "_actor_cache", + ]: + del state[k] + return state + + def __setstate__(self, state): + # Use session_str from previous checkpoint if does not exist + session_str = state.pop("_session_str") + self.__dict__.setdefault("_session_str", session_str) + # Use start_time from previous checkpoint if does not exist + start_time = state.pop("_start_time") + self.__dict__.setdefault("_start_time", start_time) + + self.__dict__.update(state) + self._checkpoint_manager = self._create_checkpoint_manager() + + +class _TrialExecutorWrapper: + """Wraps around TrialExecutor class, intercepts API calls and warns users + of restricted API access. + + This is meant to facilitate restricting + the current API exposure of TrialExecutor by TrialScheduler. + """ + + def __init__( + self, + trial_executor: "_FakeRayTrialExecutor", + whitelist_attr: Optional[set] = None, + ): + self._trial_executor = trial_executor + self._whitelist_attr = whitelist_attr or set() + + for attr in self._whitelist_attr: + assert hasattr(self._trial_executor, attr) + + def __getattr__(self, attr): + if attr not in self._whitelist_attr: + if log_once("restrict_accessing_trial_executor"): + logger.warning( + f"You are trying to access {attr} interface of " + f"TrialExecutor in TrialScheduler, which is being " + f"restricted. If you believe it is reasonable for " + f"your scheduler to access this TrialExecutor API, " + f"please reach out to Ray team on GitHub. A more " + f"strict API access pattern would be enforced " + f"starting 1.12.0" + ) + return getattr(self._trial_executor, attr) + + +@DeveloperAPI +class TrialRunnerWrapper: + """Wraps around TrialRunner class, intercepts API calls and warns users + of restricted API access. + + This is meant to facilitate restricting + the current API exposure of TrialRunner by TrialScheduler. + """ + + _EXECUTOR_ATTR = "trial_executor" + + def __init__( + self, + tune_controller: TuneController, + trial_executor: Any, + runner_whitelist_attr: Optional[set] = None, + executor_whitelist_attr: Optional[set] = None, + ): + self._tune_controller = tune_controller + self._trial_executor = _TrialExecutorWrapper( + trial_executor, executor_whitelist_attr + ) + self._runner_whitelist_attr = runner_whitelist_attr or set() + + for attr in self._runner_whitelist_attr: + assert hasattr(self, attr) + + def __getattr__(self, attr): + if attr == self._EXECUTOR_ATTR: + return self._trial_executor + if attr not in self._runner_whitelist_attr: + if log_once("restrict_accessing_tune_controller"): + logger.warning( + f"You are trying to access {attr} interface of " + f"TrialRunner in TrialScheduler, which is being " + f"restricted. If you believe it is reasonable for " + f"your scheduler to access this TrialRunner API, " + f"please reach out to Ray team on GitHub. A more " + f"strict API access pattern would be enforced " + f"starting 1.12s.0" + ) + return getattr(self._tune_controller, attr) + + +def _get_max_pending_trials(search_alg: SearchAlgorithm) -> int: + max_pending_trials = os.getenv("TUNE_MAX_PENDING_TRIALS_PG", "auto") + + if max_pending_trials != "auto": + return int(max_pending_trials) + + # Else, auto detect. + + # Only BasicVariantGenerator supports > 1 pending trials. + # This is because we don't want to generate too many trials + # before we fit the searcher model. + if not isinstance(search_alg, BasicVariantGenerator): + return 1 + + # Allow up to at least 200 pending trials to trigger fast autoscaling + min_autoscaling_rate = 200 + + # Allow more pending trials for larger clusters (based on number of CPUs) + cluster_cpus = ray.cluster_resources().get("CPU", 1.0) + max_pending_trials = max(min_autoscaling_rate, int(cluster_cpus * 1.1)) + + if max_pending_trials > min_autoscaling_rate: + logger.warning( + f"The maximum number of pending trials has been " + f"automatically set to the number of available " + f"cluster CPUs, which is high " + f"({max_pending_trials} CPUs/pending trials). " + f"If you're running an experiment with a large number " + f"of trials, this could lead to scheduling overhead. " + f"In this case, consider setting the " + f"`TUNE_MAX_PENDING_TRIALS_PG` environment variable " + f"to the desired maximum number of concurrent pending trials." + ) + + return max_pending_trials + + +class _FakeRayTrialExecutor: + """The TuneController does not use a RayTrialExecutor anymore. + + Instead, we pass this fake executor for searchers/schedulers to use + as an interface. + + In the future, we should have the searchers/schedulers either interact with + the tune controller, or define a different API for more fine-grained scheduler + control. + """ + + def __init__(self, tune_controller: TuneController): + self._tune_controller = tune_controller + + def pause_trial(self, trial: Trial, should_checkpoint: bool = True): + return self._tune_controller._schedule_trial_pause( + trial, should_checkpoint=should_checkpoint + ) + + def save( + self, + trial: Trial, + result: Optional[Dict] = None, + ) -> Optional[_FutureTrainingResult]: + return self._tune_controller._schedule_trial_save(trial=trial, result=result) + + def has_resources_for_trial(self, trial: Trial): + return True + + @property + def _resource_updater(self): + return self._tune_controller._resource_updater + + def force_reconcilation_on_next_step_end(self): + pass diff --git a/lib/python3.12/site-packages/ray/tune/search/__init__.py b/lib/python3.12/site-packages/ray/tune/search/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..39afca111dceaf388d3bffba2a501c5498d598e4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/__init__.py @@ -0,0 +1,153 @@ +from ray._common.utils import get_function_args +from ray.tune.search.basic_variant import BasicVariantGenerator +from ray.tune.search.concurrency_limiter import ConcurrencyLimiter +from ray.tune.search.repeater import Repeater +from ray.tune.search.search_algorithm import SearchAlgorithm +from ray.tune.search.search_generator import SearchGenerator +from ray.tune.search.searcher import Searcher +from ray.tune.search.variant_generator import grid_search +from ray.util import PublicAPI + + +def _import_variant_generator(): + return BasicVariantGenerator + + +def _import_ax_search(): + from ray.tune.search.ax.ax_search import AxSearch + + return AxSearch + + +def _import_hyperopt_search(): + from ray.tune.search.hyperopt.hyperopt_search import HyperOptSearch + + return HyperOptSearch + + +def _import_bayesopt_search(): + from ray.tune.search.bayesopt.bayesopt_search import BayesOptSearch + + return BayesOptSearch + + +def _import_bohb_search(): + from ray.tune.search.bohb.bohb_search import TuneBOHB + + return TuneBOHB + + +def _import_nevergrad_search(): + from ray.tune.search.nevergrad.nevergrad_search import NevergradSearch + + return NevergradSearch + + +def _import_optuna_search(): + from ray.tune.search.optuna.optuna_search import OptunaSearch + + return OptunaSearch + + +def _import_zoopt_search(): + from ray.tune.search.zoopt.zoopt_search import ZOOptSearch + + return ZOOptSearch + + +def _import_hebo_search(): + from ray.tune.search.hebo.hebo_search import HEBOSearch + + return HEBOSearch + + +SEARCH_ALG_IMPORT = { + "variant_generator": _import_variant_generator, + "random": _import_variant_generator, + "ax": _import_ax_search, + "hyperopt": _import_hyperopt_search, + "bayesopt": _import_bayesopt_search, + "bohb": _import_bohb_search, + "nevergrad": _import_nevergrad_search, + "optuna": _import_optuna_search, + "zoopt": _import_zoopt_search, + "hebo": _import_hebo_search, +} + + +@PublicAPI(stability="beta") +def create_searcher( + search_alg, + **kwargs, +): + """Instantiate a search algorithm based on the given string. + + This is useful for swapping between different search algorithms. + + Args: + search_alg: The search algorithm to use. + metric: The training result objective value attribute. Stopping + procedures will use this attribute. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + **kwargs: Additional parameters. + These keyword arguments will be passed to the initialization + function of the chosen class. + Returns: + ray.tune.search.Searcher: The search algorithm. + Example: + >>> from ray import tune # doctest: +SKIP + >>> search_alg = tune.create_searcher('ax') # doctest: +SKIP + """ + + search_alg = search_alg.lower() + if search_alg not in SEARCH_ALG_IMPORT: + raise ValueError( + f"The `search_alg` argument must be one of " + f"{list(SEARCH_ALG_IMPORT)}. " + f"Got: {search_alg}" + ) + + SearcherClass = SEARCH_ALG_IMPORT[search_alg]() + + search_alg_args = get_function_args(SearcherClass) + trimmed_kwargs = {k: v for k, v in kwargs.items() if k in search_alg_args} + + return SearcherClass(**trimmed_kwargs) + + +UNRESOLVED_SEARCH_SPACE = str( + "You passed a `{par}` parameter to {cls} that contained unresolved search " + "space definitions. {cls} should however be instantiated with fully " + "configured search spaces only. To use Ray Tune's automatic search space " + "conversion, pass the space definition as part of the `param_space` argument " + "to `tune.Tuner()` instead." +) + +UNDEFINED_SEARCH_SPACE = str( + "Trying to sample a configuration from {cls}, but no search " + "space has been defined. Either pass the `{space}` argument when " + "instantiating the search algorithm, or pass a `param_space` to " + "`tune.Tuner()`." +) + +UNDEFINED_METRIC_MODE = str( + "Trying to sample a configuration from {cls}, but the `metric` " + "({metric}) or `mode` ({mode}) parameters have not been set. " + "Either pass these arguments when instantiating the search algorithm, " + "or pass them to `tune.TuneConfig()`." +) + + +__all__ = [ + "SearchAlgorithm", + "Searcher", + "ConcurrencyLimiter", + "Repeater", + "BasicVariantGenerator", + "grid_search", + "SearchGenerator", + "UNRESOLVED_SEARCH_SPACE", + "UNDEFINED_SEARCH_SPACE", + "UNDEFINED_METRIC_MODE", +] diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..415dbf725ba615b35fe7d7df139c38170ff346d4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/_mock.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/_mock.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7170d5a4dde3dfb1f72b85d77820d3f2149759e3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/_mock.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/basic_variant.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/basic_variant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0cb40cd1b7711a8b4ea8219801f72b6f0a6f6bf Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/basic_variant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/concurrency_limiter.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/concurrency_limiter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0367db8df83e10c45feeb2ff91e2895a5473b241 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/concurrency_limiter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/repeater.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/repeater.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6356aec8d2c181ebe1c5eac3ad194e9e8d1ba486 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/repeater.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/sample.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/sample.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92be589086da61e1ea1ac68159c251c882269454 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/sample.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_algorithm.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_algorithm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5671a392a88df84afdc1dd76b2c7cfa88d734288 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_algorithm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_generator.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01103a57d26ec3141bd5941a4bd6322362b74f83 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/search_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/searcher.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/searcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a1ef3f4c3a162b421d338179d6f171cefdea8fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/searcher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44afacd767bfb1c70db847199b82ad31e5ef7073 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/__pycache__/variant_generator.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/__pycache__/variant_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49d933b20756807da0a52fc20462eb44d4e9d393 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/__pycache__/variant_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/_mock.py b/lib/python3.12/site-packages/ray/tune/search/_mock.py new file mode 100644 index 0000000000000000000000000000000000000000..f1e8eb0b2140ebcebf4f6528642ed9ce7a1afd18 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/_mock.py @@ -0,0 +1,55 @@ +from typing import Dict, List, Optional + +from ray.tune.experiment import Trial +from ray.tune.search import ConcurrencyLimiter, Searcher +from ray.tune.search.search_generator import SearchGenerator + + +class _MockSearcher(Searcher): + def __init__(self, **kwargs): + self.live_trials = {} + self.counter = {"result": 0, "complete": 0} + self.final_results = [] + self.stall = False + self.results = [] + super(_MockSearcher, self).__init__(**kwargs) + + def suggest(self, trial_id: str): + if not self.stall: + self.live_trials[trial_id] = 1 + return {"test_variable": 2} + return None + + def on_trial_result(self, trial_id: str, result: Dict): + self.counter["result"] += 1 + self.results += [result] + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + self.counter["complete"] += 1 + if result: + self._process_result(result) + if trial_id in self.live_trials: + del self.live_trials[trial_id] + + def _process_result(self, result: Dict): + self.final_results += [result] + + +class _MockSuggestionAlgorithm(SearchGenerator): + def __init__(self, max_concurrent: Optional[int] = None, **kwargs): + self.searcher = _MockSearcher(**kwargs) + if max_concurrent: + self.searcher = ConcurrencyLimiter( + self.searcher, max_concurrent=max_concurrent + ) + super(_MockSuggestionAlgorithm, self).__init__(self.searcher) + + @property + def live_trials(self) -> List[Trial]: + return self.searcher.live_trials + + @property + def results(self) -> List[Dict]: + return self.searcher.results diff --git a/lib/python3.12/site-packages/ray/tune/search/ax/__init__.py b/lib/python3.12/site-packages/ray/tune/search/ax/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d2cc1df85b97152386a332358657e13e2dd06ede --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/ax/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.ax.ax_search import AxSearch + +__all__ = ["AxSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7497b1cfe4e8d54821f95dcfb2c55b2d6190c93 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/ax_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/ax_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0824dd8ba0ea8b7144e3000c88b1086469f96e9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/ax/__pycache__/ax_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/ax/ax_search.py b/lib/python3.12/site-packages/ray/tune/search/ax/ax_search.py new file mode 100644 index 0000000000000000000000000000000000000000..5cc408a89c73104f62f3fdb531c2c01be406686f --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/ax/ax_search.py @@ -0,0 +1,432 @@ +import copy +import logging +from typing import Dict, List, Optional, Union + +import numpy as np + +from ray import cloudpickle +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Float, + Integer, + LogUniform, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import flatten_dict, unflatten_list_dict + +try: + import ax + from ax.service.ax_client import AxClient +except ImportError: + ax = AxClient = None + +# This exception only exists in newer Ax releases for python 3.7 +try: + from ax.exceptions.core import DataRequiredError + from ax.exceptions.generation_strategy import MaxParallelismReachedException +except ImportError: + MaxParallelismReachedException = DataRequiredError = Exception + + +logger = logging.getLogger(__name__) + + +class AxSearch(Searcher): + """Uses `Ax `_ to optimize hyperparameters. + + Ax is a platform for understanding, managing, deploying, and + automating adaptive experiments. Ax provides an easy to use + interface with BoTorch, a flexible, modern library for Bayesian + optimization in PyTorch. More information can be found in https://ax.dev/. + + To use this search algorithm, you must install Ax: + + .. code-block:: bash + + $ pip install ax-platform + + Parameters: + space: Parameters in the experiment search space. + Required elements in the dictionaries are: "name" (name of + this parameter, string), "type" (type of the parameter: "range", + "fixed", or "choice", string), "bounds" for range parameters + (list of two values, lower bound first), "values" for choice + parameters (list of values), and "value" for fixed parameters + (single value). + metric: Name of the metric used as objective in this + experiment. This metric must be present in `raw_data` argument + to `log_data`. This metric must also be present in the dict + reported/returned by the Trainable. If None but a mode was passed, + the `ray.tune.result.DEFAULT_METRIC` will be used per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. Defaults to "max". + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + parameter_constraints: Parameter constraints, such as + "x3 >= x4" or "x3 + x4 >= 2". + outcome_constraints: Outcome constraints of form + "metric_name >= bound", like "m1 <= 3." + ax_client: Optional AxClient instance. If this is set, do + not pass any values to these parameters: `space`, `metric`, + `parameter_constraints`, `outcome_constraints`. + **ax_kwargs: Passed to AxClient instance. Ignored if `AxClient` is not + None. + + Tune automatically converts search spaces to Ax's format: + + .. code-block:: python + + from ray import tune + from ray.tune.search.ax import AxSearch + + config = { + "x1": tune.uniform(0.0, 1.0), + "x2": tune.uniform(0.0, 1.0) + } + + def easy_objective(config): + for i in range(100): + intermediate_result = config["x1"] + config["x2"] * i + tune.report({"score": intermediate_result}) + + ax_search = AxSearch() + tuner = tune.Tuner( + easy_objective, + tune_config=tune.TuneConfig( + search_alg=ax_search, + metric="score", + mode="max", + ), + param_space=config, + ) + tuner.fit() + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + from ray import tune + from ray.tune.search.ax import AxSearch + + parameters = [ + {"name": "x1", "type": "range", "bounds": [0.0, 1.0]}, + {"name": "x2", "type": "range", "bounds": [0.0, 1.0]}, + ] + + def easy_objective(config): + for i in range(100): + intermediate_result = config["x1"] + config["x2"] * i + tune.report({"score": intermediate_result}) + + ax_search = AxSearch(space=parameters, metric="score", mode="max") + tuner = tune.Tuner( + easy_objective, + tune_config=tune.TuneConfig( + search_alg=ax_search, + ), + ) + tuner.fit() + + """ + + def __init__( + self, + space: Optional[Union[Dict, List[Dict]]] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + parameter_constraints: Optional[List] = None, + outcome_constraints: Optional[List] = None, + ax_client: Optional[AxClient] = None, + **ax_kwargs, + ): + assert ( + ax is not None + ), """Ax must be installed! + You can install AxSearch with the command: + `pip install ax-platform`.""" + + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + + super(AxSearch, self).__init__( + metric=metric, + mode=mode, + ) + + self._ax = ax_client + self._ax_kwargs = ax_kwargs or {} + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space) + + self._space = space + self._parameter_constraints = parameter_constraints + self._outcome_constraints = outcome_constraints + + self._points_to_evaluate = copy.deepcopy(points_to_evaluate) + + self._parameters = [] + self._live_trial_mapping = {} + + if self._ax or self._space: + self._setup_experiment() + + def _setup_experiment(self): + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + if not self._ax: + self._ax = AxClient(**self._ax_kwargs) + + try: + exp = self._ax.experiment + has_experiment = True + except ValueError: + has_experiment = False + + if not has_experiment: + if not self._space: + raise ValueError( + "You have to create an Ax experiment by calling " + "`AxClient.create_experiment()`, or you should pass an " + "Ax search space as the `space` parameter to `AxSearch`, " + "or pass a `param_space` dict to `tune.Tuner()`." + ) + if self._mode not in ["min", "max"]: + raise ValueError( + "Please specify the `mode` argument when initializing " + "the `AxSearch` object or pass it to `tune.TuneConfig()`." + ) + self._ax.create_experiment( + parameters=self._space, + objective_name=self._metric, + parameter_constraints=self._parameter_constraints, + outcome_constraints=self._outcome_constraints, + minimize=self._mode != "max", + ) + else: + if any( + [ + self._space, + self._parameter_constraints, + self._outcome_constraints, + self._mode, + self._metric, + ] + ): + raise ValueError( + "If you create the Ax experiment yourself, do not pass " + "values for these parameters to `AxSearch`: {}.".format( + [ + "space", + "parameter_constraints", + "outcome_constraints", + "mode", + "metric", + ] + ) + ) + + exp = self._ax.experiment + + # Update mode and metric from experiment if it has been passed + self._mode = "min" if exp.optimization_config.objective.minimize else "max" + self._metric = exp.optimization_config.objective.metric.name + + self._parameters = list(exp.parameters) + + if self._ax._enforce_sequential_optimization: + logger.warning( + "Detected sequential enforcement. Be sure to use " + "a ConcurrencyLimiter." + ) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ): + if self._ax: + return False + space = self.convert_search_space(config) + self._space = space + if metric: + self._metric = metric + if mode: + self._mode = mode + + self._setup_experiment() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._ax: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + if self._points_to_evaluate: + config = self._points_to_evaluate.pop(0) + parameters, trial_index = self._ax.attach_trial(config) + else: + try: + parameters, trial_index = self._ax.get_next_trial() + except (MaxParallelismReachedException, DataRequiredError): + return None + + self._live_trial_mapping[trial_id] = trial_index + try: + suggested_config = unflatten_list_dict(parameters) + except AssertionError: + # Fails to unflatten if keys are out of order, which only happens + # if search space includes a list with both constants and + # tunable hyperparameters: + # Ex: "a": [1, tune.uniform(2, 3), 4] + suggested_config = unflatten_list_dict( + {k: parameters[k] for k in sorted(parameters.keys())} + ) + return suggested_config + + def on_trial_complete(self, trial_id, result=None, error=False): + """Notification for the completion of trial. + + Data of form key value dictionary of metric names and values. + """ + if result: + self._process_result(trial_id, result) + self._live_trial_mapping.pop(trial_id) + + def _process_result(self, trial_id, result): + ax_trial_index = self._live_trial_mapping[trial_id] + metrics_to_include = [self._metric] + [ + oc.metric.name + for oc in self._ax.experiment.optimization_config.outcome_constraints + ] + metric_dict = {} + for key in metrics_to_include: + val = result[key] + if np.isnan(val) or np.isinf(val): + # Don't report trials with NaN metrics to Ax + self._ax.abandon_trial( + trial_index=ax_trial_index, + reason=f"nan/inf metrics reported by {trial_id}", + ) + return + metric_dict[key] = (val, None) + self._ax.complete_trial(trial_index=ax_trial_index, raw_data=metric_dict) + + @staticmethod + def convert_search_space(spec: Dict): + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to an Ax search space." + ) + + # Flatten and resolve again after checking for grid search. + spec = flatten_dict(spec, prevent_delimiter=True) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + def resolve_value(par, domain): + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + logger.warning( + "AxSearch does not support quantization. Dropped quantization." + ) + sampler = sampler.sampler + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + return { + "name": par, + "type": "range", + "bounds": [domain.lower, domain.upper], + "value_type": "float", + "log_scale": True, + } + elif isinstance(sampler, Uniform): + return { + "name": par, + "type": "range", + "bounds": [domain.lower, domain.upper], + "value_type": "float", + "log_scale": False, + } + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + return { + "name": par, + "type": "range", + "bounds": [domain.lower, domain.upper - 1], + "value_type": "int", + "log_scale": True, + } + elif isinstance(sampler, Uniform): + return { + "name": par, + "type": "range", + "bounds": [domain.lower, domain.upper - 1], + "value_type": "int", + "log_scale": False, + } + elif isinstance(domain, Categorical): + if isinstance(sampler, Uniform): + return {"name": par, "type": "choice", "values": domain.categories} + + raise ValueError( + "AxSearch does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + # Parameter name is e.g. "a/b/c" for nested dicts, + # "a/d/0", "a/d/1" for nested lists (using the index in the list) + fixed_values = [ + {"name": "/".join(str(p) for p in path), "type": "fixed", "value": val} + for path, val in resolved_vars + ] + resolved_values = [ + resolve_value("/".join(str(p) for p in path), domain) + for path, domain in domain_vars + ] + + return fixed_values + resolved_values + + def save(self, checkpoint_path: str): + save_object = self.__dict__ + with open(checkpoint_path, "wb") as outputFile: + cloudpickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path: str): + with open(checkpoint_path, "rb") as inputFile: + save_object = cloudpickle.load(inputFile) + self.__dict__.update(save_object) diff --git a/lib/python3.12/site-packages/ray/tune/search/basic_variant.py b/lib/python3.12/site-packages/ray/tune/search/basic_variant.py new file mode 100644 index 0000000000000000000000000000000000000000..73bdaa31c5e8cba26c322d1d60b09bcbb4af9f53 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/basic_variant.py @@ -0,0 +1,421 @@ +import copy +import itertools +import os +import uuid +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +import numpy as np + +from ray.air._internal.usage import tag_searcher +from ray.tune.error import TuneError +from ray.tune.experiment.config_parser import _create_trial_from_spec, _make_parser +from ray.tune.search.sample import _BackwardsCompatibleNumpyRng, np_random_generator +from ray.tune.search.search_algorithm import SearchAlgorithm +from ray.tune.search.variant_generator import ( + _count_spec_samples, + _count_variants, + _flatten_resolved_vars, + _get_preset_variants, + format_vars, + generate_variants, +) +from ray.tune.utils.util import _atomic_save, _load_newest_checkpoint +from ray.util import PublicAPI + +if TYPE_CHECKING: + from ray.tune.experiment import Experiment + +SERIALIZATION_THRESHOLD = 1e6 + + +class _VariantIterator: + """Iterates over generated variants from the search space. + + This object also toggles between lazy evaluation and + eager evaluation of samples. If lazy evaluation is enabled, + this object cannot be serialized. + """ + + def __init__(self, iterable, lazy_eval=False): + self.lazy_eval = lazy_eval + self.iterable = iterable + self._has_next = True + if lazy_eval: + self._load_value() + else: + self.iterable = list(iterable) + self._has_next = bool(self.iterable) + + def _load_value(self): + try: + self.next_value = next(self.iterable) + except StopIteration: + self._has_next = False + + def has_next(self): + return self._has_next + + def __next__(self): + if self.lazy_eval: + current_value = self.next_value + self._load_value() + return current_value + current_value = self.iterable.pop(0) + self._has_next = bool(self.iterable) + return current_value + + +class _TrialIterator: + """Generates trials from the spec. + + Args: + uuid_prefix: Used in creating the trial name. + num_samples: Number of samples from distribution + (same as tune.TuneConfig). + unresolved_spec: Experiment specification + that might have unresolved distributions. + constant_grid_search: Should random variables be sampled + first before iterating over grid variants (True) or not (False). + points_to_evaluate: Configurations that will be tried out without sampling. + lazy_eval: Whether variants should be generated + lazily or eagerly. This is toggled depending + on the size of the grid search. + start: index at which to start counting trials. + random_state (int | np.random.Generator | np.random.RandomState): + Seed or numpy random generator to use for reproducible results. + If None (default), will use the global numpy random generator + (``np.random``). Please note that full reproducibility cannot + be guaranteed in a distributed environment. + """ + + def __init__( + self, + uuid_prefix: str, + num_samples: int, + unresolved_spec: dict, + constant_grid_search: bool = False, + points_to_evaluate: Optional[List] = None, + lazy_eval: bool = False, + start: int = 0, + random_state: Optional[ + Union[int, "np_random_generator", np.random.RandomState] + ] = None, + ): + self.parser = _make_parser() + self.num_samples = num_samples + self.uuid_prefix = uuid_prefix + self.num_samples_left = num_samples + self.unresolved_spec = unresolved_spec + self.constant_grid_search = constant_grid_search + self.points_to_evaluate = points_to_evaluate or [] + self.num_points_to_evaluate = len(self.points_to_evaluate) + self.counter = start + self.lazy_eval = lazy_eval + self.variants = None + self.random_state = random_state + + def create_trial(self, resolved_vars, spec): + trial_id = self.uuid_prefix + ("%05d" % self.counter) + experiment_tag = str(self.counter) + # Always append resolved vars to experiment tag? + if resolved_vars: + experiment_tag += "_{}".format(format_vars(resolved_vars)) + self.counter += 1 + return _create_trial_from_spec( + spec, + self.parser, + evaluated_params=_flatten_resolved_vars(resolved_vars), + trial_id=trial_id, + experiment_tag=experiment_tag, + ) + + def __next__(self): + """Generates Trial objects with the variant generation process. + + Uses a fixed point iteration to resolve variants. All trials + should be able to be generated at once. + + See also: `ray.tune.search.variant_generator`. + + Returns: + Trial object + """ + + if "run" not in self.unresolved_spec: + raise TuneError("Must specify `run` in {}".format(self.unresolved_spec)) + + if self.variants and self.variants.has_next(): + # This block will be skipped upon instantiation. + # `variants` will be set later after the first loop. + resolved_vars, spec = next(self.variants) + return self.create_trial(resolved_vars, spec) + + if self.points_to_evaluate: + config = self.points_to_evaluate.pop(0) + self.num_samples_left -= 1 + self.variants = _VariantIterator( + _get_preset_variants( + self.unresolved_spec, + config, + constant_grid_search=self.constant_grid_search, + random_state=self.random_state, + ), + lazy_eval=self.lazy_eval, + ) + resolved_vars, spec = next(self.variants) + return self.create_trial(resolved_vars, spec) + elif self.num_samples_left > 0: + self.variants = _VariantIterator( + generate_variants( + self.unresolved_spec, + constant_grid_search=self.constant_grid_search, + random_state=self.random_state, + ), + lazy_eval=self.lazy_eval, + ) + self.num_samples_left -= 1 + resolved_vars, spec = next(self.variants) + return self.create_trial(resolved_vars, spec) + else: + raise StopIteration + + def __iter__(self): + return self + + +@PublicAPI +class BasicVariantGenerator(SearchAlgorithm): + """Uses Tune's variant generation for resolving variables. + + This is the default search algorithm used if no other search algorithm + is specified. + + + Args: + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + max_concurrent: Maximum number of concurrently running trials. + If 0 (default), no maximum is enforced. + constant_grid_search: If this is set to ``True``, Ray Tune will + *first* try to sample random values and keep them constant over + grid search parameters. If this is set to ``False`` (default), + Ray Tune will sample new random parameters in each grid search + condition. + random_state: + Seed or numpy random generator to use for reproducible results. + If None (default), will use the global numpy random generator + (``np.random``). Please note that full reproducibility cannot + be guaranteed in a distributed environment. + + + Example: + + .. code-block:: python + + from ray import tune + + # This will automatically use the `BasicVariantGenerator` + tuner = tune.Tuner( + lambda config: config["a"] + config["b"], + tune_config=tune.TuneConfig( + num_samples=4 + ), + param_space={ + "a": tune.grid_search([1, 2]), + "b": tune.randint(0, 3) + }, + ) + tuner.fit() + + In the example above, 8 trials will be generated: For each sample + (``4``), each of the grid search variants for ``a`` will be sampled + once. The ``b`` parameter will be sampled randomly. + + The generator accepts a pre-set list of points that should be evaluated. + The points will replace the first samples of each experiment passed to + the ``BasicVariantGenerator``. + + Each point will replace one sample of the specified ``num_samples``. If + grid search variables are overwritten with the values specified in the + presets, the number of samples will thus be reduced. + + Example: + + .. code-block:: python + + from ray import tune + from ray.tune.search.basic_variant import BasicVariantGenerator + + tuner = tune.Tuner( + lambda config: config["a"] + config["b"], + tune_config=tune.TuneConfig( + search_alg=BasicVariantGenerator(points_to_evaluate=[ + {"a": 2, "b": 2}, + {"a": 1}, + {"b": 2} + ]), + num_samples=4 + ), + param_space={ + "a": tune.grid_search([1, 2]), + "b": tune.randint(0, 3) + }, + ) + tuner.fit() + + The example above will produce six trials via four samples: + + - The first sample will produce one trial with ``a=2`` and ``b=2``. + - The second sample will produce one trial with ``a=1`` and ``b`` sampled + randomly + - The third sample will produce two trials, one for each grid search + value of ``a``. It will be ``b=2`` for both of these trials. + - The fourth sample will produce two trials, one for each grid search + value of ``a``. ``b`` will be sampled randomly and independently for + both of these trials. + + """ + + CKPT_FILE_TMPL = "basic-variant-state-{}.json" + + def __init__( + self, + points_to_evaluate: Optional[List[Dict]] = None, + max_concurrent: int = 0, + constant_grid_search: bool = False, + random_state: Optional[ + Union[int, "np_random_generator", np.random.RandomState] + ] = None, + ): + tag_searcher(self) + self._trial_generator = [] + self._iterators = [] + self._trial_iter = None + self._finished = False + self._random_state = _BackwardsCompatibleNumpyRng(random_state) + + self._points_to_evaluate = points_to_evaluate or [] + + # Unique prefix for all trials generated, e.g., trial ids start as + # 2f1e_00001, 2f1ef_00002, 2f1ef_0003, etc. Overridable for testing. + force_test_uuid = os.environ.get("_TEST_TUNE_TRIAL_UUID") + if force_test_uuid: + self._uuid_prefix = force_test_uuid + "_" + else: + self._uuid_prefix = str(uuid.uuid1().hex)[:5] + "_" + + self._total_samples = 0 + self.max_concurrent = max_concurrent + self._constant_grid_search = constant_grid_search + self._live_trials = set() + + @property + def total_samples(self): + return self._total_samples + + def add_configurations( + self, experiments: Union["Experiment", List["Experiment"], Dict[str, Dict]] + ): + """Chains generator given experiment specifications. + + Arguments: + experiments: Experiments to run. + """ + from ray.tune.experiment import _convert_to_experiment_list + + experiment_list = _convert_to_experiment_list(experiments) + + for experiment in experiment_list: + grid_vals = _count_spec_samples(experiment.spec, num_samples=1) + lazy_eval = grid_vals > SERIALIZATION_THRESHOLD + if lazy_eval: + warnings.warn( + f"The number of pre-generated samples ({grid_vals}) " + "exceeds the serialization threshold " + f"({int(SERIALIZATION_THRESHOLD)}). Resume ability is " + "disabled. To fix this, reduce the number of " + "dimensions/size of the provided grid search." + ) + + previous_samples = self._total_samples + points_to_evaluate = copy.deepcopy(self._points_to_evaluate) + self._total_samples += _count_variants(experiment.spec, points_to_evaluate) + iterator = _TrialIterator( + uuid_prefix=self._uuid_prefix, + num_samples=experiment.spec.get("num_samples", 1), + unresolved_spec=experiment.spec, + constant_grid_search=self._constant_grid_search, + points_to_evaluate=points_to_evaluate, + lazy_eval=lazy_eval, + start=previous_samples, + random_state=self._random_state, + ) + self._iterators.append(iterator) + self._trial_generator = itertools.chain(self._trial_generator, iterator) + + def next_trial(self): + """Provides one Trial object to be queued into the TrialRunner. + + Returns: + Trial: Returns a single trial. + """ + if self.is_finished(): + return None + if self.max_concurrent > 0 and len(self._live_trials) >= self.max_concurrent: + return None + if not self._trial_iter: + self._trial_iter = iter(self._trial_generator) + try: + trial = next(self._trial_iter) + self._live_trials.add(trial.trial_id) + return trial + except StopIteration: + self._trial_generator = [] + self._trial_iter = None + self.set_finished() + return None + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + if trial_id in self._live_trials: + self._live_trials.remove(trial_id) + + def get_state(self): + if any(iterator.lazy_eval for iterator in self._iterators): + return False + state = self.__dict__.copy() + del state["_trial_generator"] + return state + + def set_state(self, state): + self.__dict__.update(state) + for iterator in self._iterators: + self._trial_generator = itertools.chain(self._trial_generator, iterator) + + def save_to_dir(self, dirpath, session_str): + if any(iterator.lazy_eval for iterator in self._iterators): + return False + state_dict = self.get_state() + _atomic_save( + state=state_dict, + checkpoint_dir=dirpath, + file_name=self.CKPT_FILE_TMPL.format(session_str), + tmp_file_name=".tmp_generator", + ) + + def has_checkpoint(self, dirpath: str): + """Whether a checkpoint file exists within dirpath.""" + return any(Path(dirpath).glob(self.CKPT_FILE_TMPL.format("*"))) + + def restore_from_dir(self, dirpath: str): + """Restores self + searcher + search wrappers from dirpath.""" + state_dict = _load_newest_checkpoint(dirpath, self.CKPT_FILE_TMPL.format("*")) + if not state_dict: + raise RuntimeError("Unable to find checkpoint in {}.".format(dirpath)) + self.set_state(state_dict) diff --git a/lib/python3.12/site-packages/ray/tune/search/bayesopt/__init__.py b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2569dcd5849a69f9818e90ac6db9cc4891c2da9c --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.bayesopt.bayesopt_search import BayesOptSearch + +__all__ = ["BayesOptSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..006779d442f630889b93916140009073afc92e85 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/bayesopt_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/bayesopt_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77e25253c9c78a5fa5ae8dcbcf929cf87308af4c Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/bayesopt/__pycache__/bayesopt_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/bayesopt/bayesopt_search.py b/lib/python3.12/site-packages/ray/tune/search/bayesopt/bayesopt_search.py new file mode 100644 index 0000000000000000000000000000000000000000..97b67192bcc88aafe81a9dd5d361b03884d84688 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/bayesopt/bayesopt_search.py @@ -0,0 +1,449 @@ +import json +import logging +import pickle +from collections import defaultdict +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import Domain, Float, Quantized, Uniform +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils import flatten_dict +from ray.tune.utils.util import is_nan_or_inf, unflatten_dict + +try: # Python 3 only -- needed for lint test. + import bayes_opt as byo +except ImportError: + byo = None + + +if TYPE_CHECKING: + from ray.tune import ExperimentAnalysis + +logger = logging.getLogger(__name__) + + +def _dict_hash(config, precision): + flatconfig = flatten_dict(config) + for param, value in flatconfig.items(): + if isinstance(value, float): + flatconfig[param] = "{:.{digits}f}".format(value, digits=precision) + + hashed = json.dumps(flatconfig, sort_keys=True, default=str) + return hashed + + +class BayesOptSearch(Searcher): + """Uses bayesian-optimization/BayesianOptimization to optimize hyperparameters. + + bayesian-optimization/BayesianOptimization is a library for Bayesian Optimization. More + info can be found here: https://github.com/bayesian-optimization/BayesianOptimization. + + This searcher will automatically filter out any NaN, inf or -inf + results. + + You will need to install bayesian-optimization/BayesianOptimization via the following: + + .. code-block:: bash + + pip install bayesian-optimization==1.4.3 + + Initializing this search algorithm with a ``space`` requires that it's + in the ``BayesianOptimization`` search space format. Otherwise, you + should instead pass in a Tune search space into ``Tuner(param_space=...)``, + and the search space will be automatically converted for you. + + See this ``BayesianOptimization`` example notebook + `_ + for an example. + + Args: + space: Continuous search space. Parameters will be sampled from + this space which will be used to run trials. + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + utility_kwargs: Parameters to define the utility function. + The default value is a dictionary with three keys: + - kind: ucb (Upper Confidence Bound) + - kappa: 2.576 + - xi: 0.0 + random_state: Used to initialize BayesOpt. + random_search_steps: Number of initial random searches. + This is necessary to avoid initial local overfitting + of the Bayesian process. + verbose: Sets verbosity level for BayesOpt packages. + patience: If patience is set and we've repeated a trial numerous times, + we terminate the experiment. + skip_duplicate: skip duplicate config + analysis: Optionally, the previous analysis to integrate. + + Tune automatically converts search spaces to BayesOptSearch's format: + + .. code-block:: python + + from ray import tune + from ray.tune.search.bayesopt import BayesOptSearch + + config = { + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100) + } + + bayesopt = BayesOptSearch(metric="mean_loss", mode="min") + tuner = tune.Tuner( + my_func, + tune_config=tune.TuneConfig( + search_alg=baysopt, + ), + param_space=config, + ) + tuner.fit() + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + from ray import tune + from ray.tune.search.bayesopt import BayesOptSearch + + space = { + 'width': (0, 20), + 'height': (-100, 100), + } + bayesopt = BayesOptSearch(space, metric="mean_loss", mode="min") + tuner = tune.Tuner( + my_func, + tune_config=tune.TuneConfig( + search_alg=bayesopt, + ), + ) + tuner.fit() + + """ + + # bayes_opt.BayesianOptimization: Optimization object + optimizer = None + + def __init__( + self, + space: Optional[Dict] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + utility_kwargs: Optional[Dict] = None, + random_state: int = 42, + random_search_steps: int = 10, + verbose: int = 0, + patience: int = 5, + skip_duplicate: bool = True, + analysis: Optional["ExperimentAnalysis"] = None, + ): + assert byo is not None, ( + "BayesOpt must be installed!. You can install BayesOpt with" + " the command: `pip install bayesian-optimization==1.4.3`." + ) + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + self._config_counter = defaultdict(int) + self._patience = patience + # int: Precision at which to hash values. + self.repeat_float_precision = 5 + if self._patience <= 0: + raise ValueError("patience must be set to a value greater than 0!") + self._skip_duplicate = skip_duplicate + super(BayesOptSearch, self).__init__( + metric=metric, + mode=mode, + ) + + if utility_kwargs is None: + # The defaults arguments are the same + # as in the package BayesianOptimization + utility_kwargs = dict( + kind="ucb", + kappa=2.576, + xi=0.0, + ) + + if mode == "max": + self._metric_op = 1.0 + elif mode == "min": + self._metric_op = -1.0 + + self._points_to_evaluate = points_to_evaluate + + self._live_trial_mapping = {} + self._buffered_trial_results = [] + self.random_search_trials = random_search_steps + self._total_random_search_trials = 0 + + self.utility = byo.UtilityFunction(**utility_kwargs) + + self._analysis = analysis + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space, join=True) + + self._space = space + self._verbose = verbose + self._random_state = random_state + + self.optimizer = None + if space: + self._setup_optimizer() + + def _setup_optimizer(self): + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + self.optimizer = byo.BayesianOptimization( + f=None, + pbounds=self._space, + verbose=self._verbose, + random_state=self._random_state, + ) + + # Registering the provided analysis, if given + if self._analysis is not None: + self.register_analysis(self._analysis) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self.optimizer: + return False + space = self.convert_search_space(config) + self._space = space + if metric: + self._metric = metric + if mode: + self._mode = mode + + if self._mode == "max": + self._metric_op = 1.0 + elif self._mode == "min": + self._metric_op = -1.0 + + self._setup_optimizer() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + """Return new point to be explored by black box function. + + Args: + trial_id: Id of the trial. + This is a short alphanumerical string. + + Returns: + Either a dictionary describing the new point to explore or + None, when no new point is to be explored for the time being. + """ + if not self.optimizer: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + if self._points_to_evaluate: + config = self._points_to_evaluate.pop(0) + else: + # We compute the new point to explore + config = self.optimizer.suggest(self.utility) + + config_hash = _dict_hash(config, self.repeat_float_precision) + # Check if already computed + already_seen = config_hash in self._config_counter + self._config_counter[config_hash] += 1 + top_repeats = max(self._config_counter.values()) + + # If patience is set and we've repeated a trial numerous times, + # we terminate the experiment. + if self._patience is not None and top_repeats > self._patience: + return Searcher.FINISHED + # If we have seen a value before, we'll skip it. + if already_seen and self._skip_duplicate: + logger.info("Skipping duplicated config: {}.".format(config)) + return None + + # If we are still in the random search part and we are waiting for + # trials to complete + if len(self._buffered_trial_results) < self.random_search_trials: + # We check if we have already maxed out the number of requested + # random search trials + if self._total_random_search_trials == self.random_search_trials: + # If so we stop the suggestion and return None + return None + # Otherwise we increase the total number of rndom search trials + if config: + self._total_random_search_trials += 1 + + # Save the new trial to the trial mapping + self._live_trial_mapping[trial_id] = config + + # Return a deep copy of the mapping + return unflatten_dict(config) + + def register_analysis(self, analysis: "ExperimentAnalysis"): + """Integrate the given analysis into the gaussian process. + + Args: + analysis: Optionally, the previous analysis + to integrate. + """ + for (_, report), params in zip( + analysis.dataframe(metric=self._metric, mode=self._mode).iterrows(), + analysis.get_all_configs().values(), + ): + # We add the obtained results to the + # gaussian process optimizer + self._register_result(params, report) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + """Notification for the completion of trial. + + Args: + trial_id: Id of the trial. + This is a short alphanumerical string. + result: Dictionary of result. + May be none when some error occurs. + error: Boolean representing a previous error state. + The result should be None when error is True. + """ + # We try to get the parameters used for this trial + params = self._live_trial_mapping.pop(trial_id, None) + + # The results may be None if some exception is raised during the trial. + # Also, if the parameters are None (were already processed) + # we interrupt the following procedure. + # Additionally, if somehow the error is True but + # the remaining values are not we also block the method + if result is None or params is None or error: + return + + # If we don't have to execute some random search steps + if len(self._buffered_trial_results) >= self.random_search_trials: + # we simply register the obtained result + self._register_result(params, result) + return + + # We store the results into a temporary cache + self._buffered_trial_results.append((params, result)) + + # If the random search finished, + # we update the BO with all the computer points. + if len(self._buffered_trial_results) == self.random_search_trials: + for params, result in self._buffered_trial_results: + self._register_result(params, result) + + def _register_result(self, params: Tuple[str], result: Dict): + """Register given tuple of params and results.""" + if is_nan_or_inf(result[self.metric]): + return + self.optimizer.register(params, self._metric_op * result[self.metric]) + + def get_state(self) -> Dict[str, Any]: + state = self.__dict__.copy() + return state + + def set_state(self, state: Dict[str, Any]): + self.__dict__.update(state) + + def save(self, checkpoint_path: str): + """Storing current optimizer state.""" + save_object = self.get_state() + with open(checkpoint_path, "wb") as f: + pickle.dump(save_object, f) + + def restore(self, checkpoint_path: str): + """Restoring current optimizer state.""" + with open(checkpoint_path, "rb") as f: + save_object = pickle.load(f) + + if isinstance(save_object, dict): + self.set_state(save_object) + else: + # Backwards compatibility + ( + self.optimizer, + self._buffered_trial_results, + self._total_random_search_trials, + self._config_counter, + self._points_to_evaluate, + ) = save_object + + @staticmethod + def convert_search_space(spec: Dict, join: bool = False) -> Dict: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a BayesOpt search space." + ) + + # Flatten and resolve again after checking for grid search. + spec = flatten_dict(spec, prevent_delimiter=True) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + def resolve_value(domain: Domain) -> Tuple[float, float]: + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + logger.warning( + "BayesOpt search does not support quantization. " + "Dropped quantization." + ) + sampler = sampler.get_sampler() + + if isinstance(domain, Float): + if domain.sampler is not None and not isinstance( + domain.sampler, Uniform + ): + logger.warning( + "BayesOpt does not support specific sampling methods. " + "The {} sampler will be dropped.".format(sampler) + ) + return (domain.lower, domain.upper) + + raise ValueError( + "BayesOpt does not support parameters of type " + "`{}`".format(type(domain).__name__) + ) + + # Parameter name is e.g. "a/b/c" for nested dicts + bounds = {"/".join(path): resolve_value(domain) for path, domain in domain_vars} + + if join: + spec.update(bounds) + bounds = spec + + return bounds diff --git a/lib/python3.12/site-packages/ray/tune/search/bohb/__init__.py b/lib/python3.12/site-packages/ray/tune/search/bohb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..40068b1562adcd327627c8b8f641079c95364ccc --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/bohb/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.bohb.bohb_search import BOHB, TuneBOHB + +__all__ = ["BOHB", "TuneBOHB"] diff --git a/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac247500139ab9239387acf55612dba505f3270d Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/bohb_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/bohb_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..289f3c01f9f46c96a3645142d4c82df023e69d29 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/bohb/__pycache__/bohb_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/bohb/bohb_search.py b/lib/python3.12/site-packages/ray/tune/search/bohb/bohb_search.py new file mode 100644 index 0000000000000000000000000000000000000000..5fa701b700f7f5adf1eed308364725c779c5466d --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/bohb/bohb_search.py @@ -0,0 +1,383 @@ +"""BOHB (Bayesian Optimization with HyperBand)""" + +import copy +import logging +import math +from typing import Dict, List, Optional, Union + +# use cloudpickle instead of pickle to make BOHB obj +# pickleable +from ray import cloudpickle +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + LogUniform, + Normal, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import flatten_dict, unflatten_list_dict + +try: + import ConfigSpace + from hpbandster.optimizers.config_generators.bohb import BOHB +except ImportError: + BOHB = ConfigSpace = None + +logger = logging.getLogger(__name__) + + +class _BOHBJobWrapper: + """Mock object for HpBandSter to process.""" + + def __init__(self, loss: float, budget: float, config: Dict): + self.result = {"loss": loss} + self.kwargs = {"budget": budget, "config": config.copy()} + self.exception = None + + +class TuneBOHB(Searcher): + """BOHB suggestion component. + + + Requires HpBandSter and ConfigSpace to be installed. You can install + HpBandSter and ConfigSpace with: ``pip install hpbandster ConfigSpace``. + + This should be used in conjunction with HyperBandForBOHB. + + Args: + space: Continuous ConfigSpace search space. + Parameters will be sampled from this space which will be used + to run trials. + bohb_config: configuration for HpBandSter BOHB algorithm + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + seed: Optional random seed to initialize the random number + generator. Setting this should lead to identical initial + configurations at each run. + max_concurrent: Number of maximum concurrent trials. + If this Searcher is used in a ``ConcurrencyLimiter``, the + ``max_concurrent`` value passed to it will override the + value passed here. Set to <= 0 for no limit on concurrency. + + Tune automatically converts search spaces to TuneBOHB's format: + + .. code-block:: python + + config = { + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100), + "activation": tune.choice(["relu", "tanh"]) + } + + algo = TuneBOHB(metric="mean_loss", mode="min") + bohb = HyperBandForBOHB( + time_attr="training_iteration", + metric="mean_loss", + mode="min", + max_t=100) + run(my_trainable, config=config, scheduler=bohb, search_alg=algo) + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + import ConfigSpace as CS + + config_space = CS.ConfigurationSpace() + config_space.add_hyperparameter( + CS.UniformFloatHyperparameter("width", lower=0, upper=20)) + config_space.add_hyperparameter( + CS.UniformFloatHyperparameter("height", lower=-100, upper=100)) + config_space.add_hyperparameter( + CS.CategoricalHyperparameter( + name="activation", choices=["relu", "tanh"])) + + algo = TuneBOHB( + config_space, metric="mean_loss", mode="min") + bohb = HyperBandForBOHB( + time_attr="training_iteration", + metric="mean_loss", + mode="min", + max_t=100) + run(my_trainable, scheduler=bohb, search_alg=algo) + + """ + + def __init__( + self, + space: Optional[Union[Dict, "ConfigSpace.ConfigurationSpace"]] = None, + bohb_config: Optional[Dict] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + seed: Optional[int] = None, + max_concurrent: int = 0, + ): + assert ( + BOHB is not None + ), """HpBandSter must be installed! + You can install HpBandSter with the command: + `pip install hpbandster ConfigSpace`.""" + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + self.trial_to_params = {} + self._metric = metric + + self._bohb_config = bohb_config + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space) + + self._space = space + self._seed = seed + + self.running = set() + self.paused = set() + + self._max_concurrent = max_concurrent + self._points_to_evaluate = points_to_evaluate + + super(TuneBOHB, self).__init__( + metric=self._metric, + mode=mode, + ) + + if self._space: + self._setup_bohb() + + def set_max_concurrency(self, max_concurrent: int) -> bool: + self._max_concurrent = max_concurrent + return True + + def _setup_bohb(self): + from hpbandster.optimizers.config_generators.bohb import BOHB + + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + if self._mode == "max": + self._metric_op = -1.0 + elif self._mode == "min": + self._metric_op = 1.0 + + if self._seed is not None: + self._space.seed(self._seed) + + self.running = set() + self.paused = set() + + bohb_config = self._bohb_config or {} + self.bohber = BOHB(self._space, **bohb_config) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self._space: + return False + space = self.convert_search_space(config) + self._space = space + + if metric: + self._metric = metric + if mode: + self._mode = mode + + self._setup_bohb() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._space: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + max_concurrent = ( + self._max_concurrent if self._max_concurrent > 0 else float("inf") + ) + if len(self.running) >= max_concurrent: + return None + + if self._points_to_evaluate: + config = self._points_to_evaluate.pop(0) + else: + # This parameter is not used in hpbandster implementation. + config, _ = self.bohber.get_config(None) + self.trial_to_params[trial_id] = copy.deepcopy(config) + self.running.add(trial_id) + return unflatten_list_dict(config) + + def on_trial_result(self, trial_id: str, result: Dict): + if trial_id not in self.paused: + self.running.add(trial_id) + if "hyperband_info" not in result: + logger.warning( + "BOHB Info not detected in result. Are you using " + "HyperBandForBOHB as a scheduler?" + ) + elif "budget" in result.get("hyperband_info", {}): + hbs_wrapper = self.to_wrapper(trial_id, result) + self.bohber.new_result(hbs_wrapper) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + del self.trial_to_params[trial_id] + self.paused.discard(trial_id) + self.running.discard(trial_id) + + def to_wrapper(self, trial_id: str, result: Dict) -> _BOHBJobWrapper: + return _BOHBJobWrapper( + self._metric_op * result[self.metric], + result["hyperband_info"]["budget"], + self.trial_to_params[trial_id], + ) + + # BOHB Specific. + # TODO(team-ml): Refactor alongside HyperBandForBOHB + def on_pause(self, trial_id: str): + self.paused.add(trial_id) + self.running.discard(trial_id) + + def on_unpause(self, trial_id: str): + self.paused.discard(trial_id) + self.running.add(trial_id) + + @staticmethod + def convert_search_space(spec: Dict) -> "ConfigSpace.ConfigurationSpace": + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a TuneBOHB search space." + ) + + # Flatten and resolve again after checking for grid search. + spec = flatten_dict(spec, prevent_delimiter=True) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + def resolve_value( + par: str, domain: Domain + ) -> ConfigSpace.hyperparameters.Hyperparameter: + quantize = None + + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + quantize = sampler.q + sampler = sampler.sampler + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + lower = domain.lower + upper = domain.upper + if quantize: + lower = math.ceil(domain.lower / quantize) * quantize + upper = math.floor(domain.upper / quantize) * quantize + return ConfigSpace.UniformFloatHyperparameter( + par, lower=lower, upper=upper, q=quantize, log=True + ) + elif isinstance(sampler, Uniform): + lower = domain.lower + upper = domain.upper + if quantize: + lower = math.ceil(domain.lower / quantize) * quantize + upper = math.floor(domain.upper / quantize) * quantize + return ConfigSpace.UniformFloatHyperparameter( + par, lower=lower, upper=upper, q=quantize, log=False + ) + elif isinstance(sampler, Normal): + return ConfigSpace.hyperparameters.NormalFloatHyperparameter( + par, mu=sampler.mean, sigma=sampler.sd, q=quantize, log=False + ) + + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + lower = domain.lower + upper = domain.upper + if quantize: + lower = math.ceil(domain.lower / quantize) * quantize + upper = math.floor(domain.upper / quantize) * quantize + else: + # Tune search space integers are exclusive + upper -= 1 + return ConfigSpace.UniformIntegerHyperparameter( + par, lower=lower, upper=upper, q=quantize, log=True + ) + elif isinstance(sampler, Uniform): + lower = domain.lower + upper = domain.upper + if quantize: + lower = math.ceil(domain.lower / quantize) * quantize + upper = math.floor(domain.upper / quantize) * quantize + else: + # Tune search space integers are exclusive + upper -= 1 + return ConfigSpace.UniformIntegerHyperparameter( + par, lower=lower, upper=upper, q=quantize, log=False + ) + + elif isinstance(domain, Categorical): + if isinstance(sampler, Uniform): + return ConfigSpace.CategoricalHyperparameter( + par, choices=domain.categories + ) + + raise ValueError( + "TuneBOHB does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + cs = ConfigSpace.ConfigurationSpace() + for path, domain in domain_vars: + par = "/".join(str(p) for p in path) + value = resolve_value(par, domain) + cs.add_hyperparameter(value) + + return cs + + def save(self, checkpoint_path: str): + save_object = self.__dict__ + with open(checkpoint_path, "wb") as outputFile: + cloudpickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path: str): + with open(checkpoint_path, "rb") as inputFile: + save_object = cloudpickle.load(inputFile) + self.__dict__.update(save_object) diff --git a/lib/python3.12/site-packages/ray/tune/search/concurrency_limiter.py b/lib/python3.12/site-packages/ray/tune/search/concurrency_limiter.py new file mode 100644 index 0000000000000000000000000000000000000000..847db982ea881a0079dd78cdf8b4689ac1c02401 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/concurrency_limiter.py @@ -0,0 +1,176 @@ +import copy +import logging +from typing import Dict, List, Optional + +from ray.tune.search.searcher import Searcher +from ray.tune.search.util import _set_search_properties_backwards_compatible +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI +class ConcurrencyLimiter(Searcher): + """A wrapper algorithm for limiting the number of concurrent trials. + + Certain Searchers have their own internal logic for limiting + the number of concurrent trials. If such a Searcher is passed to a + ``ConcurrencyLimiter``, the ``max_concurrent`` of the + ``ConcurrencyLimiter`` will override the ``max_concurrent`` value + of the Searcher. The ``ConcurrencyLimiter`` will then let the + Searcher's internal logic take over. + + Args: + searcher: Searcher object that the + ConcurrencyLimiter will manage. + max_concurrent: Maximum concurrent samples from the underlying + searcher. + batch: Whether to wait for all concurrent samples + to finish before updating the underlying searcher. + + Example: + + .. code-block:: python + + from ray.tune.search import ConcurrencyLimiter + search_alg = HyperOptSearch(metric="accuracy") + search_alg = ConcurrencyLimiter(search_alg, max_concurrent=2) + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=search_alg + ), + ) + tuner.fit() + """ + + def __init__(self, searcher: Searcher, max_concurrent: int, batch: bool = False): + assert type(max_concurrent) is int and max_concurrent > 0 + self.searcher = searcher + self.max_concurrent = max_concurrent + self.batch = batch + self.live_trials = set() + self.num_unfinished_live_trials = 0 + self.cached_results = {} + self._limit_concurrency = True + + if not isinstance(searcher, Searcher): + raise RuntimeError( + f"The `ConcurrencyLimiter` only works with `Searcher` " + f"objects (got {type(searcher)}). Please try to pass " + f"`max_concurrent` to the search generator directly." + ) + + self._set_searcher_max_concurrency() + + super(ConcurrencyLimiter, self).__init__( + metric=self.searcher.metric, mode=self.searcher.mode + ) + + def _set_searcher_max_concurrency(self): + # If the searcher has special logic for handling max concurrency, + # we do not do anything inside the ConcurrencyLimiter + self._limit_concurrency = not self.searcher.set_max_concurrency( + self.max_concurrent + ) + + def set_max_concurrency(self, max_concurrent: int) -> bool: + # Determine if this behavior is acceptable, or if it should + # raise an exception. + self.max_concurrent = max_concurrent + return True + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + self._set_searcher_max_concurrency() + return _set_search_properties_backwards_compatible( + self.searcher.set_search_properties, metric, mode, config, **spec + ) + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._limit_concurrency: + return self.searcher.suggest(trial_id) + + assert ( + trial_id not in self.live_trials + ), f"Trial ID {trial_id} must be unique: already found in set." + if len(self.live_trials) >= self.max_concurrent: + logger.debug( + f"Not providing a suggestion for {trial_id} due to " + "concurrency limit: %s/%s.", + len(self.live_trials), + self.max_concurrent, + ) + return + + suggestion = self.searcher.suggest(trial_id) + if suggestion not in (None, Searcher.FINISHED): + self.live_trials.add(trial_id) + self.num_unfinished_live_trials += 1 + return suggestion + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + if not self._limit_concurrency: + return self.searcher.on_trial_complete(trial_id, result=result, error=error) + + if trial_id not in self.live_trials: + return + elif self.batch: + self.cached_results[trial_id] = (result, error) + self.num_unfinished_live_trials -= 1 + if self.num_unfinished_live_trials <= 0: + # Update the underlying searcher once the + # full batch is completed. + for trial_id, (result, error) in self.cached_results.items(): + self.searcher.on_trial_complete( + trial_id, result=result, error=error + ) + self.live_trials.remove(trial_id) + self.cached_results = {} + self.num_unfinished_live_trials = 0 + else: + return + else: + self.searcher.on_trial_complete(trial_id, result=result, error=error) + self.live_trials.remove(trial_id) + self.num_unfinished_live_trials -= 1 + + def on_trial_result(self, trial_id: str, result: Dict) -> None: + self.searcher.on_trial_result(trial_id, result) + + def add_evaluated_point( + self, + parameters: Dict, + value: float, + error: bool = False, + pruned: bool = False, + intermediate_values: Optional[List[float]] = None, + ): + return self.searcher.add_evaluated_point( + parameters, value, error, pruned, intermediate_values + ) + + def get_state(self) -> Dict: + state = self.__dict__.copy() + del state["searcher"] + return copy.deepcopy(state) + + def set_state(self, state: Dict): + self.__dict__.update(state) + + def save(self, checkpoint_path: str): + self.searcher.save(checkpoint_path) + + def restore(self, checkpoint_path: str): + self.searcher.restore(checkpoint_path) + + # BOHB Specific. + # TODO(team-ml): Refactor alongside HyperBandForBOHB + def on_pause(self, trial_id: str): + self.searcher.on_pause(trial_id) + + def on_unpause(self, trial_id: str): + self.searcher.on_unpause(trial_id) diff --git a/lib/python3.12/site-packages/ray/tune/search/hebo/__init__.py b/lib/python3.12/site-packages/ray/tune/search/hebo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66e5c5675bb22e2d013a89b994146192763ed1b3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/hebo/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.hebo.hebo_search import HEBOSearch + +__all__ = ["HEBOSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51c405934b052380b428f9cce2419ddd73bffdf2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/hebo_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/hebo_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9a74502c2e463c4466bfc22b603bc0c652cc056 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/hebo/__pycache__/hebo_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/hebo/hebo_search.py b/lib/python3.12/site-packages/ray/tune/search/hebo/hebo_search.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f130c5d514461364ed566ec5ed2004abb01975 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/hebo/hebo_search.py @@ -0,0 +1,466 @@ +import logging +import pickle +from typing import Dict, List, Optional, Union + +import numpy as np +import pandas as pd + +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + LogUniform, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import is_nan_or_inf, unflatten_dict, validate_warmstart + +try: # Python 3 only -- needed for lint test. + import hebo + import torch # hebo has torch as a dependency +except ImportError: + hebo = None + +logger = logging.getLogger(__name__) + +SPACE_ERROR_MESSAGE = ( + "Space must be either a HEBO DesignSpace object" + "or a dictionary with ONLY tune search spaces." +) + + +class HEBOSearch(Searcher): + """Uses HEBO (Heteroscedastic Evolutionary Bayesian Optimization) + to optimize hyperparameters. + + HEBO is a cutting edge black-box optimization framework created + by Huawei's Noah Ark. More info can be found here: + https://github.com/huawei-noah/HEBO/tree/master/HEBO. + + `space` can either be a HEBO's `DesignSpace` object or a dict of Tune + search spaces. + + Please note that the first few trials will be random and used + to kickstart the search process. In order to achieve good results, + we recommend setting the number of trials to at least 16. + + Maximum number of concurrent trials is determined by ``max_concurrent`` + argument. Trials will be done in batches of ``max_concurrent`` trials. + If this Searcher is used in a ``ConcurrencyLimiter``, the + ``max_concurrent`` value passed to it will override the value passed + here. + + Args: + space: A dict mapping parameter names to Tune search spaces or a + HEBO DesignSpace object. + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + evaluated_rewards: If you have previously evaluated the + parameters passed in as points_to_evaluate you can avoid + re-running those trials by passing in the reward attributes + as a list so the optimiser can be told the results without + needing to re-compute the trial. Must be the same length as + points_to_evaluate. + random_state_seed: Seed for reproducible + results. Defaults to None. Please note that setting this to a value + will change global random states for `numpy` and `torch` + on initalization and loading from checkpoint. + max_concurrent: Number of maximum concurrent trials. + If this Searcher is used in a ``ConcurrencyLimiter``, the + ``max_concurrent`` value passed to it will override the + value passed here. + **kwargs: The keyword arguments will be passed to `HEBO()``. + + Tune automatically converts search spaces to HEBO's format: + + .. code-block:: python + + from ray import tune + from ray.tune.search.hebo import HEBOSearch + + config = { + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100) + } + + hebo = HEBOSearch(metric="mean_loss", mode="min") + tuner = tune.Tuner( + trainable_function, + tune_config=tune.TuneConfig( + search_alg=hebo + ), + param_space=config + ) + tuner.fit() + + Alternatively, you can pass a HEBO `DesignSpace` object manually to the + Searcher: + + .. code-block:: python + + from ray import tune + from ray.tune.search.hebo import HEBOSearch + from hebo.design_space.design_space import DesignSpace + + space_config = [ + {'name' : 'width', 'type' : 'num', 'lb' : 0, 'ub' : 20}, + {'name' : 'height', 'type' : 'num', 'lb' : -100, 'ub' : 100}, + ] + space = DesignSpace().parse(space_config) + + hebo = HEBOSearch(space, metric="mean_loss", mode="min") + tuner = tune.Tuner( + trainable_function, + tune_config=tune.TuneConfig( + search_alg=hebo + ) + ) + tuner.fit() + + """ + + def __init__( + self, + space: Optional[ + Union[Dict, "hebo.design_space.design_space.DesignSpace"] + ] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + evaluated_rewards: Optional[List] = None, + random_state_seed: Optional[int] = None, + max_concurrent: int = 8, + **kwargs, + ): + assert hebo is not None, ( + "HEBO must be installed! You can install HEBO with" + " the command: `pip install 'HEBO>=0.2.0'`." + "This error may also be caused if HEBO" + " dependencies have bad versions. Try updating HEBO" + " first." + ) + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + assert ( + isinstance(max_concurrent, int) and max_concurrent >= 1 + ), "`max_concurrent` must be an integer and at least 1." + if random_state_seed is not None: + assert isinstance( + random_state_seed, int + ), "random_state_seed must be None or int, got '{}'.".format( + type(random_state_seed) + ) + super(HEBOSearch, self).__init__(metric=metric, mode=mode) + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if resolved_vars: + raise TypeError(SPACE_ERROR_MESSAGE) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space) + elif space is not None and not isinstance( + space, hebo.design_space.design_space.DesignSpace + ): + raise TypeError(SPACE_ERROR_MESSAGE + " Got {}.".format(type(space))) + + self._hebo_config = kwargs + self._random_state_seed = random_state_seed + self._space = space + self._points_to_evaluate = points_to_evaluate + self._evaluated_rewards = evaluated_rewards + self._initial_points = [] + self._live_trial_mapping = {} + + self._max_concurrent = max_concurrent + self._suggestions_cache = [] + self._batch_filled = False + + self._opt = None + if space: + self._setup_optimizer() + + def set_max_concurrency(self, max_concurrent: int) -> bool: + self._max_concurrent = max_concurrent + return True + + def _setup_optimizer(self): + # HEBO internally minimizes, so "max" => -1 + if self._mode == "max": + self._metric_op = -1.0 + elif self._mode == "min": + self._metric_op = 1.0 + + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + if not isinstance(self._space, hebo.design_space.design_space.DesignSpace): + raise ValueError( + f"Invalid search space: {type(self._space)}. Either pass a " + f"valid search space to the `HEBOSearch` class or pass " + f"a `param_space` parameter to `tune.Tuner()`" + ) + + if self._space.num_paras <= 0: + raise ValueError( + "Got empty search space. Please make sure to pass " + "a valid search space with at least one parameter to " + "`HEBOSearch`" + ) + + if self._random_state_seed is not None: + np.random.seed(self._random_state_seed) + torch.random.manual_seed(self._random_state_seed) + + self._opt = hebo.optimizers.hebo.HEBO(space=self._space, **self._hebo_config) + + if self._points_to_evaluate: + validate_warmstart( + self._space.para_names, + self._points_to_evaluate, + self._evaluated_rewards, + ) + if self._evaluated_rewards: + self._opt.observe( + pd.DataFrame(self._points_to_evaluate), + np.array(self._evaluated_rewards) * self._metric_op, + ) + else: + self._initial_points = self._points_to_evaluate + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self._opt: + return False + space = self.convert_search_space(config) + self._space = space + + if metric: + self._metric = metric + if mode: + self._mode = mode + + self._setup_optimizer() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._opt: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + if not self._live_trial_mapping: + self._batch_filled = False + + if self._initial_points: + params = self._initial_points.pop(0) + suggestion = pd.DataFrame([params], index=[0]) + else: + if ( + self._batch_filled + or len(self._live_trial_mapping) >= self._max_concurrent + ): + return None + if not self._suggestions_cache: + suggestion = self._opt.suggest(n_suggestions=self._max_concurrent) + self._suggestions_cache = suggestion.to_dict("records") + params = self._suggestions_cache.pop(0) + suggestion = pd.DataFrame([params], index=[0]) + self._live_trial_mapping[trial_id] = suggestion + if len(self._live_trial_mapping) >= self._max_concurrent: + self._batch_filled = True + return unflatten_dict(params) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + """Notification for the completion of trial. + + HEBO always minimizes.""" + + if result: + self._process_result(trial_id, result) + self._live_trial_mapping.pop(trial_id) + + def _process_result(self, trial_id: str, result: Dict): + trial_info = self._live_trial_mapping[trial_id] + if result and not is_nan_or_inf(result[self._metric]): + self._opt.observe( + trial_info, np.array([self._metric_op * result[self._metric]]) + ) + + def add_evaluated_point( + self, + parameters: Dict, + value: float, + error: bool = False, + pruned: bool = False, + intermediate_values: Optional[List[float]] = None, + ): + if intermediate_values: + logger.warning("HEBO doesn't use intermediate_values. Ignoring.") + if not error and not pruned: + self._opt.observe( + pd.DataFrame( + [ + { + k: v + for k, v in parameters.items() + if k in self._opt.space.para_names + } + ] + ), + np.array([value]) * self._metric_op, + ) + else: + logger.warning( + "Only non errored and non pruned points can be added to HEBO." + ) + + def save(self, checkpoint_path: str): + """Storing current optimizer state.""" + if self._random_state_seed is not None: + numpy_random_state = np.random.get_state() + torch_random_state = torch.get_rng_state() + else: + numpy_random_state = None + torch_random_state = None + save_object = self.__dict__.copy() + save_object["__numpy_random_state"] = numpy_random_state + save_object["__torch_random_state"] = torch_random_state + with open(checkpoint_path, "wb") as f: + pickle.dump(save_object, f) + + def restore(self, checkpoint_path: str): + """Restoring current optimizer state.""" + with open(checkpoint_path, "rb") as f: + save_object = pickle.load(f) + + if isinstance(save_object, dict): + numpy_random_state = save_object.pop("__numpy_random_state", None) + torch_random_state = save_object.pop("__torch_random_state", None) + self.__dict__.update(save_object) + else: + # Backwards compatibility + ( + self._opt, + self._initial_points, + numpy_random_state, + torch_random_state, + self._live_trial_mapping, + self._max_concurrent, + self._suggestions_cache, + self._space, + self._hebo_config, + self._batch_filled, + ) = save_object + if numpy_random_state is not None: + np.random.set_state(numpy_random_state) + if torch_random_state is not None: + torch.random.set_rng_state(torch_random_state) + + @staticmethod + def convert_search_space(spec: Dict, prefix: str = "") -> Dict: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + params = [] + + if not domain_vars and not grid_vars: + return {} + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a HEBO search space." + ) + + def resolve_value(par: str, domain: Domain): + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + logger.warning( + "HEBO search does not support quantization. " + "Dropped quantization." + ) + sampler = sampler.get_sampler() + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + return { + "name": par, + "type": "pow", + "lb": domain.lower, + "ub": domain.upper, + } + elif isinstance(sampler, Uniform): + return { + "name": par, + "type": "num", + "lb": domain.lower, + "ub": domain.upper, + } + + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + return { + "name": par, + "type": "pow_int", + "lb": domain.lower, + "ub": domain.upper - 1, # Upper bound exclusive + } + elif isinstance(sampler, Uniform): + return { + "name": par, + "type": "int", + "lb": domain.lower, + "ub": domain.upper - 1, # Upper bound exclusive + } + elif isinstance(domain, Categorical): + return { + "name": par, + "type": "cat", + "categories": list(domain.categories), + } + + raise ValueError( + "HEBO does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + for path, domain in domain_vars: + par = "/".join([str(p) for p in ((prefix,) + path if prefix else path)]) + value = resolve_value(par, domain) + params.append(value) + + return hebo.design_space.design_space.DesignSpace().parse(params) diff --git a/lib/python3.12/site-packages/ray/tune/search/hyperopt/__init__.py b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3b142015f04bf2eec6b1be7e54d545621dfe2953 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.hyperopt.hyperopt_search import HyperOptSearch + +__all__ = ["HyperOptSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32ff56f83b7c719a2089abe08969be7656950d11 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/hyperopt_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/hyperopt_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5275926fefd70c598520392df0b6e138314a9af Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/hyperopt/__pycache__/hyperopt_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/hyperopt/hyperopt_search.py b/lib/python3.12/site-packages/ray/tune/search/hyperopt/hyperopt_search.py new file mode 100644 index 0000000000000000000000000000000000000000..d1d3f964acb7117fde438195ea5ba092ad1f4802 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/hyperopt/hyperopt_search.py @@ -0,0 +1,559 @@ +import copy +import logging +from functools import partial +from typing import Any, Dict, List, Optional + +import numpy as np + +# Use cloudpickle instead of pickle to make lambda funcs in HyperOpt pickleable +from ray import cloudpickle +from ray.tune.error import TuneError +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + LogUniform, + Normal, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import assign_value, parse_spec_vars +from ray.tune.utils import flatten_dict + +try: + hyperopt_logger = logging.getLogger("hyperopt") + hyperopt_logger.setLevel(logging.WARNING) + import hyperopt as hpo + from hyperopt.pyll import Apply +except ImportError: + hpo = None + Apply = None + + +logger = logging.getLogger(__name__) + + +HYPEROPT_UNDEFINED_DETAILS = ( + " This issue can also come up with HyperOpt if your search space only " + "contains constant variables, which is not supported by HyperOpt. In that case, " + "don't pass any searcher or add sample variables to the search space." +) + + +class HyperOptSearch(Searcher): + """A wrapper around HyperOpt to provide trial suggestions. + + HyperOpt a Python library for serial and parallel optimization + over awkward search spaces, which may include real-valued, discrete, + and conditional dimensions. More info can be found at + http://hyperopt.github.io/hyperopt. + + HyperOptSearch uses the Tree-structured Parzen Estimators algorithm, + though it can be trivially extended to support any algorithm HyperOpt + supports. + + To use this search algorithm, you will need to install HyperOpt: + + .. code-block:: bash + + pip install -U hyperopt + + + Parameters: + space: HyperOpt configuration. Parameters will be sampled + from this configuration and will be used to override + parameters generated in the variant generation process. + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + n_initial_points: number of random evaluations of the + objective function before starting to approximate it with + tree parzen estimators. Defaults to 20. + random_state_seed: seed for reproducible + results. Defaults to None. + gamma: parameter governing the tree parzen + estimators suggestion algorithm. Defaults to 0.25. + + Tune automatically converts search spaces to HyperOpt's format: + + .. code-block:: python + + config = { + 'width': tune.uniform(0, 20), + 'height': tune.uniform(-100, 100), + 'activation': tune.choice(["relu", "tanh"]) + } + + current_best_params = [{ + 'width': 10, + 'height': 0, + 'activation': "relu", + }] + + hyperopt_search = HyperOptSearch( + metric="mean_loss", mode="min", + points_to_evaluate=current_best_params) + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=hyperopt_search + ), + param_space=config + ) + tuner.fit() + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + space = { + 'width': hp.uniform('width', 0, 20), + 'height': hp.uniform('height', -100, 100), + 'activation': hp.choice("activation", ["relu", "tanh"]) + } + + current_best_params = [{ + 'width': 10, + 'height': 0, + 'activation': "relu", + }] + + hyperopt_search = HyperOptSearch( + space, metric="mean_loss", mode="min", + points_to_evaluate=current_best_params) + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=hyperopt_search + ), + ) + tuner.fit() + + """ + + def __init__( + self, + space: Optional[Dict] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + n_initial_points: int = 20, + random_state_seed: Optional[int] = None, + gamma: float = 0.25, + ): + assert ( + hpo is not None + ), "HyperOpt must be installed! Run `pip install hyperopt`." + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + super(HyperOptSearch, self).__init__( + metric=metric, + mode=mode, + ) + # hyperopt internally minimizes, so "max" => -1 + if mode == "max": + self.metric_op = -1.0 + elif mode == "min": + self.metric_op = 1.0 + + if n_initial_points is None: + self.algo = hpo.tpe.suggest + else: + self.algo = partial(hpo.tpe.suggest, n_startup_jobs=n_initial_points) + if gamma is not None: + self.algo = partial(self.algo, gamma=gamma) + + self._points_to_evaluate = copy.deepcopy(points_to_evaluate) + + self._live_trial_mapping = {} + self.rstate = np.random.RandomState(random_state_seed) + + self.domain = None + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space) + self._space = space + self._setup_hyperopt() + + def _setup_hyperopt(self) -> None: + from hyperopt.fmin import generate_trials_to_calculate + + if not self._space: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + + HYPEROPT_UNDEFINED_DETAILS + ) + + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + if self._points_to_evaluate is None: + self._hpopt_trials = hpo.Trials() + self._points_to_evaluate = 0 + else: + assert isinstance(self._points_to_evaluate, (list, tuple)) + + for i in range(len(self._points_to_evaluate)): + config = self._points_to_evaluate[i] + self._convert_categories_to_indices(config) + # HyperOpt treats initial points as LIFO, reverse to get FIFO + self._points_to_evaluate = list(reversed(self._points_to_evaluate)) + self._hpopt_trials = generate_trials_to_calculate(self._points_to_evaluate) + self._hpopt_trials.refresh() + self._points_to_evaluate = len(self._points_to_evaluate) + + self.domain = hpo.Domain(lambda spc: spc, self._space) + + def _convert_categories_to_indices(self, config) -> None: + """Convert config parameters for categories into hyperopt-compatible + representations where instead the index of the category is expected.""" + + def _lookup(config_dict, space_dict, key): + if isinstance(config_dict[key], dict): + for k in config_dict[key]: + _lookup(config_dict[key], space_dict[key], k) + else: + if ( + key in space_dict + and isinstance(space_dict[key], hpo.base.pyll.Apply) + and space_dict[key].name == "switch" + ): + if len(space_dict[key].pos_args) > 0: + categories = [ + a.obj + for a in space_dict[key].pos_args[1:] + if a.name == "literal" + ] + try: + idx = categories.index(config_dict[key]) + except ValueError as exc: + msg = ( + f"Did not find category with value " + f"`{config_dict[key]}` in " + f"hyperopt parameter `{key}`. " + ) + + if isinstance(config_dict[key], int): + msg += ( + "In previous versions, a numerical " + "index was expected for categorical " + "values of `points_to_evaluate`, " + "but in ray>=1.2.0, the categorical " + "value is expected to be directly " + "provided. " + ) + + msg += "Please make sure the specified category is valid." + raise ValueError(msg) from exc + config_dict[key] = idx + + for k in config: + _lookup(config, self._space, k) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self.domain: + return False + space = self.convert_search_space(config) + self._space = space + + if metric: + self._metric = metric + if mode: + self._mode = mode + + self.metric_op = -1.0 if self._mode == "max" else 1.0 + + self._setup_hyperopt() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self.domain: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + + HYPEROPT_UNDEFINED_DETAILS + ) + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + if self._points_to_evaluate > 0: + using_point_to_evaluate = True + new_trial = self._hpopt_trials.trials[self._points_to_evaluate - 1] + self._points_to_evaluate -= 1 + else: + using_point_to_evaluate = False + new_ids = self._hpopt_trials.new_trial_ids(1) + self._hpopt_trials.refresh() + + # Get new suggestion from Hyperopt + new_trials = self.algo( + new_ids, + self.domain, + self._hpopt_trials, + self.rstate.randint(2**31 - 1), + ) + self._hpopt_trials.insert_trial_docs(new_trials) + self._hpopt_trials.refresh() + new_trial = new_trials[0] + self._live_trial_mapping[trial_id] = (new_trial["tid"], new_trial) + + # Taken from HyperOpt.base.evaluate + config = hpo.base.spec_from_misc(new_trial["misc"]) + + # We have to flatten nested spaces here so parameter names match + config = flatten_dict(config, flatten_list=True) + + ctrl = hpo.base.Ctrl(self._hpopt_trials, current_trial=new_trial) + memo = self.domain.memo_from_config(config) + hpo.utils.use_obj_for_literal_in_memo( + self.domain.expr, ctrl, hpo.base.Ctrl, memo + ) + + try: + suggested_config = hpo.pyll.rec_eval( + self.domain.expr, + memo=memo, + print_node_on_error=self.domain.rec_eval_print_node_on_error, + ) + except (AssertionError, TypeError) as e: + if using_point_to_evaluate and ( + isinstance(e, AssertionError) or "GarbageCollected" in str(e) + ): + raise ValueError( + "HyperOpt encountered a GarbageCollected switch argument. " + "Usually this is caused by a config in " + "`points_to_evaluate` " + "missing a key present in `space`. Ensure that " + "`points_to_evaluate` contains " + "all non-constant keys from `space`.\n" + "Config from `points_to_evaluate`: " + f"{config}\n" + "HyperOpt search space: " + f"{self._space}" + ) from e + raise e + return copy.deepcopy(suggested_config) + + def on_trial_result(self, trial_id: str, result: Dict) -> None: + ho_trial = self._get_hyperopt_trial(trial_id) + if ho_trial is None: + return + now = hpo.utils.coarse_utcnow() + ho_trial["book_time"] = now + ho_trial["refresh_time"] = now + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ) -> None: + """Notification for the completion of trial. + + The result is internally negated when interacting with HyperOpt + so that HyperOpt can "maximize" this value, as it minimizes on default. + """ + ho_trial = self._get_hyperopt_trial(trial_id) + if ho_trial is None: + return + ho_trial["refresh_time"] = hpo.utils.coarse_utcnow() + if error: + ho_trial["state"] = hpo.base.JOB_STATE_ERROR + ho_trial["misc"]["error"] = (str(TuneError), "Tune Error") + self._hpopt_trials.refresh() + elif result: + self._process_result(trial_id, result) + del self._live_trial_mapping[trial_id] + + def _process_result(self, trial_id: str, result: Dict) -> None: + ho_trial = self._get_hyperopt_trial(trial_id) + if not ho_trial: + return + ho_trial["refresh_time"] = hpo.utils.coarse_utcnow() + + ho_trial["state"] = hpo.base.JOB_STATE_DONE + hp_result = self._to_hyperopt_result(result) + ho_trial["result"] = hp_result + self._hpopt_trials.refresh() + + def _to_hyperopt_result(self, result: Dict) -> Dict: + try: + return {"loss": self.metric_op * result[self.metric], "status": "ok"} + except KeyError as e: + raise RuntimeError( + f"Hyperopt expected to see the metric `{self.metric}` in the " + f"last result, but it was not found. To fix this, make " + f"sure your call to `tune.report` or your return value of " + f"your trainable class `step()` contains the above metric " + f"as a key." + ) from e + + def _get_hyperopt_trial(self, trial_id: str) -> Optional[Dict]: + if trial_id not in self._live_trial_mapping: + return + hyperopt_tid = self._live_trial_mapping[trial_id][0] + return [t for t in self._hpopt_trials.trials if t["tid"] == hyperopt_tid][0] + + def get_state(self) -> Dict: + return { + "hyperopt_trials": self._hpopt_trials, + "rstate": self.rstate.get_state(), + } + + def set_state(self, state: Dict) -> None: + self._hpopt_trials = state["hyperopt_trials"] + self.rstate.set_state(state["rstate"]) + + def save(self, checkpoint_path: str) -> None: + save_object = self.__dict__.copy() + save_object["__rstate"] = self.rstate.get_state() + with open(checkpoint_path, "wb") as f: + cloudpickle.dump(save_object, f) + + def restore(self, checkpoint_path: str) -> None: + with open(checkpoint_path, "rb") as f: + save_object = cloudpickle.load(f) + + if "__rstate" not in save_object: + # Backwards compatibility + self.set_state(save_object) + else: + self.rstate.set_state(save_object.pop("__rstate")) + self.__dict__.update(save_object) + + @staticmethod + def convert_search_space(spec: Dict, prefix: str = "") -> Dict: + spec = copy.deepcopy(spec) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if not domain_vars and not grid_vars: + return {} + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a HyperOpt search space." + ) + + def resolve_value(par: str, domain: Domain) -> Any: + quantize = None + + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + quantize = sampler.q + sampler = sampler.sampler + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + if quantize: + return hpo.hp.qloguniform( + par, np.log(domain.lower), np.log(domain.upper), quantize + ) + return hpo.hp.loguniform( + par, np.log(domain.lower), np.log(domain.upper) + ) + elif isinstance(sampler, Uniform): + if quantize: + return hpo.hp.quniform( + par, domain.lower, domain.upper, quantize + ) + return hpo.hp.uniform(par, domain.lower, domain.upper) + elif isinstance(sampler, Normal): + if quantize: + return hpo.hp.qnormal(par, sampler.mean, sampler.sd, quantize) + return hpo.hp.normal(par, sampler.mean, sampler.sd) + + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + if quantize: + return hpo.base.pyll.scope.int( + hpo.hp.qloguniform( + par, + np.log(domain.lower), + np.log(domain.upper), + quantize, + ) + ) + return hpo.base.pyll.scope.int( + hpo.hp.qloguniform( + par, np.log(domain.lower), np.log(domain.upper - 1), 1.0 + ) + ) + elif isinstance(sampler, Uniform): + if quantize: + return hpo.base.pyll.scope.int( + hpo.hp.quniform( + par, domain.lower, domain.upper - 1, quantize + ) + ) + return hpo.hp.uniformint(par, domain.lower, high=domain.upper - 1) + elif isinstance(domain, Categorical): + if isinstance(sampler, Uniform): + return hpo.hp.choice( + par, + [ + ( + HyperOptSearch.convert_search_space( + category, prefix=par + ) + if isinstance(category, dict) + else ( + HyperOptSearch.convert_search_space( + dict(enumerate(category)), prefix=f"{par}/{i}" + ) + if isinstance(category, list) + and len(category) > 0 + and isinstance(category[0], Domain) + else ( + resolve_value(f"{par}/{i}", category) + if isinstance(category, Domain) + else category + ) + ) + ) + for i, category in enumerate(domain.categories) + ], + ) + + raise ValueError( + "HyperOpt does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + for path, domain in domain_vars: + par = "/".join([str(p) for p in ((prefix,) + path if prefix else path)]) + value = resolve_value(par, domain) + assign_value(spec, path, value) + + return spec diff --git a/lib/python3.12/site-packages/ray/tune/search/nevergrad/__init__.py b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..da69e05d07ddabd6dba42739a3173edb1befa752 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.nevergrad.nevergrad_search import NevergradSearch + +__all__ = ["NevergradSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b622ec74054cae7d0508615b76be4bc838545388 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/nevergrad_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/nevergrad_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1a4d17c57328dd6d41ab73f9de0e44c32ea1325 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/nevergrad/__pycache__/nevergrad_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/nevergrad/nevergrad_search.py b/lib/python3.12/site-packages/ray/tune/search/nevergrad/nevergrad_search.py new file mode 100644 index 0000000000000000000000000000000000000000..e26c9de2fe25cf2b0c6f0c60ef0641aed04adc39 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/nevergrad/nevergrad_search.py @@ -0,0 +1,374 @@ +import inspect +import logging +import math +import pickle +from typing import Dict, List, Optional, Sequence, Type, Union + +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + LogUniform, + Quantized, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import flatten_dict, unflatten_dict + +try: + import nevergrad as ng + from nevergrad.optimization import Optimizer + from nevergrad.optimization.base import ConfiguredOptimizer + + Parameter = ng.p.Parameter +except ImportError: + ng = None + Optimizer = None + ConfiguredOptimizer = None + Parameter = None + +logger = logging.getLogger(__name__) + + +class NevergradSearch(Searcher): + """Uses Nevergrad to optimize hyperparameters. + + Nevergrad is an open source tool from Facebook for derivative free + optimization. More info can be found at: + https://github.com/facebookresearch/nevergrad. + + You will need to install Nevergrad via the following command: + + .. code-block:: bash + + $ pip install nevergrad + + Parameters: + optimizer: Optimizer class provided from Nevergrad. + See here for available optimizers: + https://facebookresearch.github.io/nevergrad/optimizers_ref.html#optimizers + This can also be an instance of a `ConfiguredOptimizer`. See the + section on configured optimizers in the above link. + optimizer_kwargs: Kwargs passed in when instantiating the `optimizer` + space: Nevergrad parametrization + to be passed to optimizer on instantiation, or list of parameter + names if you passed an optimizer object. + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + + Tune automatically converts search spaces to Nevergrad's format: + + .. code-block:: python + + import nevergrad as ng + + config = { + "width": tune.uniform(0, 20), + "height": tune.uniform(-100, 100), + "activation": tune.choice(["relu", "tanh"]) + } + + current_best_params = [{ + "width": 10, + "height": 0, + "activation": relu", + }] + + ng_search = NevergradSearch( + optimizer=ng.optimizers.OnePlusOne, + metric="mean_loss", + mode="min", + points_to_evaluate=current_best_params) + + run(my_trainable, config=config, search_alg=ng_search) + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + import nevergrad as ng + + space = ng.p.Dict( + width=ng.p.Scalar(lower=0, upper=20), + height=ng.p.Scalar(lower=-100, upper=100), + activation=ng.p.Choice(choices=["relu", "tanh"]) + ) + + ng_search = NevergradSearch( + optimizer=ng.optimizers.OnePlusOne, + space=space, + metric="mean_loss", + mode="min") + + run(my_trainable, search_alg=ng_search) + + """ + + def __init__( + self, + optimizer: Optional[ + Union[Optimizer, Type[Optimizer], ConfiguredOptimizer] + ] = None, + optimizer_kwargs: Optional[Dict] = None, + space: Optional[Union[Dict, Parameter]] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + ): + assert ( + ng is not None + ), """Nevergrad must be installed! + You can install Nevergrad with the command: + `pip install nevergrad`.""" + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + + super(NevergradSearch, self).__init__(metric=metric, mode=mode) + + self._space = None + self._opt_factory = None + self._nevergrad_opt = None + self._optimizer_kwargs = optimizer_kwargs or {} + + if points_to_evaluate is None: + self._points_to_evaluate = None + elif not isinstance(points_to_evaluate, Sequence): + raise ValueError( + "Invalid object type passed for `points_to_evaluate`: " + f"{type(points_to_evaluate)}. " + "Please pass a list of points (dictionaries) instead." + ) + else: + self._points_to_evaluate = list(points_to_evaluate) + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)) + ) + space = self.convert_search_space(space) + + if isinstance(optimizer, Optimizer): + if space is not None and not isinstance(space, list): + raise ValueError( + "If you pass a configured optimizer to Nevergrad, either " + "pass a list of parameter names or None as the `space` " + "parameter." + ) + if self._optimizer_kwargs: + raise ValueError( + "If you pass in optimizer kwargs, either pass " + "an `Optimizer` subclass or an instance of " + "`ConfiguredOptimizer`." + ) + + self._parameters = space + self._nevergrad_opt = optimizer + elif ( + inspect.isclass(optimizer) and issubclass(optimizer, Optimizer) + ) or isinstance(optimizer, ConfiguredOptimizer): + self._opt_factory = optimizer + self._parameters = None + self._space = space + else: + raise ValueError( + "The `optimizer` argument passed to NevergradSearch must be " + "either an `Optimizer` or a `ConfiguredOptimizer`." + ) + + self._live_trial_mapping = {} + + if self._nevergrad_opt or self._space: + self._setup_nevergrad() + + def _setup_nevergrad(self): + if self._opt_factory: + self._nevergrad_opt = self._opt_factory( + self._space, **self._optimizer_kwargs + ) + + # nevergrad.tell internally minimizes, so "max" => -1 + if self._mode == "max": + self._metric_op = -1.0 + elif self._mode == "min": + self._metric_op = 1.0 + + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + if hasattr(self._nevergrad_opt, "instrumentation"): # added in v0.2.0 + if self._nevergrad_opt.instrumentation.kwargs: + if self._nevergrad_opt.instrumentation.args: + raise ValueError("Instrumented optimizers should use kwargs only") + if self._parameters is not None: + raise ValueError( + "Instrumented optimizers should provide " + "None as parameter_names" + ) + else: + if self._parameters is None: + raise ValueError( + "Non-instrumented optimizers should have " + "a list of parameter_names" + ) + if len(self._nevergrad_opt.instrumentation.args) != 1: + raise ValueError("Instrumented optimizers should use kwargs only") + if self._parameters is not None and self._nevergrad_opt.dimension != len( + self._parameters + ): + raise ValueError( + "len(parameters_names) must match optimizer " + "dimension for non-instrumented optimizers" + ) + + if self._points_to_evaluate: + # Nevergrad is LIFO, so we add the points to evaluate in reverse + # order. + for i in range(len(self._points_to_evaluate) - 1, -1, -1): + self._nevergrad_opt.suggest(self._points_to_evaluate[i]) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self._nevergrad_opt or self._space: + return False + space = self.convert_search_space(config) + self._space = space + + if metric: + self._metric = metric + if mode: + self._mode = mode + + self._setup_nevergrad() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._nevergrad_opt: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + suggested_config = self._nevergrad_opt.ask() + + self._live_trial_mapping[trial_id] = suggested_config + # in v0.2.0+, output of ask() is a Candidate, + # with fields args and kwargs + if not suggested_config.kwargs: + if self._parameters: + return unflatten_dict( + dict(zip(self._parameters, suggested_config.args[0])) + ) + return unflatten_dict(suggested_config.value) + else: + return unflatten_dict(suggested_config.kwargs) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + """Notification for the completion of trial. + + The result is internally negated when interacting with Nevergrad + so that Nevergrad Optimizers can "maximize" this value, + as it minimizes on default. + """ + if result: + self._process_result(trial_id, result) + + self._live_trial_mapping.pop(trial_id) + + def _process_result(self, trial_id: str, result: Dict): + ng_trial_info = self._live_trial_mapping[trial_id] + self._nevergrad_opt.tell(ng_trial_info, self._metric_op * result[self._metric]) + + def save(self, checkpoint_path: str): + save_object = self.__dict__ + with open(checkpoint_path, "wb") as outputFile: + pickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path: str): + with open(checkpoint_path, "rb") as inputFile: + save_object = pickle.load(inputFile) + self.__dict__.update(save_object) + + @staticmethod + def convert_search_space(spec: Dict) -> Parameter: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a Nevergrad search space." + ) + + # Flatten and resolve again after checking for grid search. + spec = flatten_dict(spec, prevent_delimiter=True) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + def resolve_value(domain: Domain) -> Parameter: + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + logger.warning( + "Nevergrad does not support quantization. Dropped quantization." + ) + sampler = sampler.get_sampler() + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + return ng.p.Log( + lower=domain.lower, upper=domain.upper, exponent=math.e + ) + return ng.p.Scalar(lower=domain.lower, upper=domain.upper) + + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + return ng.p.Log( + lower=domain.lower, + upper=domain.upper - 1, # Upper bound exclusive + exponent=math.e, + ).set_integer_casting() + return ng.p.Scalar( + lower=domain.lower, + upper=domain.upper - 1, # Upper bound exclusive + ).set_integer_casting() + + elif isinstance(domain, Categorical): + return ng.p.Choice(choices=domain.categories) + + raise ValueError( + "Nevergrad does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + # Parameter name is e.g. "a/b/c" for nested dicts + space = {"/".join(path): resolve_value(domain) for path, domain in domain_vars} + + return ng.p.Dict(**space) diff --git a/lib/python3.12/site-packages/ray/tune/search/optuna/__init__.py b/lib/python3.12/site-packages/ray/tune/search/optuna/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..74b450f2e684c73a0dd6178763839ba0bd3cf03a --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/optuna/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.optuna.optuna_search import OptunaSearch + +__all__ = ["OptunaSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..803bc45e2bd1f59c42e110a1e8a7751ae71bed30 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/optuna_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/optuna_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85fdf5ad4e92e6979585ee0735aeab319cb0d6e1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/optuna/__pycache__/optuna_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/optuna/optuna_search.py b/lib/python3.12/site-packages/ray/tune/search/optuna/optuna_search.py new file mode 100644 index 0000000000000000000000000000000000000000..c018301a522f8408b8180ce69b8b8ba4cd97aae2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/optuna/optuna_search.py @@ -0,0 +1,730 @@ +import functools +import logging +import pickle +import time +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from packaging import version + +from ray.air.constants import TRAINING_ITERATION +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + LogUniform, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import flatten_dict, unflatten_dict, validate_warmstart + +try: + import optuna as ot + from optuna.distributions import BaseDistribution as OptunaDistribution + from optuna.samplers import BaseSampler + from optuna.storages import BaseStorage + from optuna.trial import Trial as OptunaTrial, TrialState as OptunaTrialState +except ImportError: + ot = None + OptunaDistribution = None + BaseSampler = None + BaseStorage = None + OptunaTrialState = None + OptunaTrial = None + +logger = logging.getLogger(__name__) + +# print a warning if define by run function takes longer than this to execute +DEFINE_BY_RUN_WARN_THRESHOLD_S = 1 # 1 is arbitrary + + +class _OptunaTrialSuggestCaptor: + """Utility to capture returned values from Optuna's suggest_ methods. + + This will wrap around the ``optuna.Trial` object and decorate all + `suggest_` callables with a function capturing the returned value, + which will be saved in the ``captured_values`` dict. + """ + + def __init__(self, ot_trial: OptunaTrial) -> None: + self.ot_trial = ot_trial + self.captured_values: Dict[str, Any] = {} + + def _get_wrapper(self, func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs): + # name is always the first arg for suggest_ methods + name = kwargs.get("name", args[0]) + ret = func(*args, **kwargs) + self.captured_values[name] = ret + return ret + + return wrapper + + def __getattr__(self, item_name: str) -> Any: + item = getattr(self.ot_trial, item_name) + if item_name.startswith("suggest_") and callable(item): + return self._get_wrapper(item) + return item + + +class OptunaSearch(Searcher): + """A wrapper around Optuna to provide trial suggestions. + + `Optuna `_ is a hyperparameter optimization library. + In contrast to other libraries, it employs define-by-run style + hyperparameter definitions. + + This Searcher is a thin wrapper around Optuna's search algorithms. + You can pass any Optuna sampler, which will be used to generate + hyperparameter suggestions. + + Multi-objective optimization is supported. + + Args: + space: Hyperparameter search space definition for + Optuna's sampler. This can be either a :class:`dict` with + parameter names as keys and ``optuna.distributions`` as values, + or a Callable - in which case, it should be a define-by-run + function using ``optuna.trial`` to obtain the hyperparameter + values. The function should return either a :class:`dict` of + constant values with names as keys, or None. + For more information, see https://optuna.readthedocs.io\ +/en/stable/tutorial/10_key_features/002_configurations.html. + + .. warning:: + No actual computation should take place in the define-by-run + function. Instead, put the training logic inside the function + or class trainable passed to ``tune.Tuner()``. + + metric: The training result objective value attribute. If + None but a mode was passed, the anonymous metric ``_metric`` + will be used per default. Can be a list of metrics for + multi-objective optimization. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. Can be a list of + modes for multi-objective optimization (corresponding to + ``metric``). + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + sampler: Optuna sampler used to + draw hyperparameter configurations. Defaults to ``MOTPESampler`` + for multi-objective optimization with Optuna<2.9.0, and + ``TPESampler`` in every other case. + See https://optuna.readthedocs.io/en/stable/reference/samplers/index.html + for available Optuna samplers. + + .. warning:: + Please note that with Optuna 2.10.0 and earlier + default ``MOTPESampler``/``TPESampler`` suffer + from performance issues when dealing with a large number of + completed trials (approx. >100). This will manifest as + a delay when suggesting new configurations. + This is an Optuna issue and may be fixed in a future + Optuna release. + study_name: Optuna study name that uniquely identifies the trial + results. Defaults to ``"optuna"``. + storage: Optuna storage used for storing trial results to + storages other than in-memory storage, + for instance optuna.storages.RDBStorage. + seed: Seed to initialize sampler with. This parameter is only + used when ``sampler=None``. In all other cases, the sampler + you pass should be initialized with the seed already. + evaluated_rewards: If you have previously evaluated the + parameters passed in as points_to_evaluate you can avoid + re-running those trials by passing in the reward attributes + as a list so the optimiser can be told the results without + needing to re-compute the trial. Must be the same length as + points_to_evaluate. + + .. warning:: + When using ``evaluated_rewards``, the search space ``space`` + must be provided as a :class:`dict` with parameter names as + keys and ``optuna.distributions`` instances as values. The + define-by-run search space definition is not yet supported with + this functionality. + + Tune automatically converts search spaces to Optuna's format: + + .. code-block:: python + + from ray.tune.search.optuna import OptunaSearch + + config = { + "a": tune.uniform(6, 8) + "b": tune.loguniform(1e-4, 1e-2) + } + + optuna_search = OptunaSearch( + metric="loss", + mode="min") + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + param_space=config, + ) + tuner.fit() + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + from ray.tune.search.optuna import OptunaSearch + import optuna + + space = { + "a": optuna.distributions.FloatDistribution(6, 8), + "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True), + } + + optuna_search = OptunaSearch( + space, + metric="loss", + mode="min") + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + ) + tuner.fit() + + # Equivalent Optuna define-by-run function approach: + + def define_search_space(trial: optuna.Trial): + trial.suggest_float("a", 6, 8) + trial.suggest_float("b", 1e-4, 1e-2, log=True) + # training logic goes into trainable, this is just + # for search space definition + + optuna_search = OptunaSearch( + define_search_space, + metric="loss", + mode="min") + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + ) + tuner.fit() + + Multi-objective optimization is supported: + + .. code-block:: python + + from ray.tune.search.optuna import OptunaSearch + import optuna + + space = { + "a": optuna.distributions.FloatDistribution(6, 8), + "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True), + } + + # Note you have to specify metric and mode here instead of + # in tune.TuneConfig + optuna_search = OptunaSearch( + space, + metric=["loss1", "loss2"], + mode=["min", "max"]) + + # Do not specify metric and mode here! + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + ) + tuner.fit() + + You can pass configs that will be evaluated first using + ``points_to_evaluate``: + + .. code-block:: python + + from ray.tune.search.optuna import OptunaSearch + import optuna + + space = { + "a": optuna.distributions.FloatDistribution(6, 8), + "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True), + } + + optuna_search = OptunaSearch( + space, + points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}] + metric="loss", + mode="min") + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + ) + tuner.fit() + + Avoid re-running evaluated trials by passing the rewards together with + `points_to_evaluate`: + + .. code-block:: python + + from ray.tune.search.optuna import OptunaSearch + import optuna + + space = { + "a": optuna.distributions.FloatDistribution(6, 8), + "b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True), + } + + optuna_search = OptunaSearch( + space, + points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}] + evaluated_rewards=[0.89, 0.42] + metric="loss", + mode="min") + + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=optuna_search, + ), + ) + tuner.fit() + + .. versionadded:: 0.8.8 + + """ + + def __init__( + self, + space: Optional[ + Union[ + Dict[str, "OptunaDistribution"], + List[Tuple], + Callable[["OptunaTrial"], Optional[Dict[str, Any]]], + ] + ] = None, + metric: Optional[Union[str, List[str]]] = None, + mode: Optional[Union[str, List[str]]] = None, + points_to_evaluate: Optional[List[Dict]] = None, + sampler: Optional["BaseSampler"] = None, + study_name: Optional[str] = None, + storage: Optional["BaseStorage"] = None, + seed: Optional[int] = None, + evaluated_rewards: Optional[List] = None, + ): + assert ot is not None, "Optuna must be installed! Run `pip install optuna`." + super(OptunaSearch, self).__init__(metric=metric, mode=mode) + + if isinstance(space, dict) and space: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(space) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self).__name__) + ) + space = self.convert_search_space(space) + else: + # Flatten to support nested dicts + space = flatten_dict(space, "/") + + self._space = space + + self._points_to_evaluate = points_to_evaluate or [] + self._evaluated_rewards = evaluated_rewards + if study_name: + self._study_name = study_name + else: + self._study_name = "optuna" # Fixed study name for in-memory storage + + if sampler and seed: + logger.warning( + "You passed an initialized sampler to `OptunaSearch`. The " + "`seed` parameter has to be passed to the sampler directly " + "and will be ignored." + ) + elif sampler: + assert isinstance(sampler, BaseSampler), ( + "You can only pass an instance of " + "`optuna.samplers.BaseSampler` " + "as a sampler to `OptunaSearcher`." + ) + + self._sampler = sampler + self._seed = seed + + if storage: + assert isinstance(storage, BaseStorage), ( + "The `storage` parameter in `OptunaSearcher` must be an instance " + "of `optuna.storages.BaseStorage`." + ) + # If storage is not provided, just set self._storage to None + # so that the default in-memory storage is used. + self._storage = storage + + self._completed_trials = set() + + self._ot_trials = {} + self._ot_study = None + if self._space: + self._setup_study(mode) + + def _setup_study(self, mode: Union[str, list]): + if self._metric is None and self._mode: + if isinstance(self._mode, list): + raise ValueError( + "If ``mode`` is a list (multi-objective optimization " + "case), ``metric`` must be defined." + ) + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + pruner = ot.pruners.NopPruner() + + if self._sampler: + sampler = self._sampler + elif isinstance(mode, list) and version.parse(ot.__version__) < version.parse( + "2.9.0" + ): + # MOTPESampler deprecated in Optuna>=2.9.0 + sampler = ot.samplers.MOTPESampler(seed=self._seed) + else: + sampler = ot.samplers.TPESampler(seed=self._seed) + + if isinstance(mode, list): + study_direction_args = dict( + directions=["minimize" if m == "min" else "maximize" for m in mode], + ) + else: + study_direction_args = dict( + direction="minimize" if mode == "min" else "maximize", + ) + + self._ot_study = ot.study.create_study( + storage=self._storage, + sampler=sampler, + pruner=pruner, + study_name=self._study_name, + load_if_exists=True, + **study_direction_args, + ) + + if self._points_to_evaluate: + validate_warmstart( + self._space, + self._points_to_evaluate, + self._evaluated_rewards, + validate_point_name_lengths=not callable(self._space), + ) + if self._evaluated_rewards: + for point, reward in zip( + self._points_to_evaluate, self._evaluated_rewards + ): + self.add_evaluated_point(point, reward) + else: + for point in self._points_to_evaluate: + self._ot_study.enqueue_trial(point) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self._space: + return False + space = self.convert_search_space(config) + self._space = space + if metric: + self._metric = metric + if mode: + self._mode = mode + + self._setup_study(self._mode) + return True + + def _suggest_from_define_by_run_func( + self, + func: Callable[["OptunaTrial"], Optional[Dict[str, Any]]], + ot_trial: "OptunaTrial", + ) -> Dict: + captor = _OptunaTrialSuggestCaptor(ot_trial) + time_start = time.time() + ret = func(captor) + time_taken = time.time() - time_start + if time_taken > DEFINE_BY_RUN_WARN_THRESHOLD_S: + warnings.warn( + "Define-by-run function passed in the `space` argument " + f"took {time_taken} seconds to " + "run. Ensure that actual computation, training takes " + "place inside Tune's train functions or Trainables " + "passed to `tune.Tuner()`." + ) + if ret is not None: + if not isinstance(ret, dict): + raise TypeError( + "The return value of the define-by-run function " + "passed in the `space` argument should be " + "either None or a `dict` with `str` keys. " + f"Got {type(ret)}." + ) + if not all(isinstance(k, str) for k in ret.keys()): + raise TypeError( + "At least one of the keys in the dict returned by the " + "define-by-run function passed in the `space` argument " + "was not a `str`." + ) + return {**captor.captured_values, **ret} if ret else captor.captured_values + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._space: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + if callable(self._space): + # Define-by-run case + if trial_id not in self._ot_trials: + self._ot_trials[trial_id] = self._ot_study.ask() + + ot_trial = self._ot_trials[trial_id] + + params = self._suggest_from_define_by_run_func(self._space, ot_trial) + else: + # Use Optuna ask interface (since version 2.6.0) + if trial_id not in self._ot_trials: + self._ot_trials[trial_id] = self._ot_study.ask( + fixed_distributions=self._space + ) + ot_trial = self._ot_trials[trial_id] + params = ot_trial.params + + return unflatten_dict(params) + + def on_trial_result(self, trial_id: str, result: Dict): + if isinstance(self.metric, list): + # Optuna doesn't support incremental results + # for multi-objective optimization + return + if trial_id in self._completed_trials: + logger.warning( + f"Received additional result for trial {trial_id}, but " + f"it already finished. Result: {result}" + ) + return + metric = result[self.metric] + step = result[TRAINING_ITERATION] + ot_trial = self._ot_trials[trial_id] + ot_trial.report(metric, step) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + if trial_id in self._completed_trials: + logger.warning( + f"Received additional completion for trial {trial_id}, but " + f"it already finished. Result: {result}" + ) + return + + ot_trial = self._ot_trials[trial_id] + + if result: + if isinstance(self.metric, list): + val = [result.get(metric, None) for metric in self.metric] + else: + val = result.get(self.metric, None) + else: + val = None + ot_trial_state = OptunaTrialState.COMPLETE + if val is None: + if error: + ot_trial_state = OptunaTrialState.FAIL + else: + ot_trial_state = OptunaTrialState.PRUNED + try: + self._ot_study.tell(ot_trial, val, state=ot_trial_state) + except Exception as exc: + logger.warning(exc) # E.g. if NaN was reported + + self._completed_trials.add(trial_id) + + def add_evaluated_point( + self, + parameters: Dict, + value: float, + error: bool = False, + pruned: bool = False, + intermediate_values: Optional[List[float]] = None, + ): + if not self._space: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="space" + ) + ) + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + if callable(self._space): + raise TypeError( + "Define-by-run function passed in `space` argument is not " + "yet supported when using `evaluated_rewards`. Please provide " + "an `OptunaDistribution` dict or pass a Ray Tune " + "search space to `tune.Tuner()`." + ) + + ot_trial_state = OptunaTrialState.COMPLETE + if error: + ot_trial_state = OptunaTrialState.FAIL + elif pruned: + ot_trial_state = OptunaTrialState.PRUNED + + if intermediate_values: + intermediate_values_dict = dict(enumerate(intermediate_values)) + else: + intermediate_values_dict = None + + # If the trial state is FAILED, the value must be `None` in Optuna==4.1.0 + # Reference: https://github.com/optuna/optuna/pull/5211 + # This is a temporary fix for the issue that Optuna enforces the value + # to be `None` if the trial state is FAILED. + # TODO (hpguo): A better solution may requires us to update the base class + # to allow the `value` arg in `add_evaluated_point` being `Optional[float]`. + if ot_trial_state == OptunaTrialState.FAIL: + value = None + + trial = ot.trial.create_trial( + state=ot_trial_state, + value=value, + params=parameters, + distributions=self._space, + intermediate_values=intermediate_values_dict, + ) + + self._ot_study.add_trial(trial) + + def save(self, checkpoint_path: str): + save_object = self.__dict__.copy() + with open(checkpoint_path, "wb") as outputFile: + pickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path: str): + with open(checkpoint_path, "rb") as inputFile: + save_object = pickle.load(inputFile) + if isinstance(save_object, dict): + self.__dict__.update(save_object) + else: + # Backwards compatibility + ( + self._sampler, + self._ot_trials, + self._ot_study, + self._points_to_evaluate, + self._evaluated_rewards, + ) = save_object + + @staticmethod + def convert_search_space(spec: Dict) -> Dict[str, Any]: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if not domain_vars and not grid_vars: + return {} + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to an Optuna search space." + ) + + # Flatten and resolve again after checking for grid search. + spec = flatten_dict(spec, prevent_delimiter=True) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + def resolve_value(domain: Domain) -> ot.distributions.BaseDistribution: + quantize = None + + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + quantize = sampler.q + sampler = sampler.sampler + if isinstance(sampler, LogUniform): + logger.warning( + "Optuna does not handle quantization in loguniform " + "sampling. The parameter will be passed but it will " + "probably be ignored." + ) + + if isinstance(domain, Float): + if isinstance(sampler, LogUniform): + if quantize: + logger.warning( + "Optuna does not support both quantization and " + "sampling from LogUniform. Dropped quantization." + ) + return ot.distributions.FloatDistribution( + domain.lower, domain.upper, log=True + ) + + elif isinstance(sampler, Uniform): + if quantize: + return ot.distributions.FloatDistribution( + domain.lower, domain.upper, step=quantize + ) + return ot.distributions.FloatDistribution( + domain.lower, domain.upper + ) + + elif isinstance(domain, Integer): + if isinstance(sampler, LogUniform): + return ot.distributions.IntDistribution( + domain.lower, domain.upper - 1, step=quantize or 1, log=True + ) + elif isinstance(sampler, Uniform): + # Upper bound should be inclusive for quantization and + # exclusive otherwise + return ot.distributions.IntDistribution( + domain.lower, + domain.upper - int(bool(not quantize)), + step=quantize or 1, + ) + elif isinstance(domain, Categorical): + if isinstance(sampler, Uniform): + return ot.distributions.CategoricalDistribution(domain.categories) + + raise ValueError( + "Optuna search does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + # Parameter name is e.g. "a/b/c" for nested dicts + values = {"/".join(path): resolve_value(domain) for path, domain in domain_vars} + + return values diff --git a/lib/python3.12/site-packages/ray/tune/search/repeater.py b/lib/python3.12/site-packages/ray/tune/search/repeater.py new file mode 100644 index 0000000000000000000000000000000000000000..c9de4a156091971ba326c6e0cfca4678d6dc2261 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/repeater.py @@ -0,0 +1,199 @@ +import copy +import logging +from typing import Dict, List, Optional + +import numpy as np + +from ray.tune.search.searcher import Searcher +from ray.tune.search.util import _set_search_properties_backwards_compatible +from ray.util import PublicAPI + +logger = logging.getLogger(__name__) + +TRIAL_INDEX = "__trial_index__" +"""str: A constant value representing the repeat index of the trial.""" + + +def _warn_num_samples(searcher: Searcher, num_samples: int): + if isinstance(searcher, Repeater) and num_samples % searcher.repeat: + logger.warning( + "`num_samples` is now expected to be the total number of trials, " + "including the repeat trials. For example, set num_samples=15 if " + "you intend to obtain 3 search algorithm suggestions and repeat " + "each suggestion 5 times. Any leftover trials " + "(num_samples mod repeat) will be ignored." + ) + + +class _TrialGroup: + """Internal class for grouping trials of same parameters. + + This is used when repeating trials for reducing training variance. + + Args: + primary_trial_id: Trial ID of the "primary trial". + This trial is the one that the Searcher is aware of. + config: Suggested configuration shared across all trials + in the trial group. + max_trials: Max number of trials to execute within this group. + + """ + + def __init__(self, primary_trial_id: str, config: Dict, max_trials: int = 1): + assert type(config) is dict, "config is not a dict, got {}".format(config) + self.primary_trial_id = primary_trial_id + self.config = config + self._trials = {primary_trial_id: None} + self.max_trials = max_trials + + def add(self, trial_id: str): + assert len(self._trials) < self.max_trials + self._trials.setdefault(trial_id, None) + + def full(self) -> bool: + return len(self._trials) == self.max_trials + + def report(self, trial_id: str, score: float): + assert trial_id in self._trials + if score is None: + raise ValueError("Internal Error: Score cannot be None.") + self._trials[trial_id] = score + + def finished_reporting(self) -> bool: + return ( + None not in self._trials.values() and len(self._trials) == self.max_trials + ) + + def scores(self) -> List[Optional[float]]: + return list(self._trials.values()) + + def count(self) -> int: + return len(self._trials) + + +@PublicAPI +class Repeater(Searcher): + """A wrapper algorithm for repeating trials of same parameters. + + Set tune.TuneConfig(num_samples=...) to be a multiple of `repeat`. For example, + set num_samples=15 if you intend to obtain 3 search algorithm suggestions + and repeat each suggestion 5 times. Any leftover trials + (num_samples mod repeat) will be ignored. + + It is recommended that you do not run an early-stopping TrialScheduler + simultaneously. + + Args: + searcher: Searcher object that the + Repeater will optimize. Note that the Searcher + will only see 1 trial among multiple repeated trials. + The result/metric passed to the Searcher upon + trial completion will be averaged among all repeats. + repeat: Number of times to generate a trial with a repeated + configuration. Defaults to 1. + set_index: Sets a tune.search.repeater.TRIAL_INDEX in + Trainable/Function config which corresponds to the index of the + repeated trial. This can be used for seeds. Defaults to True. + + Example: + + .. code-block:: python + + from ray.tune.search import Repeater + + search_alg = BayesOptSearch(...) + re_search_alg = Repeater(search_alg, repeat=10) + + # Repeat 2 samples 10 times each. + tuner = tune.Tuner( + trainable, + tune_config=tune.TuneConfig( + search_alg=re_search_alg, + num_samples=20, + ), + ) + tuner.fit() + + """ + + def __init__(self, searcher: Searcher, repeat: int = 1, set_index: bool = True): + self.searcher = searcher + self.repeat = repeat + self._set_index = set_index + self._groups = [] + self._trial_id_to_group = {} + self._current_group = None + super(Repeater, self).__init__( + metric=self.searcher.metric, mode=self.searcher.mode + ) + + def suggest(self, trial_id: str) -> Optional[Dict]: + if self._current_group is None or self._current_group.full(): + config = self.searcher.suggest(trial_id) + if config is None: + return config + self._current_group = _TrialGroup( + trial_id, copy.deepcopy(config), max_trials=self.repeat + ) + self._groups.append(self._current_group) + index_in_group = 0 + else: + index_in_group = self._current_group.count() + self._current_group.add(trial_id) + + config = self._current_group.config.copy() + if self._set_index: + config[TRIAL_INDEX] = index_in_group + self._trial_id_to_group[trial_id] = self._current_group + return config + + def on_trial_complete(self, trial_id: str, result: Optional[Dict] = None, **kwargs): + """Stores the score for and keeps track of a completed trial. + + Stores the metric of a trial as nan if any of the following conditions + are met: + + 1. ``result`` is empty or not provided. + 2. ``result`` is provided but no metric was provided. + + """ + if trial_id not in self._trial_id_to_group: + logger.error( + "Trial {} not in group; cannot report score. " + "Seen trials: {}".format(trial_id, list(self._trial_id_to_group)) + ) + trial_group = self._trial_id_to_group[trial_id] + if not result or self.searcher.metric not in result: + score = np.nan + else: + score = result[self.searcher.metric] + trial_group.report(trial_id, score) + + if trial_group.finished_reporting(): + scores = trial_group.scores() + self.searcher.on_trial_complete( + trial_group.primary_trial_id, + result={self.searcher.metric: np.nanmean(scores)}, + **kwargs + ) + + def get_state(self) -> Dict: + self_state = self.__dict__.copy() + del self_state["searcher"] + return self_state + + def set_state(self, state: Dict): + self.__dict__.update(state) + + def save(self, checkpoint_path: str): + self.searcher.save(checkpoint_path) + + def restore(self, checkpoint_path: str): + self.searcher.restore(checkpoint_path) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + return _set_search_properties_backwards_compatible( + self.searcher.set_search_properties, metric, mode, config, **spec + ) diff --git a/lib/python3.12/site-packages/ray/tune/search/sample.py b/lib/python3.12/site-packages/ray/tune/search/sample.py new file mode 100644 index 0000000000000000000000000000000000000000..ccf65034a4a0edb3d0397c9fa7268536e4d03476 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/sample.py @@ -0,0 +1,767 @@ +import logging +import warnings +from copy import copy +from inspect import signature +from math import isclose +from typing import Any, Callable, Dict, List, Optional, Sequence, Union + +import numpy as np + +# Backwards compatibility +from ray.util.annotations import DeveloperAPI, PublicAPI, RayDeprecationWarning + +try: + # Added in numpy>=1.17 but we require numpy>=1.16 + np_random_generator = np.random.Generator + LEGACY_RNG = False +except AttributeError: + + class np_random_generator: + pass + + LEGACY_RNG = True + +logger = logging.getLogger(__name__) + + +_MISSING = object() # Sentinel for missing parameters. + + +def _warn_for_base() -> None: + warnings.warn( + ( + "The `base` argument is deprecated. " + "Please remove it as it is not actually needed in this method." + ), + RayDeprecationWarning, + stacklevel=2, + ) + + +class _BackwardsCompatibleNumpyRng: + """Thin wrapper to ensure backwards compatibility between + new and old numpy randomness generators. + """ + + _rng = None + + def __init__( + self, + generator_or_seed: Optional[ + Union["np_random_generator", np.random.RandomState, int] + ] = None, + ): + if generator_or_seed is None or isinstance( + generator_or_seed, (np.random.RandomState, np_random_generator) + ): + self._rng = generator_or_seed + elif LEGACY_RNG: + self._rng = np.random.RandomState(generator_or_seed) + else: + self._rng = np.random.default_rng(generator_or_seed) + + @property + def legacy_rng(self) -> bool: + return not isinstance(self._rng, np_random_generator) + + @property + def rng(self): + # don't set self._rng to np.random to avoid picking issues + return self._rng if self._rng is not None else np.random + + def __getattr__(self, name: str) -> Any: + # https://numpy.org/doc/stable/reference/random/new-or-different.html + if self.legacy_rng: + if name == "integers": + name = "randint" + elif name == "random": + name = "rand" + return getattr(self.rng, name) + + +RandomState = Union[ + None, _BackwardsCompatibleNumpyRng, np_random_generator, np.random.RandomState, int +] + + +@DeveloperAPI +class Domain: + """Base class to specify a type and valid range to sample parameters from. + + This base class is implemented by parameter spaces, like float ranges + (``Float``), integer ranges (``Integer``), or categorical variables + (``Categorical``). The ``Domain`` object contains information about + valid values (e.g. minimum and maximum values), and exposes methods that + allow specification of specific samplers (e.g. ``uniform()`` or + ``loguniform()``). + + """ + + sampler = None + default_sampler_cls = None + + def cast(self, value): + """Cast value to domain type""" + return value + + def set_sampler(self, sampler, allow_override=False): + if self.sampler and not allow_override: + raise ValueError( + "You can only choose one sampler for parameter " + "domains. Existing sampler for parameter {}: " + "{}. Tried to add {}".format( + self.__class__.__name__, self.sampler, sampler + ) + ) + self.sampler = sampler + + def get_sampler(self): + sampler = self.sampler + if not sampler: + sampler = self.default_sampler_cls() + return sampler + + def sample( + self, + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + sampler = self.get_sampler() + return sampler.sample(self, config=config, size=size, random_state=random_state) + + def is_grid(self): + return isinstance(self.sampler, Grid) + + def is_function(self): + return False + + def is_valid(self, value: Any): + """Returns True if `value` is a valid value in this domain.""" + raise NotImplementedError + + @property + def domain_str(self): + return "(unknown)" + + +@DeveloperAPI +class Sampler: + def sample( + self, + domain: Domain, + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + raise NotImplementedError + + +@DeveloperAPI +class BaseSampler(Sampler): + def __str__(self): + return "Base" + + +@DeveloperAPI +class Uniform(Sampler): + def __str__(self): + return "Uniform" + + +@DeveloperAPI +class LogUniform(Sampler): + def __init__(self, base: object = _MISSING): + if base is not _MISSING: + _warn_for_base() + + def __str__(self): + return "LogUniform" + + +@DeveloperAPI +class Normal(Sampler): + def __init__(self, mean: float = 0.0, sd: float = 0.0): + self.mean = mean + self.sd = sd + + assert self.sd > 0, "SD has to be strictly greater than 0" + + def __str__(self): + return "Normal" + + +@DeveloperAPI +class Grid(Sampler): + """Dummy sampler used for grid search""" + + def sample( + self, + domain: Domain, + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + return RuntimeError("Do not call `sample()` on grid.") + + +@DeveloperAPI +class Float(Domain): + class _Uniform(Uniform): + def sample( + self, + domain: "Float", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + assert domain.lower > float("-inf"), "Uniform needs a lower bound" + assert domain.upper < float("inf"), "Uniform needs a upper bound" + items = random_state.uniform(domain.lower, domain.upper, size=size) + return items if len(items) > 1 else domain.cast(items[0]) + + class _LogUniform(LogUniform): + def sample( + self, + domain: "Float", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + assert domain.lower > 0, "LogUniform needs a lower bound greater than 0" + assert ( + 0 < domain.upper < float("inf") + ), "LogUniform needs a upper bound greater than 0" + logmin = np.log(domain.lower) + logmax = np.log(domain.upper) + + items = np.exp(random_state.uniform(logmin, logmax, size=size)) + return items if len(items) > 1 else domain.cast(items[0]) + + class _Normal(Normal): + def sample( + self, + domain: "Float", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + assert not domain.lower or domain.lower == float( + "-inf" + ), "Normal sampling does not allow a lower value bound." + assert not domain.upper or domain.upper == float( + "inf" + ), "Normal sampling does not allow a upper value bound." + items = random_state.normal(self.mean, self.sd, size=size) + return items if len(items) > 1 else domain.cast(items[0]) + + default_sampler_cls = _Uniform + + def __init__(self, lower: Optional[float], upper: Optional[float]): + # Need to explicitly check for None + self.lower = lower if lower is not None else float("-inf") + self.upper = upper if upper is not None else float("inf") + + def cast(self, value): + return float(value) + + def uniform(self): + if not self.lower > float("-inf"): + raise ValueError( + "Uniform requires a lower bound. Make sure to set the " + "`lower` parameter of `Float()`." + ) + if not self.upper < float("inf"): + raise ValueError( + "Uniform requires a upper bound. Make sure to set the " + "`upper` parameter of `Float()`." + ) + new = copy(self) + new.set_sampler(self._Uniform()) + return new + + def loguniform(self, base: object = _MISSING): + if base is not _MISSING: + _warn_for_base() + if not self.lower > 0: + raise ValueError( + "LogUniform requires a lower bound greater than 0." + f"Got: {self.lower}. Did you pass a variable that has " + "been log-transformed? If so, pass the non-transformed value " + "instead." + ) + if not 0 < self.upper < float("inf"): + raise ValueError( + "LogUniform requires a upper bound greater than 0. " + f"Got: {self.lower}. Did you pass a variable that has " + "been log-transformed? If so, pass the non-transformed value " + "instead." + ) + new = copy(self) + new.set_sampler(self._LogUniform()) + return new + + def normal(self, mean=0.0, sd=1.0): + new = copy(self) + new.set_sampler(self._Normal(mean, sd)) + return new + + def quantized(self, q: float): + if self.lower > float("-inf") and not isclose( + self.lower / q, round(self.lower / q) + ): + raise ValueError( + f"Your lower variable bound {self.lower} is not divisible by " + f"quantization factor {q}." + ) + if self.upper < float("inf") and not isclose( + self.upper / q, round(self.upper / q) + ): + raise ValueError( + f"Your upper variable bound {self.upper} is not divisible by " + f"quantization factor {q}." + ) + + new = copy(self) + new.set_sampler(Quantized(new.get_sampler(), q), allow_override=True) + return new + + def is_valid(self, value: float): + return self.lower <= value <= self.upper + + @property + def domain_str(self): + return f"({self.lower}, {self.upper})" + + +@DeveloperAPI +class Integer(Domain): + class _Uniform(Uniform): + def sample( + self, + domain: "Integer", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + items = random_state.integers(domain.lower, domain.upper, size=size) + return items if len(items) > 1 else domain.cast(items[0]) + + class _LogUniform(LogUniform): + def sample( + self, + domain: "Integer", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + assert domain.lower > 0, "LogUniform needs a lower bound greater than 0" + assert ( + 0 < domain.upper < float("inf") + ), "LogUniform needs a upper bound greater than 0" + logmin = np.log(domain.lower) + logmax = np.log(domain.upper) + + items = np.exp(random_state.uniform(logmin, logmax, size=size)) + items = np.floor(items).astype(int) + return items if len(items) > 1 else domain.cast(items[0]) + + default_sampler_cls = _Uniform + + def __init__(self, lower, upper): + self.lower = lower + self.upper = upper + + def cast(self, value): + return int(value) + + def quantized(self, q: int): + new = copy(self) + new.set_sampler(Quantized(new.get_sampler(), q), allow_override=True) + return new + + def uniform(self): + new = copy(self) + new.set_sampler(self._Uniform()) + return new + + def loguniform(self, base: object = _MISSING): + if base is not _MISSING: + _warn_for_base() + if not self.lower > 0: + raise ValueError( + "LogUniform requires a lower bound greater than 0." + f"Got: {self.lower}. Did you pass a variable that has " + "been log-transformed? If so, pass the non-transformed value " + "instead." + ) + if not 0 < self.upper < float("inf"): + raise ValueError( + "LogUniform requires a upper bound greater than 0. " + f"Got: {self.lower}. Did you pass a variable that has " + "been log-transformed? If so, pass the non-transformed value " + "instead." + ) + new = copy(self) + new.set_sampler(self._LogUniform()) + return new + + def is_valid(self, value: int): + return self.lower <= value <= self.upper + + @property + def domain_str(self): + return f"({self.lower}, {self.upper})" + + +@DeveloperAPI +class Categorical(Domain): + class _Uniform(Uniform): + def sample( + self, + domain: "Categorical", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + # do not use .choice() directly on domain.categories + # as that will coerce them to a single dtype + indices = random_state.choice( + np.arange(0, len(domain.categories)), size=size + ) + items = [domain.categories[index] for index in indices] + return items if len(items) > 1 else domain.cast(items[0]) + + default_sampler_cls = _Uniform + + def __init__(self, categories: Sequence): + self.categories = list(categories) + + def uniform(self): + new = copy(self) + new.set_sampler(self._Uniform()) + return new + + def grid(self): + new = copy(self) + new.set_sampler(Grid()) + return new + + def __len__(self): + return len(self.categories) + + def __getitem__(self, item): + return self.categories[item] + + def is_valid(self, value: Any): + return value in self.categories + + @property + def domain_str(self): + return f"{self.categories}" + + +@DeveloperAPI +class Function(Domain): + class _CallSampler(BaseSampler): + def __try_fn(self, domain: "Function", config: Dict[str, Any]): + try: + return domain.func(config) + except (AttributeError, KeyError): + from ray.tune.search.variant_generator import _UnresolvedAccessGuard + + r = domain.func(_UnresolvedAccessGuard({"config": config})) + logger.warning( + "sample_from functions that take a spec dict are " + "deprecated. Please update your function to work with " + "the config dict directly." + ) + return r + + def sample( + self, + domain: "Function", + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + if domain.pass_config: + items = [ + ( + self.__try_fn(domain, config[i]) + if isinstance(config, list) + else self.__try_fn(domain, config) + ) + for i in range(size) + ] + else: + items = [domain.func() for i in range(size)] + + return items if len(items) > 1 else domain.cast(items[0]) + + default_sampler_cls = _CallSampler + + def __init__(self, func: Callable): + sig = signature(func) + + pass_config = True # whether we should pass `config` when calling `func` + try: + sig.bind({}) + except TypeError: + pass_config = False + + if not pass_config: + try: + sig.bind() + except TypeError as exc: + raise ValueError( + "The function passed to a `Function` parameter must be " + "callable with either 0 or 1 parameters." + ) from exc + + self.pass_config = pass_config + self.func = func + + def is_function(self): + return True + + def is_valid(self, value: Any): + return True # This is user-defined, so lets not assume anything + + @property + def domain_str(self): + return f"{self.func}()" + + +@DeveloperAPI +class Quantized(Sampler): + def __init__(self, sampler: Sampler, q: Union[float, int]): + self.sampler = sampler + self.q = q + + assert self.sampler, "Quantized() expects a sampler instance" + + def get_sampler(self): + return self.sampler + + def sample( + self, + domain: Domain, + config: Optional[Union[List[Dict], Dict]] = None, + size: int = 1, + random_state: "RandomState" = None, + ): + if not isinstance(random_state, _BackwardsCompatibleNumpyRng): + random_state = _BackwardsCompatibleNumpyRng(random_state) + + if self.q == 1: + return self.sampler.sample(domain, config, size, random_state=random_state) + + quantized_domain = copy(domain) + quantized_domain.lower = np.ceil(domain.lower / self.q) * self.q + quantized_domain.upper = np.floor(domain.upper / self.q) * self.q + values = self.sampler.sample( + quantized_domain, config, size, random_state=random_state + ) + quantized = np.round(np.divide(values, self.q)) * self.q + + if not isinstance(quantized, np.ndarray): + return domain.cast(quantized) + return list(quantized) + + +@PublicAPI +def sample_from(func: Callable[[Dict], Any]): + """Specify that tune should sample configuration values from this function. + + Arguments: + func: An callable function to draw a sample from. + """ + return Function(func) + + +@PublicAPI +def uniform(lower: float, upper: float): + """Sample a float value uniformly between ``lower`` and ``upper``. + + Sampling from ``tune.uniform(1, 10)`` is equivalent to sampling from + ``np.random.uniform(1, 10))`` + + """ + return Float(lower, upper).uniform() + + +@PublicAPI +def quniform(lower: float, upper: float, q: float): + """Sample a quantized float value uniformly between ``lower`` and ``upper``. + + Sampling from ``tune.uniform(1, 10)`` is equivalent to sampling from + ``np.random.uniform(1, 10))`` + + The value will be quantized, i.e. rounded to an integer increment of ``q``. + Quantization makes the upper bound inclusive. + + """ + return Float(lower, upper).uniform().quantized(q) + + +@PublicAPI +def loguniform(lower: float, upper: float, base: object = _MISSING): + """Sugar for sampling in different orders of magnitude. + + Args: + lower: Lower boundary of the output interval (e.g. 1e-4) + upper: Upper boundary of the output interval (e.g. 1e-2) + + """ + if base is not _MISSING: + _warn_for_base() + return Float(lower, upper).loguniform() + + +@PublicAPI +def qloguniform(lower: float, upper: float, q: float, base: object = _MISSING): + """Sugar for sampling in different orders of magnitude. + + The value will be quantized, i.e. rounded to an integer increment of ``q``. + + Quantization makes the upper bound inclusive. + + Args: + lower: Lower boundary of the output interval (e.g. 1e-4) + upper: Upper boundary of the output interval (e.g. 1e-2) + q: Quantization number. The result will be rounded to an + integer increment of this value. + + """ + if base is not _MISSING: + _warn_for_base() + return Float(lower, upper).loguniform().quantized(q) + + +@PublicAPI +def choice(categories: Sequence): + """Sample a categorical value. + + Sampling from ``tune.choice([1, 2])`` is equivalent to sampling from + ``np.random.choice([1, 2])`` + + """ + return Categorical(categories).uniform() + + +@PublicAPI +def randint(lower: int, upper: int): + """Sample an integer value uniformly between ``lower`` and ``upper``. + + ``lower`` is inclusive, ``upper`` is exclusive. + + Sampling from ``tune.randint(10)`` is equivalent to sampling from + ``np.random.randint(10)`` + + .. versionchanged:: 1.5.0 + When converting Ray Tune configs to searcher-specific search spaces, + the lower and upper limits are adjusted to keep compatibility with + the bounds stated in the docstring above. + + """ + return Integer(lower, upper).uniform() + + +@PublicAPI +def lograndint(lower: int, upper: int, base: object = _MISSING): + """Sample an integer value log-uniformly between ``lower`` and ``upper``. + + ``lower`` is inclusive, ``upper`` is exclusive. + + .. versionchanged:: 1.5.0 + When converting Ray Tune configs to searcher-specific search spaces, + the lower and upper limits are adjusted to keep compatibility with + the bounds stated in the docstring above. + + """ + if base is not _MISSING: + _warn_for_base() + return Integer(lower, upper).loguniform() + + +@PublicAPI +def qrandint(lower: int, upper: int, q: int = 1): + """Sample an integer value uniformly between ``lower`` and ``upper``. + + ``lower`` is inclusive, ``upper`` is also inclusive (!). + + The value will be quantized, i.e. rounded to an integer increment of ``q``. + Quantization makes the upper bound inclusive. + + .. versionchanged:: 1.5.0 + When converting Ray Tune configs to searcher-specific search spaces, + the lower and upper limits are adjusted to keep compatibility with + the bounds stated in the docstring above. + + """ + return Integer(lower, upper).uniform().quantized(q) + + +@PublicAPI +def qlograndint(lower: int, upper: int, q: int, base: object = _MISSING): + """Sample an integer value log-uniformly between ``lower`` and ``upper``. + + ``lower`` is inclusive, ``upper`` is also inclusive (!). + + The value will be quantized, i.e. rounded to an integer increment of ``q``. + Quantization makes the upper bound inclusive. + + .. versionchanged:: 1.5.0 + When converting Ray Tune configs to searcher-specific search spaces, + the lower and upper limits are adjusted to keep compatibility with + the bounds stated in the docstring above. + + """ + if base is not _MISSING: + _warn_for_base() + return Integer(lower, upper).loguniform().quantized(q) + + +@PublicAPI +def randn(mean: float = 0.0, sd: float = 1.0): + """Sample a float value normally with ``mean`` and ``sd``. + + Args: + mean: Mean of the normal distribution. Defaults to 0. + sd: SD of the normal distribution. Defaults to 1. + + """ + return Float(None, None).normal(mean, sd) + + +@PublicAPI +def qrandn(mean: float, sd: float, q: float): + """Sample a float value normally with ``mean`` and ``sd``. + + The value will be quantized, i.e. rounded to an integer increment of ``q``. + + Args: + mean: Mean of the normal distribution. + sd: SD of the normal distribution. + q: Quantization number. The result will be rounded to an + integer increment of this value. + + """ + return Float(None, None).normal(mean, sd).quantized(q) diff --git a/lib/python3.12/site-packages/ray/tune/search/search_algorithm.py b/lib/python3.12/site-packages/ray/tune/search/search_algorithm.py new file mode 100644 index 0000000000000000000000000000000000000000..8ae5154c976bd4e6632089d604a0080223c2292b --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/search_algorithm.py @@ -0,0 +1,127 @@ +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + from ray.tune.experiment import Experiment + + +@DeveloperAPI +class SearchAlgorithm: + """Interface of an event handler API for hyperparameter search. + + Unlike TrialSchedulers, SearchAlgorithms will not have the ability + to modify the execution (i.e., stop and pause trials). + + Trials added manually (i.e., via the Client API) will also notify + this class upon new events, so custom search algorithms should + maintain a list of trials ID generated from this class. + + See also: `ray.tune.search.BasicVariantGenerator`. + """ + + _finished = False + + _metric = None + + @property + def metric(self): + return self._metric + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + """Pass search properties to search algorithm. + + This method acts as an alternative to instantiating search algorithms + with their own specific search spaces. Instead they can accept a + Tune config through this method. + + The search algorithm will usually pass this method to their + ``Searcher`` instance. + + Args: + metric: Metric to optimize + mode: One of ["min", "max"]. Direction to optimize. + config: Tune config dict. + **spec: Any kwargs for forward compatiblity. + Info like Experiment.PUBLIC_KEYS is provided through here. + """ + if self._metric and metric: + return False + if metric: + self._metric = metric + return True + + @property + def total_samples(self): + """Get number of total trials to be generated""" + return 0 + + def add_configurations( + self, experiments: Union["Experiment", List["Experiment"], Dict[str, Dict]] + ): + """Tracks given experiment specifications. + + Arguments: + experiments: Experiments to run. + """ + raise NotImplementedError + + def next_trial(self): + """Returns single Trial object to be queued into the TrialRunner. + + Returns: + trial: Returns a Trial object. + """ + raise NotImplementedError + + def on_trial_result(self, trial_id: str, result: Dict): + """Called on each intermediate result returned by a trial. + + This will only be called when the trial is in the RUNNING state. + + Arguments: + trial_id: Identifier for the trial. + result: Result dictionary. + """ + pass + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + """Notification for the completion of trial. + + Arguments: + trial_id: Identifier for the trial. + result: Defaults to None. A dict will + be provided with this notification when the trial is in + the RUNNING state AND either completes naturally or + by manual termination. + error: Defaults to False. True if the trial is in + the RUNNING state and errors. + """ + pass + + def is_finished(self) -> bool: + """Returns True if no trials left to be queued into TrialRunner. + + Can return True before all trials have finished executing. + """ + return self._finished + + def set_finished(self): + """Marks the search algorithm as finished.""" + self._finished = True + + def has_checkpoint(self, dirpath: str) -> bool: + """Should return False if restoring is not implemented.""" + return False + + def save_to_dir(self, dirpath: str, **kwargs): + """Saves a search algorithm.""" + pass + + def restore_from_dir(self, dirpath: str): + """Restores a search algorithm along with its wrapped state.""" + pass diff --git a/lib/python3.12/site-packages/ray/tune/search/search_generator.py b/lib/python3.12/site-packages/ray/tune/search/search_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..41f267f4521447856304da745fcaab9aec896fe6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/search_generator.py @@ -0,0 +1,222 @@ +import copy +import logging +from typing import Dict, List, Optional, Union + +from ray.tune.error import TuneError +from ray.tune.experiment import Experiment, Trial, _convert_to_experiment_list +from ray.tune.experiment.config_parser import _create_trial_from_spec, _make_parser +from ray.tune.search.search_algorithm import SearchAlgorithm +from ray.tune.search.searcher import Searcher +from ray.tune.search.util import _set_search_properties_backwards_compatible +from ray.tune.search.variant_generator import _resolve_nested_dict, format_vars +from ray.tune.utils.util import ( + _atomic_save, + _load_newest_checkpoint, + flatten_dict, + merge_dicts, +) +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +def _warn_on_repeater(searcher, total_samples): + from ray.tune.search.repeater import _warn_num_samples + + _warn_num_samples(searcher, total_samples) + + +@DeveloperAPI +class SearchGenerator(SearchAlgorithm): + """Generates trials to be passed to the TrialRunner. + + Uses the provided ``searcher`` object to generate trials. This class + transparently handles repeating trials with score aggregation + without embedding logic into the Searcher. + + Args: + searcher: Search object that subclasses the Searcher base class. This + is then used for generating new hyperparameter samples. + """ + + CKPT_FILE_TMPL = "search_gen_state-{}.json" + + def __init__(self, searcher: Searcher): + assert issubclass( + type(searcher), Searcher + ), "Searcher should be subclassing Searcher." + self.searcher = searcher + self._parser = _make_parser() + self._experiment = None + self._counter = 0 # Keeps track of number of trials created. + self._total_samples = 0 # int: total samples to evaluate. + self._finished = False + + @property + def metric(self): + return self.searcher.metric + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + return _set_search_properties_backwards_compatible( + self.searcher.set_search_properties, metric, mode, config, **spec + ) + + @property + def total_samples(self): + return self._total_samples + + def add_configurations( + self, experiments: Union[Experiment, List[Experiment], Dict[str, Dict]] + ): + """Registers experiment specifications. + + Arguments: + experiments: Experiments to run. + """ + assert not self._experiment + logger.debug("added configurations") + experiment_list = _convert_to_experiment_list(experiments) + assert ( + len(experiment_list) == 1 + ), "SearchAlgorithms can only support 1 experiment at a time." + self._experiment = experiment_list[0] + experiment_spec = self._experiment.spec + self._total_samples = self._experiment.spec.get("num_samples", 1) + + _warn_on_repeater(self.searcher, self._total_samples) + if "run" not in experiment_spec: + raise TuneError("Must specify `run` in {}".format(experiment_spec)) + + def next_trial(self): + """Provides one Trial object to be queued into the TrialRunner. + + Returns: + Trial: Returns a single trial. + """ + if not self.is_finished(): + return self.create_trial_if_possible(self._experiment.spec) + return None + + def create_trial_if_possible(self, experiment_spec: Dict) -> Optional[Trial]: + logger.debug("creating trial") + trial_id = Trial.generate_id() + suggested_config = self.searcher.suggest(trial_id) + if suggested_config == Searcher.FINISHED: + self._finished = True + logger.debug("Searcher has finished.") + return + + if suggested_config is None: + return + spec = copy.deepcopy(experiment_spec) + spec["config"] = merge_dicts(spec["config"], copy.deepcopy(suggested_config)) + + # Create a new trial_id if duplicate trial is created + flattened_config = _resolve_nested_dict(spec["config"]) + self._counter += 1 + tag = "{0}_{1}".format(str(self._counter), format_vars(flattened_config)) + trial = _create_trial_from_spec( + spec, + self._parser, + evaluated_params=flatten_dict(suggested_config), + experiment_tag=tag, + trial_id=trial_id, + ) + return trial + + def on_trial_result(self, trial_id: str, result: Dict): + """Notifies the underlying searcher.""" + self.searcher.on_trial_result(trial_id, result) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + self.searcher.on_trial_complete(trial_id=trial_id, result=result, error=error) + + def is_finished(self) -> bool: + return self._counter >= self._total_samples or self._finished + + def get_state(self) -> Dict: + return { + "counter": self._counter, + "total_samples": self._total_samples, + "finished": self._finished, + "experiment": self._experiment, + } + + def set_state(self, state: Dict): + self._counter = state["counter"] + self._total_samples = state["total_samples"] + self._finished = state["finished"] + self._experiment = state["experiment"] + + def has_checkpoint(self, dirpath: str): + return bool(_load_newest_checkpoint(dirpath, self.CKPT_FILE_TMPL.format("*"))) + + def save_to_dir(self, dirpath: str, session_str: str): + """Saves self + searcher to dir. + + Separates the "searcher" from its wrappers (concurrency, repeating). + This allows the user to easily restore a given searcher. + + The save operation is atomic (write/swap). + + Args: + dirpath: Filepath to experiment dir. + session_str: Unique identifier of the current run + session. + """ + searcher = self.searcher + search_alg_state = self.get_state() + while hasattr(searcher, "searcher"): + searcher_name = type(searcher).__name__ + if searcher_name in search_alg_state: + logger.warning( + "There was a duplicate when saving {}. " + "Restore may not work properly.".format(searcher_name) + ) + else: + search_alg_state["name:" + searcher_name] = searcher.get_state() + searcher = searcher.searcher + base_searcher = searcher + # We save the base searcher separately for users to easily + # separate the searcher. + base_searcher.save_to_dir(dirpath, session_str) + _atomic_save( + state=search_alg_state, + checkpoint_dir=dirpath, + file_name=self.CKPT_FILE_TMPL.format(session_str), + tmp_file_name=".tmp_search_generator_ckpt", + ) + + def restore_from_dir(self, dirpath: str): + """Restores self + searcher + search wrappers from dirpath.""" + + searcher = self.searcher + search_alg_state = _load_newest_checkpoint( + dirpath, self.CKPT_FILE_TMPL.format("*") + ) + if not search_alg_state: + raise RuntimeError("Unable to find checkpoint in {}.".format(dirpath)) + while hasattr(searcher, "searcher"): + searcher_name = "name:" + type(searcher).__name__ + if searcher_name not in search_alg_state: + names = [ + key.split("name:")[1] + for key in search_alg_state + if key.startswith("name:") + ] + logger.warning( + "{} was not found in the experiment " + "state when restoring. Found {}.".format(searcher_name, names) + ) + else: + searcher.set_state(search_alg_state.pop(searcher_name)) + searcher = searcher.searcher + base_searcher = searcher + + logger.debug(f"searching base {base_searcher}") + base_searcher.restore_from_dir(dirpath) + self.set_state(search_alg_state) diff --git a/lib/python3.12/site-packages/ray/tune/search/searcher.py b/lib/python3.12/site-packages/ray/tune/search/searcher.py new file mode 100644 index 0000000000000000000000000000000000000000..2d0209a1df208708b26a10646bedcd00e00e3847 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/searcher.py @@ -0,0 +1,597 @@ +import copy +import glob +import logging +import os +import warnings +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from ray.air._internal.usage import tag_searcher +from ray.tune.search.util import _set_search_properties_backwards_compatible +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.tune.analysis import ExperimentAnalysis + from ray.tune.experiment import Trial + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class Searcher: + """Abstract class for wrapping suggesting algorithms. + + Custom algorithms can extend this class easily by overriding the + `suggest` method provide generated parameters for the trials. + + Any subclass that implements ``__init__`` must also call the + constructor of this class: ``super(Subclass, self).__init__(...)``. + + To track suggestions and their corresponding evaluations, the method + `suggest` will be passed a trial_id, which will be used in + subsequent notifications. + + Not all implementations support multi objectives. + + Note to Tune developers: If a new searcher is added, please update + `air/_internal/usage.py`. + + Args: + metric: The training result objective value attribute. If + list then list of training result objective value attributes + mode: If string One of {min, max}. If list then + list of max and min, determines whether objective is minimizing + or maximizing the metric attribute. Must match type of metric. + + .. code-block:: python + + class ExampleSearch(Searcher): + def __init__(self, metric="mean_loss", mode="min", **kwargs): + super(ExampleSearch, self).__init__( + metric=metric, mode=mode, **kwargs) + self.optimizer = Optimizer() + self.configurations = {} + + def suggest(self, trial_id): + configuration = self.optimizer.query() + self.configurations[trial_id] = configuration + + def on_trial_complete(self, trial_id, result, **kwargs): + configuration = self.configurations[trial_id] + if result and self.metric in result: + self.optimizer.update(configuration, result[self.metric]) + + tuner = tune.Tuner( + trainable_function, + tune_config=tune.TuneConfig( + search_alg=ExampleSearch() + ) + ) + tuner.fit() + + + """ + + FINISHED = "FINISHED" + CKPT_FILE_TMPL = "searcher-state-{}.pkl" + + def __init__( + self, + metric: Optional[str] = None, + mode: Optional[str] = None, + ): + tag_searcher(self) + self._metric = metric + self._mode = mode + + if not mode or not metric: + # Early return to avoid assertions + return + + assert isinstance( + metric, type(mode) + ), "metric and mode must be of the same type" + if isinstance(mode, str): + assert mode in ["min", "max"], "if `mode` is a str must be 'min' or 'max'!" + elif isinstance(mode, list): + assert len(mode) == len(metric), "Metric and mode must be the same length" + assert all( + mod in ["min", "max", "obs"] for mod in mode + ), "All of mode must be 'min' or 'max' or 'obs'!" + else: + raise ValueError("Mode most either be a list or string") + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + """Pass search properties to searcher. + + This method acts as an alternative to instantiating search algorithms + with their own specific search spaces. Instead they can accept a + Tune config through this method. A searcher should return ``True`` + if setting the config was successful, or ``False`` if it was + unsuccessful, e.g. when the search space has already been set. + + Args: + metric: Metric to optimize + mode: One of ["min", "max"]. Direction to optimize. + config: Tune config dict. + **spec: Any kwargs for forward compatiblity. + Info like Experiment.PUBLIC_KEYS is provided through here. + """ + return False + + def on_trial_result(self, trial_id: str, result: Dict) -> None: + """Optional notification for result during training. + + Note that by default, the result dict may include NaNs or + may not include the optimization metric. It is up to the + subclass implementation to preprocess the result to + avoid breaking the optimization process. + + Args: + trial_id: A unique string ID for the trial. + result: Dictionary of metrics for current training progress. + Note that the result dict may include NaNs or + may not include the optimization metric. It is up to the + subclass implementation to preprocess the result to + avoid breaking the optimization process. + """ + pass + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ) -> None: + """Notification for the completion of trial. + + Typically, this method is used for notifying the underlying + optimizer of the result. + + Args: + trial_id: A unique string ID for the trial. + result: Dictionary of metrics for current training progress. + Note that the result dict may include NaNs or + may not include the optimization metric. It is up to the + subclass implementation to preprocess the result to + avoid breaking the optimization process. Upon errors, this + may also be None. + error: True if the training process raised an error. + + """ + raise NotImplementedError + + def suggest(self, trial_id: str) -> Optional[Dict]: + """Queries the algorithm to retrieve the next set of parameters. + + Arguments: + trial_id: Trial ID used for subsequent notifications. + + Returns: + dict | FINISHED | None: Configuration for a trial, if possible. + If FINISHED is returned, Tune will be notified that + no more suggestions/configurations will be provided. + If None is returned, Tune will skip the querying of the + searcher for this step. + + """ + raise NotImplementedError + + def add_evaluated_point( + self, + parameters: Dict, + value: float, + error: bool = False, + pruned: bool = False, + intermediate_values: Optional[List[float]] = None, + ): + """Pass results from a point that has been evaluated separately. + + This method allows for information from outside the + suggest - on_trial_complete loop to be passed to the search + algorithm. + This functionality depends on the underlying search algorithm + and may not be always available. + + Args: + parameters: Parameters used for the trial. + value: Metric value obtained in the trial. + error: True if the training process raised an error. + pruned: True if trial was pruned. + intermediate_values: List of metric values for + intermediate iterations of the result. None if not + applicable. + + """ + raise NotImplementedError + + def add_evaluated_trials( + self, + trials_or_analysis: Union["Trial", List["Trial"], "ExperimentAnalysis"], + metric: str, + ): + """Pass results from trials that have been evaluated separately. + + This method allows for information from outside the + suggest - on_trial_complete loop to be passed to the search + algorithm. + This functionality depends on the underlying search algorithm + and may not be always available (same as ``add_evaluated_point``.) + + Args: + trials_or_analysis: Trials to pass results form to the searcher. + metric: Metric name reported by trials used for + determining the objective value. + + """ + if self.add_evaluated_point == Searcher.add_evaluated_point: + raise NotImplementedError + + # lazy imports to avoid circular dependencies + from ray.tune.analysis import ExperimentAnalysis + from ray.tune.experiment import Trial + from ray.tune.result import DONE + + if isinstance(trials_or_analysis, (list, tuple)): + trials = trials_or_analysis + elif isinstance(trials_or_analysis, Trial): + trials = [trials_or_analysis] + elif isinstance(trials_or_analysis, ExperimentAnalysis): + trials = trials_or_analysis.trials + else: + raise NotImplementedError( + "Expected input to be a `Trial`, a list of `Trial`s, or " + f"`ExperimentAnalysis`, got: {trials_or_analysis}" + ) + + any_trial_had_metric = False + + def trial_to_points(trial: Trial) -> Dict[str, Any]: + nonlocal any_trial_had_metric + has_trial_been_pruned = ( + trial.status == Trial.TERMINATED + and not trial.last_result.get(DONE, False) + ) + has_trial_finished = ( + trial.status == Trial.TERMINATED and trial.last_result.get(DONE, False) + ) + if not any_trial_had_metric: + any_trial_had_metric = ( + metric in trial.last_result and has_trial_finished + ) + if Trial.TERMINATED and metric not in trial.last_result: + return None + return dict( + parameters=trial.config, + value=trial.last_result.get(metric, None), + error=trial.status == Trial.ERROR, + pruned=has_trial_been_pruned, + intermediate_values=None, # we do not save those + ) + + for trial in trials: + kwargs = trial_to_points(trial) + if kwargs: + self.add_evaluated_point(**kwargs) + + if not any_trial_had_metric: + warnings.warn( + "No completed trial returned the specified metric. " + "Make sure the name you have passed is correct. " + ) + + def save(self, checkpoint_path: str): + """Save state to path for this search algorithm. + + Args: + checkpoint_path: File where the search algorithm + state is saved. This path should be used later when + restoring from file. + + Example: + + .. code-block:: python + + search_alg = Searcher(...) + + tuner = tune.Tuner( + cost, + tune_config=tune.TuneConfig( + search_alg=search_alg, + num_samples=5 + ), + param_space=config + ) + results = tuner.fit() + + search_alg.save("./my_favorite_path.pkl") + + .. versionchanged:: 0.8.7 + Save is automatically called by `Tuner().fit()`. You can use + `Tuner().restore()` to restore from an experiment directory + such as `~/ray_results/trainable`. + + """ + raise NotImplementedError + + def restore(self, checkpoint_path: str): + """Restore state for this search algorithm + + + Args: + checkpoint_path: File where the search algorithm + state is saved. This path should be the same + as the one provided to "save". + + Example: + + .. code-block:: python + + search_alg.save("./my_favorite_path.pkl") + + search_alg2 = Searcher(...) + search_alg2 = ConcurrencyLimiter(search_alg2, 1) + search_alg2.restore(checkpoint_path) + tuner = tune.Tuner( + cost, + tune_config=tune.TuneConfig( + search_alg=search_alg2, + num_samples=5 + ), + ) + tuner.fit() + + """ + raise NotImplementedError + + def set_max_concurrency(self, max_concurrent: int) -> bool: + """Set max concurrent trials this searcher can run. + + This method will be called on the wrapped searcher by the + ``ConcurrencyLimiter``. It is intended to allow for searchers + which have custom, internal logic handling max concurrent trials + to inherit the value passed to ``ConcurrencyLimiter``. + + If this method returns False, it signifies that no special + logic for handling this case is present in the searcher. + + Args: + max_concurrent: Number of maximum concurrent trials. + """ + return False + + def get_state(self) -> Dict: + raise NotImplementedError + + def set_state(self, state: Dict): + raise NotImplementedError + + def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"): + """Automatically saves the given searcher to the checkpoint_dir. + + This is automatically used by Tuner().fit() during a Tune job. + + Args: + checkpoint_dir: Filepath to experiment dir. + session_str: Unique identifier of the current run + session. + """ + tmp_search_ckpt_path = os.path.join(checkpoint_dir, ".tmp_searcher_ckpt") + success = True + try: + self.save(tmp_search_ckpt_path) + except NotImplementedError: + if log_once("suggest:save_to_dir"): + logger.warning("save not implemented for Searcher. Skipping save.") + success = False + + if success and os.path.exists(tmp_search_ckpt_path): + os.replace( + tmp_search_ckpt_path, + os.path.join(checkpoint_dir, self.CKPT_FILE_TMPL.format(session_str)), + ) + + def restore_from_dir(self, checkpoint_dir: str): + """Restores the state of a searcher from a given checkpoint_dir. + + Typically, you should use this function to restore from an + experiment directory such as `~/ray_results/trainable`. + + .. code-block:: python + + tuner = tune.Tuner( + cost, + run_config=tune.RunConfig( + name=self.experiment_name, + storage_path="~/my_results", + ), + tune_config=tune.TuneConfig( + search_alg=search_alg, + num_samples=5 + ), + param_space=config + ) + tuner.fit() + + search_alg2 = Searcher() + search_alg2.restore_from_dir( + os.path.join("~/my_results", self.experiment_name) + """ + + pattern = self.CKPT_FILE_TMPL.format("*") + full_paths = glob.glob(os.path.join(checkpoint_dir, pattern)) + if not full_paths: + raise RuntimeError( + "Searcher unable to find checkpoint in {}".format(checkpoint_dir) + ) # TODO + most_recent_checkpoint = max(full_paths) + self.restore(most_recent_checkpoint) + + @property + def metric(self) -> str: + """The training result objective value attribute.""" + return self._metric + + @property + def mode(self) -> str: + """Specifies if minimizing or maximizing the metric.""" + return self._mode + + +@PublicAPI +class ConcurrencyLimiter(Searcher): + """A wrapper algorithm for limiting the number of concurrent trials. + + Certain Searchers have their own internal logic for limiting + the number of concurrent trials. If such a Searcher is passed to a + ``ConcurrencyLimiter``, the ``max_concurrent`` of the + ``ConcurrencyLimiter`` will override the ``max_concurrent`` value + of the Searcher. The ``ConcurrencyLimiter`` will then let the + Searcher's internal logic take over. + + Args: + searcher: Searcher object that the + ConcurrencyLimiter will manage. + max_concurrent: Maximum concurrent samples from the underlying + searcher. + batch: Whether to wait for all concurrent samples + to finish before updating the underlying searcher. + + Example: + + .. code-block:: python + + from ray.tune.search import ConcurrencyLimiter + search_alg = HyperOptSearch(metric="accuracy") + search_alg = ConcurrencyLimiter(search_alg, max_concurrent=2) + tuner = tune.Tuner( + trainable_function, + tune_config=tune.TuneConfig( + search_alg=search_alg + ), + ) + tuner.fit() + + """ + + def __init__(self, searcher: Searcher, max_concurrent: int, batch: bool = False): + assert type(max_concurrent) is int and max_concurrent > 0 + self.searcher = searcher + self.max_concurrent = max_concurrent + self.batch = batch + self.live_trials = set() + self.num_unfinished_live_trials = 0 + self.cached_results = {} + self._limit_concurrency = True + + if not isinstance(searcher, Searcher): + raise RuntimeError( + f"The `ConcurrencyLimiter` only works with `Searcher` " + f"objects (got {type(searcher)}). Please try to pass " + f"`max_concurrent` to the search generator directly." + ) + + self._set_searcher_max_concurrency() + + super(ConcurrencyLimiter, self).__init__( + metric=self.searcher.metric, mode=self.searcher.mode + ) + + def _set_searcher_max_concurrency(self): + # If the searcher has special logic for handling max concurrency, + # we do not do anything inside the ConcurrencyLimiter + self._limit_concurrency = not self.searcher.set_max_concurrency( + self.max_concurrent + ) + + def set_max_concurrency(self, max_concurrent: int) -> bool: + # Determine if this behavior is acceptable, or if it should + # raise an exception. + self.max_concurrent = max_concurrent + return True + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + self._set_searcher_max_concurrency() + return _set_search_properties_backwards_compatible( + self.searcher.set_search_properties, metric, mode, config, **spec + ) + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._limit_concurrency: + return self.searcher.suggest(trial_id) + + assert ( + trial_id not in self.live_trials + ), f"Trial ID {trial_id} must be unique: already found in set." + if len(self.live_trials) >= self.max_concurrent: + logger.debug( + f"Not providing a suggestion for {trial_id} due to " + "concurrency limit: %s/%s.", + len(self.live_trials), + self.max_concurrent, + ) + return + + suggestion = self.searcher.suggest(trial_id) + if suggestion not in (None, Searcher.FINISHED): + self.live_trials.add(trial_id) + self.num_unfinished_live_trials += 1 + return suggestion + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + if not self._limit_concurrency: + return self.searcher.on_trial_complete(trial_id, result=result, error=error) + + if trial_id not in self.live_trials: + return + elif self.batch: + self.cached_results[trial_id] = (result, error) + self.num_unfinished_live_trials -= 1 + if self.num_unfinished_live_trials <= 0: + # Update the underlying searcher once the + # full batch is completed. + for trial_id, (result, error) in self.cached_results.items(): + self.searcher.on_trial_complete( + trial_id, result=result, error=error + ) + self.live_trials.remove(trial_id) + self.cached_results = {} + self.num_unfinished_live_trials = 0 + else: + return + else: + self.searcher.on_trial_complete(trial_id, result=result, error=error) + self.live_trials.remove(trial_id) + self.num_unfinished_live_trials -= 1 + + def on_trial_result(self, trial_id: str, result: Dict) -> None: + self.searcher.on_trial_result(trial_id, result) + + def add_evaluated_point( + self, + parameters: Dict, + value: float, + error: bool = False, + pruned: bool = False, + intermediate_values: Optional[List[float]] = None, + ): + return self.searcher.add_evaluated_point( + parameters, value, error, pruned, intermediate_values + ) + + def get_state(self) -> Dict: + state = self.__dict__.copy() + del state["searcher"] + return copy.deepcopy(state) + + def set_state(self, state: Dict): + self.__dict__.update(state) + + def save(self, checkpoint_path: str): + self.searcher.save(checkpoint_path) + + def restore(self, checkpoint_path: str): + self.searcher.restore(checkpoint_path) diff --git a/lib/python3.12/site-packages/ray/tune/search/util.py b/lib/python3.12/site-packages/ray/tune/search/util.py new file mode 100644 index 0000000000000000000000000000000000000000..9358bd32af6373aa62ee60b1d197544b10691608 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/util.py @@ -0,0 +1,31 @@ +import logging +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + + +def _set_search_properties_backwards_compatible( + set_search_properties_func, + metric: Optional[str], + mode: Optional[str], + config: Dict, + **spec +) -> bool: + """Wraps around set_search_properties() so that it is backward compatible. + + Also outputs a warning to encourage custom searchers to be updated. + """ + try: + return set_search_properties_func(metric, mode, config, **spec) + except TypeError as e: + if str(e).startswith( + "set_search_properties() got an unexpected keyword argument" + ): + logger.warning( + "Please update custom Searcher to take in function signature " + "as ``def set_search_properties(metric, mode, config, " + "**spec) -> bool``." + ) + return set_search_properties_func(metric, mode, config) + else: + raise e diff --git a/lib/python3.12/site-packages/ray/tune/search/variant_generator.py b/lib/python3.12/site-packages/ray/tune/search/variant_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5b3c7d27e6e9208e9548126c3901158b1077ac72 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/variant_generator.py @@ -0,0 +1,523 @@ +import copy +import logging +import random +import re +from collections.abc import Mapping +from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple + +import numpy + +from ray.tune.search.sample import Categorical, Domain, Function, RandomState +from ray.util.annotations import DeveloperAPI, PublicAPI + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +def generate_variants( + unresolved_spec: Dict, + constant_grid_search: bool = False, + random_state: "RandomState" = None, +) -> Generator[Tuple[Dict, Dict], None, None]: + """Generates variants from a spec (dict) with unresolved values. + + There are two types of unresolved values: + + Grid search: These define a grid search over values. For example, the + following grid search values in a spec will produce six distinct + variants in combination: + + "activation": grid_search(["relu", "tanh"]) + "learning_rate": grid_search([1e-3, 1e-4, 1e-5]) + + Lambda functions: These are evaluated to produce a concrete value, and + can express dependencies or conditional distributions between values. + They can also be used to express random search (e.g., by calling + into the `random` or `np` module). + + "cpu": lambda spec: spec.config.num_workers + "batch_size": lambda spec: random.uniform(1, 1000) + + Finally, to support defining specs in plain JSON / YAML, grid search + and lambda functions can also be defined alternatively as follows: + + "activation": {"grid_search": ["relu", "tanh"]} + "cpu": {"eval": "spec.config.num_workers"} + + Use `format_vars` to format the returned dict of hyperparameters. + + Yields: + (Dict of resolved variables, Spec object) + """ + for resolved_vars, spec in _generate_variants_internal( + unresolved_spec, + constant_grid_search=constant_grid_search, + random_state=random_state, + ): + assert not _unresolved_values(spec) + yield resolved_vars, spec + + +@PublicAPI(stability="beta") +def grid_search(values: Iterable) -> Dict[str, Iterable]: + """Specify a grid of values to search over. + + Values specified in a grid search are guaranteed to be sampled. + + If multiple grid search variables are defined, they are combined with the + combinatorial product. This means every possible combination of values will + be sampled. + + Example: + + >>> from ray import tune + >>> param_space={ + ... "x": tune.grid_search([10, 20]), + ... "y": tune.grid_search(["a", "b", "c"]) + ... } + + This will create a grid of 6 samples: + ``{"x": 10, "y": "a"}``, ``{"x": 10, "y": "b"}``, etc. + + When specifying ``num_samples`` in the + :class:`TuneConfig `, this will specify + the number of random samples per grid search combination. + + For instance, in the example above, if ``num_samples=4``, + a total of 24 trials will be started - + 4 trials for each of the 6 grid search combinations. + + Args: + values: An iterable whose parameters will be used for creating a trial grid. + + """ + return {"grid_search": values} + + +_STANDARD_IMPORTS = { + "random": random, + "np": numpy, +} + +_MAX_RESOLUTION_PASSES = 20 + + +def _resolve_nested_dict(nested_dict: Dict) -> Dict[Tuple, Any]: + """Flattens a nested dict by joining keys into tuple of paths. + + Can then be passed into `format_vars`. + """ + res = {} + for k, v in nested_dict.items(): + if isinstance(v, dict): + for k_, v_ in _resolve_nested_dict(v).items(): + res[(k,) + k_] = v_ + else: + res[(k,)] = v + return res + + +@DeveloperAPI +def format_vars(resolved_vars: Dict) -> str: + """Format variables to be used as experiment tags. + + Experiment tags are used in directory names, so this method makes sure + the resulting tags can be legally used in directory names on all systems. + + The input to this function is a dict of the form + ``{("nested", "config", "path"): "value"}``. The output will be a comma + separated string of the form ``last_key=value``, so in this example + ``path=value``. + + Note that the sanitizing implies that empty strings are possible return + values. This is expected and acceptable, as it is not a common case and + the resulting directory names will still be valid. + + Args: + resolved_vars: Dictionary mapping from config path tuples to a value. + + Returns: + Comma-separated key=value string. + """ + vars = resolved_vars.copy() + # TrialRunner already has these in the experiment_tag + for v in ["run", "env", "resources_per_trial"]: + vars.pop(v, None) + + return ",".join( + f"{_clean_value(k[-1])}={_clean_value(v)}" for k, v in sorted(vars.items()) + ) + + +def _flatten_resolved_vars(resolved_vars: Dict) -> Dict: + """Formats the resolved variable dict into a mapping of (str -> value).""" + flattened_resolved_vars_dict = {} + for pieces, value in resolved_vars.items(): + if pieces[0] == "config": + pieces = pieces[1:] + pieces = [str(piece) for piece in pieces] + flattened_resolved_vars_dict["/".join(pieces)] = value + return flattened_resolved_vars_dict + + +def _clean_value(value: Any) -> str: + """Format floats and replace invalid string characters with ``_``.""" + if isinstance(value, float): + return f"{value:.4f}" + else: + # Define an invalid alphabet, which is the inverse of the + # stated regex characters + invalid_alphabet = r"[^a-zA-Z0-9_-]+" + return re.sub(invalid_alphabet, "_", str(value)).strip("_") + + +@DeveloperAPI +def parse_spec_vars( + spec: Dict, +) -> Tuple[List[Tuple[Tuple, Any]], List[Tuple[Tuple, Any]], List[Tuple[Tuple, Any]]]: + resolved, unresolved = _split_resolved_unresolved_values(spec) + resolved_vars = list(resolved.items()) + + if not unresolved: + return resolved_vars, [], [] + + grid_vars = [] + domain_vars = [] + for path, value in unresolved.items(): + if value.is_grid(): + grid_vars.append((path, value)) + else: + domain_vars.append((path, value)) + grid_vars.sort() + + return resolved_vars, domain_vars, grid_vars + + +def _count_spec_samples(spec: Dict, num_samples=1) -> int: + """Count samples for a specific spec""" + _, domain_vars, grid_vars = parse_spec_vars(spec) + grid_count = 1 + for path, domain in grid_vars: + grid_count *= len(domain.categories) + return num_samples * grid_count + + +def _count_variants(spec: Dict, presets: Optional[List[Dict]] = None) -> int: + # Helper function: Deep update dictionary + def deep_update(d, u): + for k, v in u.items(): + if isinstance(v, Mapping): + d[k] = deep_update(d.get(k, {}), v) + else: + d[k] = v + return d + + total_samples = 0 + total_num_samples = spec.get("num_samples", 1) + # For each preset, overwrite the spec and count the samples generated + # for this preset + for preset in presets: + preset_spec = copy.deepcopy(spec) + deep_update(preset_spec["config"], preset) + total_samples += _count_spec_samples(preset_spec, 1) + total_num_samples -= 1 + + # Add the remaining samples + if total_num_samples > 0: + total_samples += _count_spec_samples(spec, total_num_samples) + return total_samples + + +def _generate_variants_internal( + spec: Dict, constant_grid_search: bool = False, random_state: "RandomState" = None +) -> Tuple[Dict, Dict]: + spec = copy.deepcopy(spec) + _, domain_vars, grid_vars = parse_spec_vars(spec) + + if not domain_vars and not grid_vars: + yield {}, spec + return + + # Variables to resolve + to_resolve = domain_vars + + all_resolved = True + if constant_grid_search: + # In this path, we first sample random variables and keep them constant + # for grid search. + # `_resolve_domain_vars` will alter `spec` directly + all_resolved, resolved_vars = _resolve_domain_vars( + spec, domain_vars, allow_fail=True, random_state=random_state + ) + if not all_resolved: + # Not all variables have been resolved, but remove those that have + # from the `to_resolve` list. + to_resolve = [(r, d) for r, d in to_resolve if r not in resolved_vars] + grid_search = _grid_search_generator(spec, grid_vars) + for resolved_spec in grid_search: + if not constant_grid_search or not all_resolved: + # In this path, we sample the remaining random variables + _, resolved_vars = _resolve_domain_vars( + resolved_spec, to_resolve, random_state=random_state + ) + + for resolved, spec in _generate_variants_internal( + resolved_spec, + constant_grid_search=constant_grid_search, + random_state=random_state, + ): + for path, value in grid_vars: + resolved_vars[path] = _get_value(spec, path) + for k, v in resolved.items(): + if ( + k in resolved_vars + and v != resolved_vars[k] + and _is_resolved(resolved_vars[k]) + ): + raise ValueError( + "The variable `{}` could not be unambiguously " + "resolved to a single value. Consider simplifying " + "your configuration.".format(k) + ) + resolved_vars[k] = v + yield resolved_vars, spec + + +def _get_preset_variants( + spec: Dict, + config: Dict, + constant_grid_search: bool = False, + random_state: "RandomState" = None, +): + """Get variants according to a spec, initialized with a config. + + Variables from the spec are overwritten by the variables in the config. + Thus, we may end up with less sampled parameters. + + This function also checks if values used to overwrite search space + parameters are valid, and logs a warning if not. + """ + spec = copy.deepcopy(spec) + + resolved, _, _ = parse_spec_vars(config) + + for path, val in resolved: + try: + domain = _get_value(spec["config"], path) + if isinstance(domain, dict): + if "grid_search" in domain: + domain = Categorical(domain["grid_search"]) + else: + # If users want to overwrite an entire subdict, + # let them do it. + domain = None + except IndexError as exc: + raise ValueError( + f"Pre-set config key `{'/'.join(path)}` does not correspond " + f"to a valid key in the search space definition. Please add " + f"this path to the `param_space` variable passed to `tune.Tuner()`." + ) from exc + + if domain: + if isinstance(domain, Domain): + if not domain.is_valid(val): + logger.warning( + f"Pre-set value `{val}` is not within valid values of " + f"parameter `{'/'.join(path)}`: {domain.domain_str}" + ) + else: + # domain is actually a fixed value + if domain != val: + logger.warning( + f"Pre-set value `{val}` is not equal to the value of " + f"parameter `{'/'.join(path)}`: {domain}" + ) + assign_value(spec["config"], path, val) + + return _generate_variants_internal( + spec, constant_grid_search=constant_grid_search, random_state=random_state + ) + + +@DeveloperAPI +def assign_value(spec: Dict, path: Tuple, value: Any): + """Assigns a value to a nested dictionary. + + Handles the special case of tuples, in which case the tuples + will be re-constructed to accommodate the updated value. + """ + parent_spec = None + parent_key = None + for k in path[:-1]: + parent_spec = spec + parent_key = k + spec = spec[k] + key = path[-1] + if not isinstance(spec, tuple): + # spec is mutable. Just assign the value. + spec[key] = value + else: + if parent_spec is None: + raise ValueError("Cannot assign value to a tuple.") + assert isinstance(key, int), "Tuple key must be an int." + # Special handling since tuples are immutable. + parent_spec[parent_key] = spec[:key] + (value,) + spec[key + 1 :] + + +def _get_value(spec: Dict, path: Tuple) -> Any: + for k in path: + spec = spec[k] + return spec + + +def _resolve_domain_vars( + spec: Dict, + domain_vars: List[Tuple[Tuple, Domain]], + allow_fail: bool = False, + random_state: "RandomState" = None, +) -> Tuple[bool, Dict]: + resolved = {} + error = True + num_passes = 0 + while error and num_passes < _MAX_RESOLUTION_PASSES: + num_passes += 1 + error = False + for path, domain in domain_vars: + if path in resolved: + continue + try: + value = domain.sample( + _UnresolvedAccessGuard(spec), random_state=random_state + ) + except RecursiveDependencyError as e: + error = e + except Exception: + raise ValueError( + "Failed to evaluate expression: {}: {}".format(path, domain) + ) + else: + assign_value(spec, path, value) + resolved[path] = value + if error: + if not allow_fail: + raise error + else: + return False, resolved + return True, resolved + + +def _grid_search_generator( + unresolved_spec: Dict, grid_vars: List +) -> Generator[Dict, None, None]: + value_indices = [0] * len(grid_vars) + + def increment(i): + value_indices[i] += 1 + if value_indices[i] >= len(grid_vars[i][1]): + value_indices[i] = 0 + if i + 1 < len(value_indices): + return increment(i + 1) + else: + return True + return False + + if not grid_vars: + yield unresolved_spec + return + + while value_indices[-1] < len(grid_vars[-1][1]): + spec = copy.deepcopy(unresolved_spec) + for i, (path, values) in enumerate(grid_vars): + assign_value(spec, path, values[value_indices[i]]) + yield spec + if grid_vars: + done = increment(0) + if done: + break + + +def _is_resolved(v) -> bool: + resolved, _ = _try_resolve(v) + return resolved + + +def _try_resolve(v) -> Tuple[bool, Any]: + if isinstance(v, Domain): + # Domain to sample from + return False, v + elif isinstance(v, dict) and len(v) == 1 and "eval" in v: + # Lambda function in eval syntax + return False, Function( + lambda spec: eval(v["eval"], _STANDARD_IMPORTS, {"spec": spec}) + ) + elif isinstance(v, dict) and len(v) == 1 and "grid_search" in v: + # Grid search values + grid_values = v["grid_search"] + return False, Categorical(grid_values).grid() + return True, v + + +def _split_resolved_unresolved_values( + spec: Dict, +) -> Tuple[Dict[Tuple, Any], Dict[Tuple, Any]]: + resolved_vars = {} + unresolved_vars = {} + for k, v in spec.items(): + resolved, v = _try_resolve(v) + if not resolved: + unresolved_vars[(k,)] = v + elif isinstance(v, dict): + # Recurse into a dict + ( + _resolved_children, + _unresolved_children, + ) = _split_resolved_unresolved_values(v) + for path, value in _resolved_children.items(): + resolved_vars[(k,) + path] = value + for path, value in _unresolved_children.items(): + unresolved_vars[(k,) + path] = value + elif isinstance(v, (list, tuple)): + # Recurse into a list + for i, elem in enumerate(v): + ( + _resolved_children, + _unresolved_children, + ) = _split_resolved_unresolved_values({i: elem}) + for path, value in _resolved_children.items(): + resolved_vars[(k,) + path] = value + for path, value in _unresolved_children.items(): + unresolved_vars[(k,) + path] = value + else: + resolved_vars[(k,)] = v + return resolved_vars, unresolved_vars + + +def _unresolved_values(spec: Dict) -> Dict[Tuple, Any]: + return _split_resolved_unresolved_values(spec)[1] + + +def _has_unresolved_values(spec: Dict) -> bool: + return True if _unresolved_values(spec) else False + + +class _UnresolvedAccessGuard(dict): + def __init__(self, *args, **kwds): + super(_UnresolvedAccessGuard, self).__init__(*args, **kwds) + self.__dict__ = self + + def __getattribute__(self, item): + value = dict.__getattribute__(self, item) + if not _is_resolved(value): + raise RecursiveDependencyError( + "`{}` recursively depends on {}".format(item, value) + ) + elif isinstance(value, dict): + return _UnresolvedAccessGuard(value) + else: + return value + + +@DeveloperAPI +class RecursiveDependencyError(Exception): + def __init__(self, msg: str): + Exception.__init__(self, msg) diff --git a/lib/python3.12/site-packages/ray/tune/search/zoopt/__init__.py b/lib/python3.12/site-packages/ray/tune/search/zoopt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..44ab8e345754da48205fbe712e4bb787a5ddfc2a --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/zoopt/__init__.py @@ -0,0 +1,3 @@ +from ray.tune.search.zoopt.zoopt_search import ZOOptSearch + +__all__ = ["ZOOptSearch"] diff --git a/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0d4bdbdb47d85932d1838b90e737a2b9dc671b7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/zoopt_search.cpython-312.pyc b/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/zoopt_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b65b8d35b23aff18d0853ff181a68f106ed4f8c4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/tune/search/zoopt/__pycache__/zoopt_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/tune/search/zoopt/zoopt_search.py b/lib/python3.12/site-packages/ray/tune/search/zoopt/zoopt_search.py new file mode 100644 index 0000000000000000000000000000000000000000..d568f94bde0fff922b0ff4b6347010e00e9f8be2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/tune/search/zoopt/zoopt_search.py @@ -0,0 +1,379 @@ +import copy +import logging +from typing import Dict, List, Optional, Tuple + +import ray +import ray.cloudpickle as pickle +from ray.tune.result import DEFAULT_METRIC +from ray.tune.search import ( + UNDEFINED_METRIC_MODE, + UNDEFINED_SEARCH_SPACE, + UNRESOLVED_SEARCH_SPACE, + Searcher, +) +from ray.tune.search.sample import ( + Categorical, + Domain, + Float, + Integer, + Quantized, + Uniform, +) +from ray.tune.search.variant_generator import parse_spec_vars +from ray.tune.utils.util import unflatten_dict + +try: + import zoopt + from zoopt import Solution, ValueType +except ImportError: + zoopt = None + Solution = ValueType = None + +logger = logging.getLogger(__name__) + + +class ZOOptSearch(Searcher): + """A wrapper around ZOOpt to provide trial suggestions. + + ZOOptSearch is a library for derivative-free optimization. It is backed by + the `ZOOpt `__ package. Currently, + Asynchronous Sequential RAndomized COordinate Shrinking (ASRacos) + is implemented in Tune. + + To use ZOOptSearch, install zoopt (>=0.4.1): ``pip install -U zoopt``. + + Tune automatically converts search spaces to ZOOpt"s format: + + .. code-block:: python + + from ray import tune + from ray.tune.search.zoopt import ZOOptSearch + + "config": { + "iterations": 10, # evaluation times + "width": tune.uniform(-10, 10), + "height": tune.uniform(-10, 10) + } + + zoopt_search_config = { + "parallel_num": 8, # how many workers to parallel + } + + zoopt_search = ZOOptSearch( + algo="Asracos", # only support Asracos currently + budget=20, # must match `num_samples` in `tune.TuneConfig()`. + dim_dict=dim_dict, + metric="mean_loss", + mode="min", + **zoopt_search_config + ) + + tuner = tune.Tuner( + my_objective, + tune_config=tune.TuneConfig( + search_alg=zoopt_search, + num_samples=20 + ), + run_config=tune.RunConfig( + name="zoopt_search", + stop={"timesteps_total": 10} + ), + param_space=config + ) + tuner.fit() + + If you would like to pass the search space manually, the code would + look like this: + + .. code-block:: python + + from ray import tune + from ray.tune.search.zoopt import ZOOptSearch + from zoopt import ValueType + + dim_dict = { + "height": (ValueType.CONTINUOUS, [-10, 10], 1e-2), + "width": (ValueType.DISCRETE, [-10, 10], False), + "layers": (ValueType.GRID, [4, 8, 16]) + } + + "config": { + "iterations": 10, # evaluation times + } + + zoopt_search_config = { + "parallel_num": 8, # how many workers to parallel + } + + zoopt_search = ZOOptSearch( + algo="Asracos", # only support Asracos currently + budget=20, # must match `num_samples` in `tune.TuneConfig()`. + dim_dict=dim_dict, + metric="mean_loss", + mode="min", + **zoopt_search_config + ) + + tuner = tune.Tuner( + my_objective, + tune_config=tune.TuneConfig( + search_alg=zoopt_search, + num_samples=20 + ), + run_config=tune.RunConfig( + name="zoopt_search", + stop={"timesteps_total": 10} + ), + ) + tuner.fit() + + Parameters: + algo: To specify an algorithm in zoopt you want to use. + Only support ASRacos currently. + budget: Number of samples. + dim_dict: Dimension dictionary. + For continuous dimensions: (continuous, search_range, precision); + For discrete dimensions: (discrete, search_range, has_order); + For grid dimensions: (grid, grid_list). + More details can be found in zoopt package. + metric: The training result objective value attribute. If None + but a mode was passed, the anonymous metric `_metric` will be used + per default. + mode: One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + points_to_evaluate: Initial parameter suggestions to be run + first. This is for when you already have some good parameters + you want to run first to help the algorithm make better suggestions + for future parameters. Needs to be a list of dicts containing the + configurations. + parallel_num: How many workers to parallel. Note that initial + phase may start less workers than this number. More details can + be found in zoopt package. + """ + + optimizer = None + + def __init__( + self, + algo: str = "asracos", + budget: Optional[int] = None, + dim_dict: Optional[Dict] = None, + metric: Optional[str] = None, + mode: Optional[str] = None, + points_to_evaluate: Optional[List[Dict]] = None, + parallel_num: int = 1, + **kwargs + ): + assert ( + zoopt is not None + ), "ZOOpt not found - please install zoopt by `pip install -U zoopt`." + assert budget is not None, "`budget` should not be None!" + if mode: + assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." + _algo = algo.lower() + assert _algo in [ + "asracos", + "sracos", + ], "`algo` must be in ['asracos', 'sracos'] currently" + + self._algo = _algo + + if isinstance(dim_dict, dict) and dim_dict: + resolved_vars, domain_vars, grid_vars = parse_spec_vars(dim_dict) + if domain_vars or grid_vars: + logger.warning( + UNRESOLVED_SEARCH_SPACE.format(par="dim_dict", cls=type(self)) + ) + dim_dict = self.convert_search_space(dim_dict, join=True) + + self._dim_dict = dim_dict + self._budget = budget + + self._metric = metric + if mode == "max": + self._metric_op = -1.0 + elif mode == "min": + self._metric_op = 1.0 + + self._points_to_evaluate = copy.deepcopy(points_to_evaluate) + + self._live_trial_mapping = {} + + self._dim_keys = [] + self.solution_dict = {} + self.best_solution_list = [] + self.optimizer = None + + self.kwargs = kwargs + + self.parallel_num = parallel_num + + super(ZOOptSearch, self).__init__(metric=self._metric, mode=mode) + + if self._dim_dict: + self._setup_zoopt() + + def _setup_zoopt(self): + if self._metric is None and self._mode: + # If only a mode was passed, use anonymous metric + self._metric = DEFAULT_METRIC + + _dim_list = [] + for k in self._dim_dict: + self._dim_keys.append(k) + _dim_list.append(self._dim_dict[k]) + + init_samples = None + if self._points_to_evaluate: + logger.warning( + "`points_to_evaluate` is ignored by ZOOpt in versions <= 0.4.1." + ) + init_samples = [ + Solution(x=tuple(point[dim] for dim in self._dim_keys)) + for point in self._points_to_evaluate + ] + dim = zoopt.Dimension2(_dim_list) + par = zoopt.Parameter(budget=self._budget, init_samples=init_samples) + if self._algo == "sracos" or self._algo == "asracos": + from zoopt.algos.opt_algorithms.racos.sracos import SRacosTune + + self.optimizer = SRacosTune( + dimension=dim, + parameter=par, + parallel_num=self.parallel_num, + **self.kwargs + ) + + def set_search_properties( + self, metric: Optional[str], mode: Optional[str], config: Dict, **spec + ) -> bool: + if self._dim_dict: + return False + space = self.convert_search_space(config) + self._dim_dict = space + + if metric: + self._metric = metric + if mode: + self._mode = mode + + if self._mode == "max": + self._metric_op = -1.0 + elif self._mode == "min": + self._metric_op = 1.0 + + self._setup_zoopt() + return True + + def suggest(self, trial_id: str) -> Optional[Dict]: + if not self._dim_dict or not self.optimizer: + raise RuntimeError( + UNDEFINED_SEARCH_SPACE.format( + cls=self.__class__.__name__, space="dim_dict" + ) + ) + if not self._metric or not self._mode: + raise RuntimeError( + UNDEFINED_METRIC_MODE.format( + cls=self.__class__.__name__, metric=self._metric, mode=self._mode + ) + ) + + _solution = self.optimizer.suggest() + + if _solution == "FINISHED": + if ray.__version__ >= "0.8.7": + return Searcher.FINISHED + else: + return None + + if _solution: + self.solution_dict[str(trial_id)] = _solution + _x = _solution.get_x() + new_trial = dict(zip(self._dim_keys, _x)) + self._live_trial_mapping[trial_id] = new_trial + return unflatten_dict(new_trial) + + def on_trial_complete( + self, trial_id: str, result: Optional[Dict] = None, error: bool = False + ): + """Notification for the completion of trial.""" + if result: + _solution = self.solution_dict[str(trial_id)] + _best_solution_so_far = self.optimizer.complete( + _solution, self._metric_op * result[self._metric] + ) + if _best_solution_so_far: + self.best_solution_list.append(_best_solution_so_far) + + del self._live_trial_mapping[trial_id] + + def save(self, checkpoint_path: str): + save_object = self.__dict__ + with open(checkpoint_path, "wb") as outputFile: + pickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path: str): + with open(checkpoint_path, "rb") as inputFile: + save_object = pickle.load(inputFile) + self.__dict__.update(save_object) + + @staticmethod + def convert_search_space(spec: Dict, join: bool = False) -> Dict[str, Tuple]: + spec = copy.deepcopy(spec) + resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) + + if not domain_vars and not grid_vars: + return {} + + if grid_vars: + raise ValueError( + "Grid search parameters cannot be automatically converted " + "to a ZOOpt search space." + ) + + def resolve_value(domain: Domain) -> Tuple: + quantize = None + + sampler = domain.get_sampler() + if isinstance(sampler, Quantized): + quantize = sampler.q + sampler = sampler.sampler + + if isinstance(domain, Float): + precision = quantize or 1e-12 + if isinstance(sampler, Uniform): + return ( + ValueType.CONTINUOUS, + [domain.lower, domain.upper], + precision, + ) + + elif isinstance(domain, Integer): + if isinstance(sampler, Uniform): + return (ValueType.DISCRETE, [domain.lower, domain.upper - 1], True) + + elif isinstance(domain, Categorical): + # Categorical variables would use ValueType.DISCRETE with + # has_partial_order=False, however, currently we do not + # keep track of category values and cannot automatically + # translate back and forth between them. + if isinstance(sampler, Uniform): + return (ValueType.GRID, domain.categories) + + raise ValueError( + "ZOOpt does not support parameters of type " + "`{}` with samplers of type `{}`".format( + type(domain).__name__, type(domain.sampler).__name__ + ) + ) + + conv_spec = { + "/".join(path): resolve_value(domain) for path, domain in domain_vars + } + + if join: + spec.update(conv_spec) + conv_spec = spec + + return conv_spec