def testHasResourcesForTrialWithCaching(self):
pgm = _PlacementGroupManager()
pgf1 = PlacementGroupFactory([{"CPU": self.head_cpus}])
pgf2 = PlacementGroupFactory([{"CPU": self.head_cpus - 1}])
executor = RayTrialExecutor(reuse_actors=True)
executor._pg_manager = pgm
executor.set_max_pending_trials(1)
def train(config):
yield 1
yield 2
yield 3
yield 4
register_trainable("resettable", train)
trial1 = Trial("resettable", placement_group_factory=pgf1)
trial2 = Trial("resettable", placement_group_factory=pgf1)
trial3 = Trial("resettable", placement_group_factory=pgf2)
assert executor.has_resources_for_trial(trial1)
assert executor.has_resources_for_trial(trial2)
assert executor.has_resources_for_trial(trial3)
executor._stage_and_update_status([trial1, trial2, trial3])
while not pgm.has_ready(trial1):
time.sleep(1)
executor._stage_and_update_status([trial1, trial2, trial3])
# Fill staging
executor._stage_and_update_status([trial1, trial2, trial3])
assert executor.has_resources_for_trial(trial1)
assert executor.has_resources_for_trial(trial2)
assert not executor.has_resources_for_trial(trial3)
executor._start_trial(trial1)
executor._stage_and_update_status([trial1, trial2, trial3])
executor.pause_trial(
trial1) # Caches the PG and removes a PG from staging
assert len(pgm._staging_futures) == 0
# This will re-schedule a placement group
pgm.reconcile_placement_groups([trial1, trial2])
assert len(pgm._staging_futures) == 1
assert not pgm.can_stage()
# We should still have resources for this trial as it has a cached PG
assert executor.has_resources_for_trial(trial1)
assert executor.has_resources_for_trial(trial2)
assert not executor.has_resources_for_trial(trial3)
class RayTrialExecutorTest(unittest.TestCase):
def setUp(self):
self.trial_executor = RayTrialExecutor()
ray.init(num_cpus=2, ignore_reinit_error=True)
_register_all() # Needed for flaky tests
def tearDown(self):
ray.shutdown()
_register_all() # re-register the evicted objects
def _simulate_starting_trial(self, trial):
future_result = self.trial_executor.get_next_executor_event(
live_trials={trial}, next_trial_exists=True
)
assert future_result.type == _ExecutorEventType.PG_READY
self.assertTrue(self.trial_executor.start_trial(trial))
self.assertEqual(Trial.RUNNING, trial.status)
def _simulate_getting_result(self, trial):
while True:
event = self.trial_executor.get_next_executor_event(
live_trials={trial}, next_trial_exists=False
)
if event.type == _ExecutorEventType.TRAINING_RESULT:
break
training_result = event.result[_ExecutorEvent.KEY_FUTURE_RESULT]
if isinstance(training_result, list):
for r in training_result:
trial.update_last_result(r)
else:
trial.update_last_result(training_result)
def _simulate_saving(self, trial):
checkpoint = self.trial_executor.save(trial, CheckpointStorage.PERSISTENT)
self.assertEqual(checkpoint, trial.saving_to)
self.assertEqual(trial.checkpoint.dir_or_data, None)
event = self.trial_executor.get_next_executor_event(
live_trials={trial}, next_trial_exists=False
)
assert event.type == _ExecutorEventType.SAVING_RESULT
self.process_trial_save(trial, event.result[_ExecutorEvent.KEY_FUTURE_RESULT])
self.assertEqual(checkpoint, trial.checkpoint)
def testStartStop(self):
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self.trial_executor.stop_trial(trial)
def testAsyncSave(self):
"""Tests that saved checkpoint value not immediately set."""
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
self._simulate_saving(trial)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def testSaveRestore(self):
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
self._simulate_saving(trial)
self.trial_executor.restore(trial)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def testPauseResume(self):
"""Tests that pausing works for trials in flight."""
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self.trial_executor.pause_trial(trial)
self.assertEqual(Trial.PAUSED, trial.status)
self._simulate_starting_trial(trial)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def testSavePauseResumeErrorRestore(self):
"""Tests that pause checkpoint does not replace restore checkpoint."""
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
# Save
self._simulate_saving(trial)
# Train
self.trial_executor.continue_training(trial)
self._simulate_getting_result(trial)
# Pause
self.trial_executor.pause_trial(trial)
self.assertEqual(Trial.PAUSED, trial.status)
self.assertEqual(trial.checkpoint.storage_mode, CheckpointStorage.MEMORY)
# Resume
self._simulate_starting_trial(trial)
# Error
trial.set_status(Trial.ERROR)
# Restore
self.trial_executor.restore(trial)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def testStartFailure(self):
_global_registry.register(TRAINABLE_CLASS, "asdf", None)
trial = Trial("asdf", resources=Resources(1, 0))
self.trial_executor.start_trial(trial)
self.assertEqual(Trial.ERROR, trial.status)
def testPauseResume2(self):
"""Tests that pausing works for trials being processed."""
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
self.trial_executor.pause_trial(trial)
self.assertEqual(Trial.PAUSED, trial.status)
self._simulate_starting_trial(trial)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def _testPauseAndStart(self, result_buffer_length):
"""Tests that unpausing works for trials being processed."""
os.environ["TUNE_RESULT_BUFFER_LENGTH"] = f"{result_buffer_length}"
os.environ["TUNE_RESULT_BUFFER_MIN_TIME_S"] = "1"
# Need a new trial executor so the ENV vars are parsed again
self.trial_executor = RayTrialExecutor()
base = max(result_buffer_length, 1)
trial = Trial("__fake")
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
self.assertEqual(trial.last_result.get(TRAINING_ITERATION), base)
self.trial_executor.pause_trial(trial)
self.assertEqual(Trial.PAUSED, trial.status)
self._simulate_starting_trial(trial)
self._simulate_getting_result(trial)
self.assertEqual(trial.last_result.get(TRAINING_ITERATION), base * 2)
self.trial_executor.stop_trial(trial)
self.assertEqual(Trial.TERMINATED, trial.status)
def testPauseAndStartNoBuffer(self):
self._testPauseAndStart(0)
def testPauseAndStartTrivialBuffer(self):
self._testPauseAndStart(1)
def testPauseAndStartActualBuffer(self):
self._testPauseAndStart(8)
def testNoResetTrial(self):
"""Tests that reset handles NotImplemented properly."""
trial = Trial("__fake")
self._simulate_starting_trial(trial)
exists = self.trial_executor.reset_trial(trial, {}, "modified_mock")
self.assertEqual(exists, False)
self.assertEqual(Trial.RUNNING, trial.status)
def testResetTrial(self):
"""Tests that reset works as expected."""
class B(Trainable):
def step(self):
return dict(timesteps_this_iter=1, done=True)
def reset_config(self, config):
self.config = config
return True
trials = self.generate_trials(
{
"run": B,
"config": {"foo": 0},
},
"grid_search",
)
trial = trials[0]
self._simulate_starting_trial(trial)
exists = self.trial_executor.reset_trial(trial, {"hi": 1}, "modified_mock")
self.assertEqual(exists, True)
self.assertEqual(trial.config.get("hi"), 1)
self.assertEqual(trial.experiment_tag, "modified_mock")
self.assertEqual(Trial.RUNNING, trial.status)
def testTrialCleanup(self):
class B(Trainable):
def step(self):
print("Step start")
time.sleep(4)
print("Step done")
return dict(my_metric=1, timesteps_this_iter=1, done=True)
def reset_config(self, config):
self.config = config
return True
def cleanup(self):
print("Cleanup start")
time.sleep(4)
print("Cleanup done")
# First check if the trials terminate gracefully by default
trials = self.generate_trials(
{
"run": B,
"config": {"foo": 0},
},
"grid_search",
)
trial = trials[0]
self._simulate_starting_trial(trial)
time.sleep(1)
print("Stop trial")
self.trial_executor.stop_trial(trial)
print("Start trial cleanup")
start = time.time()
self.trial_executor.cleanup([trial])
# 4 - 1 + 4.
self.assertGreaterEqual(time.time() - start, 6)
# Check forceful termination. It should run for much less than the
# sleep periods in the Trainable
trials = self.generate_trials(
{
"run": B,
"config": {"foo": 0},
},
"grid_search",
)
trial = trials[0]
os.environ["TUNE_FORCE_TRIAL_CLEANUP_S"] = "1"
self.trial_executor = RayTrialExecutor()
os.environ["TUNE_FORCE_TRIAL_CLEANUP_S"] = "0"
self._simulate_starting_trial(trial)
self.assertEqual(Trial.RUNNING, trial.status)
# This should be enough time for `trial._default_result_or_future`
# to return. Otherwise, PID won't show up in `trial.last_result`,
# which is asserted down below.
time.sleep(2)
print("Stop trial")
self.trial_executor.stop_trial(trial)
print("Start trial cleanup")
start = time.time()
self.trial_executor.cleanup([trial])
# less than 1 with some margin.
self.assertLess(time.time() - start, 2.0)
# also check if auto-filled metrics were returned
self.assertIn(PID, trial.last_result)
self.assertIn(TRIAL_ID, trial.last_result)
self.assertNotIn("my_metric", trial.last_result)
@staticmethod
def generate_trials(spec, name):
suggester = BasicVariantGenerator()
suggester.add_configurations({name: spec})
trials = []
while not suggester.is_finished():
trial = suggester.next_trial()
if trial:
trials.append(trial)
else:
break
return trials
def process_trial_save(self, trial, checkpoint_value):
"""Simulates trial runner save."""
checkpoint = trial.saving_to
checkpoint.dir_or_data = checkpoint_value
trial.on_checkpoint(checkpoint)