wehub-resource-sync
1081 lines
43 KiB
Python
1081 lines
43 KiB
Python
import abc
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import pathlib
|
|
import re
|
|
import tempfile
|
|
from types import MappingProxyType
|
|
from typing import Any, Collection, Dict, List, Optional, Tuple, Union
|
|
|
|
import pyarrow.fs
|
|
from packaging import version
|
|
|
|
import ray
|
|
import ray.cloudpickle as pickle
|
|
from ray.rllib.core import (
|
|
COMPONENT_LEARNER,
|
|
COMPONENT_LEARNER_GROUP,
|
|
COMPONENT_RL_MODULE,
|
|
)
|
|
from ray.rllib.utils import force_list
|
|
from ray.rllib.utils.actor_manager import FaultTolerantActorManager
|
|
from ray.rllib.utils.annotations import (
|
|
OldAPIStack,
|
|
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
|
)
|
|
from ray.rllib.utils.serialization import NOT_SERIALIZABLE, serialize_type
|
|
from ray.rllib.utils.typing import StateDict
|
|
from ray.train import Checkpoint as Checkpoint_train
|
|
from ray.tune import Checkpoint as Checkpoint_tune
|
|
from ray.tune.utils.file_transfer import sync_dir_between_nodes
|
|
from ray.util import log_once
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# The current checkpoint version used by RLlib for Algorithm and Policy checkpoints.
|
|
# History:
|
|
# 0.1: Ray 2.0.0
|
|
# A single `checkpoint-[iter num]` file for Algorithm checkpoints
|
|
# within the checkpoint directory. Policy checkpoints not supported across all
|
|
# DL frameworks.
|
|
|
|
# 1.0: Ray >=2.1.0
|
|
# An algorithm_state.pkl file for the state of the Algorithm (excluding
|
|
# individual policy states).
|
|
# One sub-dir inside the "policies" sub-dir for each policy with a
|
|
# dedicated policy_state.pkl in it for the policy state.
|
|
|
|
# 1.1: Same as 1.0, but has a new "format" field in the rllib_checkpoint.json file
|
|
# indicating, whether the checkpoint is `cloudpickle` (default) or `msgpack`.
|
|
|
|
# 1.2: Introduces the checkpoint for the new Learner API if the Learner API is enabled.
|
|
|
|
# 2.0: Introduces the Checkpointable API for all components on the new API stack
|
|
# (if the Learner-, RLModule, EnvRunner, and ConnectorV2 APIs are enabled).
|
|
|
|
CHECKPOINT_VERSION = version.Version("1.1")
|
|
CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER = version.Version("2.1")
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class Checkpointable(abc.ABC):
|
|
"""Abstract base class for a component of RLlib that can be checkpointed to disk.
|
|
|
|
Subclasses must implement the following APIs:
|
|
- save_to_path()
|
|
- restore_from_path()
|
|
- from_checkpoint()
|
|
- get_state()
|
|
- set_state()
|
|
- get_ctor_args_and_kwargs()
|
|
- get_metadata()
|
|
- get_checkpointable_components()
|
|
"""
|
|
|
|
# The state file for the implementing class.
|
|
# This file contains any state information that does NOT belong to any subcomponent
|
|
# of the implementing class (which are `Checkpointable` themselves and thus should
|
|
# have their own state- and metadata files).
|
|
# After a `save_to_path([path])` this file can be found directly in: `path/`.
|
|
STATE_FILE_NAME = "state"
|
|
|
|
# The filename of the pickle file that contains the class information of the
|
|
# Checkpointable as well as all constructor args to be passed to such a class in
|
|
# order to construct a new instance.
|
|
CLASS_AND_CTOR_ARGS_FILE_NAME = "class_and_ctor_args.pkl"
|
|
|
|
# Subclasses may set this to their own metadata filename.
|
|
# The dict returned by self.get_metadata() is stored in this JSON file.
|
|
METADATA_FILE_NAME = "metadata.json"
|
|
|
|
def save_to_path(
|
|
self,
|
|
path: Optional[Union[str, pathlib.Path]] = None,
|
|
*,
|
|
state: Optional[StateDict] = None,
|
|
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
use_msgpack: bool = False,
|
|
) -> str:
|
|
"""Saves the state of the implementing class (or `state`) to `path`.
|
|
|
|
The state of the implementing class is always saved in the following format:
|
|
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
path/
|
|
[component1]/
|
|
[component1 subcomponentA]/
|
|
...
|
|
[component1 subcomponentB]/
|
|
...
|
|
[component2]/
|
|
...
|
|
[cls.METADATA_FILE_NAME] (json)
|
|
[cls.STATE_FILE_NAME] (pkl|msgpack)
|
|
|
|
The main logic is to loop through all subcomponents of this Checkpointable
|
|
and call their respective `save_to_path` methods. Then save the remaining
|
|
(non subcomponent) state to this Checkpointable's STATE_FILE_NAME.
|
|
In the exception that a component is a FaultTolerantActorManager instance,
|
|
instead of calling `save_to_path` directly on that manager, the first healthy
|
|
actor is interpreted as the component and its `save_to_path` method is called.
|
|
Even if that actor is located on another node, the created file is automatically
|
|
synced to the local node.
|
|
|
|
Args:
|
|
path: The path to the directory to save the state of the implementing class
|
|
to. If `path` doesn't exist or is None, then a new directory will be
|
|
created (and returned).
|
|
state: An optional state dict to be used instead of getting a new state of
|
|
the implementing class through `self.get_state()`.
|
|
filesystem: PyArrow FileSystem to use to access data at the `path`.
|
|
If not specified, this is inferred from the URI scheme of `path`.
|
|
use_msgpack: Whether the state file should be written using msgpack and
|
|
msgpack_numpy (file extension is `.msgpack`), rather than pickle (file
|
|
extension is `.pkl`).
|
|
|
|
Returns:
|
|
The path (str) where the state has been saved.
|
|
"""
|
|
|
|
# If no path is given create a local temporary directory.
|
|
if path is None:
|
|
import uuid
|
|
|
|
# Get the location of the temporary directory on the OS.
|
|
tmp_dir = pathlib.Path(tempfile.gettempdir())
|
|
# Create a random directory name.
|
|
random_dir_name = str(uuid.uuid4())
|
|
# Create the path, but do not craet the directory on the
|
|
# filesystem, yet. This is done by `PyArrow`.
|
|
path = path or tmp_dir / random_dir_name
|
|
|
|
# We need a string path for `pyarrow.fs.FileSystem.from_uri`.
|
|
path = path if isinstance(path, str) else path.as_posix()
|
|
|
|
# If we have no filesystem, figure it out.
|
|
if path and not filesystem:
|
|
# Note the path needs to be a path that is relative to the
|
|
# filesystem (e.g. `gs://tmp/...` -> `tmp/...`).
|
|
filesystem, path = pyarrow.fs.FileSystem.from_uri(path)
|
|
|
|
# Make sure, path exists.
|
|
filesystem.create_dir(path, recursive=True)
|
|
|
|
# Convert to `pathlib.Path` for easy handling.
|
|
path = pathlib.Path(path)
|
|
|
|
# Write metadata file to disk.
|
|
metadata = self.get_metadata()
|
|
if "checkpoint_version" not in metadata:
|
|
metadata["checkpoint_version"] = str(
|
|
CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER
|
|
)
|
|
with filesystem.open_output_stream(
|
|
(path / self.METADATA_FILE_NAME).as_posix()
|
|
) as f:
|
|
f.write(json.dumps(metadata).encode("utf-8"))
|
|
|
|
# Write the class and constructor args information to disk. Always use pickle
|
|
# for this, because this information contains classes and maybe other
|
|
# non-serializable data.
|
|
with filesystem.open_output_stream(
|
|
(path / self.CLASS_AND_CTOR_ARGS_FILE_NAME).as_posix()
|
|
) as f:
|
|
pickle.dump(
|
|
{
|
|
"class": type(self),
|
|
"ctor_args_and_kwargs": self.get_ctor_args_and_kwargs(),
|
|
},
|
|
f,
|
|
)
|
|
|
|
# Get the entire state of this Checkpointable, or use provided `state`.
|
|
_state_provided = state is not None
|
|
# Get only the non-checkpointable components of the state. Checkpointable
|
|
# components are saved to path by their own `save_to_path` in the loop below.
|
|
state = state or self.get_state(
|
|
not_components=[c[0] for c in self.get_checkpointable_components()]
|
|
)
|
|
|
|
# Write components of `self` that themselves are `Checkpointable`.
|
|
for comp_name, comp in self.get_checkpointable_components():
|
|
# If subcomponent's name is not in `state`, ignore it and don't write this
|
|
# subcomponent's state to disk.
|
|
if _state_provided and comp_name not in state:
|
|
continue
|
|
comp_path = path / comp_name
|
|
|
|
# If component is an ActorManager, save the manager's first healthy
|
|
# actor's state to disk (even if it's on another node, in which case, we'll
|
|
# sync the generated file(s) back to this node).
|
|
if isinstance(comp, FaultTolerantActorManager):
|
|
actor_to_use = comp.healthy_actor_ids()[0]
|
|
|
|
def _get_ip(_=None):
|
|
import ray
|
|
|
|
return ray.util.get_node_ip_address()
|
|
|
|
_result = next(
|
|
iter(
|
|
comp.foreach_actor(
|
|
_get_ip,
|
|
remote_actor_ids=[actor_to_use],
|
|
)
|
|
)
|
|
)
|
|
if not _result.ok:
|
|
raise _result.get()
|
|
worker_ip_addr = _result.get()
|
|
self_ip_addr = _get_ip()
|
|
|
|
# Save the state to a temporary location on the `actor_to_use`'s
|
|
# node.
|
|
comp_state_ref = None
|
|
if _state_provided:
|
|
comp_state_ref = ray.put(state.pop(comp_name))
|
|
|
|
# If worker_addr == self_addr, save directly to the path
|
|
# provided by the user, make sure to use filesystem.
|
|
if worker_ip_addr == self_ip_addr:
|
|
comp.foreach_actor(
|
|
lambda w, _path=comp_path, _filesystem=filesystem, _state=comp_state_ref, _use_msgpack=use_msgpack: ( # noqa
|
|
w.save_to_path(
|
|
path=_path,
|
|
filesystem=_filesystem,
|
|
state=(
|
|
ray.get(_state)
|
|
if _state is not None
|
|
else w.get_state()
|
|
),
|
|
use_msgpack=_use_msgpack,
|
|
)
|
|
),
|
|
remote_actor_ids=[actor_to_use],
|
|
)
|
|
# Transfer state files from the worker node to the head node
|
|
else:
|
|
# Save the checkpoint to the temporary directory on the worker.
|
|
def _save(w, _state=comp_state_ref, _use_msgpack=use_msgpack):
|
|
import tempfile
|
|
|
|
# Create a temporary directory on the worker.
|
|
tmpdir = tempfile.mkdtemp()
|
|
w.save_to_path(
|
|
path=tmpdir,
|
|
state=(
|
|
ray.get(_state) if _state is not None else w.get_state()
|
|
),
|
|
use_msgpack=_use_msgpack,
|
|
)
|
|
return tmpdir
|
|
|
|
_result = next(
|
|
iter(comp.foreach_actor(_save, remote_actor_ids=[actor_to_use]))
|
|
)
|
|
if not _result.ok:
|
|
raise _result.get()
|
|
worker_temp_dir = _result.get()
|
|
|
|
# Sync the temporary directory from the worker to this node.
|
|
sync_dir_between_nodes(
|
|
worker_ip_addr,
|
|
worker_temp_dir,
|
|
self_ip_addr,
|
|
str(comp_path),
|
|
)
|
|
|
|
# Remove the temporary directory on the worker.
|
|
def _rmdir(_, _dir=worker_temp_dir):
|
|
import shutil
|
|
|
|
shutil.rmtree(_dir)
|
|
|
|
comp.foreach_actor(_rmdir, remote_actor_ids=[actor_to_use])
|
|
|
|
# Local component (instance stored in a property of `self`).
|
|
else:
|
|
if _state_provided:
|
|
comp_state = state.pop(comp_name)
|
|
else:
|
|
comp_state = self.get_state(components=comp_name)[comp_name]
|
|
# By providing the `state` arg, we make sure that the component does not
|
|
# have to call its own `get_state()` anymore, but uses what's provided
|
|
# here.
|
|
comp.save_to_path(
|
|
path=comp_path,
|
|
filesystem=filesystem,
|
|
state=comp_state,
|
|
use_msgpack=use_msgpack,
|
|
)
|
|
|
|
# Write all the remaining state to disk.
|
|
filename = path / (
|
|
self.STATE_FILE_NAME + (".msgpack" if use_msgpack else ".pkl")
|
|
)
|
|
with filesystem.open_output_stream(filename.as_posix()) as f:
|
|
if use_msgpack:
|
|
msgpack = try_import_msgpack(error=True)
|
|
msgpack.dump(state, f)
|
|
else:
|
|
pickle.dump(state, f)
|
|
|
|
return str(path)
|
|
|
|
def restore_from_path(
|
|
self,
|
|
path: Union[str, pathlib.Path],
|
|
*,
|
|
component: Optional[str] = None,
|
|
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
"""Restores the state of the implementing class from the given path.
|
|
|
|
If the `component` arg is provided, `path` refers to a checkpoint of a
|
|
subcomponent of `self`, thus allowing the user to load only the subcomponent's
|
|
state into `self` without affecting any of the other state information (for
|
|
example, loading only the NN state into a Checkpointable, which contains such
|
|
an NN, but also has other state information that should NOT be changed by
|
|
calling this method).
|
|
|
|
The given `path` should have the following structure and contain the following
|
|
files:
|
|
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
path/
|
|
[component1]/
|
|
[component1 subcomponentA]/
|
|
...
|
|
[component1 subcomponentB]/
|
|
...
|
|
[component2]/
|
|
...
|
|
[cls.METADATA_FILE_NAME] (json)
|
|
[cls.STATE_FILE_NAME] (pkl|msgpack)
|
|
|
|
Note that the self.METADATA_FILE_NAME file is not required to restore the state.
|
|
|
|
Args:
|
|
path: The path to load the implementing class' state from or to load the
|
|
state of only one subcomponent's state of the implementing class (if
|
|
`component` is provided).
|
|
component: If provided, `path` is interpreted as the checkpoint path of only
|
|
the subcomponent and thus, only that subcomponent's state is
|
|
restored/loaded. All other state of `self` remains unchanged in this
|
|
case.
|
|
filesystem: PyArrow FileSystem to use to access data at the `path`. If not
|
|
specified, this is inferred from the URI scheme of `path`.
|
|
**kwargs: Forward compatibility kwargs.
|
|
"""
|
|
path = path if isinstance(path, str) else path.as_posix()
|
|
|
|
if path and not filesystem:
|
|
# Note the path needs to be a path that is relative to the
|
|
# filesystem (e.g. `gs://tmp/...` -> `tmp/...`).
|
|
filesystem, path = pyarrow.fs.FileSystem.from_uri(path)
|
|
# Only here convert to a `Path` instance b/c otherwise
|
|
# cloud path gets broken (i.e. 'gs://' -> 'gs:/').
|
|
path = pathlib.Path(path)
|
|
|
|
if not _exists_at_fs_path(filesystem, path.as_posix()):
|
|
raise FileNotFoundError(f"`path` ({path}) not found!")
|
|
|
|
# Restore components of `self` that themselves are `Checkpointable`.
|
|
orig_comp_names = {c[0] for c in self.get_checkpointable_components()}
|
|
self._restore_all_subcomponents_from_path(
|
|
path=path,
|
|
filesystem=filesystem,
|
|
component=component,
|
|
**kwargs,
|
|
)
|
|
|
|
# Restore the "base" state (not individual subcomponents).
|
|
if component is None:
|
|
filename = path / self.STATE_FILE_NAME
|
|
if filename.with_suffix(".msgpack").is_file():
|
|
msgpack = try_import_msgpack(error=True)
|
|
with filesystem.open_input_stream(
|
|
filename.with_suffix(".msgpack").as_posix()
|
|
) as f:
|
|
state = msgpack.load(f, strict_map_key=False)
|
|
else:
|
|
with filesystem.open_input_stream(
|
|
filename.with_suffix(".pkl").as_posix()
|
|
) as f:
|
|
state = pickle.load(f)
|
|
self.set_state(state)
|
|
|
|
new_comp_names = {c[0] for c in self.get_checkpointable_components()}
|
|
diff_comp_names = new_comp_names - orig_comp_names
|
|
if diff_comp_names:
|
|
self._restore_all_subcomponents_from_path(
|
|
path=path,
|
|
filesystem=filesystem,
|
|
only_comp_names=diff_comp_names,
|
|
**kwargs,
|
|
)
|
|
|
|
@classmethod
|
|
def from_checkpoint(
|
|
cls,
|
|
path: Union[str, pathlib.Path],
|
|
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
**kwargs,
|
|
) -> "Checkpointable":
|
|
"""Creates a new Checkpointable instance from the given location and returns it.
|
|
|
|
Args:
|
|
path: The checkpoint path to load (a) the information on how to construct
|
|
a new instance of the implementing class and (b) the state to restore
|
|
the created instance to.
|
|
filesystem: PyArrow FileSystem to use to access data at the `path`. If not
|
|
specified, this is inferred from the URI scheme of `path`.
|
|
kwargs: Forward compatibility kwargs. Note that these kwargs are sent to
|
|
each subcomponent's `from_checkpoint()` call.
|
|
|
|
Returns:
|
|
A new instance of the implementing class, already set to the state stored
|
|
under `path`.
|
|
"""
|
|
# We need a string path for the `PyArrow` filesystem.
|
|
path = path if isinstance(path, str) else path.as_posix()
|
|
|
|
# If no filesystem is passed in create one.
|
|
if path and not filesystem:
|
|
# Note the path needs to be a path that is relative to the
|
|
# filesystem (e.g. `gs://tmp/...` -> `tmp/...`).
|
|
filesystem, path = pyarrow.fs.FileSystem.from_uri(path)
|
|
# Only here convert to a `Path` instance b/c otherwise
|
|
# cloud path gets broken (i.e. 'gs://' -> 'gs:/').
|
|
path = pathlib.Path(path)
|
|
|
|
# Get the class constructor to call and its args/kwargs.
|
|
# Try reading the pickle file first.
|
|
try:
|
|
with filesystem.open_input_stream(
|
|
(path / cls.CLASS_AND_CTOR_ARGS_FILE_NAME).as_posix()
|
|
) as f:
|
|
ctor_info = pickle.load(f)
|
|
ctor = ctor_info["class"]
|
|
ctor_args = force_list(ctor_info["ctor_args_and_kwargs"][0])
|
|
ctor_kwargs = ctor_info["ctor_args_and_kwargs"][1]
|
|
|
|
# Inspect the ctor to see, which arguments in ctor_info should be replaced
|
|
# with the user provided **kwargs.
|
|
for i, (param_name, param) in enumerate(
|
|
inspect.signature(ctor).parameters.items()
|
|
):
|
|
if param_name in kwargs:
|
|
val = kwargs.pop(param_name)
|
|
if (
|
|
param.kind == inspect._ParameterKind.POSITIONAL_OR_KEYWORD
|
|
and len(ctor_args) > i
|
|
):
|
|
ctor_args[i] = val
|
|
else:
|
|
ctor_kwargs[param_name] = val
|
|
|
|
# If the pickle file is from another python version, use provided args instead.
|
|
except (ValueError, AttributeError, ImportError, TypeError):
|
|
logger.warning(
|
|
"Could not restore original class from checkpoint at '%s' "
|
|
"(possible version mismatch), falling back to %s.",
|
|
path,
|
|
cls.__name__,
|
|
)
|
|
# Use class that this method was called on.
|
|
ctor = cls
|
|
# Use only user provided **kwargs.
|
|
ctor_args = []
|
|
ctor_kwargs = kwargs
|
|
|
|
# Check, whether the constructor actually goes together with `cls`.
|
|
if not issubclass(ctor, cls):
|
|
raise ValueError(
|
|
f"The class ({ctor}) stored in checkpoint ({path}) does not seem to be "
|
|
f"a subclass of `cls` ({cls})!"
|
|
)
|
|
elif not issubclass(ctor, Checkpointable):
|
|
raise ValueError(
|
|
f"The class ({ctor}) stored in checkpoint ({path}) does not seem to be "
|
|
"an implementer of the `Checkpointable` API!"
|
|
)
|
|
|
|
# Construct the initial object (without any particular state).
|
|
obj = ctor(*ctor_args, **ctor_kwargs)
|
|
# Restore the state of the constructed object.
|
|
obj.restore_from_path(path, filesystem=filesystem, **kwargs)
|
|
# Return the new object.
|
|
return obj
|
|
|
|
@abc.abstractmethod
|
|
def get_state(
|
|
self,
|
|
components: Optional[Union[str, Collection[str]]] = None,
|
|
*,
|
|
not_components: Optional[Union[str, Collection[str]]] = None,
|
|
**kwargs,
|
|
) -> StateDict:
|
|
"""Returns the implementing class's current state as a dict.
|
|
|
|
The returned dict must only contain msgpack-serializable data if you want to
|
|
use the `AlgorithmConfig._msgpack_checkpoints` option. Consider returning your
|
|
non msgpack-serializable data from the `Checkpointable.get_ctor_args_and_kwargs`
|
|
method, instead.
|
|
|
|
Args:
|
|
components: An optional collection of string keys to be included in the
|
|
returned state. This might be useful, if getting certain components
|
|
of the state is expensive (e.g. reading/compiling the weights of a large
|
|
NN) and at the same time, these components are not required by the
|
|
caller.
|
|
not_components: An optional list of string keys to be excluded in the
|
|
returned state, even if the same string is part of `components`.
|
|
This is useful to get the complete state of the class, except
|
|
one or a few components.
|
|
kwargs: Forward-compatibility kwargs.
|
|
|
|
Returns:
|
|
The current state of the implementing class (or only the `components`
|
|
specified, w/o those in `not_components`).
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def set_state(self, state: StateDict) -> None:
|
|
"""Sets the implementing class' state to the given state dict.
|
|
|
|
If component keys are missing in `state`, these components of the implementing
|
|
class will not be updated/set.
|
|
|
|
Args:
|
|
state: The state dict to restore the state from. Maps component keys
|
|
to the corresponding subcomponent's own state.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def get_ctor_args_and_kwargs(self) -> Tuple[Tuple, Dict[str, Any]]:
|
|
"""Returns the args/kwargs used to create `self` from its constructor.
|
|
|
|
Returns:
|
|
A tuple of the args (as a tuple) and kwargs (as a Dict[str, Any]) used to
|
|
construct `self` from its class constructor.
|
|
"""
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
def get_metadata(self) -> Dict:
|
|
"""Returns JSON writable metadata further describing the implementing class.
|
|
|
|
Note that this metadata is NOT part of any state and is thus NOT needed to
|
|
restore the state of a Checkpointable instance from a directory. Rather, the
|
|
metadata will be written into `self.METADATA_FILE_NAME` when calling
|
|
`self.save_to_path()` for the user's convenience.
|
|
|
|
Returns:
|
|
A JSON-encodable dict of metadata information.
|
|
"""
|
|
return {
|
|
"class_and_ctor_args_file": self.CLASS_AND_CTOR_ARGS_FILE_NAME,
|
|
"state_file": self.STATE_FILE_NAME,
|
|
"ray_version": ray.__version__,
|
|
"ray_commit": ray.__commit__,
|
|
}
|
|
|
|
def get_checkpointable_components(self) -> List[Tuple[str, "Checkpointable"]]:
|
|
"""Returns the implementing class's own Checkpointable subcomponents.
|
|
|
|
Returns:
|
|
A list of 2-tuples (name, subcomponent) describing the implementing class'
|
|
subcomponents, all of which have to be `Checkpointable` themselves and
|
|
whose state is therefore written into subdirectories (rather than the main
|
|
state file (self.STATE_FILE_NAME) when calling `self.save_to_path()`).
|
|
"""
|
|
return []
|
|
|
|
def _check_component(self, name, components, not_components) -> bool:
|
|
"""Returns True if a component should be checkpointed.
|
|
|
|
Args:
|
|
name: The checkpoint name.
|
|
components: A list of components that should be checkpointed.
|
|
non_components: A list of components that should not be checkpointed.
|
|
|
|
Returns:
|
|
True, if the component should be checkpointed and otherwise False.
|
|
"""
|
|
comp_list = force_list(components)
|
|
not_comp_list = force_list(not_components)
|
|
if (
|
|
components is None
|
|
or any(c.startswith(name + "/") for c in comp_list)
|
|
or name in comp_list
|
|
) and (not_components is None or name not in not_comp_list):
|
|
return True
|
|
return False
|
|
|
|
def _get_subcomponents(self, name, components):
|
|
if components is None:
|
|
return None
|
|
|
|
components = force_list(components)
|
|
subcomponents = []
|
|
for comp in components:
|
|
if comp.startswith(name + "/"):
|
|
subcomponents.append(comp[len(name) + 1 :])
|
|
|
|
return None if not subcomponents else subcomponents
|
|
|
|
def _restore_all_subcomponents_from_path(
|
|
self, path, filesystem, only_comp_names=None, component=None, **kwargs
|
|
):
|
|
for comp_name, comp in self.get_checkpointable_components():
|
|
if only_comp_names is not None and comp_name not in only_comp_names:
|
|
continue
|
|
|
|
# The value of the `component` argument for the upcoming
|
|
# `[subcomponent].restore_from_path(.., component=..)` call.
|
|
comp_arg = None
|
|
|
|
if component is None:
|
|
comp_dir = path / comp_name
|
|
# If subcomponent's dir is not in path, ignore it and don't restore this
|
|
# subcomponent's state from disk.
|
|
if not _exists_at_fs_path(filesystem, comp_dir.as_posix()):
|
|
continue
|
|
else:
|
|
comp_dir = path
|
|
|
|
# `component` is a path that starts with `comp` -> Remove the name of
|
|
# `comp` from the `component` arg in the upcoming call to `restore_..`.
|
|
if component.startswith(comp_name + "/"):
|
|
comp_arg = component[len(comp_name) + 1 :]
|
|
# `component` has nothing to do with `comp` -> Skip.
|
|
elif component != comp_name:
|
|
continue
|
|
|
|
# If component is an ActorManager, restore all the manager's healthy
|
|
# actors' states from disk (even if they are on another node, in which case,
|
|
# we'll sync checkpoint file(s) to the respective node).
|
|
if isinstance(comp, FaultTolerantActorManager):
|
|
head_node_ip = ray.util.get_node_ip_address()
|
|
all_healthy_actors = comp.healthy_actor_ids()
|
|
|
|
def _restore(
|
|
w,
|
|
_kwargs=MappingProxyType(kwargs),
|
|
_path=comp_dir,
|
|
_head_ip=head_node_ip,
|
|
_comp_arg=comp_arg,
|
|
):
|
|
import tempfile
|
|
|
|
import ray
|
|
|
|
worker_node_ip = ray.util.get_node_ip_address()
|
|
# If the worker is on the same node as the head, load the checkpoint
|
|
# directly from the path otherwise sync the checkpoint from the head
|
|
# to the worker and load it from there.
|
|
if worker_node_ip == _head_ip:
|
|
w.restore_from_path(
|
|
path=_path,
|
|
filesystem=filesystem,
|
|
component=_comp_arg,
|
|
**_kwargs,
|
|
)
|
|
else:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
sync_dir_between_nodes(
|
|
_head_ip, _path, worker_node_ip, temp_dir
|
|
)
|
|
w.restore_from_path(
|
|
temp_dir, component=_comp_arg, **_kwargs
|
|
)
|
|
|
|
comp.foreach_actor(_restore, remote_actor_ids=all_healthy_actors)
|
|
|
|
# Call `restore_from_path()` on local subcomponent, thereby passing in the
|
|
# **kwargs.
|
|
else:
|
|
comp.restore_from_path(
|
|
comp_dir, filesystem=filesystem, component=comp_arg, **kwargs
|
|
)
|
|
|
|
|
|
def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, path: str) -> bool:
|
|
"""Returns `True` if the path can be found in the filesystem."""
|
|
valid = fs.get_file_info(path)
|
|
return valid.type != pyarrow.fs.FileType.NotFound
|
|
|
|
|
|
def _is_dir(file_info: pyarrow.fs.FileInfo) -> bool:
|
|
"""Returns `True`, if the file info is from a directory."""
|
|
return file_info.type == pyarrow.fs.FileType.Directory
|
|
|
|
|
|
@OldAPIStack
|
|
def get_checkpoint_info(
|
|
checkpoint: Union[str, Checkpoint_train, Checkpoint_tune],
|
|
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Returns a dict with information about an Algorithm/Policy checkpoint.
|
|
|
|
If the given checkpoint is a >=v1.0 checkpoint directory, try reading all
|
|
information from the contained `rllib_checkpoint.json` file.
|
|
|
|
Args:
|
|
checkpoint: The checkpoint directory (str) or a Checkpoint object.
|
|
filesystem: PyArrow FileSystem to use to access data at the `checkpoint`. If not
|
|
specified, this is inferred from the URI scheme provided by `checkpoint`.
|
|
|
|
Returns:
|
|
A dict containing the keys:
|
|
"type": One of "Policy" or "Algorithm".
|
|
"checkpoint_version": A version tuple, e.g. v1.0, indicating the checkpoint
|
|
version. This will help RLlib to remain backward compatible wrt. future
|
|
Ray and checkpoint versions.
|
|
"checkpoint_dir": The directory with all the checkpoint files in it. This might
|
|
be the same as the incoming `checkpoint` arg.
|
|
"state_file": The main file with the Algorithm/Policy's state information in it.
|
|
This is usually a pickle-encoded file.
|
|
"policy_ids": An optional set of PolicyIDs in case we are dealing with an
|
|
Algorithm checkpoint. None if `checkpoint` is a Policy checkpoint.
|
|
"""
|
|
# Default checkpoint info.
|
|
info = {
|
|
"type": "Algorithm",
|
|
"format": "cloudpickle",
|
|
"checkpoint_version": CHECKPOINT_VERSION,
|
|
"checkpoint_dir": None,
|
|
"state_file": None,
|
|
"policy_ids": None,
|
|
"module_ids": None,
|
|
}
|
|
|
|
# `checkpoint` is a Checkpoint instance: Translate to directory and continue.
|
|
if isinstance(checkpoint, (Checkpoint_train, Checkpoint_tune)):
|
|
checkpoint = checkpoint.to_directory()
|
|
|
|
if checkpoint and not filesystem:
|
|
# Note the path needs to be a path that is relative to the
|
|
# filesystem (e.g. `gs://tmp/...` -> `tmp/...`).
|
|
filesystem, checkpoint = pyarrow.fs.FileSystem.from_uri(checkpoint)
|
|
# Only here convert to a `Path` instance b/c otherwise
|
|
# cloud path gets broken (i.e. 'gs://' -> 'gs:/').
|
|
checkpoint = pathlib.Path(checkpoint)
|
|
|
|
# Checkpoint is dir.
|
|
if _exists_at_fs_path(filesystem, checkpoint.as_posix()) and _is_dir(
|
|
filesystem.get_file_info(checkpoint.as_posix())
|
|
):
|
|
info.update({"checkpoint_dir": str(checkpoint)})
|
|
|
|
# Figure out whether this is an older checkpoint format
|
|
# (with a `checkpoint-\d+` file in it).
|
|
file_info_list = filesystem.get_file_info(
|
|
pyarrow.fs.FileSelector(checkpoint.as_posix(), recursive=False)
|
|
)
|
|
for file_info in file_info_list:
|
|
if file_info.is_file:
|
|
if re.match("checkpoint-\\d+", file_info.base_name):
|
|
info.update(
|
|
{
|
|
"checkpoint_version": version.Version("0.1"),
|
|
"state_file": str(file_info.base_name),
|
|
}
|
|
)
|
|
return info
|
|
|
|
# No old checkpoint file found.
|
|
|
|
# If rllib_checkpoint.json file present, read available information from it
|
|
# and then continue with the checkpoint analysis (possibly overriding further
|
|
# information).
|
|
if _exists_at_fs_path(
|
|
filesystem, (checkpoint / "rllib_checkpoint.json").as_posix()
|
|
):
|
|
# if (checkpoint / "rllib_checkpoint.json").is_file():
|
|
with filesystem.open_input_stream(
|
|
(checkpoint / "rllib_checkpoint.json").as_posix()
|
|
) as f:
|
|
# with open(checkpoint / "rllib_checkpoint.json") as f:
|
|
rllib_checkpoint_info = json.load(fp=f)
|
|
if "checkpoint_version" in rllib_checkpoint_info:
|
|
rllib_checkpoint_info["checkpoint_version"] = version.Version(
|
|
rllib_checkpoint_info["checkpoint_version"]
|
|
)
|
|
info.update(rllib_checkpoint_info)
|
|
else:
|
|
# No rllib_checkpoint.json file present: Warn and continue trying to figure
|
|
# out checkpoint info ourselves.
|
|
if log_once("no_rllib_checkpoint_json_file"):
|
|
logger.warning(
|
|
"No `rllib_checkpoint.json` file found in checkpoint directory "
|
|
f"{checkpoint}! Trying to extract checkpoint info from other files "
|
|
f"found in that dir."
|
|
)
|
|
|
|
# Policy checkpoint file found.
|
|
for extension in ["pkl", "msgpck"]:
|
|
if _exists_at_fs_path(
|
|
filesystem, (checkpoint / ("policy_state." + extension)).as_posix()
|
|
):
|
|
# if (checkpoint / ("policy_state." + extension)).is_file():
|
|
info.update(
|
|
{
|
|
"type": "Policy",
|
|
"format": "cloudpickle" if extension == "pkl" else "msgpack",
|
|
"checkpoint_version": CHECKPOINT_VERSION,
|
|
"state_file": str(checkpoint / f"policy_state.{extension}"),
|
|
}
|
|
)
|
|
return info
|
|
|
|
# Valid Algorithm checkpoint >v0 file found?
|
|
format = None
|
|
for extension in ["pkl", "msgpck", "msgpack"]:
|
|
state_file = checkpoint / f"algorithm_state.{extension}"
|
|
if (
|
|
_exists_at_fs_path(filesystem, state_file.as_posix())
|
|
and filesystem.get_file_info(state_file.as_posix()).is_file
|
|
):
|
|
format = "cloudpickle" if extension == "pkl" else "msgpack"
|
|
break
|
|
if format is None:
|
|
raise ValueError(
|
|
"Given checkpoint does not seem to be valid! No file with the name "
|
|
"`algorithm_state.[pkl|msgpack|msgpck]` (or `checkpoint-[0-9]+`) found."
|
|
)
|
|
|
|
info.update(
|
|
{
|
|
"format": format,
|
|
"state_file": str(state_file),
|
|
}
|
|
)
|
|
|
|
# Collect all policy IDs in the sub-dir "policies/".
|
|
policies_dir = checkpoint / "policies"
|
|
if _exists_at_fs_path(filesystem, policies_dir.as_posix()) and _is_dir(
|
|
filesystem.get_file_info(policies_dir.as_posix())
|
|
):
|
|
policy_ids = set()
|
|
file_info_list = filesystem.get_file_info(
|
|
pyarrow.fs.FileSelector(policies_dir.as_posix(), recursive=False)
|
|
)
|
|
for file_info in file_info_list:
|
|
policy_ids.add(file_info.base_name)
|
|
info.update({"policy_ids": policy_ids})
|
|
|
|
# Collect all module IDs in the sub-dir "learner/module_state/".
|
|
modules_dir = (
|
|
checkpoint
|
|
/ COMPONENT_LEARNER_GROUP
|
|
/ COMPONENT_LEARNER
|
|
/ COMPONENT_RL_MODULE
|
|
)
|
|
if _exists_at_fs_path(filesystem, checkpoint.as_posix()) and _is_dir(
|
|
filesystem.get_file_info(modules_dir.as_posix())
|
|
):
|
|
module_ids = set()
|
|
file_info_list = filesystem.get_file_info(
|
|
pyarrow.fs.FileSelector(modules_dir.as_posix(), recursive=False)
|
|
)
|
|
for file_info in file_info_list:
|
|
# Only add subdirs (those are the ones where the RLModule data
|
|
# is stored, not files (could be json metadata files).
|
|
module_dir = modules_dir / file_info.base_name
|
|
if _is_dir(filesystem.get_file_info(module_dir.as_posix())):
|
|
module_ids.add(file_info.base_name)
|
|
info.update({"module_ids": module_ids})
|
|
|
|
# Checkpoint is a file: Use as-is (interpreting it as old Algorithm checkpoint
|
|
# version).
|
|
elif (
|
|
_exists_at_fs_path(filesystem, checkpoint.as_posix())
|
|
and filesystem.get_file_info(checkpoint.as_posix()).is_file
|
|
):
|
|
info.update(
|
|
{
|
|
"checkpoint_version": version.Version("0.1"),
|
|
"checkpoint_dir": str(checkpoint.parent),
|
|
"state_file": str(checkpoint),
|
|
}
|
|
)
|
|
|
|
else:
|
|
raise ValueError(
|
|
f"Given checkpoint ({str(checkpoint)}) not found! Must be a "
|
|
"checkpoint directory (or a file for older checkpoint versions)."
|
|
)
|
|
|
|
return info
|
|
|
|
|
|
@OldAPIStack
|
|
def convert_to_msgpack_checkpoint(
|
|
checkpoint: Union[str, Checkpoint_train, Checkpoint_tune],
|
|
msgpack_checkpoint_dir: str,
|
|
) -> str:
|
|
"""Converts an Algorithm checkpoint (pickle based) to a msgpack based one.
|
|
|
|
Msgpack has the advantage of being python version independent.
|
|
|
|
Args:
|
|
checkpoint: The directory, in which to find the Algorithm checkpoint (pickle
|
|
based).
|
|
msgpack_checkpoint_dir: The directory, in which to create the new msgpack
|
|
based checkpoint.
|
|
|
|
Returns:
|
|
The directory in which the msgpack checkpoint has been created. Note that
|
|
this is the same as `msgpack_checkpoint_dir`.
|
|
"""
|
|
from ray.rllib.algorithms import Algorithm
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
from ray.rllib.core.rl_module import validate_module_id
|
|
|
|
# Try to import msgpack and msgpack_numpy.
|
|
msgpack = try_import_msgpack(error=True)
|
|
|
|
# Restore the Algorithm using the python version dependent checkpoint.
|
|
algo = Algorithm.from_checkpoint(checkpoint)
|
|
state = algo.__getstate__()
|
|
|
|
# Convert all code in state into serializable data.
|
|
# Serialize the algorithm class.
|
|
state["algorithm_class"] = serialize_type(state["algorithm_class"])
|
|
# Serialize the algorithm's config object.
|
|
if not isinstance(state["config"], dict):
|
|
state["config"] = state["config"].serialize()
|
|
else:
|
|
state["config"] = AlgorithmConfig._serialize_dict(state["config"])
|
|
|
|
# Extract policy states from worker state (Policies get their own
|
|
# checkpoint sub-dirs).
|
|
policy_states = {}
|
|
if "worker" in state and "policy_states" in state["worker"]:
|
|
policy_states = state["worker"].pop("policy_states", {})
|
|
|
|
# Policy mapping fn.
|
|
state["worker"]["policy_mapping_fn"] = NOT_SERIALIZABLE
|
|
# Is Policy to train function.
|
|
state["worker"]["is_policy_to_train"] = NOT_SERIALIZABLE
|
|
|
|
# Add RLlib checkpoint version (as string).
|
|
state["checkpoint_version"] = str(CHECKPOINT_VERSION)
|
|
|
|
# Write state (w/o policies) to disk.
|
|
state_file = os.path.join(msgpack_checkpoint_dir, "algorithm_state.msgpck")
|
|
with open(state_file, "wb") as f:
|
|
msgpack.dump(state, f)
|
|
|
|
# Write rllib_checkpoint.json.
|
|
with open(os.path.join(msgpack_checkpoint_dir, "rllib_checkpoint.json"), "w") as f:
|
|
json.dump(
|
|
{
|
|
"type": "Algorithm",
|
|
"checkpoint_version": state["checkpoint_version"],
|
|
"format": "msgpack",
|
|
"state_file": state_file,
|
|
"policy_ids": list(policy_states.keys()),
|
|
"ray_version": ray.__version__,
|
|
"ray_commit": ray.__commit__,
|
|
},
|
|
f,
|
|
)
|
|
|
|
# Write individual policies to disk, each in their own subdirectory.
|
|
for pid, policy_state in policy_states.items():
|
|
# From here on, disallow policyIDs that would not work as directory names.
|
|
validate_module_id(pid, error=True)
|
|
policy_dir = os.path.join(msgpack_checkpoint_dir, "policies", pid)
|
|
os.makedirs(policy_dir, exist_ok=True)
|
|
policy = algo.get_policy(pid)
|
|
policy.export_checkpoint(
|
|
policy_dir,
|
|
policy_state=policy_state,
|
|
checkpoint_format="msgpack",
|
|
)
|
|
|
|
# Release all resources used by the Algorithm.
|
|
algo.stop()
|
|
|
|
return msgpack_checkpoint_dir
|
|
|
|
|
|
@OldAPIStack
|
|
def convert_to_msgpack_policy_checkpoint(
|
|
policy_checkpoint: Union[str, Checkpoint_train, Checkpoint_tune],
|
|
msgpack_checkpoint_dir: str,
|
|
) -> str:
|
|
"""Converts a Policy checkpoint (pickle based) to a msgpack based one.
|
|
|
|
Msgpack has the advantage of being python version independent.
|
|
|
|
Args:
|
|
policy_checkpoint: The directory, in which to find the Policy checkpoint (pickle
|
|
based).
|
|
msgpack_checkpoint_dir: The directory, in which to create the new msgpack
|
|
based checkpoint.
|
|
|
|
Returns:
|
|
The directory in which the msgpack checkpoint has been created. Note that
|
|
this is the same as `msgpack_checkpoint_dir`.
|
|
"""
|
|
from ray.rllib.policy.policy import Policy
|
|
|
|
policy = Policy.from_checkpoint(policy_checkpoint)
|
|
|
|
os.makedirs(msgpack_checkpoint_dir, exist_ok=True)
|
|
policy.export_checkpoint(
|
|
msgpack_checkpoint_dir,
|
|
policy_state=policy.get_state(),
|
|
checkpoint_format="msgpack",
|
|
)
|
|
|
|
# Release all resources used by the Policy.
|
|
del policy
|
|
|
|
return msgpack_checkpoint_dir
|
|
|
|
|
|
@PublicAPI
|
|
def try_import_msgpack(error: bool = False):
|
|
"""Tries importing msgpack and msgpack_numpy and returns the patched msgpack module.
|
|
|
|
Returns None if error is False and msgpack or msgpack_numpy is not installed.
|
|
Raises an error, if error is True and the modules could not be imported.
|
|
|
|
Args:
|
|
error: Whether to raise an error if msgpack/msgpack_numpy cannot be imported.
|
|
|
|
Returns:
|
|
The `msgpack` module, with the msgpack_numpy module already patched in. This
|
|
means you can already encde and decode numpy arrays with the returned module.
|
|
|
|
Raises:
|
|
ImportError: If error=True and msgpack/msgpack_numpy is not installed.
|
|
"""
|
|
try:
|
|
import msgpack
|
|
import msgpack_numpy
|
|
|
|
# Make msgpack_numpy look like msgpack.
|
|
msgpack_numpy.patch()
|
|
|
|
return msgpack
|
|
|
|
except Exception:
|
|
if error:
|
|
raise ImportError(
|
|
"Could not import or setup msgpack and msgpack_numpy! "
|
|
"Try running `pip install msgpack msgpack_numpy` first."
|
|
)
|