mirror of
https://github.com/wassname/ray.git
synced 2026-07-23 13:10:11 +08:00
[tune] Cross-Framework Compatibility (#2646)
This commit is a first pass at restructuring the Trial execution logic to support running on multiple frameworks.
This commit is contained in:
joyyoj
+54
-170
@@ -2,16 +2,15 @@ from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import tempfile
|
||||
from collections import namedtuple
|
||||
from datetime import datetime
|
||||
import tempfile
|
||||
import time
|
||||
import traceback
|
||||
import ray
|
||||
import os
|
||||
|
||||
from ray.tune import TuneError
|
||||
from ray.tune.logger import NoopLogger, UnifiedLogger, pretty_print
|
||||
from ray.tune.logger import pretty_print, UnifiedLogger
|
||||
# NOTE(rkn): We import ray.tune.registry here instead of importing the names we
|
||||
# need because there are cyclic imports that may cause specific names to not
|
||||
# have been defined yet. See https://github.com/ray-project/ray/issues/1716.
|
||||
@@ -39,7 +38,9 @@ class Resources(
|
||||
launch additional Ray actors that use CPUs.
|
||||
extra_gpu (int): Extra GPUs to reserve in case the trial needs to
|
||||
launch additional Ray actors that use GPUs.
|
||||
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls, cpu, gpu, extra_cpu=0, extra_gpu=0):
|
||||
@@ -62,6 +63,30 @@ def has_trainable(trainable_name):
|
||||
ray.tune.registry.TRAINABLE_CLASS, trainable_name)
|
||||
|
||||
|
||||
class Checkpoint(object):
|
||||
"""Describes a checkpoint of trial state.
|
||||
|
||||
Checkpoint may be saved in different storage.
|
||||
|
||||
Attributes:
|
||||
storage (str): Storage type.
|
||||
value (str): If storage==MEMORY,value is a Python object.
|
||||
If storage==DISK,value is a path points to the checkpoint in disk.
|
||||
"""
|
||||
|
||||
MEMORY = "memory"
|
||||
DISK = "disk"
|
||||
|
||||
def __init__(self, storage, value):
|
||||
self.storage = storage
|
||||
self.value = value
|
||||
|
||||
@staticmethod
|
||||
def from_object(value=None):
|
||||
"""Creates a checkpoint from a Python object."""
|
||||
return Checkpoint(Checkpoint.MEMORY, value)
|
||||
|
||||
|
||||
class Trial(object):
|
||||
"""A trial object holds the state for one model training run.
|
||||
|
||||
@@ -110,16 +135,15 @@ class Trial(object):
|
||||
resources
|
||||
or self._get_trainable_cls().default_resource_request(self.config))
|
||||
self.stopping_criterion = stopping_criterion or {}
|
||||
self.checkpoint_freq = checkpoint_freq
|
||||
self.upload_dir = upload_dir
|
||||
self.verbose = True
|
||||
self.max_failures = max_failures
|
||||
|
||||
# Local trial state that is updated during the run
|
||||
self.last_result = None
|
||||
self._checkpoint_path = restore_path
|
||||
self._checkpoint_obj = None
|
||||
self.runner = None
|
||||
self.checkpoint_freq = checkpoint_freq
|
||||
self._checkpoint = Checkpoint(
|
||||
storage=Checkpoint.DISK, value=restore_path)
|
||||
self.status = Trial.PENDING
|
||||
self.location = None
|
||||
self.logdir = None
|
||||
@@ -136,96 +160,34 @@ class Trial(object):
|
||||
def generate_id(cls):
|
||||
return binary_to_hex(random_string())[:8]
|
||||
|
||||
def start(self, checkpoint_obj=None):
|
||||
"""Starts this trial.
|
||||
def init_logger(self):
|
||||
"""Init logger."""
|
||||
|
||||
If an error is encountered when starting the trial, an exception will
|
||||
be thrown.
|
||||
if not self.result_logger:
|
||||
if not os.path.exists(self.local_dir):
|
||||
os.makedirs(self.local_dir)
|
||||
self.logdir = tempfile.mkdtemp(
|
||||
prefix="{}_{}".format(
|
||||
str(self)[:MAX_LEN_IDENTIFIER], date_str()),
|
||||
dir=self.local_dir)
|
||||
self.result_logger = UnifiedLogger(self.config, self.logdir,
|
||||
self.upload_dir)
|
||||
|
||||
Args:
|
||||
checkpoint_obj (obj): Optional checkpoint to resume from.
|
||||
"""
|
||||
def close_logger(self):
|
||||
"""Close logger."""
|
||||
|
||||
self._setup_runner()
|
||||
if checkpoint_obj:
|
||||
self.restore_from_obj(checkpoint_obj)
|
||||
elif self._checkpoint_path:
|
||||
self.restore_from_path(self._checkpoint_path)
|
||||
elif self._checkpoint_obj:
|
||||
self.restore_from_obj(self._checkpoint_obj)
|
||||
|
||||
def stop(self, error=False, error_msg=None, stop_logger=True):
|
||||
"""Stops this trial.
|
||||
|
||||
Stops this trial, releasing all allocating resources. If stopping the
|
||||
trial fails, the run will be marked as terminated in error, but no
|
||||
exception will be thrown.
|
||||
|
||||
Args:
|
||||
error (bool): Whether to mark this trial as terminated in error.
|
||||
error_msg (str): Optional error message.
|
||||
stop_logger (bool): Whether to shut down the trial logger.
|
||||
"""
|
||||
|
||||
if error:
|
||||
self.status = Trial.ERROR
|
||||
else:
|
||||
self.status = Trial.TERMINATED
|
||||
|
||||
try:
|
||||
if error_msg and self.logdir:
|
||||
self.num_failures += 1
|
||||
error_file = os.path.join(self.logdir,
|
||||
"error_{}.txt".format(date_str()))
|
||||
with open(error_file, "w") as f:
|
||||
f.write(error_msg)
|
||||
self.error_file = error_file
|
||||
if self.runner:
|
||||
stop_tasks = []
|
||||
stop_tasks.append(self.runner.stop.remote())
|
||||
stop_tasks.append(self.runner.__ray_terminate__.remote())
|
||||
# TODO(ekl) seems like wait hangs when killing actors
|
||||
_, unfinished = ray.wait(
|
||||
stop_tasks, num_returns=2, timeout=250)
|
||||
except Exception:
|
||||
print("Error stopping runner:", traceback.format_exc())
|
||||
self.status = Trial.ERROR
|
||||
finally:
|
||||
self.runner = None
|
||||
|
||||
if stop_logger and self.result_logger:
|
||||
if self.result_logger:
|
||||
self.result_logger.close()
|
||||
self.result_logger = None
|
||||
|
||||
def pause(self):
|
||||
"""We want to release resources (specifically GPUs) when pausing an
|
||||
experiment. This results in a state similar to TERMINATED."""
|
||||
|
||||
assert self.status == Trial.RUNNING, self.status
|
||||
try:
|
||||
self.checkpoint(to_object_store=True)
|
||||
self.stop(stop_logger=False)
|
||||
self.status = Trial.PAUSED
|
||||
except Exception:
|
||||
print("Error pausing runner:", traceback.format_exc())
|
||||
self.status = Trial.ERROR
|
||||
|
||||
def unpause(self):
|
||||
"""Sets PAUSED trial to pending to allow scheduler to start."""
|
||||
assert self.status == Trial.PAUSED, self.status
|
||||
self.status = Trial.PENDING
|
||||
|
||||
def resume(self):
|
||||
"""Resume PAUSED trials. This is a blocking call."""
|
||||
|
||||
assert self.status == Trial.PAUSED, self.status
|
||||
self.start()
|
||||
|
||||
def train_remote(self):
|
||||
"""Returns Ray future for one iteration of training."""
|
||||
|
||||
assert self.status == Trial.RUNNING, self.status
|
||||
return self.runner.train.remote()
|
||||
def write_error_log(self, error_msg):
|
||||
if error_msg and self.logdir:
|
||||
self.num_failures += 1 # may be moved to outer scope?
|
||||
error_file = os.path.join(self.logdir,
|
||||
"error_{}.txt".format(date_str()))
|
||||
with open(error_file, "w") as f:
|
||||
f.write(error_msg)
|
||||
self.error_file = error_file
|
||||
|
||||
def should_stop(self, result):
|
||||
"""Whether the given result meets this trial's stopping criteria."""
|
||||
@@ -294,57 +256,7 @@ class Trial(object):
|
||||
if self.error_file else "")
|
||||
|
||||
def has_checkpoint(self):
|
||||
return self._checkpoint_path is not None or \
|
||||
self._checkpoint_obj is not None
|
||||
|
||||
def checkpoint(self, to_object_store=False):
|
||||
"""Checkpoints the state of this trial.
|
||||
|
||||
Args:
|
||||
to_object_store (bool): Whether to save to the Ray object store
|
||||
(async) vs a path on local disk (sync).
|
||||
"""
|
||||
|
||||
obj = None
|
||||
path = None
|
||||
if to_object_store:
|
||||
obj = self.runner.save_to_object.remote()
|
||||
else:
|
||||
path = ray.get(self.runner.save.remote())
|
||||
self._checkpoint_path = path
|
||||
self._checkpoint_obj = obj
|
||||
|
||||
if self.verbose:
|
||||
print("Saved checkpoint for {} to {}".format(self, path or obj))
|
||||
return path or obj
|
||||
|
||||
def restore_from_path(self, path):
|
||||
"""Restores runner state from specified path.
|
||||
|
||||
Args:
|
||||
path (str): A path where state will be restored.
|
||||
"""
|
||||
|
||||
if self.runner is None:
|
||||
print("Unable to restore - no runner")
|
||||
else:
|
||||
try:
|
||||
ray.get(self.runner.restore.remote(path))
|
||||
except Exception:
|
||||
print("Error restoring runner:", traceback.format_exc())
|
||||
self.status = Trial.ERROR
|
||||
|
||||
def restore_from_obj(self, obj):
|
||||
"""Restores runner state from the specified object."""
|
||||
|
||||
if self.runner is None:
|
||||
print("Unable to restore - no runner")
|
||||
else:
|
||||
try:
|
||||
ray.get(self.runner.restore_from_object.remote(obj))
|
||||
except Exception:
|
||||
print("Error restoring runner:", traceback.format_exc())
|
||||
self.status = Trial.ERROR
|
||||
return self._checkpoint.value is not None
|
||||
|
||||
def update_last_result(self, result, terminate=False):
|
||||
if terminate:
|
||||
@@ -357,34 +269,6 @@ class Trial(object):
|
||||
self.last_result = result
|
||||
self.result_logger.on_result(self.last_result)
|
||||
|
||||
def _setup_runner(self):
|
||||
self.status = Trial.RUNNING
|
||||
cls = ray.remote(
|
||||
num_cpus=self.resources.cpu,
|
||||
num_gpus=self.resources.gpu)(self._get_trainable_cls())
|
||||
if not self.result_logger:
|
||||
if not os.path.exists(self.local_dir):
|
||||
os.makedirs(self.local_dir)
|
||||
self.logdir = tempfile.mkdtemp(
|
||||
prefix="{}_{}".format(
|
||||
str(self)[:MAX_LEN_IDENTIFIER], date_str()),
|
||||
dir=self.local_dir)
|
||||
self.result_logger = UnifiedLogger(self.config, self.logdir,
|
||||
self.upload_dir)
|
||||
remote_logdir = self.logdir
|
||||
|
||||
def logger_creator(config):
|
||||
# Set the working dir in the remote process, for user file writes
|
||||
if not os.path.exists(remote_logdir):
|
||||
os.makedirs(remote_logdir)
|
||||
os.chdir(remote_logdir)
|
||||
return NoopLogger(config, remote_logdir)
|
||||
|
||||
# Logging for trials is handled centrally by TrialRunner, so
|
||||
# configure the remote runner to use a noop-logger.
|
||||
self.runner = cls.remote(
|
||||
config=self.config, logger_creator=logger_creator)
|
||||
|
||||
def _get_trainable_cls(self):
|
||||
return ray.tune.registry._global_registry.get(
|
||||
ray.tune.registry.TRAINABLE_CLASS, self.trainable_name)
|
||||
|
||||
Reference in New Issue
Block a user