def test_resume_from_checkpoint(ray_start_4_cpus, tmpdir): def train_func(): checkpoint = session.get_checkpoint() if checkpoint: epoch = checkpoint.to_dict()["epoch"] else: epoch = 0 for i in range(epoch, epoch + 2): session.report({"epoch": i}, checkpoint=Checkpoint.from_dict({"epoch": i})) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=scale_config ) result = trainer.fit() assert result.checkpoint.to_dict()["epoch"] == 1 # Move checkpoint to a different directory. checkpoint_dict = result.checkpoint.to_dict() checkpoint = Checkpoint.from_dict(checkpoint_dict) checkpoint_path = checkpoint.to_directory(tmpdir) resume_from = Checkpoint.from_directory(checkpoint_path) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=scale_config, resume_from_checkpoint=resume_from, ) result = trainer.fit() assert result.checkpoint.to_dict()["epoch"] == 2
def test_run(ray_start_4_cpus): """Tests that Train can be run without any specific backends.""" num_workers = 2 key = "value" value = 1 config = TestConfig() def train_func(): checkpoint = session.get_checkpoint() session.report(metrics=checkpoint.to_dict(), checkpoint=checkpoint) return checkpoint.to_dict()[key] checkpoint = Checkpoint.from_dict({ # this would be set during checkpoint saving "_current_checkpoint_id": 1, key: value, }) trainer = DataParallelTrainer( train_func, backend_config=config, resume_from_checkpoint=checkpoint, scaling_config=ScalingConfig(num_workers=num_workers), ) results = trainer.fit() assert results.checkpoint.to_dict()[key] == checkpoint.to_dict()[key]
def test_fit_train(ray_start_4_cpus): def train_func(): train.report(loss=1) trainer = DataParallelTrainer(train_loop_per_worker=train_func, scaling_config=scale_config) assert trainer.fit().metrics["loss"] == 1
def test_reuse_checkpoint(ray_start_4_cpus): def train_func(config): itr = 0 ckpt = session.get_checkpoint() if ckpt is not None: ckpt = ckpt.to_dict() itr = ckpt["iter"] + 1 for i in range(itr, config["max_iter"]): session.report( dict(test=i, training_iteration=i), checkpoint=Checkpoint.from_dict(dict(iter=i)), ) trainer = DataParallelTrainer( train_func, backend_config=TestConfig(), scaling_config=ScalingConfig(num_workers=1), ) tuner = Tuner( trainer, param_space={"train_loop_config": {"max_iter": 5}}, ) [trial] = tuner.fit()._experiment_analysis.trials checkpoint_path = trial.checkpoint.dir_or_data checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict() assert checkpoint["iter"] == 4 tuner = Tuner( trainer, param_space={"train_loop_config": {"max_iter": 10}}, ).restore(trial.local_dir) analysis = tuner.fit()._experiment_analysis trial_dfs = list(analysis.trial_dataframes.values()) assert len(trial_dfs[0]["training_iteration"]) == 5
def test_retry(ray_start_4_cpus): def train_func(): ckpt = session.get_checkpoint() restored = bool(ckpt) # Does a previous checkpoint exist? itr = 0 if ckpt: ckpt = ckpt.to_dict() itr = ckpt["iter"] + 1 for i in range(itr, 4): if i == 2 and not restored: raise Exception("try to fail me") session.report( dict(test=i, training_iteration=i), checkpoint=Checkpoint.from_dict(dict(iter=i)), ) trainer = DataParallelTrainer( train_func, backend_config=TestConfig(), scaling_config=ScalingConfig(num_workers=1), ) tuner = Tuner( trainer, run_config=RunConfig(failure_config=FailureConfig(max_failures=3)) ) analysis = tuner.fit()._experiment_analysis checkpoint_path = analysis.trials[0].checkpoint.dir_or_data checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict() assert checkpoint["iter"] == 3 trial_dfs = list(analysis.trial_dataframes.values()) assert len(trial_dfs[0]["training_iteration"]) == 4
def test_datasets(ray_start_4_cpus): num_train_data = 10 num_val_data = 6 train_dataset = ray.data.range(num_train_data) val_dataset = ray.data.range(num_val_data) def get_dataset(): # Train dataset should be sharded. train_dataset = train.get_dataset_shard("train") assert train_dataset.count( ) == num_train_data / scale_config["num_workers"] # All other datasets should not be sharded. val_dataset = train.get_dataset_shard("val") assert val_dataset.count() == num_val_data trainer = DataParallelTrainer( train_loop_per_worker=get_dataset, scaling_config=scale_config, datasets={ "train": train_dataset, "val": val_dataset }, ) trainer.fit()
def test_checkpoint(ray_start_4_cpus): def train_func(): for i in range(3): train.save_checkpoint(model=i) trainer = DataParallelTrainer(train_loop_per_worker=train_func, scaling_config=scale_config) result = trainer.fit() assert result.checkpoint.to_dict()["model"] == 2
def test_scaling_config(ray_start_4_cpus): def train_func(): assert ray.available_resources()["CPU"] == 1 train.report(loss=1) assert ray.available_resources()["CPU"] == 4 trainer = DataParallelTrainer(train_loop_per_worker=train_func, scaling_config={"num_workers": 2}) trainer.fit()
def test_checkpoint(ray_start_4_cpus): def train_func(): for i in range(3): session.report({"epoch": i}, checkpoint=Checkpoint.from_dict({"model": i})) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=scale_config ) result = trainer.fit() assert result.checkpoint.to_dict()["model"] == 2
def test_fit_train_config(ray_start_4_cpus): def train_func(config): session.report({"loss": config["x"]}) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=scale_config, train_loop_config={"x": 100}, ) assert trainer.fit().metrics["loss"] == 100
def test_scaling_config(ray_start_4_cpus): def train_func(): assert ray.available_resources()["CPU"] == 1 session.report({"loss": 1}) assert ray.available_resources()["CPU"] == 4 trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2) ) trainer.fit()
def test_tune_error(ray_start_4_cpus): def train_func(config): raise RuntimeError("Error in training function!") trainer = DataParallelTrainer( train_func, backend_config=TestConfig(), scaling_config=ScalingConfig(num_workers=1), ) tuner = Tuner( trainer, ) result_grid = tuner.fit() with pytest.raises(RuntimeError): raise result_grid[0].error
def test_preprocessor_in_checkpoint(ray_start_4_cpus): class DummyPreprocessor(Preprocessor): def __init__(self): super().__init__() self.is_same = True def train_func(): for i in range(3): session.report({"epoch": i}, checkpoint=Checkpoint.from_dict({"model": i})) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=scale_config, preprocessor=DummyPreprocessor(), ) result = trainer.fit() assert result.checkpoint.to_dict()["model"] == 2 assert result.checkpoint.to_dict()[PREPROCESSOR_KEY].is_same
def test_bad_return_in_train_loop(ray_start_4_cpus): """Test to check if returns from train loop are discarded.""" # Simulates what happens with eg. torch models class FailOnUnpickle: def __reduce__(self): raise RuntimeError("Failing") def train_loop(config): session.report({"loss": 1}) return FailOnUnpickle() trainer = DataParallelTrainer( train_loop_per_worker=train_loop, scaling_config=scale_config ) # No exception should happen here trainer.fit()
def test_tune(ray_start_4_cpus): def train_func(config): session.report({"loss": config["x"]}) trainer = DataParallelTrainer( train_loop_per_worker=train_func, train_loop_config={"x": 100}, scaling_config=scale_config, ) tuner = Tuner( trainer, param_space={"train_loop_config": {"x": tune.choice([200, 300])}}, tune_config=TuneConfig(num_samples=2), ) result_grid = tuner.fit() assert result_grid[0].metrics["loss"] in [200, 300] # Make sure original Trainer is not affected. assert trainer._train_loop_config["x"] == 100
def test_tune_checkpoint(ray_start_4_cpus): def train_func(): for i in range(9): session.report(dict(test=i)) session.report( dict(test=i + 1), checkpoint=Checkpoint.from_dict(dict(hello="world")) ) trainer = DataParallelTrainer( train_func, backend_config=TestConfig(), scaling_config=ScalingConfig(num_workers=1), ) tuner = Tuner( trainer, param_space={"train_loop_config": {"max_iter": 5}}, ) [trial] = tuner.fit()._experiment_analysis.trials checkpoint_path = trial.checkpoint.dir_or_data assert os.path.exists(checkpoint_path) checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict() assert checkpoint["hello"] == "world"
def test_trainable_name_is_overriden_data_parallel_trainer(ray_start_4_cpus): trainer = DataParallelTrainer(lambda x: x, scaling_config=ScalingConfig(num_workers=1)) _is_trainable_name_overriden(trainer)
def test_invalid_train_loop(ray_start_4_cpus): def train_loop(config, extra_arg): pass with pytest.raises(ValueError): DataParallelTrainer(train_loop_per_worker=train_loop)