chore: import upstream snapshot with attribution

python/ray/train/tests/test_result.py

@@ -0,0 +1,154 @@
+import pyarrow
+import pytest
+
+import ray
+from ray import train
+from ray.air._internal.uri_utils import URI
+from ray.air.constants import EXPR_RESULT_FILE
+from ray.train import CheckpointConfig, Result, RunConfig, ScalingConfig
+from ray.train.base_trainer import TrainingFailedError
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.torch import TorchTrainer
+from ray.tune import TuneConfig, Tuner
+
+_PARAM_SPACE = {"a": 1, "b": 2}
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def build_dummy_trainer(configs):
+    def worker_loop(_config):
+        for i in range(configs["NUM_ITERATIONS"]):
+            # Do some random reports in between checkpoints.
+            train.report({"metric_a": -100, "metric_b": -100})
+
+            if ray.train.get_context().get_world_rank() == 0:
+                with create_dict_checkpoint({"iter": i}) as checkpoint:
+                    train.report(
+                        metrics={"metric_a": i, "metric_b": -i},
+                        checkpoint=checkpoint,
+                    )
+            else:
+                train.report(metrics={"metric_a": i, "metric_b": -i})
+        raise RuntimeError()
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=worker_loop,
+        train_loop_config=_PARAM_SPACE,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
+        run_config=RunConfig(
+            name=configs["EXP_NAME"],
+            storage_path=configs["STORAGE_PATH"],
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=configs["NUM_CHECKPOINTS"],
+                checkpoint_score_attribute="metric_a",
+                checkpoint_score_order="max",
+            ),
+        ),
+    )
+    return trainer
+
+
+def build_dummy_tuner(configs):
+    return Tuner(
+        build_dummy_trainer(configs),
+        param_space={"train_loop_config": _PARAM_SPACE},
+        tune_config=TuneConfig(num_samples=1),
+    )
+
+
+@pytest.mark.parametrize("storage", ["local", "remote"])
+@pytest.mark.parametrize("mode", ["trainer", "tuner"])
+def test_result_restore(ray_start_4_cpus, tmpdir, mock_s3_bucket_uri, storage, mode):
+    NUM_ITERATIONS = 5
+    NUM_CHECKPOINTS = 3
+    if storage == "local":
+        storage_path = str(tmpdir)
+    elif storage == "remote":
+        storage_path = str(URI(mock_s3_bucket_uri))
+
+    exp_name = "test_result_restore"
+
+    configs = {
+        "EXP_NAME": exp_name,
+        "STORAGE_PATH": storage_path,
+        "NUM_ITERATIONS": NUM_ITERATIONS,
+        "NUM_CHECKPOINTS": NUM_CHECKPOINTS,
+    }
+
+    if mode == "trainer":
+        trainer = build_dummy_trainer(configs)
+        with pytest.raises(TrainingFailedError):
+            trainer.fit()
+    elif mode == "tuner":
+        tuner = build_dummy_tuner(configs)
+        tuner.fit()
+
+    # Find the trial directory to restore
+    exp_dir = str(URI(storage_path) / exp_name)
+    fs, fs_exp_dir = pyarrow.fs.FileSystem.from_uri(exp_dir)
+    for item in fs.get_file_info(pyarrow.fs.FileSelector(fs_exp_dir)):
+        if item.type == pyarrow.fs.FileType.Directory and item.base_name.startswith(
+            "TorchTrainer"
+        ):
+            trial_dir = str(URI(exp_dir) / item.base_name)
+            break
+
+    # [1] Restore from path
+    result = Result.from_path(trial_dir)
+
+    # Check if we restored all checkpoints
+    assert result.checkpoint
+    assert len(result.best_checkpoints) == NUM_CHECKPOINTS
+
+    """
+    Top-3 checkpoints with metrics:
+
+                        | iter   | metric_a    metric_b
+    checkpoint_000004        4            4          -4
+    checkpoint_000003        3            3          -3
+    checkpoint_000002        2            2          -2
+    """
+    # Check if the checkpoints bounded with correct metrics
+    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
+    assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
+
+    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
+    assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
+
+    with pytest.raises(RuntimeError, match="Invalid metric name.*"):
+        result.get_best_checkpoint(metric="invalid_metric", mode="max")
+
+    # Check if we properly restored errors
+    assert isinstance(result.error, RuntimeError)
+
+    # Check that the config is properly formatted in the result metrics
+    assert result.metrics.get("config") == {"train_loop_config": _PARAM_SPACE}
+
+    # [2] Restore from path without result.json
+    fs.delete_file((URI(trial_dir) / EXPR_RESULT_FILE).path)
+    result = Result.from_path(trial_dir)
+
+    # Do the same checks as above
+    assert result.checkpoint
+    assert len(result.best_checkpoints) == NUM_CHECKPOINTS
+
+    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
+    assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
+
+    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
+    assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
+
+    assert isinstance(result.error, RuntimeError)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))