chore: import upstream snapshot with attribution

python/ray/tune/tests/test_trial.py

@@ -0,0 +1,156 @@
+import logging
+import sys
+
+import pytest
+
+from ray.exceptions import RayActorError, RayTaskError
+from ray.tests.conftest import propagate_logs  # noqa
+from ray.train._internal.session import _TrainingResult
+from ray.train._internal.storage import StorageContext
+from ray.train.constants import RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Checkpoint
+from ray.tune.experiment import Trial
+
+
+@pytest.fixture
+def trial(tmp_path):
+    yield Trial(
+        "mock",
+        stub=True,
+        storage=mock_storage_context(storage_path=str(tmp_path)),
+    )
+
+
+@pytest.mark.parametrize("count_preemption_errors", [False, True])
+def test_handle_preemption_error(
+    trial: Trial, count_preemption_errors: bool, monkeypatch
+):
+    """Check that the Trial counts preemption errors correctly."""
+    if count_preemption_errors:
+        monkeypatch.setenv(RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE, "1")
+
+    # Case 1: Directly raised (preemption) RayActorError
+    class PreemptionRayActorError(RayActorError):
+        def preempted(self) -> bool:
+            return True
+
+    err = PreemptionRayActorError()
+    trial.handle_error(err)
+    assert trial.num_failures == (1 if count_preemption_errors else 0)
+
+    # Case 2: RayTaskError, where the cause is a (preemption) RayActorError
+    wrapped_err = RayTaskError(
+        function_name="test", traceback_str="traceback_str", cause=err
+    )
+    trial.handle_error(wrapped_err)
+    assert trial.num_failures == (2 if count_preemption_errors else 0)
+
+    # Case 3: Non-preemption error
+    non_preempted_err = RayActorError()
+    trial.handle_error(non_preempted_err)
+    assert trial.num_failures == (3 if count_preemption_errors else 1)
+
+
+def test_load_trial_from_json_state():
+    """Check that serializing a trial to a JSON string with `Trial.get_json_state`
+    and then creating a new trial using the `Trial.from_json_state` alternate
+    constructor loads the trial with equivalent state."""
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        storage=mock_storage_context(),
+    )
+    trial.create_placement_group_factory()
+    trial.init_local_path()
+    trial.status = Trial.TERMINATED
+
+    # After loading, the trial state should be the same
+    json_state, _ = trial.get_json_state()
+    new_trial = Trial.from_json_state(json_state, stub=True)
+    assert new_trial.get_json_state()[0] == json_state
+
+
+def test_set_storage(tmp_path):
+    """Test that setting the trial's storage context will update the tracked
+    checkpoint paths."""
+    original_storage = mock_storage_context()
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        storage=original_storage,
+    )
+
+    result_1 = _TrainingResult(
+        checkpoint=Checkpoint.from_directory(original_storage.checkpoint_fs_path),
+        metrics={},
+    )
+    trial.on_checkpoint(result_1)
+
+    result_2 = _TrainingResult(
+        checkpoint=Checkpoint.from_directory(original_storage.checkpoint_fs_path),
+        metrics={},
+    )
+    trial.on_checkpoint(result_2)
+
+    new_storage = StorageContext(
+        storage_path=tmp_path / "new_storage_path",
+        experiment_dir_name="new_name",
+        trial_dir_name="new_trial",
+    )
+    trial.set_storage(new_storage)
+
+    assert result_1.checkpoint.path.startswith(new_storage.trial_fs_path)
+    assert result_2.checkpoint.path.startswith(new_storage.trial_fs_path)
+
+
+def test_trial_logdir_length():
+    """Test that trial local paths with a long logdir are truncated"""
+    trial = Trial(
+        trainable_name="none",
+        stub=True,
+        config={"a" * 50: 5.0 / 7, "b" * 50: "long" * 40},
+        storage=mock_storage_context(),
+    )
+    trial.init_local_path()
+    assert len(trial.storage.trial_dir_name) < 200
+
+
+def test_should_stop(caplog, propagate_logs):  # noqa
+    """Test whether `Trial.should_stop()` works as expected given a result dict."""
+    trial = Trial(
+        "MockTrainable",
+        stub=True,
+        trial_id="abcd1234",
+        stopping_criterion={"a": 10.0, "b/c": 20.0},
+    )
+
+    # Criterion is not reached yet -> don't stop.
+    result = {"a": 9.999, "b/c": 0.0, "some_other_key": True}
+    assert not trial.should_stop(result)
+
+    # Criterion is exactly reached -> stop.
+    result = {"a": 10.0, "b/c": 0.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Criterion is exceeded -> stop.
+    result = {"a": 10000.0, "b/c": 0.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Test nested criterion.
+    result = {"a": 5.0, "b/c": 1000.0, "some_other_key": False}
+    assert trial.should_stop(result)
+
+    # Test criterion NOT found in result metrics.
+    result = {"b/c": 1000.0}
+    with caplog.at_level(logging.WARNING):
+        trial.should_stop(result)
+    assert (
+        "Stopping criterion 'a' not found in result dict! Available keys are ['b/c']."
+    ) in caplog.text
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))