chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,977 @@
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from contextlib import redirect_stderr, redirect_stdout
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
import ray.cloudpickle as ray_pickle
|
||||
from ray._common.utils import try_to_create_directory
|
||||
from ray.air._internal.util import exception_cause, skip_exceptions
|
||||
from ray.air.constants import TIME_THIS_ITER_S, TIMESTAMP, TRAINING_ITERATION
|
||||
from ray.train._internal.checkpoint_manager import _TrainingResult
|
||||
from ray.train._internal.storage import StorageContext, _exists_at_fs_path
|
||||
from ray.train.constants import DEFAULT_STORAGE_PATH, RAY_CHDIR_TO_TRIAL_DIR
|
||||
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
||||
from ray.tune.result import (
|
||||
DEBUG_METRICS,
|
||||
DONE,
|
||||
EPISODES_THIS_ITER,
|
||||
EPISODES_TOTAL,
|
||||
HOSTNAME,
|
||||
NODE_IP,
|
||||
PID,
|
||||
RESULT_DUPLICATE,
|
||||
SHOULD_CHECKPOINT,
|
||||
STDERR_FILE,
|
||||
STDOUT_FILE,
|
||||
TIME_TOTAL_S,
|
||||
TIMESTEPS_THIS_ITER,
|
||||
TIMESTEPS_TOTAL,
|
||||
TRIAL_ID,
|
||||
TRIAL_INFO,
|
||||
)
|
||||
from ray.tune.utils import UtilMonitor
|
||||
from ray.tune.utils.log import disable_ipython
|
||||
from ray.tune.utils.util import Tee
|
||||
from ray.util.annotations import DeveloperAPI, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SETUP_TIME_THRESHOLD = 10
|
||||
|
||||
# File containing dict data returned by user from `Trainable.save_checkpoint`
|
||||
_DICT_CHECKPOINT_FILE_NAME = "_dict_checkpoint.pkl"
|
||||
|
||||
|
||||
@PublicAPI
|
||||
class Trainable:
|
||||
"""Abstract class for trainable models, functions, etc.
|
||||
|
||||
A call to ``train()`` on a trainable will execute one logical iteration of
|
||||
training. As a rule of thumb, the execution time of one train call should
|
||||
be large enough to avoid overheads (i.e. more than a few seconds), but
|
||||
short enough to report progress periodically (i.e. at most a few minutes).
|
||||
|
||||
Calling ``save()`` should save the training state of a trainable to disk,
|
||||
and ``restore(path)`` should restore a trainable to the given state.
|
||||
|
||||
Generally you only need to implement ``setup``, ``step``,
|
||||
``save_checkpoint``, and ``load_checkpoint`` when subclassing Trainable.
|
||||
|
||||
Other implementation methods that may be helpful to override are
|
||||
``log_result``, ``reset_config``, ``cleanup``, and ``_export_model``.
|
||||
|
||||
Tune will convert this class into a Ray actor, which runs on a separate process.
|
||||
By default, Tune will also change the current working directory of this process to
|
||||
its corresponding trial-level log directory ``self.logdir``.
|
||||
This is designed so that different trials that run on the same physical node won't
|
||||
accidentally write to the same location and overstep each other.
|
||||
|
||||
The behavior of changing the working directory can be disabled by setting the
|
||||
`RAY_CHDIR_TO_TRIAL_DIR=0` environment variable. This allows access to files
|
||||
in the original working directory, but relative paths should be used for read only
|
||||
purposes, and you must make sure that the directory is synced on all nodes if
|
||||
running on multiple machines.
|
||||
|
||||
The `TUNE_ORIG_WORKING_DIR` environment variable was the original workaround for
|
||||
accessing paths relative to the original working directory. This environment
|
||||
variable is deprecated, and the `RAY_CHDIR_TO_TRIAL_DIR` environment variable
|
||||
described above should be used instead.
|
||||
|
||||
This class supports checkpointing to and restoring from remote storage.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Dict[str, Any] = None,
|
||||
storage: Optional[StorageContext] = None,
|
||||
):
|
||||
"""Initialize a Trainable.
|
||||
|
||||
Sets up logging and points ``self.logdir`` to a directory in which
|
||||
training outputs should be placed.
|
||||
|
||||
Subclasses should prefer defining ``setup()`` instead of overriding
|
||||
``__init__()`` directly.
|
||||
|
||||
Args:
|
||||
config: Trainable-specific configuration data. By default
|
||||
will be saved as ``self.config``.
|
||||
storage: StorageContext object that contains persistent storage paths
|
||||
"""
|
||||
|
||||
self.config = config or {}
|
||||
trial_info = self.config.pop(TRIAL_INFO, None)
|
||||
|
||||
if self.is_actor():
|
||||
disable_ipython()
|
||||
|
||||
self._storage = storage
|
||||
if storage:
|
||||
assert storage.trial_fs_path
|
||||
logger.debug(f"StorageContext on the TRAINABLE:\n{storage}")
|
||||
|
||||
# TODO(justinvyu): Rename/remove logdir.
|
||||
self._logdir = None
|
||||
self._setup_logdir()
|
||||
|
||||
self._stdout_context = self._stdout_fp = self._stdout_stream = None
|
||||
self._stderr_context = self._stderr_fp = self._stderr_stream = None
|
||||
self._stderr_logging_handler = None
|
||||
|
||||
stdout_file = self.config.pop(STDOUT_FILE, None)
|
||||
stderr_file = self.config.pop(STDERR_FILE, None)
|
||||
|
||||
self._iteration = 0
|
||||
self._time_total = 0.0
|
||||
self._timesteps_total = None
|
||||
self._episodes_total = None
|
||||
self._time_since_restore = 0.0
|
||||
self._timesteps_since_restore = 0
|
||||
self._iterations_since_restore = 0
|
||||
self._last_result = None
|
||||
self._restored = False
|
||||
self._trial_info = trial_info
|
||||
self._stdout_file = stdout_file
|
||||
self._stderr_file = stderr_file
|
||||
|
||||
self._start_time = time.time()
|
||||
self._local_ip = ray.util.get_node_ip_address()
|
||||
|
||||
self._open_logfiles(stdout_file, stderr_file)
|
||||
|
||||
self.setup(copy.deepcopy(self.config))
|
||||
setup_time = time.time() - self._start_time
|
||||
if setup_time > SETUP_TIME_THRESHOLD:
|
||||
logger.info(
|
||||
"Trainable.setup took {:.3f} seconds. If your "
|
||||
"trainable is slow to initialize, consider setting "
|
||||
"reuse_actors=True to reduce actor creation "
|
||||
"overheads.".format(setup_time)
|
||||
)
|
||||
log_sys_usage = self.config.get("log_sys_usage", False)
|
||||
self._monitor = UtilMonitor(start=log_sys_usage)
|
||||
|
||||
@classmethod
|
||||
def default_resource_request(
|
||||
cls, config: Dict[str, Any]
|
||||
) -> Optional[PlacementGroupFactory]:
|
||||
"""Provides a static resource requirement for the given configuration.
|
||||
|
||||
This can be overridden by sub-classes to set the correct trial resource
|
||||
allocation, so the user does not need to.
|
||||
|
||||
.. testcode::
|
||||
|
||||
@classmethod
|
||||
def default_resource_request(cls, config):
|
||||
return PlacementGroupFactory([{"CPU": 1}, {"CPU": 1}])
|
||||
|
||||
|
||||
Args:
|
||||
config: The Trainable's config dict.
|
||||
|
||||
Returns:
|
||||
PlacementGroupFactory: A PlacementGroupFactory consumed by Tune
|
||||
for queueing.
|
||||
"""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def resource_help(cls, config: Dict):
|
||||
"""Returns a help string for configuring this trainable's resources.
|
||||
|
||||
Args:
|
||||
config: The Trainer's config dict.
|
||||
|
||||
Returns:
|
||||
A help string describing the resources required by the trainable.
|
||||
"""
|
||||
return ""
|
||||
|
||||
def get_current_ip_pid(self):
|
||||
return self._local_ip, os.getpid()
|
||||
|
||||
def get_auto_filled_metrics(
|
||||
self,
|
||||
now: Optional[datetime] = None,
|
||||
time_this_iter: Optional[float] = None,
|
||||
timestamp: Optional[int] = None,
|
||||
debug_metrics_only: bool = False,
|
||||
) -> dict:
|
||||
"""Return a dict with metrics auto-filled by the trainable.
|
||||
|
||||
If ``debug_metrics_only`` is True, only metrics that don't
|
||||
require at least one iteration will be returned
|
||||
(``ray.tune.result.DEBUG_METRICS``).
|
||||
"""
|
||||
if now is None:
|
||||
now = datetime.today()
|
||||
autofilled = {
|
||||
TRIAL_ID: self.trial_id,
|
||||
"date": now.strftime("%Y-%m-%d_%H-%M-%S"),
|
||||
"timestamp": timestamp if timestamp else int(time.mktime(now.timetuple())),
|
||||
TIME_THIS_ITER_S: time_this_iter,
|
||||
TIME_TOTAL_S: self._time_total,
|
||||
PID: os.getpid(),
|
||||
HOSTNAME: platform.node(),
|
||||
NODE_IP: self._local_ip,
|
||||
"config": self.config,
|
||||
"time_since_restore": self._time_since_restore,
|
||||
"iterations_since_restore": self._iterations_since_restore,
|
||||
}
|
||||
if self._timesteps_since_restore:
|
||||
autofilled["timesteps_since_restore"] = self._timesteps_since_restore
|
||||
|
||||
if debug_metrics_only:
|
||||
autofilled = {k: v for k, v in autofilled.items() if k in DEBUG_METRICS}
|
||||
return autofilled
|
||||
|
||||
def is_actor(self):
|
||||
try:
|
||||
actor_id = ray._private.worker.global_worker.actor_id
|
||||
return actor_id != actor_id.nil()
|
||||
except Exception:
|
||||
# If global_worker is not instantiated, we're not in an actor
|
||||
return False
|
||||
|
||||
def train_buffered(self, buffer_time_s: float, max_buffer_length: int = 1000):
|
||||
"""Runs multiple iterations of training.
|
||||
|
||||
Calls ``train()`` internally. Collects and combines multiple results.
|
||||
This function will run ``self.train()`` repeatedly until one of
|
||||
the following conditions is met: 1) the maximum buffer length is
|
||||
reached, 2) the maximum buffer time is reached, or 3) a checkpoint
|
||||
was created. Even if the maximum time is reached, it will always
|
||||
block until at least one result is received.
|
||||
|
||||
Args:
|
||||
buffer_time_s: Maximum time to buffer. The next result
|
||||
received after this amount of time has passed will return
|
||||
the whole buffer.
|
||||
max_buffer_length: Maximum number of results to buffer.
|
||||
|
||||
Returns:
|
||||
A list of result dicts collected from each call to ``train()``.
|
||||
"""
|
||||
results = []
|
||||
|
||||
now = time.time()
|
||||
send_buffer_at = now + buffer_time_s
|
||||
while now < send_buffer_at or not results: # At least one result
|
||||
result = self.train()
|
||||
results.append(result)
|
||||
if result.get(DONE, False):
|
||||
# If the trial is done, return
|
||||
break
|
||||
elif result.get(SHOULD_CHECKPOINT, False):
|
||||
# If a checkpoint was created, return
|
||||
break
|
||||
elif result.get(RESULT_DUPLICATE):
|
||||
# If the function API trainable completed, return
|
||||
break
|
||||
elif len(results) >= max_buffer_length:
|
||||
# If the buffer is full, return
|
||||
break
|
||||
now = time.time()
|
||||
|
||||
return results
|
||||
|
||||
def train(self):
|
||||
"""Runs one logical iteration of training.
|
||||
|
||||
Calls ``step()`` internally. Subclasses should override ``step()``
|
||||
instead to return results.
|
||||
This method automatically fills the following fields in the result:
|
||||
|
||||
`done` (bool): training is terminated. Filled only if not provided.
|
||||
|
||||
`time_this_iter_s` (float): Time in seconds this iteration
|
||||
took to run. This may be overridden in order to override the
|
||||
system-computed time difference.
|
||||
|
||||
`time_total_s` (float): Accumulated time in seconds for this
|
||||
entire experiment.
|
||||
|
||||
`training_iteration` (int): The index of this
|
||||
training iteration, e.g. call to train(). This is incremented
|
||||
after `step()` is called.
|
||||
|
||||
`pid` (str): The pid of the training process.
|
||||
|
||||
`date` (str): A formatted date of when the result was processed.
|
||||
|
||||
`timestamp` (str): A UNIX timestamp of when the result
|
||||
was processed. This may be overridden.
|
||||
|
||||
`hostname` (str): Hostname of the machine hosting the training
|
||||
process.
|
||||
|
||||
`node_ip` (str): Node ip of the machine hosting the training
|
||||
process.
|
||||
|
||||
Returns:
|
||||
A dict that describes training progress.
|
||||
"""
|
||||
start = time.time()
|
||||
try:
|
||||
result = self.step()
|
||||
except Exception as e:
|
||||
skipped = skip_exceptions(e)
|
||||
raise skipped from exception_cause(skipped)
|
||||
|
||||
assert isinstance(result, dict), "step() needs to return a dict."
|
||||
|
||||
# We do not modify internal state nor update this result if duplicate.
|
||||
if RESULT_DUPLICATE in result:
|
||||
return result
|
||||
|
||||
result = result.copy()
|
||||
|
||||
self._iteration += 1
|
||||
self._iterations_since_restore += 1
|
||||
|
||||
if result.get(TIME_THIS_ITER_S) is not None:
|
||||
time_this_iter = result[TIME_THIS_ITER_S]
|
||||
else:
|
||||
time_this_iter = time.time() - start
|
||||
self._time_total += time_this_iter
|
||||
self._time_since_restore += time_this_iter
|
||||
|
||||
result_timestamp = result.get(TIMESTAMP, None)
|
||||
|
||||
result.setdefault(DONE, False)
|
||||
|
||||
# self._timesteps_total should only be tracked if increments are provided
|
||||
if result.get(TIMESTEPS_THIS_ITER) is not None:
|
||||
if self._timesteps_total is None:
|
||||
self._timesteps_total = 0
|
||||
self._timesteps_total += result[TIMESTEPS_THIS_ITER]
|
||||
self._timesteps_since_restore += result[TIMESTEPS_THIS_ITER]
|
||||
|
||||
# self._episodes_total should only be tracked if increments provided
|
||||
if result.get(EPISODES_THIS_ITER) is not None:
|
||||
if self._episodes_total is None:
|
||||
self._episodes_total = 0
|
||||
self._episodes_total += result[EPISODES_THIS_ITER]
|
||||
|
||||
# self._timesteps_total should not override user-provided total
|
||||
if self._timesteps_total is not None:
|
||||
result.setdefault(TIMESTEPS_TOTAL, self._timesteps_total)
|
||||
if self._episodes_total is not None:
|
||||
result.setdefault(EPISODES_TOTAL, self._episodes_total)
|
||||
result.setdefault(TRAINING_ITERATION, self._iteration)
|
||||
|
||||
now = datetime.today()
|
||||
result.update(
|
||||
self.get_auto_filled_metrics(
|
||||
now=now, time_this_iter=time_this_iter, timestamp=result_timestamp
|
||||
)
|
||||
)
|
||||
|
||||
monitor_data = self._monitor.get_data()
|
||||
if monitor_data:
|
||||
result.update(monitor_data)
|
||||
|
||||
self.log_result(result)
|
||||
|
||||
if self._stdout_context:
|
||||
self._stdout_stream.flush()
|
||||
if self._stderr_context:
|
||||
self._stderr_stream.flush()
|
||||
|
||||
self._last_result = result
|
||||
|
||||
if self._storage:
|
||||
# Launch background tasks to sync artifacts at some specified frequency.
|
||||
self._storage.persist_artifacts()
|
||||
|
||||
return result
|
||||
|
||||
def get_state(self):
|
||||
return {
|
||||
"iteration": self._iteration,
|
||||
"timesteps_total": self._timesteps_total,
|
||||
"time_total": self._time_total,
|
||||
"episodes_total": self._episodes_total,
|
||||
"last_result": self._last_result,
|
||||
"ray_version": ray.__version__,
|
||||
}
|
||||
|
||||
def _report_class_trainable_checkpoint(
|
||||
self, checkpoint_dir: str, checkpoint_dict_or_path: Union[str, Dict]
|
||||
) -> _TrainingResult:
|
||||
"""Report a checkpoint saved via Trainable.save_checkpoint.
|
||||
|
||||
Need to handle both dict or path checkpoint returned by the user's
|
||||
`save_checkpoint` method.
|
||||
|
||||
This is to get class trainables to work with storage backend used by
|
||||
function trainables.
|
||||
This basically re-implements `tune.report` for class trainables,
|
||||
making sure to persist the checkpoint to storage.
|
||||
"""
|
||||
if isinstance(checkpoint_dict_or_path, dict):
|
||||
with Path(checkpoint_dir, _DICT_CHECKPOINT_FILE_NAME).open("wb") as f:
|
||||
ray_pickle.dump(checkpoint_dict_or_path, f)
|
||||
elif isinstance(checkpoint_dict_or_path, str):
|
||||
if checkpoint_dict_or_path != checkpoint_dir:
|
||||
raise ValueError(
|
||||
"The returned checkpoint path from `save_checkpoint` "
|
||||
"must be None or the same as the provided path argument."
|
||||
f"Got {checkpoint_dict_or_path} != {checkpoint_dir}"
|
||||
)
|
||||
|
||||
local_checkpoint = ray.tune.Checkpoint.from_directory(checkpoint_dir)
|
||||
|
||||
metrics = self._last_result.copy() if self._last_result else {}
|
||||
|
||||
if self._storage:
|
||||
# The checkpoint index is updated with the current result.
|
||||
# NOTE: This is no longer using "iteration" as the folder indexing
|
||||
# to be consistent with fn trainables.
|
||||
self._storage._update_checkpoint_index(metrics)
|
||||
|
||||
persisted_checkpoint = self._storage.persist_current_checkpoint(
|
||||
local_checkpoint
|
||||
)
|
||||
|
||||
checkpoint_result = _TrainingResult(
|
||||
checkpoint=persisted_checkpoint, metrics=metrics
|
||||
)
|
||||
# Persist trial artifacts to storage.
|
||||
self._storage.persist_artifacts(
|
||||
force=self._storage.sync_config.sync_artifacts_on_checkpoint
|
||||
)
|
||||
else:
|
||||
# `storage=None` only happens when initializing the
|
||||
# Trainable manually, outside of Tune/Train.
|
||||
# In this case, no storage is set, so the default behavior
|
||||
# is to just not upload anything and report a local checkpoint.
|
||||
# This is fine for the main use case of local debugging.
|
||||
checkpoint_result = _TrainingResult(
|
||||
checkpoint=local_checkpoint, metrics=metrics
|
||||
)
|
||||
return checkpoint_result
|
||||
|
||||
@DeveloperAPI
|
||||
def save(self, checkpoint_dir: Optional[str] = None) -> _TrainingResult:
|
||||
"""Saves the current model state to a checkpoint.
|
||||
|
||||
Subclasses should override ``save_checkpoint()`` instead to save state.
|
||||
|
||||
Args:
|
||||
checkpoint_dir: Optional dir to place the checkpoint.
|
||||
|
||||
Returns:
|
||||
The given or created checkpoint directory.
|
||||
|
||||
Note the return value matches up with what is expected of `restore()`.
|
||||
"""
|
||||
if not isinstance(self, ray.tune.trainable.FunctionTrainable):
|
||||
# Use a temporary directory if no checkpoint_dir is provided.
|
||||
use_temp_dir = not checkpoint_dir
|
||||
checkpoint_dir = checkpoint_dir or tempfile.mkdtemp()
|
||||
os.makedirs(checkpoint_dir, exist_ok=True)
|
||||
|
||||
checkpoint_dict_or_path = self.save_checkpoint(checkpoint_dir)
|
||||
checkpoint_result = self._report_class_trainable_checkpoint(
|
||||
checkpoint_dir, checkpoint_dict_or_path
|
||||
)
|
||||
|
||||
# Clean up the temporary directory, since it's already been
|
||||
# reported + persisted to storage. If no storage is set, the user is
|
||||
# running the Trainable locally and is responsible for cleaning
|
||||
# up the checkpoint directory themselves.
|
||||
if use_temp_dir and self._storage:
|
||||
shutil.rmtree(checkpoint_dir, ignore_errors=True)
|
||||
else:
|
||||
checkpoint_result: _TrainingResult = self.save_checkpoint(None)
|
||||
assert isinstance(checkpoint_result, _TrainingResult)
|
||||
assert self._last_result
|
||||
# Update the checkpoint result to include auto-filled metrics.
|
||||
checkpoint_result.metrics.update(self._last_result)
|
||||
|
||||
return checkpoint_result
|
||||
|
||||
@DeveloperAPI
|
||||
def restore(
|
||||
self, checkpoint_path: Union[str, "ray.tune.Checkpoint", _TrainingResult]
|
||||
):
|
||||
"""Restores training state from a given model checkpoint.
|
||||
|
||||
These checkpoints are returned from calls to save().
|
||||
|
||||
Subclasses should override ``load_checkpoint()`` instead to
|
||||
restore state.
|
||||
This method restores additional metadata saved with the checkpoint.
|
||||
|
||||
`checkpoint_path` should match with the return from ``save()``.
|
||||
|
||||
Args:
|
||||
checkpoint_path: training result that was returned by a
|
||||
previous call to `save()`.
|
||||
"""
|
||||
# TODO(justinvyu): This also supports restoring from a Checkpoint object
|
||||
# or a path, which are legacy APIs that RLlib depends on.
|
||||
# RLlib should remove this dependency since `restore` is a DeveloperAPI.
|
||||
if isinstance(checkpoint_path, str):
|
||||
checkpoint_path = ray.tune.Checkpoint.from_directory(checkpoint_path)
|
||||
if isinstance(checkpoint_path, ray.tune.Checkpoint):
|
||||
checkpoint_result = _TrainingResult(checkpoint=checkpoint_path, metrics={})
|
||||
else:
|
||||
checkpoint_result: _TrainingResult = checkpoint_path
|
||||
|
||||
assert isinstance(checkpoint_result, _TrainingResult), type(checkpoint_result)
|
||||
checkpoint = checkpoint_result.checkpoint
|
||||
checkpoint_metrics = checkpoint_result.metrics
|
||||
self._iteration = checkpoint_metrics.get(TRAINING_ITERATION, 0)
|
||||
self._time_total = checkpoint_metrics.get(TIME_TOTAL_S, 0)
|
||||
self._time_since_restore = 0.0
|
||||
self._iterations_since_restore = 0
|
||||
|
||||
# TODO(justinvyu): This stuff should be moved to rllib.
|
||||
self._timesteps_total = checkpoint_metrics.get(TIMESTEPS_TOTAL)
|
||||
self._timesteps_since_restore = 0
|
||||
self._episodes_total = checkpoint_metrics.get(EPISODES_TOTAL)
|
||||
|
||||
if not _exists_at_fs_path(checkpoint.filesystem, checkpoint.path):
|
||||
raise ValueError(
|
||||
f"Could not recover from checkpoint as it does not exist on "
|
||||
f"storage anymore. "
|
||||
f"Got storage fs type `{checkpoint.filesystem.type_name}` and "
|
||||
f"path: {checkpoint.path}"
|
||||
)
|
||||
|
||||
# TODO(justinvyu): [cls_trainable_support]
|
||||
# This is to conform to the public class Trainable `load_checkpoint` API.
|
||||
if not isinstance(self, ray.tune.trainable.FunctionTrainable):
|
||||
# Need to convert Checkpoint -> local path or dict
|
||||
# (depending on what the output of save_checkpoint was)
|
||||
with checkpoint.as_directory() as checkpoint_dir:
|
||||
checkpoint_path = Path(checkpoint_dir)
|
||||
dict_checkpoint_file = checkpoint_path / _DICT_CHECKPOINT_FILE_NAME
|
||||
if dict_checkpoint_file.exists():
|
||||
# If this was a dict checkpoint, load it as a dict
|
||||
with open(dict_checkpoint_file, "rb") as f:
|
||||
checkpoint_dict = ray_pickle.load(f)
|
||||
self.load_checkpoint(checkpoint_dict)
|
||||
else:
|
||||
self.load_checkpoint(checkpoint_dir)
|
||||
else:
|
||||
# TODO(justinvyu): The Function Trainable case doesn't conform
|
||||
# to the load_checkpoint API at the moment.
|
||||
self.load_checkpoint(checkpoint_result)
|
||||
|
||||
self._restored = True
|
||||
|
||||
logger.info(f"Restored on {self._local_ip} from checkpoint: {checkpoint}")
|
||||
|
||||
def export_model(
|
||||
self, export_formats: Union[List[str], str], export_dir: Optional[str] = None
|
||||
):
|
||||
"""Exports model based on export_formats.
|
||||
|
||||
Subclasses should override _export_model() to actually
|
||||
export model to local directory.
|
||||
|
||||
Args:
|
||||
export_formats: Format or list of (str) formats
|
||||
that should be exported.
|
||||
export_dir: Optional dir to place the exported model.
|
||||
Defaults to self.logdir.
|
||||
|
||||
Returns:
|
||||
A dict that maps ExportFormats to successfully exported models.
|
||||
"""
|
||||
if isinstance(export_formats, str):
|
||||
export_formats = [export_formats]
|
||||
export_dir = export_dir or self.logdir
|
||||
return self._export_model(export_formats, export_dir)
|
||||
|
||||
def reset(self, new_config, storage=None):
|
||||
"""Resets trial for use with new config.
|
||||
|
||||
Subclasses should override reset_config() to actually
|
||||
reset actor behavior for the new config."""
|
||||
self.config = new_config
|
||||
|
||||
self._storage = storage
|
||||
|
||||
trial_info = new_config.pop(TRIAL_INFO, None)
|
||||
if trial_info:
|
||||
self._trial_info = trial_info
|
||||
|
||||
self._setup_logdir()
|
||||
|
||||
stdout_file = new_config.pop(STDOUT_FILE, None)
|
||||
stderr_file = new_config.pop(STDERR_FILE, None)
|
||||
|
||||
self._close_logfiles()
|
||||
self._open_logfiles(stdout_file, stderr_file)
|
||||
|
||||
success = self.reset_config(new_config)
|
||||
if not success:
|
||||
return False
|
||||
|
||||
# Reset attributes. Will be overwritten by `restore` if a checkpoint
|
||||
# is provided.
|
||||
self._iteration = 0
|
||||
self._time_total = 0.0
|
||||
self._timesteps_total = None
|
||||
self._episodes_total = None
|
||||
self._time_since_restore = 0.0
|
||||
self._timesteps_since_restore = 0
|
||||
self._iterations_since_restore = 0
|
||||
self._restored = False
|
||||
|
||||
return True
|
||||
|
||||
def reset_config(self, new_config: Dict) -> bool:
|
||||
"""Resets configuration without restarting the trial.
|
||||
|
||||
This method is optional, but can be implemented to speed up algorithms
|
||||
such as PBT, and to allow performance optimizations such as running
|
||||
experiments with reuse_actors=True.
|
||||
|
||||
Args:
|
||||
new_config: Updated hyperparameter configuration
|
||||
for the trainable.
|
||||
|
||||
Returns:
|
||||
True if reset was successful else False.
|
||||
"""
|
||||
return False
|
||||
|
||||
def _setup_logdir(self):
|
||||
"""Set up the trial log directory.
|
||||
|
||||
Sets _logdir and changes the working directory to the trial directory
|
||||
on the worker process when running with Tune.
|
||||
|
||||
`_logdir` is the **per trial** directory for the Trainable.
|
||||
"""
|
||||
if self._storage:
|
||||
self._logdir = self._storage.trial_working_directory
|
||||
else:
|
||||
logdir_prefix = datetime.today().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
try_to_create_directory(DEFAULT_STORAGE_PATH)
|
||||
self._logdir = tempfile.mkdtemp(
|
||||
prefix=logdir_prefix, dir=DEFAULT_STORAGE_PATH
|
||||
)
|
||||
|
||||
os.makedirs(self._logdir, exist_ok=True)
|
||||
|
||||
if self._storage:
|
||||
os.environ.setdefault("TUNE_ORIG_WORKING_DIR", os.getcwd())
|
||||
|
||||
if bool(int(os.environ.get(RAY_CHDIR_TO_TRIAL_DIR, "1"))):
|
||||
os.chdir(self._logdir)
|
||||
|
||||
def _open_logfiles(self, stdout_file, stderr_file):
|
||||
"""Create loggers. Open stdout and stderr logfiles."""
|
||||
if stdout_file:
|
||||
stdout_path = (Path(self._logdir) / stdout_file).expanduser().as_posix()
|
||||
self._stdout_fp = open(stdout_path, "a+")
|
||||
self._stdout_stream = Tee(sys.stdout, self._stdout_fp)
|
||||
self._stdout_context = redirect_stdout(self._stdout_stream)
|
||||
self._stdout_context.__enter__()
|
||||
|
||||
if stderr_file:
|
||||
stderr_path = (Path(self._logdir) / stderr_file).expanduser().as_posix()
|
||||
self._stderr_fp = open(stderr_path, "a+")
|
||||
self._stderr_stream = Tee(sys.stderr, self._stderr_fp)
|
||||
self._stderr_context = redirect_stderr(self._stderr_stream)
|
||||
self._stderr_context.__enter__()
|
||||
|
||||
# Add logging handler to root ray logger
|
||||
formatter = logging.Formatter(
|
||||
"[%(levelname)s %(asctime)s] "
|
||||
"%(filename)s: %(lineno)d "
|
||||
"%(message)s"
|
||||
)
|
||||
self._stderr_logging_handler = logging.StreamHandler(self._stderr_fp)
|
||||
self._stderr_logging_handler.setFormatter(formatter)
|
||||
ray.logger.addHandler(self._stderr_logging_handler)
|
||||
|
||||
def _close_logfiles(self):
|
||||
"""Close stdout and stderr logfiles."""
|
||||
if self._stderr_logging_handler:
|
||||
ray.logger.removeHandler(self._stderr_logging_handler)
|
||||
|
||||
if self._stdout_context:
|
||||
self._stdout_stream.flush()
|
||||
self._stdout_context.__exit__(None, None, None)
|
||||
self._stdout_fp.close()
|
||||
self._stdout_context = None
|
||||
if self._stderr_context:
|
||||
self._stderr_stream.flush()
|
||||
self._stderr_context.__exit__(None, None, None)
|
||||
self._stderr_fp.close()
|
||||
self._stderr_context = None
|
||||
|
||||
def stop(self):
|
||||
"""Releases all resources used by this trainable.
|
||||
|
||||
Calls ``Trainable.cleanup`` internally. Subclasses should override
|
||||
``Trainable.cleanup`` for custom cleanup procedures.
|
||||
"""
|
||||
if self._monitor.is_alive():
|
||||
self._monitor.stop()
|
||||
self._monitor.join()
|
||||
self.cleanup()
|
||||
|
||||
self._close_logfiles()
|
||||
|
||||
@property
|
||||
def logdir(self):
|
||||
"""Directory of the results and checkpoints for this Trainable.
|
||||
|
||||
Note that the current working directory will also be changed to this.
|
||||
"""
|
||||
return self._logdir
|
||||
|
||||
@property
|
||||
def trial_name(self):
|
||||
"""Trial name for the corresponding trial of this Trainable.
|
||||
|
||||
This is not set if not using Tune.
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray.tune import Trainable
|
||||
|
||||
name = Trainable().trial_name
|
||||
"""
|
||||
if self._trial_info:
|
||||
return self._trial_info.trial_name
|
||||
else:
|
||||
return "default"
|
||||
|
||||
@property
|
||||
def trial_id(self):
|
||||
"""Trial ID for the corresponding trial of this Trainable.
|
||||
|
||||
This is not set if not using Tune.
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray.tune import Trainable
|
||||
|
||||
trial_id = Trainable().trial_id
|
||||
"""
|
||||
if self._trial_info:
|
||||
return self._trial_info.trial_id
|
||||
else:
|
||||
return "default"
|
||||
|
||||
@property
|
||||
def trial_resources(self) -> Optional[PlacementGroupFactory]:
|
||||
"""Resources currently assigned to the trial of this Trainable.
|
||||
|
||||
This is not set if not using Tune.
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray.tune import Trainable
|
||||
|
||||
trial_resources = Trainable().trial_resources
|
||||
"""
|
||||
if self._trial_info:
|
||||
return self._trial_info.trial_resources
|
||||
else:
|
||||
return None
|
||||
|
||||
@property
|
||||
def iteration(self):
|
||||
"""Current training iteration.
|
||||
|
||||
This value is automatically incremented every time `train()` is called
|
||||
and is automatically inserted into the training result dict.
|
||||
|
||||
"""
|
||||
return self._iteration
|
||||
|
||||
@property
|
||||
def training_iteration(self):
|
||||
"""Current training iteration (same as `self.iteration`).
|
||||
|
||||
This value is automatically incremented every time `train()` is called
|
||||
and is automatically inserted into the training result dict.
|
||||
|
||||
"""
|
||||
return self._iteration
|
||||
|
||||
def get_config(self):
|
||||
"""Returns configuration passed in by Tune."""
|
||||
return self.config
|
||||
|
||||
def step(self) -> Dict:
|
||||
"""Subclasses should override this to implement train().
|
||||
|
||||
The return value will be automatically passed to the loggers. Users
|
||||
can also return `tune.result.DONE` or `tune.result.SHOULD_CHECKPOINT`
|
||||
as a key to manually trigger termination or checkpointing of this
|
||||
trial. Note that manual checkpointing only works when subclassing
|
||||
Trainables.
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
|
||||
Returns:
|
||||
A dict that describes training progress.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def save_checkpoint(self, checkpoint_dir: str) -> Optional[Dict]:
|
||||
"""Subclasses should override this to implement ``save()``.
|
||||
|
||||
Warning:
|
||||
Do not rely on absolute paths in the implementation of
|
||||
``Trainable.save_checkpoint`` and ``Trainable.load_checkpoint``.
|
||||
|
||||
Use ``validate_save_restore`` to catch ``Trainable.save_checkpoint``/
|
||||
``Trainable.load_checkpoint`` errors before execution.
|
||||
|
||||
>>> from ray.tune.utils import validate_save_restore
|
||||
>>> MyTrainableClass = ... # doctest: +SKIP
|
||||
>>> validate_save_restore(MyTrainableClass) # doctest: +SKIP
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
|
||||
Args:
|
||||
checkpoint_dir: The directory where the checkpoint
|
||||
file must be stored. In a Tune run, if the trial is paused,
|
||||
the provided path may be temporary and moved.
|
||||
|
||||
Returns:
|
||||
A dict or None. If dict, the return value will
|
||||
be automatically serialized by Tune. In that case,
|
||||
``Trainable.load_checkpoint()`` will receive the dict upon restore.
|
||||
|
||||
Example:
|
||||
>>> trainable, trainable1, trainable2 = ... # doctest: +SKIP
|
||||
>>> print(trainable1.save_checkpoint("/tmp/checkpoint_1")) # doctest: +SKIP
|
||||
"/tmp/checkpoint_1"
|
||||
>>> print(trainable2.save_checkpoint("/tmp/checkpoint_2")) # doctest: +SKIP
|
||||
{"some": "data"}
|
||||
>>> trainable.save_checkpoint("/tmp/bad_example") # doctest: +SKIP
|
||||
"/tmp/NEW_CHECKPOINT_PATH/my_checkpoint_file" # This will error.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def load_checkpoint(self, checkpoint: Optional[Dict]):
|
||||
"""Subclasses should override this to implement restore().
|
||||
|
||||
Warning:
|
||||
In this method, do not rely on absolute paths. The absolute
|
||||
path of the checkpoint_dir used in ``Trainable.save_checkpoint``
|
||||
may be changed.
|
||||
|
||||
If ``Trainable.save_checkpoint`` returned a prefixed string, the
|
||||
prefix of the checkpoint string returned by
|
||||
``Trainable.save_checkpoint`` may be changed.
|
||||
This is because trial pausing depends on temporary directories.
|
||||
|
||||
The directory structure under the checkpoint_dir provided to
|
||||
``Trainable.save_checkpoint`` is preserved.
|
||||
|
||||
See the examples below.
|
||||
|
||||
Example:
|
||||
>>> import os
|
||||
>>> from ray.tune.trainable import Trainable
|
||||
>>> class Example(Trainable):
|
||||
... def save_checkpoint(self, checkpoint_path):
|
||||
... my_checkpoint_path = os.path.join(checkpoint_path, "my/path")
|
||||
... return my_checkpoint_path
|
||||
... def load_checkpoint(self, my_checkpoint_path):
|
||||
... print(my_checkpoint_path)
|
||||
>>> trainer = Example()
|
||||
>>> # This is used when PAUSED.
|
||||
>>> checkpoint_result = trainer.save() # doctest: +SKIP
|
||||
>>> trainer.restore(checkpoint_result) # doctest: +SKIP
|
||||
|
||||
If `Trainable.save_checkpoint` returned a dict, then Tune will directly pass
|
||||
the dict data as the argument to this method.
|
||||
|
||||
Example:
|
||||
>>> from ray.tune.trainable import Trainable
|
||||
>>> class Example(Trainable):
|
||||
... def save_checkpoint(self, checkpoint_path):
|
||||
... return {"my_data": 1}
|
||||
... def load_checkpoint(self, checkpoint_dict):
|
||||
... print(checkpoint_dict["my_data"])
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
|
||||
Args:
|
||||
checkpoint: If dict, the return value is as
|
||||
returned by ``save_checkpoint``. Otherwise, the directory
|
||||
the checkpoint was stored in.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def setup(self, config: Dict):
|
||||
"""Subclasses should override this for custom initialization.
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
|
||||
Args:
|
||||
config: Hyperparameters and other configs given.
|
||||
Copy of `self.config`.
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
def log_result(self, result: Dict):
|
||||
"""Subclasses can optionally override this to customize logging.
|
||||
|
||||
The logging here is done on the worker process rather than
|
||||
the driver.
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
|
||||
Args:
|
||||
result: Training result returned by step().
|
||||
"""
|
||||
pass
|
||||
|
||||
def cleanup(self):
|
||||
"""Subclasses should override this for any cleanup on stop.
|
||||
|
||||
If any Ray actors are launched in the Trainable (i.e., with a RLlib
|
||||
trainer), be sure to kill the Ray actor process here.
|
||||
|
||||
This process should be lightweight. Per default,
|
||||
|
||||
You can kill a Ray actor by calling `ray.kill(actor)`
|
||||
on the actor or removing all references to it and waiting for garbage
|
||||
collection
|
||||
|
||||
.. versionadded:: 0.8.7
|
||||
"""
|
||||
pass
|
||||
|
||||
def _export_model(self, export_formats: List[str], export_dir: str):
|
||||
"""Subclasses should override this to export model.
|
||||
|
||||
Args:
|
||||
export_formats: List of formats that should be exported.
|
||||
export_dir: Directory to place exported models.
|
||||
|
||||
Returns:
|
||||
A dict that maps ExportFormats to successfully exported models.
|
||||
"""
|
||||
return {}
|
||||
|
||||
def _implements_method(self, key):
|
||||
return hasattr(self, key) and callable(getattr(self, key))
|
||||
Reference in New Issue
Block a user