mirror of
https://github.com/wassname/ray.git
synced 2026-07-28 11:25:04 +08:00
[tune] logger refactor part 2: Add SyncerCallback (#11748)
Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
This commit is contained in:
co-authored by
Richard Liaw
parent
05c4e3fb2a
commit
007634fd1b
@@ -1,12 +1,15 @@
|
||||
import inspect
|
||||
import time
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from typing import Callable, Union
|
||||
|
||||
import ray
|
||||
from ray import tune
|
||||
from ray.rllib import _register_all
|
||||
@@ -18,7 +21,7 @@ from ray.tune.error import TuneError
|
||||
from ray.tune.ray_trial_executor import RayTrialExecutor
|
||||
from ray.tune.resources import Resources
|
||||
from ray.tune.suggest import BasicVariantGenerator
|
||||
from ray.tune.syncer import CloudSyncer
|
||||
from ray.tune.syncer import CloudSyncer, SyncerCallback, get_node_syncer
|
||||
from ray.tune.utils.trainable import TrainableUtil
|
||||
from ray.tune.trial import Trial
|
||||
from ray.tune.trial_runner import TrialRunner
|
||||
@@ -55,6 +58,19 @@ def _start_new_cluster():
|
||||
return cluster
|
||||
|
||||
|
||||
class _PerTrialSyncerCallback(SyncerCallback):
|
||||
def __init__(
|
||||
self,
|
||||
get_sync_fn: Callable[["Trial"], Union[None, bool, Callable]]):
|
||||
self._get_sync_fn = get_sync_fn
|
||||
super(_PerTrialSyncerCallback, self).__init__(None)
|
||||
|
||||
def _create_trial_syncer(self, trial: "Trial"):
|
||||
sync_fn = self._get_sync_fn(trial)
|
||||
return get_node_syncer(
|
||||
trial.logdir, remote_dir=trial.logdir, sync_function=sync_fn)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def start_connected_cluster():
|
||||
# Start the Ray processes.
|
||||
@@ -255,7 +271,9 @@ def test_trial_migration(start_connected_emptyhead_cluster, trainable_id):
|
||||
node = cluster.add_node(num_cpus=1)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
runner = TrialRunner(BasicVariantGenerator())
|
||||
syncer_callback = _PerTrialSyncerCallback(
|
||||
lambda trial: trial.trainable_name == "__fake")
|
||||
runner = TrialRunner(BasicVariantGenerator(), callbacks=[syncer_callback])
|
||||
kwargs = {
|
||||
"stopping_criterion": {
|
||||
"training_iteration": 4
|
||||
@@ -263,7 +281,6 @@ def test_trial_migration(start_connected_emptyhead_cluster, trainable_id):
|
||||
"checkpoint_freq": 2,
|
||||
"max_failures": 2,
|
||||
"remote_checkpoint_dir": MOCK_REMOTE_DIR,
|
||||
"sync_to_driver_fn": trainable_id == "__fake",
|
||||
}
|
||||
|
||||
# Test recovery of trial that hasn't been checkpointed
|
||||
@@ -316,7 +333,6 @@ def test_trial_migration(start_connected_emptyhead_cluster, trainable_id):
|
||||
"training_iteration": 3
|
||||
},
|
||||
"remote_checkpoint_dir": MOCK_REMOTE_DIR,
|
||||
"sync_to_driver_fn": trainable_id == "__fake",
|
||||
}
|
||||
t3 = Trial(trainable_id, **kwargs)
|
||||
runner.add_trial(t3)
|
||||
@@ -341,7 +357,9 @@ def test_trial_requeue(start_connected_emptyhead_cluster, trainable_id):
|
||||
node = cluster.add_node(num_cpus=1)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
runner = TrialRunner(BasicVariantGenerator())
|
||||
syncer_callback = _PerTrialSyncerCallback(
|
||||
lambda trial: trial.trainable_name == "__fake")
|
||||
runner = TrialRunner(BasicVariantGenerator(), callbacks=[syncer_callback])
|
||||
kwargs = {
|
||||
"stopping_criterion": {
|
||||
"training_iteration": 5
|
||||
@@ -349,7 +367,6 @@ def test_trial_requeue(start_connected_emptyhead_cluster, trainable_id):
|
||||
"checkpoint_freq": 1,
|
||||
"max_failures": 1,
|
||||
"remote_checkpoint_dir": MOCK_REMOTE_DIR,
|
||||
"sync_to_driver_fn": trainable_id == "__fake",
|
||||
}
|
||||
|
||||
trials = [Trial(trainable_id, **kwargs), Trial(trainable_id, **kwargs)]
|
||||
@@ -382,7 +399,13 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
|
||||
node = cluster.add_node(num_cpus=1)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
runner = TrialRunner(BasicVariantGenerator())
|
||||
class _SyncerCallback(SyncerCallback):
|
||||
def _create_trial_syncer(self, trial: "Trial"):
|
||||
client = mock_storage_client()
|
||||
return MockNodeSyncer(trial.logdir, trial.logdir, client)
|
||||
|
||||
syncer_callback = _SyncerCallback(None)
|
||||
runner = TrialRunner(BasicVariantGenerator(), callbacks=[syncer_callback])
|
||||
kwargs = {
|
||||
"stopping_criterion": {
|
||||
"training_iteration": 4
|
||||
@@ -390,7 +413,6 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
|
||||
"checkpoint_freq": 2,
|
||||
"max_failures": 2,
|
||||
"remote_checkpoint_dir": MOCK_REMOTE_DIR,
|
||||
"sync_to_driver_fn": trainable_id == "__fake_remote",
|
||||
}
|
||||
|
||||
# The following patches only affect __fake_remote.
|
||||
@@ -415,33 +437,26 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
|
||||
# TrainableUtil will not check this path unless we mock it.
|
||||
mock_find.side_effect = hide_remote_path(find_func)
|
||||
mock_pkl_ckpt.side_effect = hide_remote_path(pickle_func)
|
||||
with patch("ray.tune.logger.get_node_syncer") as mock_get_node_syncer:
|
||||
|
||||
def mock_get_syncer_fn(local_dir, remote_dir, sync_function):
|
||||
client = mock_storage_client()
|
||||
return MockNodeSyncer(local_dir, remote_dir, client)
|
||||
# Test recovery of trial that has been checkpointed
|
||||
t1 = Trial(trainable_id, **kwargs)
|
||||
runner.add_trial(t1)
|
||||
|
||||
mock_get_node_syncer.side_effect = mock_get_syncer_fn
|
||||
# Start trial, process result (x2), process save
|
||||
for _ in range(4):
|
||||
runner.step()
|
||||
assert t1.has_checkpoint()
|
||||
|
||||
# Test recovery of trial that has been checkpointed
|
||||
t1 = Trial(trainable_id, **kwargs)
|
||||
runner.add_trial(t1)
|
||||
|
||||
# Start trial, process result (x2), process save
|
||||
for _ in range(4):
|
||||
cluster.add_node(num_cpus=1)
|
||||
cluster.remove_node(node)
|
||||
cluster.wait_for_nodes()
|
||||
shutil.rmtree(os.path.dirname(t1.checkpoint.value))
|
||||
runner.step() # Collect result 3, kick off + fail result 4
|
||||
runner.step() # Dispatch restore
|
||||
runner.step() # Process restore + step 4
|
||||
for _ in range(3):
|
||||
if t1.status != Trial.TERMINATED:
|
||||
runner.step()
|
||||
assert t1.has_checkpoint()
|
||||
|
||||
cluster.add_node(num_cpus=1)
|
||||
cluster.remove_node(node)
|
||||
cluster.wait_for_nodes()
|
||||
shutil.rmtree(os.path.dirname(t1.checkpoint.value))
|
||||
runner.step() # Collect result 3, kick off + fail result 4
|
||||
runner.step() # Dispatch restore
|
||||
runner.step() # Process restore + step 4
|
||||
for _ in range(3):
|
||||
if t1.status != Trial.TERMINATED:
|
||||
runner.step()
|
||||
assert t1.status == Trial.TERMINATED, runner.debug_string()
|
||||
|
||||
|
||||
@@ -454,7 +469,12 @@ def test_cluster_down_simple(start_connected_cluster, tmpdir, trainable_id):
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
dirpath = str(tmpdir)
|
||||
runner = TrialRunner(local_checkpoint_dir=dirpath, checkpoint_period=0)
|
||||
syncer_callback = _PerTrialSyncerCallback(
|
||||
lambda trial: trial.trainable_name == "__fake")
|
||||
runner = TrialRunner(
|
||||
local_checkpoint_dir=dirpath,
|
||||
checkpoint_period=0,
|
||||
callbacks=[syncer_callback])
|
||||
kwargs = {
|
||||
"stopping_criterion": {
|
||||
"training_iteration": 2
|
||||
@@ -462,7 +482,6 @@ def test_cluster_down_simple(start_connected_cluster, tmpdir, trainable_id):
|
||||
"checkpoint_freq": 1,
|
||||
"max_failures": 1,
|
||||
"remote_checkpoint_dir": MOCK_REMOTE_DIR,
|
||||
"sync_to_driver_fn": trainable_id == "__fake",
|
||||
}
|
||||
trials = [Trial(trainable_id, **kwargs), Trial(trainable_id, **kwargs)]
|
||||
for t in trials:
|
||||
|
||||
@@ -11,7 +11,6 @@ import ray
|
||||
from ray.rllib import _register_all
|
||||
|
||||
from ray import tune
|
||||
from ray.tune import TuneError
|
||||
from ray.tune.syncer import CommandBasedClient
|
||||
|
||||
|
||||
@@ -87,8 +86,8 @@ class TestSyncFunctionality(unittest.TestCase):
|
||||
|
||||
def testClusterProperString(self):
|
||||
"""Tests that invalid commands throw.."""
|
||||
with self.assertRaises(TuneError):
|
||||
# This raises TuneError because logger is init in safe zone.
|
||||
with self.assertRaises(ValueError):
|
||||
# This raises ValueError because logger is init in safe zone.
|
||||
sync_config = tune.SyncConfig(sync_to_driver="ls {target}")
|
||||
[trial] = tune.run(
|
||||
"__fake",
|
||||
@@ -100,8 +99,8 @@ class TestSyncFunctionality(unittest.TestCase):
|
||||
sync_config=sync_config,
|
||||
).trials
|
||||
|
||||
with self.assertRaises(TuneError):
|
||||
# This raises TuneError because logger is init in safe zone.
|
||||
with self.assertRaises(ValueError):
|
||||
# This raises ValueError because logger is init in safe zone.
|
||||
sync_config = tune.SyncConfig(sync_to_driver="ls {source}")
|
||||
[trial] = tune.run(
|
||||
"__fake",
|
||||
|
||||
@@ -124,11 +124,7 @@ class TuneExampleTest(unittest.TestCase):
|
||||
class AutoInitTest(unittest.TestCase):
|
||||
def testTuneRestore(self):
|
||||
self.assertFalse(ray.is_initialized())
|
||||
tune.run(
|
||||
"__fake",
|
||||
name="TestAutoInit",
|
||||
stop={"training_iteration": 1},
|
||||
ray_auto_init=True)
|
||||
tune.run("__fake", name="TestAutoInit", stop={"training_iteration": 1})
|
||||
self.assertTrue(ray.is_initialized())
|
||||
|
||||
def tearDown(self):
|
||||
|
||||
Reference in New Issue
Block a user