[tune] logger refactor part 2: Add SyncerCallback (#11748)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
This commit is contained in:
Kai Fricke
2020-11-03 21:04:40 -08:00
committed by GitHub
co-authored by Richard Liaw
parent 05c4e3fb2a
commit 007634fd1b
10 changed files with 211 additions and 193 deletions
+52 -33
View File
@@ -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:
+4 -5
View File
@@ -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",
+1 -5
View File
@@ -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):