chore: import upstream snapshot with attribution

python/ray/tune/tests/execution/test_controller_control_integration.py

@@ -0,0 +1,147 @@
+import sys
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import PlacementGroupFactory, register_trainable
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+STORAGE = mock_storage_context()
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_stop_trial(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """Stopping a trial while RUNNING or PENDING should work.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testStopTrial
+    """
+
+    register_mock_trainable()
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
+    )
+    kwargs = {
+        "stopping_criterion": {"training_iteration": 10},
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 1}]),
+        "config": {"sleep": 1},
+        "storage": STORAGE,
+    }
+    trials = [
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+        Trial(MOCK_TRAINABLE_NAME, **kwargs),
+    ]
+    for t in trials:
+        runner.add_trial(t)
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until 2 trials started
+    while counter.get("RUNNING", 0) != 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("PENDING", 0) == 2
+
+    # Stop trial that is running
+    for trial in trials:
+        if trial.status == Trial.RUNNING:
+            runner._schedule_trial_stop(trial)
+            break
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until the next trial started
+    while counter.get("RUNNING", 0) < 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("TERMINATED", 0) == 1
+    assert counter.get("PENDING", 0) == 1
+
+    # Stop trial that is pending
+    for trial in trials:
+        if trial.status == Trial.PENDING:
+            runner._schedule_trial_stop(trial)
+            break
+
+    counter = Counter(t.status for t in trials)
+
+    # Wait until 2 trials are running again
+    while counter.get("RUNNING", 0) < 2:
+        runner.step()
+        counter = Counter(t.status for t in trials)
+
+    assert counter.get("RUNNING", 0) == 2
+    assert counter.get("TERMINATED", 0) == 2
+    assert counter.get("PENDING", 0) == 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_remove_actor_tracking(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
+    """When we reuse actors, actors that have been requested but not started
+    should not be tracked in ``_stopping_actors``.
+
+    When actors are re-used, we cancel original actor requests for the trial.
+    If these actors haven't been alive, there won't be a stop future to be resolved,
+    and thus they would remain in ``TuneController._stopping_actors`` until they
+    get cleaned up after 600 seconds.
+
+    This test asserts that these actors are not tracked in
+    ``TuneController._stopping_actors`` at all.
+
+    We start 4 actors, and one can run at a time. Actors are re-used across trials.
+    When the experiment ends, we expect that only one actor is left to track
+    in ``self._stopping_trials``.
+    """
+    runner = TuneController(
+        resource_manager_factory=lambda: resource_manager_cls(),
+        reuse_actors=True,
+        storage=STORAGE,
+    )
+
+    def train_fn(config):
+        return 1
+
+    register_trainable("test_remove_actor_tracking", train_fn)
+
+    kwargs = {
+        "placement_group_factory": PlacementGroupFactory([{"CPU": 4, "GPU": 2}]),
+        "storage": STORAGE,
+    }
+    trials = [Trial("test_remove_actor_tracking", **kwargs) for i in range(4)]
+    for t in trials:
+        runner.add_trial(t)
+
+    while not runner.is_finished():
+        runner.step()
+
+    # Only one actor should be left to stop
+    assert len(runner._stopping_actors) == 1
+
+    runner.cleanup()
+
+    assert len(runner._stopping_actors) == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))