chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,10 @@
|
||||
from ray.tune.trainable.function_trainable import FunctionTrainable, wrap_function
|
||||
from ray.tune.trainable.trainable import Trainable
|
||||
from ray.tune.trainable.util import with_parameters
|
||||
|
||||
__all__ = [
|
||||
"Trainable",
|
||||
"FunctionTrainable",
|
||||
"with_parameters",
|
||||
"wrap_function",
|
||||
]
|
||||
@@ -0,0 +1,266 @@
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
from functools import partial
|
||||
from numbers import Number
|
||||
from typing import Any, Callable, Dict, Optional, Type
|
||||
|
||||
from ray.air._internal.util import RunnerThread, StartTraceback
|
||||
from ray.air.constants import _ERROR_FETCH_TIMEOUT
|
||||
from ray.train._internal.checkpoint_manager import _TrainingResult
|
||||
from ray.train._internal.session import (
|
||||
TrialInfo,
|
||||
_TrainSession,
|
||||
get_session,
|
||||
init_session,
|
||||
shutdown_session,
|
||||
)
|
||||
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
||||
from ray.tune.result import DEFAULT_METRIC, RESULT_DUPLICATE, SHOULD_CHECKPOINT
|
||||
from ray.tune.trainable.trainable import Trainable
|
||||
from ray.tune.utils import _detect_config_single
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Time between FunctionTrainable checks when fetching
|
||||
# new results after signaling the reporter to continue
|
||||
|
||||
NULL_MARKER = ".null_marker"
|
||||
TEMP_MARKER = ".temp_marker"
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class FunctionTrainable(Trainable):
|
||||
"""Trainable that runs a user function reporting results.
|
||||
|
||||
This mode of execution does not support checkpoint/restore."""
|
||||
|
||||
_name = "func"
|
||||
|
||||
def setup(self, config):
|
||||
init_session(
|
||||
training_func=lambda: self._trainable_func(self.config),
|
||||
trial_info=TrialInfo(
|
||||
name=self.trial_name,
|
||||
id=self.trial_id,
|
||||
resources=self.trial_resources,
|
||||
logdir=self._storage.trial_driver_staging_path,
|
||||
driver_ip=None,
|
||||
driver_node_id=None,
|
||||
experiment_name=self._storage.experiment_dir_name,
|
||||
),
|
||||
storage=self._storage,
|
||||
synchronous_result_reporting=True,
|
||||
# Set all Train-specific properties to None.
|
||||
world_rank=None,
|
||||
local_rank=None,
|
||||
node_rank=None,
|
||||
local_world_size=None,
|
||||
world_size=None,
|
||||
dataset_shard=None,
|
||||
checkpoint=None,
|
||||
)
|
||||
self._last_training_result: Optional[_TrainingResult] = None
|
||||
|
||||
def _trainable_func(self, config: Dict[str, Any]):
|
||||
"""Subclasses can override this to set the trainable func."""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def _start(self):
|
||||
def entrypoint():
|
||||
try:
|
||||
return self._trainable_func(self.config)
|
||||
except Exception as e:
|
||||
raise StartTraceback from e
|
||||
|
||||
# the runner thread is not started until the first call to _train
|
||||
self._runner = RunnerThread(
|
||||
target=entrypoint, error_queue=self._error_queue, daemon=True
|
||||
)
|
||||
# if not alive, try to start
|
||||
self._status_reporter._start()
|
||||
try:
|
||||
self._runner.start()
|
||||
except RuntimeError:
|
||||
# If this is reached, it means the thread was started and is
|
||||
# now done or has raised an exception.
|
||||
pass
|
||||
|
||||
def step(self):
|
||||
"""Implements train() for a Function API.
|
||||
|
||||
If the RunnerThread finishes without reporting "done",
|
||||
Tune will automatically provide a magic keyword __duplicate__
|
||||
along with a result with "done=True". The TrialRunner will handle the
|
||||
result accordingly (see tune/tune_controller.py).
|
||||
"""
|
||||
session: _TrainSession = get_session()
|
||||
if not session.training_started:
|
||||
session.start()
|
||||
|
||||
training_result: Optional[_TrainingResult] = session.get_next()
|
||||
|
||||
if not training_result:
|
||||
# The `RESULT_DUPLICATE` result should have been the last
|
||||
# result reported by the session, which triggers cleanup.
|
||||
raise RuntimeError(
|
||||
"Should not have reached here. The TuneController should not "
|
||||
"have scheduled another `train` remote call."
|
||||
"It should have scheduled a `stop` instead "
|
||||
"after the training function exits."
|
||||
)
|
||||
|
||||
metrics = training_result.metrics
|
||||
# This keyword appears if the train_func using the Function API
|
||||
# finishes without "done=True". This duplicates the last result, but
|
||||
# the TuneController will not log this result again.
|
||||
# TuneController will also inject done=True to the result,
|
||||
# and proceed to queue up a STOP decision for the trial.
|
||||
if RESULT_DUPLICATE in metrics:
|
||||
metrics[SHOULD_CHECKPOINT] = False
|
||||
|
||||
self._last_training_result = training_result
|
||||
if training_result.checkpoint is not None:
|
||||
# TODO(justinvyu): Result/checkpoint reporting can be combined.
|
||||
# For now, since result/checkpoint reporting is separate, this
|
||||
# special key will tell Tune to pull the checkpoint from
|
||||
# the `last_training_result`.
|
||||
metrics[SHOULD_CHECKPOINT] = True
|
||||
return metrics
|
||||
|
||||
def execute(self, fn):
|
||||
return fn(self)
|
||||
|
||||
def save_checkpoint(self, checkpoint_dir: str = ""):
|
||||
if checkpoint_dir:
|
||||
raise ValueError("Checkpoint dir should not be used with function API.")
|
||||
|
||||
# TODO(justinvyu): This currently breaks the `save_checkpoint` interface.
|
||||
# TRAIN -> SAVE remote calls get processed sequentially,
|
||||
# so `_last_training_result.checkpoint` holds onto the latest ckpt.
|
||||
return self._last_training_result
|
||||
|
||||
def load_checkpoint(self, checkpoint_result: _TrainingResult):
|
||||
# TODO(justinvyu): This currently breaks the `load_checkpoint` interface.
|
||||
session = get_session()
|
||||
session.loaded_checkpoint = checkpoint_result.checkpoint
|
||||
|
||||
def cleanup(self):
|
||||
session = get_session()
|
||||
try:
|
||||
# session.finish raises any Exceptions from training.
|
||||
# Do not wait for thread termination here (timeout=0).
|
||||
session.finish(timeout=0)
|
||||
finally:
|
||||
# Check for any errors that might have been missed.
|
||||
session._report_thread_runner_error()
|
||||
# Shutdown session even if session.finish() raises an Exception.
|
||||
shutdown_session()
|
||||
|
||||
def reset_config(self, new_config):
|
||||
session = get_session()
|
||||
|
||||
# Wait for thread termination so it is save to re-use the same actor.
|
||||
thread_timeout = int(os.environ.get("TUNE_FUNCTION_THREAD_TIMEOUT_S", 2))
|
||||
session.finish(timeout=thread_timeout)
|
||||
if session.training_thread.is_alive():
|
||||
# Did not finish within timeout, reset unsuccessful.
|
||||
return False
|
||||
|
||||
session.reset(
|
||||
training_func=lambda: self._trainable_func(self.config),
|
||||
trial_info=TrialInfo(
|
||||
name=self.trial_name,
|
||||
id=self.trial_id,
|
||||
resources=self.trial_resources,
|
||||
logdir=self._storage.trial_working_directory,
|
||||
driver_ip=None,
|
||||
driver_node_id=None,
|
||||
experiment_name=self._storage.experiment_dir_name,
|
||||
),
|
||||
storage=self._storage,
|
||||
)
|
||||
|
||||
self._last_result = {}
|
||||
return True
|
||||
|
||||
def _report_thread_runner_error(self, block=False):
|
||||
try:
|
||||
e = self._error_queue.get(block=block, timeout=_ERROR_FETCH_TIMEOUT)
|
||||
raise StartTraceback from e
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
def wrap_function(
|
||||
train_func: Callable[[Any], Any], name: Optional[str] = None
|
||||
) -> Type["FunctionTrainable"]:
|
||||
inherit_from = (FunctionTrainable,)
|
||||
|
||||
if hasattr(train_func, "__mixins__"):
|
||||
inherit_from = train_func.__mixins__ + inherit_from
|
||||
|
||||
func_args = inspect.getfullargspec(train_func).args
|
||||
use_config_single = _detect_config_single(train_func)
|
||||
|
||||
if not use_config_single:
|
||||
raise ValueError(
|
||||
"Unknown argument found in the Trainable function. "
|
||||
"The function args must include a single 'config' positional parameter.\n"
|
||||
"Found: {}".format(func_args)
|
||||
)
|
||||
|
||||
resources = getattr(train_func, "_resources", None)
|
||||
|
||||
class ImplicitFunc(*inherit_from):
|
||||
_name = name or (
|
||||
train_func.__name__ if hasattr(train_func, "__name__") else "func"
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return self._name
|
||||
|
||||
def _trainable_func(self, config):
|
||||
fn = partial(train_func, config)
|
||||
|
||||
def handle_output(output):
|
||||
if not output:
|
||||
return
|
||||
elif isinstance(output, dict):
|
||||
get_session().report(output)
|
||||
elif isinstance(output, Number):
|
||||
get_session().report({DEFAULT_METRIC: output})
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid return or yield value. Either return/yield "
|
||||
"a single number or a dictionary object in your "
|
||||
"trainable function."
|
||||
)
|
||||
|
||||
output = None
|
||||
if inspect.isgeneratorfunction(train_func):
|
||||
for output in fn():
|
||||
handle_output(output)
|
||||
else:
|
||||
output = fn()
|
||||
handle_output(output)
|
||||
|
||||
# If train_func returns, we need to notify the main event loop
|
||||
# of the last result while avoiding double logging. This is done
|
||||
# with the keyword RESULT_DUPLICATE -- see tune/tune_controller.py.
|
||||
get_session().report({RESULT_DUPLICATE: True})
|
||||
return output
|
||||
|
||||
@classmethod
|
||||
def default_resource_request(
|
||||
cls, config: Dict[str, Any]
|
||||
) -> Optional[PlacementGroupFactory]:
|
||||
if not isinstance(resources, PlacementGroupFactory) and callable(resources):
|
||||
return resources(config)
|
||||
return resources
|
||||
|
||||
return ImplicitFunc
|
||||
@@ -0,0 +1,105 @@
|
||||
import json
|
||||
from collections import deque
|
||||
from numbers import Number
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from ray.train._internal.checkpoint_manager import _CheckpointManager
|
||||
from ray.tune.utils.serialization import (
|
||||
TuneFunctionEncoder,
|
||||
_loads_with_cloudpickle,
|
||||
)
|
||||
|
||||
|
||||
class _TrainingRunMetadata:
|
||||
"""Serializable struct for holding runtime trial metadata.
|
||||
|
||||
Runtime metadata is data that changes and is updated on runtime. This includes
|
||||
e.g. the last result, the currently available checkpoints, and the number
|
||||
of errors encountered for a trial.
|
||||
"""
|
||||
|
||||
def __init__(self, n_steps: Tuple[int] = (5, 10)):
|
||||
# General metadata
|
||||
self.start_time = None
|
||||
|
||||
# Errors
|
||||
self.num_failures = 0
|
||||
self.num_failures_after_restore = 0
|
||||
|
||||
self.error_filename = None
|
||||
self.pickled_error_filename = None
|
||||
|
||||
# Results and metrics
|
||||
self.last_result = {}
|
||||
self.last_result_time = -float("inf")
|
||||
|
||||
# stores in memory max/min/avg/last-n-avg/last result for each
|
||||
# metric by trial
|
||||
self.metric_analysis = {}
|
||||
self._n_steps = n_steps
|
||||
self.metric_n_steps = {}
|
||||
|
||||
# Checkpoints
|
||||
self.checkpoint_manager: Optional[_CheckpointManager] = None
|
||||
|
||||
self._cached_json = None
|
||||
|
||||
def invalidate_cache(self):
|
||||
self._cached_json = None
|
||||
|
||||
def update_metric(self, metric: str, value: Number, step: Optional[int] = 1):
|
||||
if metric not in self.metric_analysis:
|
||||
self.metric_analysis[metric] = {
|
||||
"max": value,
|
||||
"min": value,
|
||||
"avg": value,
|
||||
"last": value,
|
||||
}
|
||||
self.metric_n_steps[metric] = {}
|
||||
for n in self._n_steps:
|
||||
key = "last-{:d}-avg".format(n)
|
||||
self.metric_analysis[metric][key] = value
|
||||
# Store n as string for correct restore.
|
||||
self.metric_n_steps[metric][str(n)] = deque([value], maxlen=n)
|
||||
else:
|
||||
step = step or 1
|
||||
self.metric_analysis[metric]["max"] = max(
|
||||
value, self.metric_analysis[metric]["max"]
|
||||
)
|
||||
self.metric_analysis[metric]["min"] = min(
|
||||
value, self.metric_analysis[metric]["min"]
|
||||
)
|
||||
self.metric_analysis[metric]["avg"] = (
|
||||
1 / step * (value + (step - 1) * self.metric_analysis[metric]["avg"])
|
||||
)
|
||||
self.metric_analysis[metric]["last"] = value
|
||||
|
||||
for n in self._n_steps:
|
||||
key = "last-{:d}-avg".format(n)
|
||||
self.metric_n_steps[metric][str(n)].append(value)
|
||||
self.metric_analysis[metric][key] = sum(
|
||||
self.metric_n_steps[metric][str(n)]
|
||||
) / len(self.metric_n_steps[metric][str(n)])
|
||||
self.invalidate_cache()
|
||||
|
||||
def __setattr__(self, key, value):
|
||||
super().__setattr__(key, value)
|
||||
if key not in {"_cached_json"}:
|
||||
self.invalidate_cache()
|
||||
|
||||
def get_json_state(self) -> str:
|
||||
if self._cached_json is None:
|
||||
data = self.__dict__
|
||||
data.pop("_cached_json", None)
|
||||
self._cached_json = json.dumps(data, indent=2, cls=TuneFunctionEncoder)
|
||||
|
||||
return self._cached_json
|
||||
|
||||
@classmethod
|
||||
def from_json_state(cls, json_state: str) -> "_TrainingRunMetadata":
|
||||
state = _loads_with_cloudpickle(json_state)
|
||||
|
||||
run_metadata = cls()
|
||||
run_metadata.__dict__.update(state)
|
||||
|
||||
return run_metadata
|
||||
@@ -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))
|
||||
@@ -0,0 +1,60 @@
|
||||
from typing import Dict, Optional
|
||||
|
||||
from ray.train._checkpoint import Checkpoint as TrainCheckpoint
|
||||
from ray.train._internal.session import _warn_session_misuse, get_session
|
||||
from ray.train.constants import (
|
||||
V2_MIGRATION_GUIDE_MESSAGE,
|
||||
_v2_migration_warnings_enabled,
|
||||
)
|
||||
from ray.train.utils import _copy_doc, _log_deprecation_warning
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
@_copy_doc(TrainCheckpoint)
|
||||
class Checkpoint(TrainCheckpoint):
|
||||
# NOTE: This is just a pass-through wrapper around `ray.train.Checkpoint`
|
||||
# in order to detect whether the import module was correct `ray.tune.Checkpoint`.
|
||||
pass
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@_warn_session_misuse()
|
||||
def report(metrics: Dict, *, checkpoint: Optional[Checkpoint] = None) -> None:
|
||||
"""Report metrics and optionally save and register a checkpoint to Ray Tune.
|
||||
|
||||
If a checkpoint is provided, it will be
|
||||
:ref:`persisted to storage <persistent-storage-guide>`.
|
||||
|
||||
.. note::
|
||||
|
||||
Each invocation of this method will automatically increment the underlying
|
||||
``training_iteration`` number. The physical meaning of this "iteration" is
|
||||
defined by user depending on how often they call ``report``.
|
||||
It does not necessarily map to one epoch.
|
||||
|
||||
Args:
|
||||
metrics: The metrics you want to report.
|
||||
checkpoint: The optional checkpoint you want to report.
|
||||
"""
|
||||
if checkpoint and not isinstance(checkpoint, Checkpoint):
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(
|
||||
"The `Checkpoint` class should be imported from `ray.tune` "
|
||||
"when passing it to `ray.tune.report` in a Tune function. "
|
||||
"Please update your imports. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
get_session().report(metrics, checkpoint=checkpoint)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@_warn_session_misuse()
|
||||
def get_checkpoint() -> Optional[Checkpoint]:
|
||||
"""Access the latest reported checkpoint to resume from if one exists."""
|
||||
|
||||
return get_session().loaded_checkpoint
|
||||
|
||||
|
||||
def _in_tune_session() -> bool:
|
||||
return get_session() and get_session().world_rank is None
|
||||
@@ -0,0 +1,250 @@
|
||||
import inspect
|
||||
import logging
|
||||
import types
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, Union
|
||||
|
||||
import ray
|
||||
from ray.tune.execution.placement_groups import (
|
||||
PlacementGroupFactory,
|
||||
resource_dict_to_pg_factory,
|
||||
)
|
||||
from ray.tune.registry import _ParameterRegistry
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.tune.trainable import Trainable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
def with_parameters(trainable: Union[Type["Trainable"], Callable], **kwargs):
|
||||
"""Wrapper for trainables to pass arbitrary large data objects.
|
||||
|
||||
This wrapper function will store all passed parameters in the Ray
|
||||
object store and retrieve them when calling the function. It can thus
|
||||
be used to pass arbitrary data, even datasets, to Tune trainables.
|
||||
|
||||
This can also be used as an alternative to ``functools.partial`` to pass
|
||||
default arguments to trainables.
|
||||
|
||||
When used with the function API, the trainable function is called with
|
||||
the passed parameters as keyword arguments. When used with the class API,
|
||||
the ``Trainable.setup()`` method is called with the respective kwargs.
|
||||
|
||||
If the data already exists in the object store (are instances of
|
||||
ObjectRef), using ``tune.with_parameters()`` is not necessary. You can
|
||||
instead pass the object refs to the training function via the ``config``
|
||||
or use Python partials.
|
||||
|
||||
Args:
|
||||
trainable: Trainable to wrap.
|
||||
**kwargs: parameters to store in object store.
|
||||
|
||||
Function API example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray import tune
|
||||
|
||||
def train_fn(config, data=None):
|
||||
for sample in data:
|
||||
loss = update_model(sample)
|
||||
tune.report(dict(loss=loss))
|
||||
|
||||
data = HugeDataset(download=True)
|
||||
|
||||
tuner = Tuner(
|
||||
tune.with_parameters(train_fn, data=data),
|
||||
# ...
|
||||
)
|
||||
tuner.fit()
|
||||
|
||||
Class API example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray import tune
|
||||
|
||||
class MyTrainable(tune.Trainable):
|
||||
def setup(self, config, data=None):
|
||||
self.data = data
|
||||
self.iter = iter(self.data)
|
||||
self.next_sample = next(self.iter)
|
||||
|
||||
def step(self):
|
||||
loss = update_model(self.next_sample)
|
||||
try:
|
||||
self.next_sample = next(self.iter)
|
||||
except StopIteration:
|
||||
return {"loss": loss, done: True}
|
||||
return {"loss": loss}
|
||||
|
||||
data = HugeDataset(download=True)
|
||||
|
||||
tuner = Tuner(
|
||||
tune.with_parameters(MyTrainable, data=data),
|
||||
# ...
|
||||
)
|
||||
|
||||
Returns:
|
||||
A wrapped trainable that has the provided ``kwargs`` injected via the
|
||||
Ray object store at call time.
|
||||
"""
|
||||
from ray.tune.trainable import Trainable
|
||||
|
||||
if not callable(trainable) or (
|
||||
inspect.isclass(trainable) and not issubclass(trainable, Trainable)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`tune.with_parameters() only works with function trainables "
|
||||
f"or classes that inherit from `tune.Trainable()`. Got type: "
|
||||
f"{type(trainable)}."
|
||||
)
|
||||
|
||||
parameter_registry = _ParameterRegistry()
|
||||
ray._private.worker._post_init_hooks.append(parameter_registry.flush)
|
||||
|
||||
# Objects are moved into the object store
|
||||
prefix = f"{str(trainable)}_"
|
||||
for k, v in kwargs.items():
|
||||
parameter_registry.put(prefix + k, v)
|
||||
|
||||
trainable_name = getattr(trainable, "__name__", "tune_with_parameters")
|
||||
keys = set(kwargs.keys())
|
||||
|
||||
if inspect.isclass(trainable):
|
||||
# Class trainable
|
||||
|
||||
class _Inner(trainable):
|
||||
def setup(self, config):
|
||||
setup_kwargs = {}
|
||||
for k in keys:
|
||||
setup_kwargs[k] = parameter_registry.get(prefix + k)
|
||||
super(_Inner, self).setup(config, **setup_kwargs)
|
||||
|
||||
trainable_with_params = _Inner
|
||||
else:
|
||||
# Function trainable
|
||||
|
||||
def inner(config):
|
||||
fn_kwargs = {}
|
||||
for k in keys:
|
||||
fn_kwargs[k] = parameter_registry.get(prefix + k)
|
||||
return trainable(config, **fn_kwargs)
|
||||
|
||||
trainable_with_params = inner
|
||||
|
||||
if hasattr(trainable, "__mixins__"):
|
||||
trainable_with_params.__mixins__ = trainable.__mixins__
|
||||
|
||||
# If the trainable has been wrapped with `tune.with_resources`, we should
|
||||
# keep the `_resources` attribute around
|
||||
if hasattr(trainable, "_resources"):
|
||||
trainable_with_params._resources = trainable._resources
|
||||
|
||||
trainable_with_params.__name__ = trainable_name
|
||||
return trainable_with_params
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
def with_resources(
|
||||
trainable: Union[Type["Trainable"], Callable],
|
||||
resources: Union[
|
||||
Dict[str, float],
|
||||
PlacementGroupFactory,
|
||||
Callable[[dict], PlacementGroupFactory],
|
||||
],
|
||||
):
|
||||
"""Wrapper for trainables to specify resource requests.
|
||||
|
||||
This wrapper allows specification of resource requirements for a specific
|
||||
trainable. It will override potential existing resource requests (use
|
||||
with caution!).
|
||||
|
||||
The main use case is to request resources for function trainables when used
|
||||
with the Tuner() API.
|
||||
|
||||
Class trainables should usually just implement the ``default_resource_request()``
|
||||
method.
|
||||
|
||||
Args:
|
||||
trainable: Trainable to wrap.
|
||||
resources: Resource dict, placement group factory, or callable that takes
|
||||
in a config dict and returns a placement group factory.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray import tune
|
||||
from ray.tune.tuner import Tuner
|
||||
|
||||
def train_fn(config):
|
||||
return len(ray.get_gpu_ids()) # Returns 2
|
||||
|
||||
tuner = Tuner(
|
||||
tune.with_resources(train_fn, resources={"gpu": 2}),
|
||||
# ...
|
||||
)
|
||||
results = tuner.fit()
|
||||
|
||||
Returns:
|
||||
A trainable annotated with the requested resources so that Tune can
|
||||
schedule trials accordingly.
|
||||
"""
|
||||
from ray.tune.trainable import Trainable
|
||||
|
||||
if not callable(trainable) or (
|
||||
inspect.isclass(trainable) and not issubclass(trainable, Trainable)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`tune.with_resources() only works with function trainables "
|
||||
f"or classes that inherit from `tune.Trainable()`. Got type: "
|
||||
f"{type(trainable)}."
|
||||
)
|
||||
|
||||
if isinstance(resources, PlacementGroupFactory):
|
||||
pgf = resources
|
||||
elif isinstance(resources, dict):
|
||||
pgf = resource_dict_to_pg_factory(resources)
|
||||
elif callable(resources):
|
||||
pgf = resources
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid resource type for `with_resources()`: {type(resources)}"
|
||||
)
|
||||
|
||||
if not inspect.isclass(trainable):
|
||||
if isinstance(trainable, types.MethodType):
|
||||
# Methods cannot set arbitrary attributes, so we have to wrap them
|
||||
def _trainable(config):
|
||||
return trainable(config)
|
||||
|
||||
_trainable._resources = pgf
|
||||
return _trainable
|
||||
|
||||
# Just set an attribute. This will be resolved later in `wrap_function()`.
|
||||
try:
|
||||
trainable._resources = pgf
|
||||
except AttributeError as e:
|
||||
raise RuntimeError(
|
||||
"Could not use `tune.with_resources()` on the supplied trainable. "
|
||||
"Wrap your trainable in a regular function before passing it "
|
||||
"to Ray Tune."
|
||||
) from e
|
||||
else:
|
||||
|
||||
class ResourceTrainable(trainable):
|
||||
@classmethod
|
||||
def default_resource_request(
|
||||
cls, config: Dict[str, Any]
|
||||
) -> Optional[PlacementGroupFactory]:
|
||||
if not isinstance(pgf, PlacementGroupFactory) and callable(pgf):
|
||||
return pgf(config)
|
||||
return pgf
|
||||
|
||||
ResourceTrainable.__name__ = trainable.__name__
|
||||
trainable = ResourceTrainable
|
||||
|
||||
return trainable
|
||||
Reference in New Issue
Block a user