chore: import upstream snapshot with attribution

python/ray/air/tests/_test_experiment_restore_run.py

@@ -0,0 +1,180 @@
+import collections
+import json
+import os
+import random
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import ray
+from ray import train, tune
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.tune.experiment import Trial
+
+RUNNER_TYPE = os.environ.get("RUNNER_TYPE", "trainer")
+STORAGE_PATH = os.environ.get("STORAGE_PATH", "/tmp/ray_results")
+EXP_NAME = os.environ.get("EXP_NAME", "restore_integration_test")
+CALLBACK_DUMP_FILE = os.environ.get(
+    "CALLBACK_DUMP_FILE", "/tmp/callback_dump_file.json"
+)
+CSV_DATA_FILE = os.environ.get("CSV_DATA_FILE", "/tmp/dummy.csv")
+
+TIME_PER_ITER_S = float(os.environ.get("TIME_PER_ITER_S", "0.5"))
+NUM_TRIALS = int(os.environ.get("NUM_TRIALS", "1"))
+MAX_CONCURRENT_TRIALS = int(os.environ.get("MAX_CONCURRENT_TRIALS", "2"))
+ITERATIONS_PER_TRIAL = int(os.environ.get("ITERATIONS_PER_TRIAL", "64"))
+
+
+class StatefulCallback(tune.Callback):
+    def __init__(self):
+        self._trial_iterations = collections.defaultdict(list)
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        result: Dict,
+        **info,
+    ):
+        self._trial_iterations[trial.trial_id].append(result["training_iteration"])
+
+    def on_experiment_end(self, trials: List["Trial"], **info):
+        # Save callback contents to file
+        with open(CALLBACK_DUMP_FILE, "w") as f:
+            json.dump(self.get_state(), f, indent=2)
+
+    def get_state(self) -> Optional[Dict]:
+        return {"trial_iters": self._trial_iterations.copy()}
+
+    def set_state(self, state: Dict):
+        self._trial_iterations = state["trial_iters"]
+
+
+class StatefulSearcher(tune.search.Searcher):
+    def __init__(
+        self,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+    ):
+        super().__init__(metric=metric, mode=mode)
+        self._trial_count = 0
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        self._trial_count += 1
+        return {"id": self._trial_count}
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ) -> None:
+        pass
+
+    def save(self, checkpoint_path: str):
+        with open(checkpoint_path, "w") as f:
+            json.dump({"trial_count": self._trial_count}, f)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "r") as f:
+            state = json.load(f)
+        self._trial_count = state["trial_count"]
+
+
+def train_fn(config: dict, data: Optional[dict] = None):
+    checkpoint = train.get_checkpoint()
+    start = load_dict_checkpoint(checkpoint)["iteration"] + 1 if checkpoint else 1
+
+    training_started_marker = Path(
+        os.environ.get("RUN_STARTED_MARKER", "/tmp/does-not-exist")
+    )
+    if training_started_marker.exists():
+        # Multiple workers may be trying to delete the same marker
+        try:
+            training_started_marker.unlink()
+        except FileNotFoundError:
+            pass
+
+    for iteration in range(start, ITERATIONS_PER_TRIAL + 1):
+        time.sleep(TIME_PER_ITER_S)
+
+        with create_dict_checkpoint({"iteration": iteration}) as checkpoint:
+            train.report({"score": random.random()}, checkpoint=checkpoint)
+
+
+def tuner(experiment_path: str, run_config: tune.RunConfig) -> tune.ResultGrid:
+    trainable = tune.with_resources(train_fn, resources={"CPU": 1})
+    trainable = tune.with_parameters(trainable, data={"dummy_data": [1, 2, 3]})
+
+    if tune.Tuner.can_restore(experiment_path):
+        tuner = tune.Tuner.restore(
+            experiment_path, trainable=trainable, resume_errored=True
+        )
+    else:
+        tuner = tune.Tuner(
+            trainable,
+            run_config=run_config,
+            tune_config=tune.TuneConfig(
+                num_samples=8,
+                max_concurrent_trials=2,
+                search_alg=StatefulSearcher(),
+            ),
+        )
+
+    result_grid = tuner.fit()
+    return result_grid
+
+
+def trainer(experiment_path: str, run_config: train.RunConfig) -> train.Result:
+    dataset_size = 128
+    num_workers = 4
+
+    def train_loop_per_worker(config):
+        # Wrap the other train_fn with a check for the dataset.
+        assert train.get_dataset_shard("train")
+        train_fn(config)
+
+    datasets = {
+        "train": ray.data.range(dataset_size),
+        "valid": ray.data.read_csv(CSV_DATA_FILE),
+    }
+
+    if DataParallelTrainer.can_restore(experiment_path):
+        trainer = DataParallelTrainer.restore(
+            experiment_path,
+            datasets=datasets,
+            train_loop_per_worker=train_loop_per_worker,
+        )
+    else:
+        trainer = DataParallelTrainer(
+            train_loop_per_worker,
+            datasets=datasets,
+            scaling_config=train.ScalingConfig(
+                num_workers=num_workers, trainer_resources={"CPU": 0}
+            ),
+            run_config=run_config,
+        )
+
+    result = trainer.fit()
+    return result
+
+
+if __name__ == "__main__":
+    experiment_path = os.path.join(STORAGE_PATH, EXP_NAME)
+
+    ray.init()
+
+    run_config = train.RunConfig(
+        storage_path=STORAGE_PATH,
+        name=EXP_NAME,
+        checkpoint_config=train.CheckpointConfig(num_to_keep=1),
+        callbacks=[StatefulCallback()],
+    )
+
+    if RUNNER_TYPE == "tuner":
+        tuner(experiment_path, run_config)
+    elif RUNNER_TYPE == "trainer":
+        trainer(experiment_path, run_config)
+    else:
+        raise NotImplementedError(
+            "`RUNNER_TYPE` environment var must be one of ['tuner', 'trainer']"
+        )

python/ray/air/tests/conftest.py

@@ -0,0 +1,23 @@
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+import copy
+
+import pytest
+
+import ray
+from ray.tests.conftest import pytest_runtest_makereport  # noqa
+
+
+@pytest.fixture
+def restore_data_context(request):
+    """Restore any DataContext changes after the test runs"""
+    original = copy.deepcopy(ray.data.context.DataContext.get_current())
+    yield
+    ray.data.context.DataContext._set_current(original)
+
+
+@pytest.fixture
+def disable_fallback_to_object_extension(request, restore_data_context):
+    """Disables fallback to ArrowPythonObjectType"""
+    ray.data.context.DataContext.get_current().enable_fallback_to_arrow_object_ext_type = (
+        False
+    )

python/ray/air/tests/execution/test_barrier.py

@@ -0,0 +1,54 @@
+from typing import Optional, Type
+
+import pytest
+
+from ray.air.execution._internal.barrier import Barrier
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+def test_barrier_max_results():
+    """Test the `max_results` attribute.
+
+    - Set max_results=10
+    - Assert that the barrier completion callback is not invoked with num_results<10
+    - Assert that callback is invoked with num_results=10
+    - Assert that callback is not invoked again when more events arrive
+    - Assert that more events can arrive without triggering the callback after resetting
+    """
+    barrier = Barrier(max_results=10, on_completion=_raise(AssertionError))
+
+    for i in range(9):
+        barrier.arrive(i)
+
+    assert not barrier.completed
+
+    # Will trigger the on_completion callback
+    with pytest.raises(AssertionError):
+        barrier.arrive(10)
+
+    assert barrier.completed
+
+    assert barrier.num_results == 10
+
+    # Further events will not trigger callback again
+    barrier.arrive(11)
+
+    barrier.reset()
+
+    assert not barrier.completed
+
+    # After flushing more events can arrive
+    barrier.arrive(12)
+    assert barrier.num_results == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_e2e_train_flow.py

@@ -0,0 +1,266 @@
+import random
+from typing import Any, List, Optional
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in these ways:
+
+    1. On init if ``actor_init_kill`` is passed as a kwarg
+    2. On setup_1() if ``actor_setup_kill`` is passed as a kwarg (RayActorError)
+    3. On setup_1() if ``actor_setup_fail`` is passed as a kwarg (RayTaskError)
+    4. On train() if ``actor_train_kill`` is passed as a kwarg (RayTaskError)
+    5. On train() if ``actor_train_fail`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_init_kill"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def setup_1(self):
+        if self.kwargs.get("actor_setup_kill"):
+            raise SystemExit
+
+        if self.kwargs.get("actor_setup_fail"):
+            raise RuntimeError("Setup")
+
+        return True
+
+    def setup_2(self):
+        return True
+
+    def train(self, value: float) -> float:
+        if value == 4:
+            if self.kwargs.get("actor_train_kill"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("actor_train_fail"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TrainFlow:
+    """This is a Ray Train-like execution flow.
+
+    - We want to run 4 actors in total ("trials")
+    - Each actor runs two init functions
+    - We train all actors in parallel for 10 iterations
+    - Errors can come up on actor construction, in the init functions,
+      or during training
+    - When an actor fails, restart that actor
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 4
+        self._tracked_actors = []
+        self._actors_stopped = 0
+
+        self._actors_to_replace = set()
+
+        self._ready_actors = set()
+        self._training_barrier = Barrier(
+            max_results=self._actors_to_run,
+            on_completion=self.training_barrier_completed,
+        )
+        self._restart_training = None
+
+        self._training_iter = 0
+        self._results = []
+
+        self._errors = errors
+
+    def setup_actors(self):
+        for actor_id in range(self._actors_to_run):
+            error_kwargs = {}
+            if self._errors:
+                error = random.choice(self._errors)
+                error_kwargs[error] = True
+
+            print("Actor", actor_id, "will be failing with", error_kwargs)
+
+            tracked_actor = self._actor_manager.add_actor(
+                cls=Actor,
+                kwargs={"id": actor_id, **error_kwargs},
+                resource_request=ResourceRequest([{"CPU": 1}]),
+                on_start=self.actor_started,
+                on_stop=self.actor_stopped,
+                on_error=self.actor_error,
+            )
+            self._tracked_actors.append(tracked_actor)
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_1",
+            on_error=self.setup_error,
+            on_result=self.setup_1_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._ready_actors.discard(tracked_actor)
+
+        if tracked_actor in self._actors_to_replace:
+            self._replace_actor(tracked_actor=tracked_actor)
+        else:
+            self._actors_stopped += 1
+            self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._ready_actors.discard(tracked_actor)
+        self._replace_actor(tracked_actor=tracked_actor)
+
+    def _replace_actor(self, tracked_actor: TrackedActor):
+        actor_index = self._tracked_actors.index(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_index},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._tracked_actors[actor_index] = replacement_actor
+
+    def setup_1_result(self, tracked_actor: TrackedActor, result: Any):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_2",
+            on_error=self.setup_error,
+            on_result=self.setup_2_result,
+        )
+
+    def setup_2_result(self, tracked_actor: TrackedActor, result: Any):
+        self._ready_actors.add(tracked_actor)
+
+        if len(self._ready_actors) == self._actors_to_run:
+            self.continue_training()
+
+    def setup_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def continue_training(self):
+        if self._restart_training:
+            self._training_iter = self._restart_training
+        else:
+            self._training_iter += 1
+
+        self._training_barrier.reset()
+        self._actor_manager.schedule_actor_tasks(
+            self._tracked_actors,
+            "train",
+            args=(self._training_iter,),
+            on_result=self._training_barrier.arrive,
+            on_error=self.training_error,
+        )
+
+    def training_barrier_completed(self, barrier: Barrier):
+        self._results.append([res for _, res in barrier.get_results()])
+        self._restart_training = None
+
+        # If less than 10 epochs, continue training
+        if self._training_iter < 10:
+            return self.continue_training()
+
+        # Else, training finished
+        for tracked_actor in self._tracked_actors:
+            self._actor_manager.remove_actor(tracked_actor)
+
+    def training_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._restart_training = self._training_iter
+
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._ready_actors.discard(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def run(self):
+        self.setup_actors()
+
+        while not self._finished:
+            self._actor_manager.next()
+
+    def get_results(self) -> List[List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_init_kill",
+        "actor_setup_kill",
+        "actor_setup_fail",
+        "actor_train_kill",
+        "actor_train_fail",
+        # Chaos - every actor fails somehow, but in different ways
+        [
+            "actor_init_kill",
+            "actor_setup_kill",
+            "actor_setup_fail",
+            "actor_train_kill",
+            "actor_train_fail",
+        ],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TrainFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert results == [[i] * 4 for i in range(1, 11)], results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_e2e_tune_flow.py

@@ -0,0 +1,227 @@
+import random
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in three ways:
+
+    1. On init if ``actor_error_init`` is passed as a kwarg
+    2. On run() if ``actor_error_task`` is passed as a kwarg (RayActorError)
+    3. On run() if ``task_error`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_error_init"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def run(self, value: float) -> float:
+        if value == 2:
+            if self.kwargs.get("actor_error_task"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("task_error"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TuneFlow:
+    """This is a Ray Tune-like execution flow.
+
+    - We want to run 10 actors in total ("trials")
+    - Each actor collects 11 results sequentially
+    - We schedule up to 6 actors at the same time
+    - Every step, we see if we should add any new actors
+    - Otherwise, we just yield control to the event manager and process events one
+      by one
+    - When an actor is started, start training flow
+    - When a result comes in, schedule next future
+      - If this is the 11th result, stop actor
+    - When the last actor is stopped, set state to finished
+
+    - When an actor fails, restart
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 10
+        self._actors_started = 0
+        self._actors_stopped = 0
+        self._max_pending = 6
+
+        self._actor_to_id = {}
+        self._results = defaultdict(list)
+
+        self._errors = errors
+
+    def maybe_add_actors(self):
+        if self._actors_started >= self._actors_to_run:
+            return
+
+        if self._actor_manager.num_pending_actors >= self._max_pending:
+            return
+
+        error_kwargs = {}
+        if self._errors:
+            error = random.choice(self._errors)
+            error_kwargs[error] = True
+
+        actor_id = self._actors_started
+
+        print("Actor", actor_id, "will be failing with", error_kwargs)
+
+        tracked_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_id, **error_kwargs},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[tracked_actor] = actor_id
+
+        self._actors_started += 1
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "run",
+            kwargs={"value": 0},
+            on_error=self.task_error,
+            on_result=self.task_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._actors_stopped += 1
+        self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def task_result(self, tracked_actor: TrackedActor, result: float):
+        actor_id = self._actor_to_id[tracked_actor]
+        self._results[actor_id].append(result)
+
+        if result == 10:
+            self._actor_manager.remove_actor(tracked_actor)
+        else:
+            self._actor_manager.schedule_actor_task(
+                tracked_actor,
+                "run",
+                kwargs={"value": result + 1},
+                on_result=self.task_result,
+                on_error=self.task_error,
+            )
+
+    def task_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_stopped -= 1  # account for extra stop
+        self._actor_manager.remove_actor(tracked_actor)
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def run(self):
+        while not self._finished:
+            self.maybe_add_actors()
+            self._actor_manager.next(timeout=1)
+
+    def get_results(self) -> Dict[int, List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_error_init",
+        "actor_error_task",
+        "task_error",
+        # Chaos - every actor fails somehow, but in different ways
+        ["actor_error_init", "actor_error_task", "task_error"],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TuneFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert all(res[-1] == 10 for res in results.values()), results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_event_manager.py

@@ -0,0 +1,224 @@
+import time
+from typing import Any, Type
+
+import pytest
+
+import ray
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.event_manager import RayEventManager
+from ray.exceptions import RayTaskError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+def succeeding(ret: Any = None) -> Any:
+    return ret
+
+
+@ray.remote
+def failing(exc: Type[Exception], *args) -> None:
+    raise exc(*args)
+
+
+@ray.remote
+def sleeping(seconds: int, result: Any) -> Any:
+    time.sleep(seconds)
+    return result
+
+
+def test_track_future_success(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that the on_result callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_future(succeeding.remote("a"), on_result=on_result)
+
+    event_manager.wait()
+    assert "a" in seen
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_success_no_callback(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that passing no callback still succeeds.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(succeeding.remote("a"))
+
+    event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that the on_error callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    class CustomError(RuntimeError):
+        pass
+
+    def on_error(exception: Exception):
+        seen.add(exception)
+
+    event_manager.track_future(failing.remote(CustomError), on_error=on_error)
+
+    event_manager.wait()
+    assert isinstance(seen.pop(), CustomError)
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error_no_callback(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that passing no callback raises the original error.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(failing.remote(RuntimeError))
+
+    with pytest.raises(RuntimeError):
+        event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+@pytest.mark.parametrize("results_per_wait", [None, 1, 5, 10, 100])
+def test_many_futures(ray_start_4_cpus, results_per_wait):
+    """Schedule 500 succeeding and failing futures.
+
+    Check that the callbacks get triggered correctly, independent of the number
+    of results we await per call to RayEventManager.wait().
+    """
+    num_futures = 500
+
+    event_manager = RayEventManager()
+
+    seen_results = set()
+    seen_errors = set()
+
+    def on_result(result: Any):
+        seen_results.add(result)
+
+    def on_error(exception: RayTaskError):
+        seen_errors.add(exception.cause.args[0])
+
+    for i in range(num_futures):
+        event_manager.track_futures(
+            [
+                succeeding.remote("a" + str(i)),
+                failing.remote(RuntimeError, "b" + str(i)),
+            ],
+            on_result=on_result,
+            on_error=on_error,
+        )
+
+    while event_manager.num_futures > 0:
+        event_manager.wait(num_results=results_per_wait)
+
+    for i in range(num_futures):
+        assert "a" + str(i) in seen_results
+        assert "b" + str(i) in seen_errors
+
+
+def test_timeout(ray_start_4_cpus):
+    """Test the timeout parameter.
+
+    Start 4 tasks: Two succeed immediately, two after 1 second.
+
+    After waiting for 0.5 seconds, the first two tasks should have returned.
+    After waiting for up to 5 seconds, the other two tasks should have returned.
+    But because the tasks take only 0.5 seconds to run, we should have waited
+    way less than 5 seconds.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            sleeping.remote(1, "c"),
+            sleeping.remote(1, "d"),
+        ],
+        on_result=on_result,
+    )
+
+    start = time.monotonic()
+    event_manager.wait(num_results=None, timeout=0.5)
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" not in seen
+    assert "d" not in seen
+
+    event_manager.wait(num_results=None, timeout=5)
+    taken = time.monotonic() - start
+
+    assert "c" in seen
+    assert "d" in seen
+
+    # Should have returned much earlier than after 5 seconds
+    assert taken < 3
+
+    assert not event_manager._tracked_futures
+
+
+def test_task_barrier(ray_start_4_cpus):
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_completion(barrier: Barrier):
+        seen.update(barrier.get_results())
+
+    barrier = Barrier(max_results=4, on_completion=on_completion)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            succeeding.remote("c"),
+            succeeding.remote("d"),
+            sleeping.remote(2, "e"),
+        ],
+        on_result=barrier.arrive,
+    )
+
+    event_manager.wait(num_results=4)
+
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" in seen
+    assert "d" in seen
+    assert "e" not in seen
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_manager_fixed.py

@@ -0,0 +1,178 @@
+import pytest
+
+import ray
+from ray.air.execution.resources.fixed import FixedResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_4_CPU = ResourceRequest([{"CPU": 4}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture(scope="module")
+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 test_acquire_return_resources(ray_start_4_cpus):
+    manager = FixedResourceManager(total_resources={"CPU": 4})
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_4_CPU)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+    ready_2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.free_resources(ready_2)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+
+def test_numerical_error(ray_start_4_cpus):
+    """Make sure we don't run into numerical errors when using fractional resources.
+
+    Legacy test: test_trial_runner::TrialRunnerTest::testResourceNumericalError
+    """
+    manager = FixedResourceManager(
+        total_resources={"CPU": 0.99, "GPU": 0.99, "a": 0.99}
+    )
+    resource_request = ResourceRequest([{"CPU": 0.33, "GPU": 0.33, "a": 0.33}])
+
+    for i in range(3):
+        manager.request_resources(resource_request)
+        assert manager.acquire_resources(
+            resource_request=resource_request
+        ), manager._available_resources
+
+    assert manager._available_resources["CPU"] == 0
+    assert manager._available_resources["GPU"] == 0
+    assert manager._available_resources["a"] == 0
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request resources with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    manager.request_resources(REQUEST_1_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request resources with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT placement strategies."""
+    manager = FixedResourceManager()
+
+    req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+    if strategy.startswith("STRICT_"):
+        with pytest.raises(RuntimeError):
+            manager.request_resources(req)
+    else:
+        manager.request_resources(req)
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy_nested(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT_SPREAD within a PG."""
+
+    @ray.remote
+    def nested_test():
+        manager = FixedResourceManager()
+
+        req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+        if strategy == "STRICT_SPREAD":
+            with pytest.raises(RuntimeError):
+                manager.request_resources(req)
+        else:
+            manager.request_resources(req)
+
+    pg = ray.util.placement_group([{"CPU": 2}])
+    ray.wait([pg.ready()])
+
+    try:
+        ray.get(
+            nested_test.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg, placement_group_capture_child_tasks=True
+                )
+            ).remote()
+        )
+    finally:
+        ray.util.remove_placement_group(pg)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_manager_placement_group.py

@@ -0,0 +1,409 @@
+import time
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 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 _count_pg_states():
+    counter = Counter()
+    for _, pg_info in ray.util.placement_group_table().items():
+        counter[pg_info["state"]] += 1
+    return counter
+
+
+def test_request_cancel_resources(ray_start_4_cpus):
+    """Test that canceling a resource request clears the PG futures.
+
+    - Create request
+    - Assert actual PG is created
+    - Cancel request
+    - Assert staging future is removed
+    - Assert actual PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Could be pending or created
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    assert manager.get_resource_futures()
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    assert not manager.get_resource_futures()
+
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    """Tests that acquiring and returning resources works.
+
+    - At the start, no resources should be ready (no PG scheduled)
+    - Request resources for 2 CPUs
+    - (wait until they are ready)
+    - Assert that these 2 CPUs are available to be acquired
+    - Acquire
+    - Assert that there are no 2 CPU resources available anymore
+    - Free resources
+    - Assert that the 2 CPU resources are still not available (no new request)
+        - This is also tested in includes test_request_cancel_resources
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Request PG
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    # Acquire PG
+    acquired = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Free resources
+    manager.free_resources(acquired)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG still exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_request_pending(ray_start_4_cpus):
+    """Test that requesting too many resources leads to pending PGs.
+
+    - Cluster of 4 CPUs
+    - Request 3 PGs a 2 CPUs
+    - Acquire 2 PGs
+    - Assert no resources are available anymore
+    - Return both PGs
+    - Assert resources are available again
+    - Cancel request
+    - Assert no resources are available again
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until some are ready
+    ray.wait(manager.get_resource_futures(), num_returns=2)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert len(manager.get_resource_futures()) == 1
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 2
+    assert pg_states["PENDING"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    acq1 = manager.acquire_resources(REQUEST_2_CPU)
+    acq2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.free_resources(acq1)
+    manager.free_resources(acq2)
+
+    # Third PG becomes ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 2
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 3
+
+
+def test_acquire_unavailable(ray_start_4_cpus):
+    """Test that acquiring resources that are not available returns None.
+
+    - Try to acquire
+    - Assert this does not work
+    - Request resources
+    - Wait until ready
+    - Acquire
+    - Assert this did work
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.acquire_resources(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.acquire_resources(REQUEST_2_CPU)
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {"CPU": 1}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {"CPU": 1}
+    assert res2 == {"CPU": 2}
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request PG with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {}
+    assert res2 == {"CPU": 2}
+
+
+def test_capture_child_tasks(ray_start_4_cpus):
+    """Test that child tasks are captured when creating placement groups.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind a remote task that needs 2 CPUs to run
+    - Assert that it can be scheduled from within the first bundle
+
+    This is only the case if child tasks are captured in the placement groups, as
+    there is only 1 CPU available outside (on a 4 CPU cluster). The 2 CPUs
+    thus have to come from the placement group.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def needs_cpus():
+        return "Ok"
+
+    @ray.remote
+    def spawn_child_task(num_cpus: int):
+        return ray.get(needs_cpus.options(num_cpus=num_cpus).remote())
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([spawn_child_task])
+
+    res = ray.get(av1.remote(2), timeout=2.0)
+
+    assert res
+
+
+def test_clear_state(ray_start_4_cpus):
+    """Test that clearing state will remove existing placement groups.
+
+    - Create resource request
+    - Wait until PG is scheduled
+    - Assert that Ray PG is created
+    - Call `mgr.clear()`
+    - Assert that resources are not ready anymore
+    - Assert that Ray PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 0
+
+    manager.clear()
+
+    assert not manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_internal_state(ray_start_4_cpus):
+    """Test internal state mappings of the placement group manager.
+
+    This test makes assumptions and assertions around the internal state transition
+    of private properties of the placement group resource manager.
+
+    If you change internal handling logic of the manager, you may need to change this
+    test as well.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+
+    assert manager.update_interval_s == 0
+
+    manager.has_resources_ready(REQUEST_2_CPU)
+
+    # The key may exist but the set should be empty
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+
+    ####
+    # 1. Request, wait until ready, cancel
+
+    # Request resources
+    manager.request_resources(REQUEST_2_CPU)
+
+    # PG should be staged
+    assert manager._request_to_staged_pgs[REQUEST_2_CPU]
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    assert manager._pg_to_request[pg] == REQUEST_2_CPU
+
+    # Staging future should exist
+    assert manager._pg_to_staging_future[pg]
+    fut = manager._pg_to_staging_future[pg]
+    assert manager._staging_future_to_pg[fut] == pg
+
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+
+    # PG should now be ready
+    assert manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # PG should not be staged anymore
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    # Staging future should not exist anymore
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+
+    # Cancel request
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    # PG should not be ready anymore
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # All PGs should be fully removed
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, cancel while staging
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    # Cancel request before it's ready
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    # Assert no leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, acquire, free
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+    # Acquire
+    acquired_resources = manager.acquire_resources(resource_request=REQUEST_2_CPU)
+    # Assert no staging/ready leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # We still retain this mapping
+    assert manager._pg_to_request
+    # And we keep track of acquired PGs
+    assert pg in manager._acquired_pgs
+
+    # Free PG
+    manager.free_resources(acquired_resources)
+    # State should be cleared now
+    assert not manager._pg_to_request
+    assert not manager._acquired_pgs
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_request.py

@@ -0,0 +1,38 @@
+import pytest
+
+from ray.air.execution.resources.request import ResourceRequest
+
+
+def test_request_same():
+    """Test that resource requests are the same if they share the same properties."""
+
+    assert ResourceRequest([{"CPU": 1}]) == ResourceRequest([{"CPU": 1}])
+
+    # multiple bundles work
+    assert ResourceRequest([{"CPU": 1}, {"CPU": 2}]) == ResourceRequest(
+        [{"CPU": 1}, {"CPU": 2}]
+    )
+
+    # multiple resources work
+    assert ResourceRequest([{"CPU": 1, "GPU": 1}]) == ResourceRequest(
+        [{"CPU": 1, "GPU": 1}]
+    )
+
+    # 0 resources are ignored
+    assert ResourceRequest([{"CPU": 0, "GPU": 1}]) == ResourceRequest([{"GPU": 1}])
+
+    # PACK is implicit
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") == ResourceRequest(
+        [{"CPU": 1}]
+    )
+
+    # Non match: different strategy
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") != ResourceRequest(
+        [{"CPU": 1}], strategy="SPREAD"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_tracked_actor.py

@@ -0,0 +1,366 @@
+import gc
+import threading
+import time
+from collections import Counter
+from typing import Any, Optional, Type
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+class Started(RuntimeError):
+    pass
+
+
+class Stopped(RuntimeError):
+    pass
+
+
+class Failed(RuntimeError):
+    pass
+
+
+class Result(RuntimeError):
+    pass
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def cleanup():
+    # Garbage collect at the start
+    # This ensures that all resources are freed up for the upcoming test.
+    gc.collect()
+    yield
+
+
+class Actor:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def task(self, value: Any):
+        return value
+
+
+@ray.remote(num_cpus=4)
+def fn():
+    return True
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("actor_cls", [Actor, ray.remote(Actor)])
+@pytest.mark.parametrize("kill", [False, True])
+def test_start_stop_actor(ray_start_4_cpus, resource_manager_cls, actor_cls, kill):
+    """Test that starting and stopping actors work and invokes a callback.
+
+    - Start an actor
+    - Starting should trigger start callback
+    - Schedule actor task, which should resolve (meaning actor successfully started)
+    - Stop actor, which should resolve and trigger stop callback
+    - Schedule remote fn that takes up all cluster resources. This should resolve,
+      meaning that the actor was stopped successfully.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # Start actor, set callbacks
+    tracked_actor = actor_manager.add_actor(
+        cls=actor_cls,
+        kwargs={"key": "val"},
+        resource_request=ResourceRequest([{"CPU": 4}]),
+        on_start=_raise(Started),
+        on_stop=_raise(Stopped),
+        on_error=_raise(Failed),
+    )
+
+    # Actor should be started
+    with pytest.raises(Started):
+        actor_manager.next()
+
+    # Schedule task on actor which should resolve (actor successfully started)
+    actor_manager.schedule_actor_task(
+        tracked_actor, "task", (1,), on_result=_raise(Result)
+    )
+
+    with pytest.raises(Result):
+        actor_manager.next()
+
+    # Now we can assert that there are no CPUS resources available anymore.
+    # Note that actor starting is asynchronous, so we can't assert this right away
+    # - that's why we wait for the actor task to resolve first.
+    assert ray.available_resources().get("CPU", 0.0) == 0, ray.available_resources()
+
+    # Stop actor
+    actor_manager.remove_actor(tracked_actor, kill=kill)
+
+    with pytest.raises(Stopped):
+        actor_manager.next()
+
+    # This task takes up all the cluster resources. It should resolve now that
+    # the actor was terminated.
+    assert ray.get(fn.remote(), timeout=5)
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_start_many_actors(ray_start_4_cpus, resource_manager_cls):
+    """Test that starting more actors than fit onto the cluster works.
+
+    - Request 10 actors
+    - 4 can be started. Assert they are started
+    - Stop 2
+    - Assert 2 are stopped and 2 new ones are started
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    running_actors = []
+    # stats keeps track of started/stopped actors
+    stats = Counter()
+
+    def start_callback(tracked_actor):
+        running_actors.append(tracked_actor)
+        stats["started"] += 1
+
+    def stop_callback(tracked_actor):
+        running_actors.remove(tracked_actor)
+        stats["stopped"] += 1
+
+    # start 10 actors
+    expected_actors = []
+    for i in range(10):
+        tracked_actor = actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"key": "val"},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=start_callback,
+            on_stop=stop_callback,
+            on_error=_raise(Failed),
+        )
+        expected_actors.append(tracked_actor)
+
+    # wait for some actor starts
+    for i in range(4):
+        actor_manager.next()
+
+    # we should now have 4 started actors
+    assert stats["started"] == 4
+    assert stats["stopped"] == 0
+    assert len(running_actors) == 4
+    assert set(running_actors) == set(expected_actors[:4])
+
+    # stop 2 actors
+    actor_manager.remove_actor(running_actors[0])
+    actor_manager.remove_actor(running_actors[1])
+
+    # Wait four times, twice for termination, twice for start
+    for i in range(4):
+        actor_manager.next()
+
+    # we should have 4 running actors, 6 started and 2 stopped
+    assert stats["started"] == 6
+    assert stats["stopped"] == 2
+    assert len(running_actors) == 4
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("where", ["init", "fn"])
+def test_actor_fail(ray_start_4_cpus, cleanup, resource_manager_cls, where):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # keep track of failed tasks and actors
+    stats = Counter()
+
+    @ray.remote
+    class FailingActor:
+        def __init__(self, where):
+            self._where = where
+            if self._where == "init":
+                raise RuntimeError("INIT")
+
+        def fn(self):
+            if self._where == "fn":
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+            return True
+
+    def fail_callback_actor(tracked_actor, exception):
+        stats["failed_actor"] += 1
+
+    def fail_callback_task(tracked_actor, exception):
+        stats["failed_task"] += 1
+
+    # Start actor
+    tracked_actor = actor_manager.add_actor(
+        cls=FailingActor,
+        kwargs={"where": where},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_error=fail_callback_actor,
+    )
+
+    if where != "init":
+        # Wait until it is started. This won't invoke any callback, yet
+        actor_manager.next()
+
+        assert stats["failed_actor"] == 0
+        assert stats["failed_task"] == 0
+
+        # Schedule task
+        actor_manager.schedule_actor_task(
+            tracked_actor, "fn", on_error=fail_callback_task
+        )
+
+    # Yield control and wait for task resolution. This will invoke the callback.
+    actor_manager.next()
+
+    assert stats["failed_actor"] == 1
+    assert stats["failed_task"] == bool(where != "init")
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_stop_actor_before_start(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls
+):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    @ray.remote
+    class HangingActor:
+        def __init__(self):
+            while not hang_marker.exists():
+                time.sleep(0.05)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=_raise(RuntimeError, "Should not have started"),
+        on_stop=_raise(RuntimeError, "Should not have stopped"),
+    )
+    while not actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    # Actor started but hasn't triggered on_start, yet
+    actor_manager.remove_actor(tracked_actor)
+    hang_marker.write_text("")
+    while actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    assert actor_manager.num_live_actors == 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("start_thread", [False, True])
+def test_stop_actor_custom_future(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls, start_thread
+):
+    """If we pass a custom stop future, the actor should still be shutdown by GC.
+
+    This should also be the case when we start a thread in the background, as we
+    do e.g. in Ray Tune's function runner.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    actor_name = f"stopping_actor_{resource_manager_cls.__name__}_{start_thread}"
+
+    @ray.remote(name=actor_name)
+    class HangingStopActor:
+        def __init__(self):
+            self._thread = None
+            self._stop_event = threading.Event()
+            if start_thread:
+
+                def entrypoint():
+                    while True:
+                        print("Thread!")
+                        time.sleep(1)
+                        if self._stop_event.is_set():
+                            sys.exit(0)
+
+                self._thread = threading.Thread(target=entrypoint)
+                self._thread.start()
+
+        def stop(self):
+            print("Waiting")
+            while not hang_marker.exists():
+                time.sleep(0.05)
+            self._stop_event.set()
+            print("stopped")
+
+    start_barrier = Barrier(max_results=1)
+    stop_barrier = Barrier(max_results=1)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingStopActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=start_barrier.arrive,
+        on_stop=stop_barrier.arrive,
+    )
+    while not start_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor is alive
+    assert ray.get_actor(actor_name)
+
+    stop_future = actor_manager.schedule_actor_task(tracked_actor, "stop")
+    actor_manager.remove_actor(tracked_actor, kill=False, stop_future=stop_future)
+
+    assert not stop_barrier.completed
+
+    hang_marker.write_text("!")
+
+    while not stop_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor should have stopped now and should get cleaned up
+    with pytest.raises(ValueError):
+        ray.get_actor(actor_name)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_tracked_actor_task.py

@@ -0,0 +1,122 @@
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+RESOURCE_MANAGERS = [FixedResourceManager, PlacementGroupResourceManager]
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    def foo(self, val, error: bool = False):
+        if error:
+            raise RuntimeError
+        return val
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_resolve(ray_start_4_cpus, resource_manager_cls):
+    """Test that the `on_result` callback is invoked when a task completes.
+
+    - Instantiate global data object
+    - Schedule task that returns a value
+    - The callback writes the returned value to the global data object
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] = result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (4, False), on_result=result_callback
+    )
+    actor_manager.next()
+    actor_manager.next()
+
+    assert seen["data"] == 4
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+@pytest.mark.parametrize("num_tasks", [1, 10, 100])
+def test_resolve_many(ray_start_4_cpus, resource_manager_cls, num_tasks):
+    """Schedule ``num_tasks`` tasks and wait until ``wait_for_events`` of them resolve.
+
+    Every resolved task will increase a counter by its return value (1).
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] += result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.next()
+
+    for i in range(num_tasks):
+        actor_manager.schedule_actor_task(
+            tracked_actor, "foo", (1, False), on_result=result_callback
+        )
+
+    for i in range(num_tasks):
+        actor_manager.next()
+        assert seen["data"] == i + 1
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_noop(ray_start_4_cpus, resource_manager_cls):
+    """When no `on_error` callback is specified, errors should be ignored."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(tracked_actor, "foo", (1, True))
+    actor_manager.next()
+    actor_manager.next()
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_custom(ray_start_4_cpus, resource_manager_cls):
+    """When an `on_error` callback is specified, it is invoked."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    stats = Counter()
+
+    def error_callback(tracked_actor, exception):
+        stats["exception"] += 1
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (1, True), on_error=error_callback
+    )
+
+    actor_manager.next()
+    actor_manager.next()
+    assert stats["exception"] == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/mocked_wandb_integration.py

@@ -0,0 +1,136 @@
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import Dict
+from unittest.mock import Mock
+
+from wandb.util import json_dumps_safer
+
+import ray
+from ray.air.integrations.wandb import WandbLoggerCallback, _WandbLoggingActor
+
+
+class Trial(
+    namedtuple(
+        "MockTrial",
+        [
+            "config",
+            "trial_id",
+            "trial_name",
+            "experiment_dir_name",
+            "placement_group_factory",
+            "local_path",
+        ],
+    )
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+@dataclass
+class LoggingActorState:
+    args: list
+    kwargs: dict
+    exclude: list
+    logs: list
+    config: dict
+
+
+class _FakeConfig:
+    def __init__(self):
+        self.config = {}
+
+    def update(self, config, *args, **kwargs):
+        self.config.update(config)
+
+
+class _MockWandbAPI:
+    """Thread-safe.
+
+    Note: Not implemented to mock re-init behavior properly. Proceed with caution."""
+
+    def __init__(self):
+        self.logs = []
+        self.config = _FakeConfig()
+
+    def init(self, *args, **kwargs):
+        mock = Mock()
+        mock.args = args
+        mock.kwargs = kwargs
+
+        if "config" in kwargs:
+            self.config.update(kwargs["config"])
+
+        return mock
+
+    def log(self, data, step=None):
+        try:
+            json_dumps_safer(data)
+        except Exception:
+            self.logs.append("serialization error")
+        else:
+            self.logs.append(data)
+
+    def finish(self):
+        pass
+
+    def get_logs(self):
+        return self.logs
+
+    def get_config(self):
+        return self.config.config
+
+
+class _MockWandbLoggingActor(_WandbLoggingActor):
+    _mock_wandb_api_cls = _MockWandbAPI
+
+    def __init__(self, logdir, queue, exclude, to_config, *args, **kwargs):
+        super(_MockWandbLoggingActor, self).__init__(
+            logdir, queue, exclude, to_config, *args, **kwargs
+        )
+        self._wandb = self._mock_wandb_api_cls()
+
+    def get_state(self):
+        return LoggingActorState(
+            args=self.args,
+            kwargs=self.kwargs,
+            exclude=self._exclude,
+            logs=self._wandb.get_logs(),
+            config=self._wandb.get_config(),
+        )
+
+
+class WandbTestExperimentLogger(WandbLoggerCallback):
+    """Wandb logger with mocked Wandb API gateway (one per trial)."""
+
+    _logger_actor_cls = _MockWandbLoggingActor
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._saved_actor_states: Dict["Trial", LoggingActorState] = {}
+
+    def _cleanup_logging_actor(self, trial: "Trial", **kwargs):
+        logging_actor_state: LoggingActorState = ray.get(
+            self._trial_logging_actors[trial].get_state.remote()
+        )
+        self._saved_actor_states[trial] = logging_actor_state
+        super()._cleanup_logging_actor(trial, **kwargs)
+
+    @property
+    def trial_logging_actor_states(self) -> Dict["Trial", LoggingActorState]:
+        return self._saved_actor_states
+
+
+def get_mock_wandb_logger(mock_api_cls=_MockWandbAPI, **kwargs):
+    class MockWandbLoggingActor(_MockWandbLoggingActor):
+        _mock_wandb_api_cls = mock_api_cls
+
+    logger = WandbTestExperimentLogger(
+        project="test_project",
+        api_key="1234",
+        **kwargs,
+    )
+    logger._logger_actor_cls = MockWandbLoggingActor
+    return logger

python/ray/air/tests/test_air_usage.py

@@ -0,0 +1,210 @@
+"""Unit tests for AIR telemetry."""
+
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pyarrow.fs
+import pytest
+from packaging.version import Version
+
+import ray
+from ray import train, tune
+from ray._common.usage.usage_lib import TagKey
+from ray.air._internal import usage as air_usage
+from ray.air._internal.usage import AirEntrypoint
+from ray.air.integrations import comet, mlflow, wandb
+from ray.train._internal.storage import StorageContext
+from ray.tune.callback import Callback
+from ray.tune.experiment.experiment import Experiment
+from ray.tune.logger import LoggerCallback
+from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES
+
+
+def _mock_record_from_module(module, monkeypatch):
+    recorded = {}
+
+    def mock_record_extra_usage_tag(key: TagKey, value: str):
+        recorded[key] = value
+
+    monkeypatch.setattr(
+        module,
+        "record_extra_usage_tag",
+        mock_record_extra_usage_tag,
+    )
+    return recorded
+
+
+@pytest.fixture
+def mock_record(monkeypatch):
+    import ray.air._internal.usage
+
+    yield _mock_record_from_module(ray.air._internal.usage, monkeypatch=monkeypatch)
+
+
+def train_fn(config):
+    train.report({"score": 1})
+
+
+@pytest.fixture
+def tuner(tmp_path):
+    yield tune.Tuner(train_fn, run_config=tune.RunConfig(storage_path=str(tmp_path)))
+
+
+@pytest.fixture
+def trainer(tmp_path):
+    from ray.train.data_parallel_trainer import DataParallelTrainer
+
+    yield DataParallelTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(storage_path=str(tmp_path)),
+    )
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "storage_path_filesystem_expected",
+    [
+        ("/tmp/test", None, "local"),
+        ("s3://", None, "s3"),
+        ("gs://test", None, "gcs"),
+        ("mock://test", None, "mock"),
+        ("test", pyarrow.fs.LocalFileSystem(), "custom"),
+    ],
+)
+def test_tag_storage_type(storage_path_filesystem_expected, mock_record, monkeypatch):
+    # Don't write anything to storage for the test.
+    monkeypatch.setattr(StorageContext, "_create_validation_file", lambda _: None)
+    monkeypatch.setattr(StorageContext, "_check_validation_file", lambda _: None)
+
+    storage_path, storage_filesystem, expected = storage_path_filesystem_expected
+
+    if Version(pyarrow.__version__) < Version("17.0.0") and storage_path.startswith(
+        "gs://"
+    ):
+        pytest.skip("GCS support requires pyarrow >= 17.0.0")
+
+    storage = StorageContext(
+        storage_path=storage_path,
+        experiment_dir_name="test",
+        storage_filesystem=storage_filesystem,
+    )
+    air_usage.tag_storage_type(storage)
+    assert mock_record[TagKey.AIR_STORAGE_CONFIGURATION] == expected
+
+
+class _CustomLoggerCallback(LoggerCallback):
+    pass
+
+
+class _CustomCallback(Callback):
+    pass
+
+
+_TEST_CALLBACKS = [
+    wandb.WandbLoggerCallback,
+    mlflow.MLflowLoggerCallback,
+    comet.CometLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomCallback,
+]
+
+
+def test_tag_setup_wandb(mock_record):
+    from ray.air.integrations.wandb import _setup_wandb
+
+    with patch.dict(os.environ, {wandb.WANDB_MODE_ENV_VAR: "disabled"}):
+        _setup_wandb(trial_id="a", trial_name="b", config={}, _wandb=MagicMock())
+    assert mock_record[TagKey.AIR_SETUP_WANDB_INTEGRATION_USED] == "1"
+
+
+def test_tag_setup_mlflow(mock_record, monkeypatch):
+    from ray.air.integrations.mlflow import setup_mlflow
+
+    monkeypatch.setattr(ray.air.integrations.mlflow, "_MLflowLoggerUtil", MagicMock())
+    setup_mlflow()
+    assert mock_record[TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED] == "1"
+
+
+@pytest.mark.parametrize(
+    "callback_classes_expected",
+    [
+        (None, None),
+        ([], None),
+        ([lambda: None], None),
+        (
+            DEFAULT_CALLBACK_CLASSES,
+            {cls.__name__: 1 for cls in DEFAULT_CALLBACK_CLASSES},
+        ),
+        (
+            _TEST_CALLBACKS,
+            {
+                "WandbLoggerCallback": 1,
+                "MLflowLoggerCallback": 1,
+                "CometLoggerCallback": 1,
+                "CustomLoggerCallback": 2,
+                "CustomCallback": 1,
+            },
+        ),
+    ],
+)
+def test_tag_callbacks(mock_record, callback_classes_expected):
+    callback_classes, expected = callback_classes_expected
+
+    callbacks = (
+        [callback_cls() for callback_cls in callback_classes]
+        if callback_classes
+        else None
+    )
+
+    air_usage.tag_callbacks(callbacks)
+
+    callback_usage_str = mock_record.pop(TagKey.AIR_CALLBACKS, None)
+    callback_counts = json.loads(callback_usage_str) if callback_usage_str else None
+    assert callback_counts == expected
+
+
+def test_tag_env_vars(ray_start_4_cpus, mock_record, tuner):
+    """Test that env vars are recorded properly, and arbitrary user environment
+    variables are ignored."""
+    env_vars_to_record = {
+        "TUNE_GLOBAL_CHECKPOINT_S": "20",
+        "TUNE_MAX_PENDING_TRIALS_PG": "1",
+    }
+    untracked_env_vars = {"RANDOM_USER_ENV_VAR": "asdf"}
+
+    with patch.dict(os.environ, {**env_vars_to_record, **untracked_env_vars}):
+        tuner.fit()
+
+    recorded_env_vars = json.loads(mock_record[TagKey.AIR_ENV_VARS])
+    assert sorted(env_vars_to_record) == sorted(recorded_env_vars)
+
+
+@pytest.mark.parametrize("entrypoint", list(AirEntrypoint))
+def test_tag_air_entrypoint(ray_start_4_cpus, mock_record, entrypoint, tuner, trainer):
+    if entrypoint == AirEntrypoint.TUNE_RUN:
+        tune.run(train_fn)
+    elif entrypoint == AirEntrypoint.TUNE_RUN_EXPERIMENTS:
+        experiment_spec = Experiment("experiment", train_fn)
+        tune.run_experiments(experiments=experiment_spec)
+    elif entrypoint == AirEntrypoint.TUNER:
+        tuner.fit()
+    elif entrypoint == AirEntrypoint.TRAINER:
+        trainer.fit()
+
+    assert mock_record[TagKey.AIR_ENTRYPOINT] == entrypoint.value
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_errors.py

@@ -0,0 +1,284 @@
+"""
+This test suite covers error handling and propagation in Ray Train/Tune.
+
+There are two main error types to test:
+1. Trainable errors: These happen in the remote actor itself.
+    -> Within this, we should test:
+        - fail_fast=True/False/'raise'
+        - AIR Trainer w/o Tuner, AIR Trainer w/ Tuner, Tuner w/ function trainable
+2. Tune driver errors: These happen in the Tune event-handling loop.
+    -> Within this, we should test:
+        - Errors occurring at different points in the Tune loop
+          (on_trial_result, on_checkpoint, on_step_begin, etc.)
+
+These tests should:
+- Assert how errors from the trainable/Trainer get propagated to the user.
+- Assert how errors from the Tune driver get propagated to the user.
+"""
+
+import gc
+import threading
+import time
+from tempfile import TemporaryDirectory
+
+import pytest
+
+import ray
+from ray import train, tune
+from ray._common.test_utils import wait_for_condition
+from ray._raylet import GcsClient
+from ray.cluster_utils import Cluster
+from ray.core.generated import autoscaler_pb2
+from ray.tests.conftest import *  # noqa
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.trainer import BaseTrainer, TrainingFailedError
+from ray.tune import TuneError, Tuner
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4, configure_logging=False)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def gc_collect():
+    # Make sure to cleanup as much as possible between
+    # unit tests that share a Ray session
+    yield
+    gc.collect()
+
+
+@pytest.fixture
+def cluster_setup(ray_start_cluster_head: Cluster):
+    # Sets up a cluster with 3 nodes: head node + 2 workers
+    cluster = ray_start_cluster_head
+    nodes = []
+    nodes.append(cluster.add_node(resources={"worker1": 1, "cpu": 1, "coordinator": 1}))
+    nodes.append(cluster.add_node(resources={"worker2": 1, "cpu": 1}))
+    cluster.wait_for_nodes()
+
+    @ray.remote
+    def get_node_id():
+        return ray.get_runtime_context().get_node_id()
+
+    worker1_node_id = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
+    worker2_node_id = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
+    wait_for_condition(
+        lambda: len({node["NodeID"] for node in ray.nodes() if (node["Alive"])}) == 3
+    )
+
+    yield cluster, nodes, [
+        worker1_node_id,
+        worker2_node_id,
+    ]
+
+
+class _TestSpecificError(RuntimeError):
+    pass
+
+
+class FailingCallback(tune.Callback):
+    def __init__(self, error_on: str):
+        self.error_on = error_on
+
+    def on_trial_result(self, *args, **kwargs):
+        if self.error_on == "on_trial_result":
+            raise _TestSpecificError(f"Failing on {self.error_on}!")
+
+
+class FailingTrainer(BaseTrainer):
+    def training_loop(self) -> None:
+        raise _TestSpecificError("There is an error in trainer!")
+
+
+def passing_fn(config):
+    # Trigger all the driver events (on_checkpoint, on_trial_save, etc.)
+    with TemporaryDirectory() as tmpdir:
+        train.report({"score": 1}, checkpoint=train.Checkpoint.from_directory(tmpdir))
+
+
+def failing_fn(config):
+    raise _TestSpecificError("Failing!")
+
+
+trainable_map = {
+    "function": failing_fn,
+    "trainer": FailingTrainer(),
+}
+
+
+@pytest.mark.parametrize("fail_fast", [False, True, "raise"])
+def test_trainable_error_with_tuner(ray_start_4_cpus, fail_fast):
+    tuner = Tuner(
+        trainable=failing_fn,
+        run_config=tune.RunConfig(
+            name=f"tuner_errors-fail_fast={fail_fast}",
+            failure_config=tune.FailureConfig(fail_fast=fail_fast),
+        ),
+        tune_config=tune.TuneConfig(num_samples=2),
+    )
+
+    if fail_fast is False:
+        # Both trials should complete with an error.
+        results = tuner.fit()
+        assert len(results) == 2
+        for i in range(2):
+            assert results[i].error
+    elif fail_fast is True:
+        # The first trial errors -> the experiment finishes immediately.
+        results = tuner.fit()
+        errors = [result.error for result in results if result.error]
+        assert len(errors) == 1
+    elif fail_fast == "raise":
+        # The original error gets raised to the user
+        with pytest.raises(_TestSpecificError):
+            tuner.fit()
+
+
+@pytest.mark.parametrize("fail_fast", [False, True, "raise"])
+def test_trainable_error_with_trainer(ray_start_4_cpus, tmp_path, fail_fast):
+    name = f"test_trainer_errors-fail_fast={fail_fast}"
+    trainer = FailingTrainer(
+        run_config=train.RunConfig(
+            storage_path=str(tmp_path),
+            name=name,
+            failure_config=train.FailureConfig(fail_fast=fail_fast),
+        ),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+
+    if fail_fast in [False, True]:
+        # There is only 1 "trial" for a Trainer,
+        # so fail_fast = True/False doesn't change the behavior
+        # In both cases, the error should get wrapped and raised.
+        with pytest.raises(TrainingFailedError) as exc_info:
+            trainer.fit()
+
+        # The cause of the error should be the trainable error
+        assert isinstance(exc_info.value.__cause__, _TestSpecificError)
+
+        assert TrainingFailedError._RESTORE_MSG.format(
+            trainer_cls_name="FailingTrainer", path=str(tmp_path / name)
+        ) in str(exc_info.value)
+        assert TrainingFailedError._FAILURE_CONFIG_MSG in str(exc_info.value)
+
+    elif fail_fast == "raise":
+        # The original error gets raised to the user
+        with pytest.raises(_TestSpecificError):
+            trainer.fit()
+
+
+# TODO(ml-team): Test all the driver hooks once driver error propagation is fixed
+
+
+@pytest.mark.parametrize("error_on", ["on_trial_result"])
+def test_driver_error_with_tuner(ray_start_4_cpus, error_on):
+    tuner = Tuner(
+        trainable=passing_fn,
+        run_config=tune.RunConfig(
+            name=f"test_driver_errors_with_tuner-error_on={error_on}",
+            callbacks=[FailingCallback(error_on=error_on)],
+        ),
+    )
+
+    # All driver errors should get propagated to the user in the same way
+    with pytest.raises(TuneError) as exc_info:
+        tuner.fit()
+
+    # TODO(ml-team): Assert the cause error type once driver error propagation is fixed
+    assert "_TestSpecificError" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("error_at_level", ["worker", "coordinator"])
+def test_preemption_handling(
+    cluster_setup,
+    tmp_path,
+    error_at_level: str,
+):
+    """Integration test for node preemption handling in Ray Train/Tune.
+    Even though `max_failures=0`, preemption errors should still be retried."""
+    cluster, nodes, node_ids = cluster_setup
+    # node 1 = coordinator and worker, node 2 = worker
+    coordinator_node, worker_node = nodes
+    coordinator_node_id, worker_node_id = node_ids
+
+    num_workers = 2
+    tmp_path.joinpath("markers").mkdir()
+
+    def train_fn(config):
+        checkpoint = train.get_checkpoint()
+        start_iter = 0
+        if checkpoint:
+            start_iter = load_dict_checkpoint(checkpoint)["iter"] + 1
+            print(f"Restored at iter = {start_iter}")
+
+        for iter in range(start_iter, 6):
+            with create_dict_checkpoint({"iter": iter}) as checkpoint:
+                ray.train.report({"iter": iter}, checkpoint=checkpoint)
+
+            if iter == 2:
+                # Write a "done marker" to tell the driver to simulate a preemption.
+                tmp_path.joinpath(
+                    "markers", str(ray.train.get_context().get_world_rank())
+                ).touch()
+                # Await execution.
+                time.sleep(120)
+
+    def launch_training():
+        trainer = DataParallelTrainer(
+            train_loop_per_worker=train_fn,
+            scaling_config=train.ScalingConfig(
+                num_workers=num_workers,
+                trainer_resources={"coordinator": 1},
+                resources_per_worker={"cpu": 1},  # worker2 and worker3
+            ),
+            run_config=train.RunConfig(
+                storage_path=str(tmp_path),
+                name="test_preemption_error",
+                failure_config=train.FailureConfig(fail_fast=False, max_failures=0),
+            ),
+        )
+        result = trainer.fit()
+        assert result.metrics["iter"] == 5
+
+    t = threading.Thread(target=launch_training)
+    t.start()
+
+    # Wait until the workers are ready for preemption (after a few checkpoints).
+    while len(list(tmp_path.joinpath("markers").glob("*"))) < num_workers:
+        time.sleep(0.5)
+
+    if error_at_level == "coordinator":
+        node, node_id = coordinator_node, coordinator_node_id
+    elif error_at_level == "worker":
+        node, node_id = worker_node, worker_node_id
+    else:
+        raise NotImplementedError(f"Invalid error_at_level = {error_at_level}")
+
+    # Preempt a node.
+    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
+    print("Draining node...")
+    is_accepted, _ = gcs_client.drain_node(
+        node_id,
+        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
+        "preemption",
+        0,
+    )
+    assert is_accepted
+    print("Killing node...")
+    cluster.remove_node(node, allow_graceful=True)
+    print("Adding new node..")  # so that the job can be rescheduled
+    # New node can replace a preempted coordinator or worker
+    # NOTE: `cluster.add_node` only works in the main thread.
+    cluster.add_node(resources={"coordinator": 1, "cpu": 1})
+    t.join()  # Assert no errors during training.
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/air/tests/test_experiment_restore.py

@@ -0,0 +1,232 @@
+import json
+import os
+import shutil
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from ray.tune.analysis import ExperimentAnalysis
+from ray.tune.result_grid import ResultGrid
+
+_RUN_SCRIPT_FILENAME = "_test_experiment_restore_run.py"
+
+
+def _kill_process_if_needed(
+    process: subprocess.Popen, timeout_s: float = 10, poll_interval_s: float = 1.0
+):
+    """Kills a process if it hasn't finished in `timeout_s` seconds.
+    Polls every `poll_interval_s` seconds to check if the process is still running."""
+    kill_timeout = time.monotonic() + timeout_s
+    while process.poll() is None and time.monotonic() < kill_timeout:
+        time.sleep(poll_interval_s)
+    if process.poll() is None:
+        process.terminate()
+
+
+def _print_message(message):
+    sep = "=" * 50
+    print(f"\n{sep}\n{message}\n{sep}\n")
+
+
+@pytest.mark.parametrize("runner_type", ["tuner", "trainer"])
+def test_experiment_restore(tmp_path, runner_type):
+    """
+    This is an integration stress test for experiment restoration.
+
+
+    Test setup:
+
+    - For Tuner.restore:
+        - 8 trials, with a max of 2 running concurrently (--> 4 rounds of trials)
+        - Each iteration takes 0.5 seconds
+        - Each trial runs for 8 iterations --> 4 seconds
+        - Each round of 2 trials should take 4 seconds
+        - Without any interrupts/restoration:
+            - Minimum runtime: 4 rounds * 4 seconds / round = 16 seconds
+        - The test will stop the script with a SIGINT at a random time between
+        6-10 iterations each restore.
+
+    - For Trainer.restore:
+        - 1 trial with 4 workers
+        - Each iteration takes 0.5 seconds
+        - Runs for 32 iterations --> Minimum runtime = 16 seconds
+        - The test will stop the script with a SIGINT at a random time between
+        6-10 iterations after each restore.
+
+    Requirements:
+    - Req 1: Training progress persisted
+        - The experiment should progress monotonically.
+        (The training iteration shouldn't go backward at any point)
+        - Trials shouldn't start from scratch.
+    - Req 2: Searcher state saved/restored correctly
+    - Req 3: Callback state saved/restored correctly
+    """
+
+    np.random.seed(2023)
+
+    script_path = Path(__file__).parent / _RUN_SCRIPT_FILENAME
+
+    # Args to pass into the script as environment variables
+    exp_name = f"{runner_type}_restore_integration_test"
+    callback_dump_file = tmp_path / f"{runner_type}-callback_dump_file.json"
+    storage_path = tmp_path / "ray_results"
+    if storage_path.exists():
+        shutil.rmtree(storage_path)
+
+    csv_file = str(tmp_path / "dummy_data.csv")
+    dummy_df = pd.DataFrame({"x": np.arange(128), "y": 2 * np.arange(128)})
+    dummy_df.to_csv(csv_file)
+
+    run_started_marker = tmp_path / "run_started_marker"
+
+    time_per_iter_s = 0.5
+    max_concurrent = 2
+
+    if runner_type == "tuner":
+        iters_per_trial = 8
+        num_trials = 8
+    elif runner_type == "trainer":
+        iters_per_trial = 32
+        num_trials = 1
+
+    total_iters = iters_per_trial * num_trials
+
+    env = os.environ.copy()
+    env.update(
+        {
+            "RUNNER_TYPE": runner_type,
+            "STORAGE_PATH": str(storage_path),
+            "EXP_NAME": exp_name,
+            "CALLBACK_DUMP_FILE": str(callback_dump_file),
+            "RUN_STARTED_MARKER": str(run_started_marker),
+            "TIME_PER_ITER_S": str(time_per_iter_s),
+            "ITERATIONS_PER_TRIAL": str(iters_per_trial),
+            "NUM_TRIALS": str(num_trials),
+            "MAX_CONCURRENT_TRIALS": str(max_concurrent),
+            "CSV_DATA_FILE": csv_file,
+        }
+    )
+
+    # Variables used in the loop
+    return_code = None
+    total_runtime = 0
+    run_iter = 0
+    progress = 0
+    progress_history = []
+
+    poll_interval_s = 0.1
+    test_start_time = time.monotonic()
+
+    while True:
+        run_started_marker.write_text("", encoding="utf-8")
+
+        run = subprocess.Popen([sys.executable, script_path], env=env)
+        run_iter += 1
+
+        _print_message(f"Started run #{run_iter} w/ PID = {run.pid}")
+
+        # Start the timer after the first trial has entered its training loop.
+        while run.poll() is None and run_started_marker.exists():
+            time.sleep(poll_interval_s)
+
+        # If the run already finished, then exit immediately.
+        if run.poll() is not None:
+            return_code = run.poll()
+            break
+
+        timeout_s = np.random.uniform(6 * time_per_iter_s, 10 * time_per_iter_s)
+
+        _print_message(
+            "Training has started...\n"
+            f"Interrupting after {timeout_s:.2f} seconds\n"
+            f"Currently at {total_runtime:.2f} seconds"
+        )
+
+        # Sleep for a random amount of time, then stop the run.
+        start_time = time.monotonic()
+        time.sleep(timeout_s)
+        total_runtime += time.monotonic() - start_time
+
+        return_code = run.poll()
+        if return_code is None:
+            # Send "SIGINT" to stop the run
+            _print_message(f"Sending SIGUSR1 to run #{run_iter} w/ PID = {run.pid}")
+            run.send_signal(signal.SIGUSR1)
+
+            # Make sure the process is stopped forcefully after a timeout.
+            _kill_process_if_needed(run)
+        else:
+            _print_message("Run has already terminated!")
+            break
+
+        # Check up on the results.
+        results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
+        iters = [result.metrics.get("training_iteration", 0) for result in results]
+        progress = sum(iters) / total_iters
+        progress_history.append(progress)
+        _print_message(
+            f"Number of trials = {len(results)}\n"
+            f"% completion = {progress} ({sum(iters)} iters / {total_iters})\n"
+            f"Currently at {total_runtime:.2f} seconds"
+        )
+
+    _print_message(
+        f"Total number of restorations = {run_iter}\n"
+        f"Total runtime = {total_runtime:.2f}\n"
+        f"Return code = {return_code}"
+    )
+    test_end_time = time.monotonic()
+
+    assert progress == 1.0
+
+    # The script shouldn't have errored. (It should have finished by this point.)
+    assert return_code == 0, (
+        f"The script errored with return code: {return_code}.\n"
+        f"Check the `{_RUN_SCRIPT_FILENAME}` script for any issues. "
+    )
+
+    # Req 1: training progress persisted
+    # Check that progress increases monotonically (we never go backwards/start from 0)
+    assert np.all(np.diff(progress_history) >= 0), (
+        "Expected progress to increase monotonically. Instead, got:\n"
+        "{progress_history}"
+    )
+
+    # Req 2: searcher state
+    results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
+    # Check that all trials have unique ids assigned by the searcher (if applicable)
+    ids = [result.config.get("id", -1) for result in results]
+    ids = [id for id in ids if id >= 0]
+    if ids:
+        assert sorted(ids) == list(range(1, num_trials + 1)), (
+            "Expected the searcher to assign increasing id for each trial, but got:"
+            f"{ids}"
+        )
+
+    # Req 3: callback state
+    with open(callback_dump_file, "r") as f:
+        callback_state = json.load(f)
+
+    trial_iters = callback_state["trial_iters"]
+    for iters in trial_iters.values():
+        # Check that the callback has data for each trial, for all iters
+        # NOTE: There may be some duplicate data, due to the fact that
+        # the callback will be updated on every `on_trial_result` hook,
+        # but the trial may crash before the corresponding checkpoint gets processed.
+        assert sorted(set(iters)) == list(
+            range(1, iters_per_trial + 1)
+        ), f"Expected data from all iterations, but got: {iters}"
+
+    _print_message(f"Success! Test took {test_end_time - test_start_time:.2f} seconds.")
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_comet.py

@@ -0,0 +1,366 @@
+import unittest
+from collections import namedtuple
+from unittest.mock import patch
+
+from ray.air.integrations.comet import CometLoggerCallback
+
+
+class MockTrial(
+    namedtuple("MockTrial", ["config", "trial_name", "trial_id", "logdir"])
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+class InitializationTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+
+    def test_class_variable_to_instance(self):
+        """Test that class variables get properly assigned to instance
+        variables.
+        """
+        logger = self.logger
+        self.assertEqual(logger._to_exclude, logger._exclude_results)
+        self.assertEqual(logger._to_system, logger._system_results)
+        self.assertEqual(logger._to_other, logger._other_results)
+        self.assertEqual(logger._to_episodes, logger._episode_results)
+
+    def test_configure_experiment_defaults(self):
+        """Test CometLoggerCallback._configure_experiment_defaults."""
+        logger = self.logger
+
+        # Test that autologging features are properly disabled
+        exclude = CometLoggerCallback._exclude_autolog
+        for option in exclude:
+            self.assertFalse(logger.experiment_kwargs.get(option))
+        del logger
+
+        # Don't disable logging if user overwrites defaults by passing in args
+        for include_option in exclude:
+            # This unpacks to become e.g. CometLoggerCallback(log_env_cpu=True)
+            logger = CometLoggerCallback(**{include_option: True})
+            for option in exclude:
+                if option == include_option:
+                    self.assertTrue(logger.experiment_kwargs.get(option))
+                else:
+                    self.assertFalse(logger.experiment_kwargs.get(option))
+
+
+class HelperMethodTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+
+    def test_check_key_name(self):
+
+        logger = self.logger
+        # Return True when key == item
+        self.assertTrue(logger._check_key_name("name", "name"))
+        # Return True when key.startswith(item + "/")
+        self.assertTrue(logger._check_key_name("name/", "name"))
+        # Return False when item.startswith(key + "/")
+        self.assertFalse(logger._check_key_name("name", "name/"))
+        # Return False when key != item and not key.startswith(item."/")
+        self.assertFalse(logger._check_key_name("name", "x"))
+
+
+@patch("comet_ml.OfflineExperiment")
+@patch("comet_ml.Experiment")
+class OnlineVsOfflineTests(unittest.TestCase):
+    def setUp(self):
+        self.loggers = {
+            "online": CometLoggerCallback(),
+            "offline": CometLoggerCallback(online=False),
+        }
+
+        self.trial = MockTrial({"p1": 1}, "trial_1", 1, "artifact")
+
+    def test_online_dispatch(self, experiment, offline_experiment):
+
+        # To start, there should be no experiments
+        experiment.assert_not_called()
+        offline_experiment.assert_not_called()
+
+        # Start online experiment
+        logger = self.loggers["online"]
+        logger.log_trial_start(self.trial)
+
+        # Check that Experiment was called and OfflineExperiment was not
+        experiment.assert_called_once()
+        offline_experiment.assert_not_called()
+
+    def test_offline_dispatch(self, experiment, offline_experiment):
+
+        # To start, there should be no experiments
+        experiment.assert_not_called()
+        offline_experiment.assert_not_called()
+
+        # Start online experiment
+        logger = self.loggers["offline"]
+        logger.log_trial_start(self.trial)
+
+        # Check that Experiment was called and OfflineExperiment was not
+        experiment.assert_not_called()
+        offline_experiment.assert_called_once()
+
+
+@patch("comet_ml.OfflineExperiment")
+@patch("comet_ml.Experiment")
+class LogTrialStartTest(unittest.TestCase):
+    def setUp(self):
+        self.loggers = {
+            "online": CometLoggerCallback(),
+            "offline": CometLoggerCallback(online=False),
+        }
+
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 1, "artifact"),
+        ]
+
+    def test_existing_trialexperiment(self, experiment, offline_experiment):
+
+        mocks = {"online": experiment, "offline": offline_experiment}
+        for option in ["online", "offline"]:
+            logger = self.loggers[option]
+            mock = mocks[option]
+
+            # This should create an experiment
+            logger.log_trial_start(self.trials[0])
+            mock.assert_called_once()
+
+            # This should NOT create an experiment because it's the same trial
+            logger.log_trial_start(self.trials[0])
+            mock.assert_called_once()
+
+            # This should create another new experiment
+            logger.log_trial_start(self.trials[1])
+
+            # Number of times the mock was called
+            num_calls = len(mock.call_args_list)
+
+            # Assert that Experiment/OfflineExperiment was called twice
+            self.assertEqual(num_calls, 2)
+
+    def test_set_global_experiment(self, experiment, offline_experiment):
+        for option in ["online", "offline"]:
+            logger = self.loggers[option]
+            with patch("comet_ml.config.set_global_experiment") as mock:
+                logger.log_trial_start(self.trials[0])
+                mock.assert_called_with(None)
+                mock.assert_called_once()
+                mock.reset_mock()
+
+    def test_experiment_addtags(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        logger.log_trial_start(self.trials[0])
+        experiment.return_value.add_tags.assert_called_with(logger.tags)
+
+    def test_experiment_setname(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        trial = self.trials[0]
+        logger.log_trial_start(trial)
+        experiment.return_value.set_name.assert_called_with(trial.trial_name)
+
+    def test_experiment_logparams(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        trial = self.trials[0]
+        logger.log_trial_start(trial)
+        config = trial.config.copy()
+        config.pop("callbacks", None)
+        experiment.return_value.log_parameters.assert_called_with(config)
+
+
+class ExperimentKwargsTest(unittest.TestCase):
+    @patch("comet_ml.Experiment")
+    def test_kwargs_passthrough(self, experiment):
+        """Test that additional keyword arguments to CometLoggerCallback get
+        passed through to comet_ml.Experiment on log_trial_start
+        """
+        experiment_kwargs = {"kwarg_1": "val_1"}
+        logger = CometLoggerCallback(**experiment_kwargs)
+        trial = MockTrial({"parameter": 1}, "trial2", 1, "artifact")
+        logger.log_trial_start(trial)
+
+        # These are the default kwargs that get passed to create the experiment
+        expected_kwargs = {kwarg: False for kwarg in logger._exclude_autolog}
+        expected_kwargs.update(experiment_kwargs)
+
+        experiment.assert_called_with(**expected_kwargs)
+
+
+@patch("comet_ml.Experiment")
+class LogTrialResultTests(unittest.TestCase):
+    """
+    * test log_others logs
+    * test log_system logs
+    * test log_curve logs
+    """
+
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 1, "artifact"),
+        ]
+        self.result = {
+            "config": {"p1": 1},
+            "node_ip": "0.0.0.0",
+            "hostname": "hostname_val",
+            "pid": "1234",
+            "date": "2000-01-01",
+            "experiment_id": "1234",
+            "trial_id": 1,
+            "experiment_tag": "tag1",
+            "hist_stats/episode_reward": [1, 0, 1, -1, 0, 1],
+            "hist_stats/episode_lengths": [1, 2, 3, 4, 5, 6],
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+
+    def test_log_parameters(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        # Check parameters are logged properly.
+        logger.log_trial_result(1, trial, self.result)
+
+        config_update = result.copy().pop("config", {})
+        config_update.pop("callbacks", None)  # Remove callbacks
+
+        experiment.return_value.log_parameters.assert_any_call(config_update)
+
+    def test_log_metrics(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+        step = result["training_iteration"]
+
+        logger.log_trial_result(1, trial, self.result)
+        result_metrics = {
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+
+        method = experiment.return_value.log_metrics
+        method.assert_any_call(result_metrics, step=step)
+
+    def test_log_other(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        logger.log_trial_result(1, trial, result)
+        result_other = {
+            "experiment_id": "1234",
+            "trial_id": 1,
+            "experiment_tag": "tag1",
+        }
+        method = experiment.return_value.log_others
+
+        # for k,v in result_other.items():
+        method.assert_any_call(result_other)
+
+    def test_log_system(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        logger.log_trial_result(1, trial, result)
+        result_system = {
+            "node_ip": "0.0.0.0",
+            "hostname": "hostname_val",
+            "pid": "1234",
+            "date": "2000-01-01",
+        }
+        method = experiment.return_value.log_system_info
+        for k, v in result_system.items():
+            method.assert_any_call(k, v)
+
+    def test_log_curve(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+
+        # Check parameters are logged properly.
+        result = self.result
+        step = result["training_iteration"]
+        logger.log_trial_result(1, trial, result)
+
+        results_curve = {
+            "hist_stats/episode_reward": [1, 0, 1, -1, 0, 1],
+            "hist_stats/episode_lengths": [1, 2, 3, 4, 5, 6],
+        }
+
+        method = experiment.return_value.log_curve
+        print(method.call_args_list)
+        for k, v in results_curve.items():
+
+            method.assert_any_call(k, x=range(len(v)), y=v, step=step)
+
+
+@patch("comet_ml.Experiment")
+class LogTrialEndTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 2, "artifact"),
+            MockTrial({"p1": 2}, "trial_3", 3, "artifact"),
+        ]
+
+    def test_not_started_exception(self, experiment):
+        logger = self.logger
+        with self.assertRaises(KeyError):
+            logger.log_trial_end(self.trials[0])
+
+    def test_repeat_throws_error(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        with self.assertRaises(KeyError):
+            logger.log_trial_end(trial)
+
+    def test_log_trial_end(self, experiment):
+        logger = self.logger
+        trials = self.trials
+        method = experiment.return_value.end
+
+        # Should not have ended yet
+        method.assert_not_called()
+
+        for trial in trials:
+            logger.log_trial_start(trial)
+            logger.log_trial_end(trial)
+
+        self.assertEqual(len(method.call_args_list), len(trials))
+
+    def test_del(self, experiment):
+        logger = self.logger
+
+        for trial in self.trials:
+            logger.log_trial_start(trial)
+
+        end = experiment.return_value.end
+        end.assert_not_called()
+
+        logger.__del__()
+
+        self.assertEqual(len(end.call_args_list), len(self.trials))
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_mlflow.py

@@ -0,0 +1,442 @@
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from collections import namedtuple
+from unittest.mock import MagicMock, patch
+
+import pytest
+from mlflow.tracking import MlflowClient
+
+import ray
+from ray._private.dict import flatten_dict
+from ray.air._internal.mlflow import _MLflowLoggerUtil
+from ray.air.integrations.mlflow import MLflowLoggerCallback, _NoopModule, setup_mlflow
+from ray.train.torch import TorchTrainer
+from ray.tune import Tuner
+
+
+class MockTrial(
+    namedtuple("MockTrial", ["config", "trial_name", "trial_id", "local_path"])
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+class Mock_MLflowLoggerUtil(_MLflowLoggerUtil):
+    def save_artifacts(self, dir, run_id):
+        self.artifact_saved = True
+        self.artifact_info = {"dir": dir, "run_id": run_id}
+
+
+def clear_env_vars():
+    os.environ.pop("MLFLOW_EXPERIMENT_NAME", None)
+    os.environ.pop("MLFLOW_EXPERIMENT_ID", None)
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    """Automatically start and stop Ray for each test."""
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+def test_setup_mlflow_in_train_worker(ray_start_4_cpus):
+    """Test that setup_mlflow works in a Train worker."""
+
+    def train_func(config):
+        setup_mlflow(
+            experiment_name="test_exp",
+            create_experiment_if_not_exists=True,
+        )
+
+    trainer = TorchTrainer(train_func)
+    trainer.fit()
+
+
+def test_setup_mlflow_in_tune_trial(ray_start_4_cpus):
+    """Test that setup_mlflow works in a Tune trial."""
+
+    def train_func(config):
+        setup_mlflow(
+            experiment_name="test_exp",
+            create_experiment_if_not_exists=True,
+        )
+
+    tuner = Tuner(train_func)
+    result_grid = tuner.fit()
+
+    assert all(res.error is None for res in result_grid)
+
+
+class MLflowTest(unittest.TestCase):
+    def setUp(self):
+        self.tracking_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
+        self.registry_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
+
+        client = MlflowClient(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        client.create_experiment(name="existing_experiment")
+        # Mlflow > 2 creates a "Default" experiment which has ID 0, so we start our
+        # test with ID 1.
+        assert client.get_experiment_by_name("existing_experiment").experiment_id == "1"
+
+    def tearDown(self) -> None:
+        pass
+
+    def testMlFlowLoggerCallbackConfig(self):
+        # Explicitly pass in all args.
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test_exp",
+        )
+        logger.setup()
+        self.assertEqual(
+            logger.mlflow_util._mlflow.get_tracking_uri(), self.tracking_uri
+        )
+        self.assertEqual(
+            logger.mlflow_util._mlflow.get_registry_uri(), self.registry_uri
+        )
+        self.assertListEqual(
+            [e.name for e in logger.mlflow_util._mlflow.search_experiments()],
+            ["test_exp", "existing_experiment", "Default"],
+        )
+        self.assertEqual(logger.mlflow_util.experiment_id, "2")
+
+        # Check if client recognizes already existing experiment.
+        logger = MLflowLoggerCallback(
+            experiment_name="existing_experiment",
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in experiment name as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "2")
+
+        # Pass in existing experiment name as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in existing experiment id as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in non existing experiment id as env var.
+        # This should create a new experiment.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "500"
+        with self.assertRaises(ValueError):
+            logger = MLflowLoggerCallback(
+                tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+            )
+            logger.setup()
+
+        # Experiment id env var should take precedence over name env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Using tags
+        tags = {"user_name": "John", "git_commit_hash": "abc123"}
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_tags"
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri, tags=tags
+        )
+        logger.setup()
+        self.assertEqual(logger.tags, tags)
+
+    @patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
+    def testMlFlowLoggerLogging(self):
+        clear_env_vars()
+        trial_config = {"par1": "a", "par2": "b"}
+        trial = MockTrial(trial_config, "trial1", 0, "artifact")
+
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test1",
+            save_artifact=True,
+            tags={"hello": "world"},
+        )
+        logger.setup()
+
+        # Check if run is created with proper tags.
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
+        self.assertEqual(len(all_runs), 1)
+        # all_runs is a pandas dataframe.
+        all_runs = all_runs.to_dict(orient="records")
+        run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
+        self.assertDictEqual(
+            run.data.tags,
+            {"hello": "world", "trial_name": "trial1", "mlflow.runName": "trial1"},
+        )
+        self.assertEqual(logger._trial_runs[trial], run.info.run_id)
+        # Params should be logged.
+        self.assertDictEqual(run.data.params, trial_config)
+
+        # When same trial is started again, new run should not be created.
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
+        self.assertEqual(len(all_runs), 1)
+
+        # Check metrics are logged properly.
+        result = {
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+        logger.on_trial_result(0, [], trial, result)
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        # metric3 is not logged since it cannot be converted to float.
+        self.assertDictEqual(
+            run.data.metrics, {"metric1": 0.8, "metric2": 1.0, "training_iteration": 0}
+        )
+
+        # Check that artifact is logged on termination.
+        logger.on_trial_complete(0, [], trial)
+        self.assertTrue(logger.mlflow_util.artifact_saved)
+        self.assertDictEqual(
+            logger.mlflow_util.artifact_info,
+            {"dir": "artifact", "run_id": run.info.run_id},
+        )
+
+        # Check if params are logged at the end.
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        self.assertDictEqual(run.data.params, trial_config)
+
+    @patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
+    def testMlFlowLoggerLogging_logAtEnd(self):
+        clear_env_vars()
+        trial_config = {"par1": "a", "par2": "b"}
+        trial = MockTrial(trial_config, "trial1", 0, "artifact")
+
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test_log_at_end",
+            tags={"hello": "world"},
+            log_params_on_trial_end=True,
+        )
+        logger.setup()
+        exp_id = logger.mlflow_util.experiment_id
+
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=[exp_id])
+        self.assertEqual(len(all_runs), 1)
+        # all_runs is a pandas dataframe.
+        all_runs = all_runs.to_dict(orient="records")
+        run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
+
+        # Params should NOT be logged at start.
+        self.assertDictEqual(run.data.params, {})
+
+        # Check that params are logged at the end.
+        logger.on_trial_complete(0, [], trial)
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        self.assertDictEqual(run.data.params, trial_config)
+
+    def testMlFlowSetupExplicit(self):
+        clear_env_vars()
+        trial_config = {"par1": 4, "par2": 9.0}
+
+        # No MLflow config passed in.
+        with self.assertRaises(ValueError):
+            setup_mlflow(trial_config)
+
+        # Invalid experiment-id
+        with self.assertRaises(ValueError):
+            setup_mlflow(trial_config, experiment_id="500")
+
+        # Set to experiment that does not already exist.
+        with self.assertRaises(ValueError):
+            setup_mlflow(
+                trial_config,
+                experiment_id="500",
+                experiment_name="new_experiment",
+                tracking_uri=self.tracking_uri,
+            )
+
+        mlflow = setup_mlflow(
+            trial_config,
+            experiment_id="500",
+            experiment_name="existing_experiment",
+            tracking_uri=self.tracking_uri,
+        )
+        mlflow.end_run()
+
+    @patch("ray.train.get_context")
+    def testMlFlowSetupRankNonRankZero(self, mock_get_context):
+        """Assert that non-rank-0 workers get a noop module"""
+        mock_context = MagicMock()
+        mock_context.get_world_rank.return_value = 1
+
+        mock_get_context.return_value = mock_context
+
+        mlflow = setup_mlflow({})
+        assert isinstance(mlflow, _NoopModule)
+
+        mlflow.log_metrics()
+        mlflow.sklearn.save_model(None, "model_directory")
+
+
+class MLflowUtilTest(unittest.TestCase):
+    def setUp(self):
+        self.dirpath = tempfile.mkdtemp()
+        import mlflow
+
+        mlflow.set_tracking_uri("sqlite:///" + self.dirpath + "/mlflow.sqlite")
+        mlflow.create_experiment(name="existing_experiment")
+
+        self.mlflow_util = _MLflowLoggerUtil()
+        self.tracking_uri = mlflow.get_tracking_uri()
+
+    def tearDown(self):
+        shutil.rmtree(self.dirpath)
+
+    def test_experiment_id(self):
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri, experiment_id="0")
+        assert self.mlflow_util.experiment_id == "0"
+
+    def test_experiment_id_env_var(self):
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
+        assert self.mlflow_util.experiment_id == "0"
+        del os.environ["MLFLOW_EXPERIMENT_ID"]
+
+    def test_experiment_name(self):
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="existing_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "1"
+
+    def test_run_started_with_correct_experiment(self):
+        experiment_name = "my_experiment_name"
+        # Make sure run is started under the correct experiment.
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name=experiment_name
+        )
+        run = self.mlflow_util.start_run(set_active=True)
+        assert (
+            run.info.experiment_id
+            == self.mlflow_util._mlflow.get_experiment_by_name(
+                experiment_name
+            ).experiment_id
+        )
+
+        self.mlflow_util.end_run()
+
+    def test_experiment_name_env_var(self):
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
+        assert self.mlflow_util.experiment_id == "1"
+        del os.environ["MLFLOW_EXPERIMENT_NAME"]
+
+    def test_id_precedence(self):
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "0"
+        del os.environ["MLFLOW_EXPERIMENT_ID"]
+
+    def test_new_experiment(self):
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "2"
+
+    def test_setup_fail(self):
+        with self.assertRaises(ValueError):
+            self.mlflow_util.setup_mlflow(
+                tracking_uri=self.tracking_uri,
+                experiment_name="new_experiment2",
+                create_experiment_if_not_exists=False,
+            )
+
+    def test_log_params(self):
+        params = {"a": "a", "x": {"y": "z"}}
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        run = self.mlflow_util.start_run()
+        run_id = run.info.run_id
+        self.mlflow_util.log_params(params_to_log=params, run_id=run_id)
+
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.params == flatten_dict(params)
+
+        params2 = {"b": "b"}
+        self.mlflow_util.start_run(set_active=True)
+        self.mlflow_util.log_params(params_to_log=params2, run_id=run_id)
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.params == flatten_dict(
+            {
+                **params,
+                **params2,
+            }
+        )
+
+        self.mlflow_util.end_run()
+
+    def test_log_metrics(self):
+        metrics = {"a": 1.0, "x": {"y": 2.0}}
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        run = self.mlflow_util.start_run()
+        run_id = run.info.run_id
+        self.mlflow_util.log_metrics(metrics_to_log=metrics, run_id=run_id, step=0)
+
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.metrics == flatten_dict(metrics)
+
+        metrics2 = {"b": 1.0}
+        self.mlflow_util.start_run(set_active=True)
+        self.mlflow_util.log_metrics(metrics_to_log=metrics2, run_id=run_id, step=0)
+        assert self.mlflow_util._mlflow.get_run(
+            run_id=run_id
+        ).data.metrics == flatten_dict(
+            {
+                **metrics,
+                **metrics2,
+            }
+        )
+        self.mlflow_util.end_run()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_wandb.py

@@ -0,0 +1,591 @@
+"""Tests for wandb integration.
+
+Note: These tests use a set of mocked APIs:
+- _MockWandbAPI: Mocks wandb API calls (ex: wandb.init).
+- _MockWandbLoggingActor: The same as the regular _WandbLoggingActor,
+    except using the mocked wandb API
+- WandbTestExperimentLogger: Thin subclass of `WandbLoggerCallback` to use for testing.
+    Provides a helper `trial_logging_actors` property that can be used to
+    access attributes of the remote actors for assertions.
+- Use the `get_mock_wandb_logger` helper method to create a logger with
+    a custom mock wandb API class. (Ex: If you want to override some wandb API methods.)
+
+Template for testing with these mocks:
+
+    wandb_logger_kwargs = {}
+    logger = get_mock_wandb_logger(mock_api_cls=_MockWandbAPI, **wandb_logger_kwargs)
+    logger.setup()
+
+    # From now on, the API key is in the env variable.
+    # Start the remote logging actor
+    logger.on_trial_start(0, [], trial)
+    # Log some results
+    result = {}
+    logger.on_trial_result(0, [], trial, result)
+    # Send a STOP signal to the logging actor
+    logger.on_trial_complete(0, [], trial)
+    # This will wait for the logging actor to finish + cleanup
+    logger.on_experiment_end(trials=[trial])
+
+    # Now, we can access properties of the logging actors
+    # (must happen after `on_trial_end` and `on_experiment_end`)
+    logger_state = logger.trial_logging_actor_states[trial]
+    # logger_state.logs, logger_state.config, logger_state.kwargs, ...
+"""
+
+import gc
+import os
+import tempfile
+import time
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+
+import ray
+from ray.air.integrations.wandb import (
+    WANDB_ENV_VAR,
+    WANDB_GROUP_ENV_VAR,
+    WANDB_POPULATE_RUN_LOCATION_HOOK,
+    WANDB_PROJECT_ENV_VAR,
+    WANDB_SETUP_API_KEY_HOOK,
+    RunDisabled,
+    WandbLoggerCallback,
+    _QueueItem,
+    _WandbLoggingActor,
+    setup_wandb,
+)
+from ray.air.tests.mocked_wandb_integration import (
+    Trial,
+    WandbTestExperimentLogger,
+    _MockWandbAPI,
+    _MockWandbLoggingActor,
+    get_mock_wandb_logger,
+)
+from ray.exceptions import RayActorError
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def trial():
+    trial_config = {"par1": 4, "par2": 9.12345678}
+    trial = Trial(
+        trial_config,
+        0,
+        "trial_0",
+        "trainable",
+        PlacementGroupFactory([{"CPU": 1}]),
+        "/tmp",
+    )
+    yield trial
+
+
+@pytest.fixture(autouse=True)
+def wandb_env():
+    """Clean up W&B env var before and after each test.
+
+    Even if we use monkeypatch in the test, this is useful to remove environment
+    variables that are set on the laptop when running tests locally.
+    """
+    if WANDB_ENV_VAR in os.environ:
+        del os.environ[WANDB_ENV_VAR]
+    yield
+    if WANDB_ENV_VAR in os.environ:
+        del os.environ[WANDB_ENV_VAR]
+
+
+def fake_wandb_populate_run_location_hook():
+    """Fake user-provided hook to populate W&B environment variables."""
+    os.environ[WANDB_PROJECT_ENV_VAR] = "test_project"
+    os.environ[WANDB_GROUP_ENV_VAR] = "test_group"
+
+
+FAKE_WANDB_POPULATE_RUN_LOCATION_HOOK_IMPORT_PATH = (
+    "ray.air.tests.test_integration_wandb.fake_wandb_populate_run_location_hook"
+)
+
+
+class TestWandbLogger:
+    def test_wandb_logger_project_group(self, monkeypatch):
+        monkeypatch.setenv(WANDB_PROJECT_ENV_VAR, "test_project_from_env_var")
+        monkeypatch.setenv(WANDB_GROUP_ENV_VAR, "test_group_from_env_var")
+        # Read project and group name from environment variable
+        logger = WandbTestExperimentLogger(api_key="1234")
+        logger.setup()
+        assert logger.project == "test_project_from_env_var"
+        assert logger.group == "test_group_from_env_var"
+
+    def test_wandb_logger_api_key_config(self, monkeypatch):
+        # No API key
+        with pytest.raises(ValueError):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+
+        # Fetch API key from argument even if external hook and WANDB_ENV_VAR set
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "abcde",
+        )
+        # API Key in config
+        logger = WandbTestExperimentLogger(project="test_project", api_key="1234")
+        logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "1234"
+
+    def test_wandb_logger_api_key_file(self, monkeypatch):
+        # Fetch API key from file even if external hook and WANDB_ENV_VAR set
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "abcde",
+        )
+        # API Key file
+        with tempfile.NamedTemporaryFile("wt") as fp:
+            fp.write("5678")
+            fp.flush()
+
+            logger = WandbTestExperimentLogger(
+                project="test_project", api_key_file=fp.name
+            )
+            logger.setup()
+            assert os.environ[WANDB_ENV_VAR] == "5678"
+
+    def test_wandb_logger_api_key_env_var(self, monkeypatch):
+        # API Key from env var takes precedence over external hook and
+        # logged in W&B API key
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "1234",
+        )
+        mock_wandb = Mock(api=Mock(api_key="efgh"))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "1234"
+        mock_wandb.ensure_configured.assert_not_called()
+
+    def test_wandb_logger_api_key_external_hook(self, monkeypatch):
+        # API Key from external hook if API key not provided through
+        # argument or WANDB_ENV_VAR and user not already logged in to W&B
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+
+        mock_wandb = Mock(api=Mock(api_key=None))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "abcd"
+        mock_wandb.ensure_configured.assert_called_once()
+
+        mock_wandb = Mock(ensure_configured=Mock(side_effect=AttributeError()))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "abcd"
+
+    def test_wandb_logger_api_key_from_wandb_login(self, monkeypatch):
+        # No API key should get set if user is already logged in to W&B
+        # and they didn't pass API key through argument or env var.
+        # External hook should not be called because user already logged
+        # in takes precedence.
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        mock_wandb = Mock()
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ.get(WANDB_ENV_VAR) is None
+        mock_wandb.ensure_configured.assert_called_once()
+
+    def test_wandb_logger_run_location_external_hook(self, monkeypatch):
+        with patch.dict(os.environ):
+            # No project
+            with pytest.raises(ValueError):
+                logger = WandbTestExperimentLogger(api_key="1234")
+                logger.setup()
+
+            # Project and group env vars from external hook
+            monkeypatch.setenv(
+                WANDB_POPULATE_RUN_LOCATION_HOOK,
+                FAKE_WANDB_POPULATE_RUN_LOCATION_HOOK_IMPORT_PATH,
+            )
+            logger = WandbTestExperimentLogger(api_key="1234")
+            logger.setup()
+            assert os.environ[WANDB_PROJECT_ENV_VAR] == "test_project"
+            assert os.environ[WANDB_GROUP_ENV_VAR] == "test_group"
+
+    def test_wandb_logger_start(self, monkeypatch, trial):
+        monkeypatch.setenv(WANDB_ENV_VAR, "9012")
+        # API Key in env
+        logger = WandbTestExperimentLogger(project="test_project")
+        logger.setup()
+        # From now on, the API key is in the env variable.
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert logger_state.kwargs["project"] == "test_project"
+        assert logger_state.kwargs["id"] == trial.trial_id
+        assert logger_state.kwargs["name"] == trial.trial_name
+        assert logger_state.kwargs["group"] == trial.experiment_dir_name
+        assert "config" in logger_state.exclude
+
+        del logger
+
+        # log config.
+        logger = WandbTestExperimentLogger(project="test_project", log_config=True)
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert "config" not in logger_state.exclude
+        assert "metric" not in logger_state.exclude
+
+        del logger
+
+        # Exclude metric.
+        logger = WandbTestExperimentLogger(project="test_project", excludes=["metric"])
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert "config" in logger_state.exclude
+        assert "metric" in logger_state.exclude
+
+        del logger
+
+    def test_wandb_logger_reporting(self, trial):
+        logger = WandbTestExperimentLogger(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger.on_trial_start(0, [], trial)
+        r1 = {
+            "metric1": 0.8,
+            "metric2": 1.4,
+            "metric3": np.asarray(32.0),
+            "metric4": np.float32(32.0),
+            "const": "text",
+            "config": trial.config,
+        }
+        logger.on_trial_result(0, [], trial, r1)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        logged = logger.trial_logging_actor_states[trial].logs[0]
+        assert "metric1" in logged
+        assert "metric2" not in logged
+        assert "metric3" in logged
+        assert "metric4" in logged
+        assert "const" not in logged
+        assert "config" not in logged
+
+    def test_wandb_logger_auto_config_keys(self, trial):
+        logger = WandbTestExperimentLogger(project="test_project", api_key="1234")
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        result = {key: 0 for key in WandbLoggerCallback.AUTO_CONFIG_KEYS}
+        logger.on_trial_result(0, [], trial, result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        config = logger.trial_logging_actor_states[trial].config
+        # The results in `AUTO_CONFIG_KEYS` should be saved as training configuration
+        # instead of output metrics.
+        assert set(WandbLoggerCallback.AUTO_CONFIG_KEYS) < set(config)
+
+    def test_wandb_logger_exclude_config(self):
+        trial = Trial(
+            config={"param1": 0, "param2": 0},
+            trial_id=0,
+            trial_name="trial_0",
+            experiment_dir_name="trainable",
+            placement_group_factory=PlacementGroupFactory([{"CPU": 1}]),
+            local_path=tempfile.gettempdir(),
+        )
+        logger = WandbTestExperimentLogger(
+            project="test_project",
+            api_key="1234",
+            excludes=(["param2"] + WandbLoggerCallback.AUTO_CONFIG_KEYS),
+        )
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+
+        # We need to test that `excludes` also applies to `AUTO_CONFIG_KEYS`.
+        result = {key: 0 for key in WandbLoggerCallback.AUTO_CONFIG_KEYS}
+        logger.on_trial_result(0, [], trial, result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+
+        config = logger.trial_logging_actor_states[trial].config
+        assert set(config) == {"param1"}
+
+    def test_set_serializability_result(self, trial):
+        """Tests that objects that contain sets can be serialized by wandb."""
+        logger = WandbTestExperimentLogger(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger.on_trial_start(0, [], trial)
+
+        # Testing for https://github.com/ray-project/ray/issues/28541
+        rllib_result = {
+            "env": "simple_spread",
+            "framework": "torch",
+            "num_gpus": 1,
+            "num_workers": 20,
+            "num_envs_per_env_runner": 1,
+            "compress_observations": True,
+            "lambda": 0.99,
+            "train_batch_size": 512,
+            "sgd_minibatch_size": 32,
+            "num_sgd_iter": 5,
+            "batch_mode": "truncate_episodes",
+            "entropy_coeff": 0.01,
+            "lr": 2e-05,
+            "multiagent": {
+                "policies": {"shared_policy"},
+                "policy_mapping_fn": lambda x: x,
+            },
+        }
+        logger.on_trial_result(0, [], trial, rllib_result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        logged = logger.trial_logging_actor_states[trial].logs[0]
+        assert logged != "serialization error"
+
+    def test_wandb_logging_actor_api_key(self, trial, monkeypatch):
+        """Tests that the wandb API key get propagated as an environment variable to
+        the remote logging actors."""
+
+        def mock_run(actor_cls):
+            return os.environ.get(WANDB_ENV_VAR)
+
+        monkeypatch.setattr(_MockWandbLoggingActor, "run", mock_run)
+
+        logger = WandbLoggerCallback(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger._logger_actor_cls = _MockWandbLoggingActor
+        logger.setup()
+        logger.log_trial_start(trial)
+        actor_env_var = ray.get(logger._trial_logging_futures[trial])
+        assert actor_env_var == "1234"
+
+    def test_wandb_finish(self, trial, tmp_path):
+        """Test that logging actors are cleaned up upon experiment completion."""
+        marker = tmp_path / "hang_marker"
+        marker.write_text("")
+
+        class HangingFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                while marker.exists():
+                    time.sleep(0.1)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=HangingFinishMockWandbAPI,
+            upload_timeout=1.0,
+        )
+        logger.setup()
+        logger.on_trial_start(0, [], trial)
+        logger.on_trial_complete(0, [], trial)
+        # Signalling stop will not cleanup fully due to the hanging finish
+        assert logger._trial_logging_actors
+        marker.unlink()
+        # wandb.finish has ended -> experiment end hook should cleanup actors fully
+        logger.on_experiment_end(trials=[trial])
+        assert not logger._trial_logging_actors
+
+    def test_wandb_kill_hanging_actor(self, trial):
+        """Test that logging actors are killed if exceeding the upload timeout
+        upon experiment completion."""
+
+        class HangingFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                time.sleep(5)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=HangingFinishMockWandbAPI,
+            upload_timeout=0.1,
+        )
+        logger.setup()
+        logger.on_trial_start(0, [], trial)
+        logger.on_trial_complete(0, [], trial)
+        # Signalling stop will not cleanup fully due to the hanging finish
+        assert logger._trial_logging_actors
+        actor = logger._trial_logging_actors[trial]
+        # Experiment end hook should kill actors since upload_timeout < 5
+        logger.on_experiment_end(trials=[trial])
+        assert not logger._trial_logging_actors
+        gc.collect()
+        with pytest.raises(RayActorError):
+            ray.get(actor.get_state.remote())
+
+    def test_wandb_destructor(self, trial):
+        """Test that the WandbLoggerCallback destructor forcefully cleans up
+        logging actors."""
+
+        class SlowFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                time.sleep(5)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=SlowFinishMockWandbAPI,
+            upload_timeout=1.0,
+        )
+
+        logger.setup()
+        # Triggers logging actor run loop
+        logger.on_trial_start(0, [], trial)
+        actor = logger._trial_logging_actors[trial]
+        del logger
+        gc.collect()
+        with pytest.raises(RayActorError):
+            ray.get(actor.get_state.remote())
+
+    def test_wandb_logging_actor_fault_tolerance(self, trial):
+        """Tests that failing wandb logging actors are restarted"""
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            fail_marker = Path(tempdir) / "fail_marker"
+
+            class _FailingWandbLoggingActor(_MockWandbLoggingActor):
+                def _handle_result(self, result):
+                    if (
+                        result.get("training_iteration") == 3
+                        and not fail_marker.exists()
+                    ):
+                        fail_marker.write_text("Ok")
+                        raise SystemExit
+
+                    return super()._handle_result(result)
+
+            logger = WandbLoggerCallback(
+                project="test_project", api_key="1234", excludes=["metric2"]
+            )
+            logger._logger_actor_cls = _FailingWandbLoggingActor
+            logger.setup()
+            logger.log_trial_start(trial)
+
+            actor = logger._trial_logging_actors[trial]
+            queue = logger._trial_queues[trial]
+
+            logger.log_trial_result(1, trial, result={"training_iteration": 1})
+            logger.log_trial_result(2, trial, result={"training_iteration": 2})
+            logger.log_trial_result(3, trial, result={"training_iteration": 3})
+
+            logger.log_trial_result(4, trial, result={"training_iteration": 4})
+            logger.log_trial_result(5, trial, result={"training_iteration": 5})
+
+            queue.put((_QueueItem.END, None))
+
+            # Wait for the actor's run method to complete
+            ray.get(logger._trial_logging_futures[trial])
+
+            state = ray.get(actor.get_state.remote())
+            assert [metrics["training_iteration"] for metrics in state.logs] == [4, 5]
+
+    def test_wandb_restart(self, trial):
+        """Test that the WandbLoggerCallback reuses actors for trial restarts."""
+
+        logger = WandbLoggerCallback(project="test_project", api_key="1234")
+        logger._logger_actor_cls = _MockWandbLoggingActor
+        logger.setup()
+
+        assert len(logger._trial_logging_futures) == 0
+        assert len(logger._logging_future_to_trial) == 0
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
+
+def test_wandb_logging_process_run_info_hook(monkeypatch):
+    """
+    Test WANDB_PROCESS_RUN_INFO_HOOK in _WandbLoggingActor is
+    correctly called by calling _WandbLoggingActor.run() mocking
+    out calls to wandb.
+    """
+    mock_queue = Mock(get=Mock(return_value=(_QueueItem.END, None)))
+    monkeypatch.setenv(
+        "WANDB_PROCESS_RUN_INFO_HOOK", "mock_wandb_process_run_info_hook"
+    )
+
+    with patch.object(ray.air.integrations.wandb, "load_class") as mock_load_class:
+        logging_process = _WandbLoggingActor(
+            logdir="/tmp", queue=mock_queue, exclude=[], to_config=[]
+        )
+        logging_process._wandb = Mock()
+        logging_process.run()
+
+    logging_process._wandb.init.assert_called_once()
+    run = logging_process._wandb.init.return_value
+    mock_load_class.assert_called_once_with("mock_wandb_process_run_info_hook")
+    external_hook = mock_load_class.return_value
+    external_hook.assert_called_once_with(run)
+    logging_process._wandb.finish.assert_called_once()
+
+
+def test_wandb_logger_rank_zero_only(trial, monkeypatch):
+    """Test that logging is disabled for non-rank-0 workers when rank_zero_only is True."""
+
+    monkeypatch.setenv(
+        WANDB_ENV_VAR,
+        "abcde",
+    )
+
+    mock_session = Mock()
+    mock_session.experiment_name = "test_project"
+    mock_session.trial_name = "trial_0"
+    mock_session.trial_id = "trial_0"
+
+    # Test case 1: rank_zero_only=True, rank 0
+    mock_session.world_rank = 0
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+    # Test case 2: rank_zero_only=True, non-rank-0
+    mock_session.world_rank = 1
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert isinstance(run, RunDisabled)
+
+    # Test case 3: rank_zero_only=False, any rank
+    mock_session.world_rank = 1
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=False, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+    # Test case 4: rank_zero_only=True, no session
+    with patch("ray.air.integrations.wandb.get_session", return_value=None):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_keras_callback.py

@@ -0,0 +1,223 @@
+import os
+import sys
+from typing import Dict, Tuple
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.air.integrations.keras import ReportCheckpointCallback
+    from ray.train.tensorflow import TensorflowTrainer
+
+
+class TestReportCheckpointCallback:
+    @pytest.fixture(name="model")
+    def model_fixture(self):
+        model = tf.keras.Sequential(
+            [tf.keras.layers.InputLayer(input_shape=(1,)), tf.keras.layers.Dense(1)]
+        )
+        model.compile(
+            optimizer="sgd",
+            loss="mean_squared_error",
+            metrics=["accuracy"],
+        )
+        return model
+
+    @patch("ray.train.report")
+    @pytest.mark.parametrize(
+        "metrics, expected_metrics_keys",
+        [
+            (None, {"loss", "accuracy", "val_loss", "val_accuracy"}),
+            ("loss", {"loss"}),
+            (["loss", "accuracy"], {"loss", "accuracy"}),
+            ({"spam": "loss"}, {"spam"}),
+        ],
+    )
+    def test_reported_metrics_contain_expected_keys(
+        self, mock_report, metrics, expected_metrics_keys, model
+    ):
+        # Reported metrics contain different keys depending on the value passed to the
+        # `metrics` parameter. This test varies the value of `metrics` and asserts that
+        # the reported keys are correct.
+        model.fit(
+            x=np.zeros((1, 1)),
+            y=np.zeros((1, 1)),
+            validation_data=(np.zeros((1, 1)), np.zeros((1, 1))),
+            callbacks=[ReportCheckpointCallback(metrics=metrics)],
+        )
+
+        for (metrics,), _ in ray.train.report.call_args_list:
+            assert metrics.keys() == expected_metrics_keys
+
+    @patch("ray.train.report")
+    def test_report_with_default_arguments(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` with default arguments. The test
+        # simulates the end of an epoch, and asserts that a metric and checkpoint are
+        # reported.
+        callback = ReportCheckpointCallback()
+        callback.set_model(model)
+
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 1
+        metrics, checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert metrics == {"loss": 0}
+        assert checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_checkpoint_on_list(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `checkpoint_on` is a `list`. The
+        # test simulates each event in `checkpoint_on`, and asserts that a checkpoint
+        # is reported for each event.
+        callback = ReportCheckpointCallback(
+            checkpoint_on=["epoch_end", "train_batch_end"]
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 2
+        _, first_checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert first_checkpoint is not None
+        _, second_checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert second_checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_report_metrics_on_list(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `report_metrics_on` is a `list`.
+        # The test simulates each event in `report_metrics_on`, and asserts that metrics
+        # are reported for each event.
+        callback = ReportCheckpointCallback(
+            report_metrics_on=["epoch_end", "train_batch_end"]
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 1})
+
+        assert len(ray.train.report.call_args_list) == 2
+        first_metric, _ = self.parse_call(ray.train.report.call_args_list[0])
+        assert first_metric == {"loss": 0}
+        second_metric, _ = self.parse_call(ray.train.report.call_args_list[1])
+        assert second_metric == {"loss": 1}
+
+    @patch("ray.train.report")
+    def test_report_and_checkpoint_on_different_events(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `report_metrics_on` and
+        # `checkpoint_on` are different. The test asserts that:
+        # 1. Checkpoints are reported on `checkpoint_on`
+        # 2. Metrics are reported on `report_metrics_on`
+        # 3. Metrics are reported with checkpoints
+        callback = ReportCheckpointCallback(
+            report_metrics_on="train_batch_end", checkpoint_on="epoch_end"
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 1})
+
+        assert len(ray.train.report.call_args_list) == 2
+        first_metric, first_checkpoint = self.parse_call(
+            ray.train.report.call_args_list[0]
+        )
+        assert first_metric == {"loss": 0}
+        assert first_checkpoint is None
+        second_metric, second_checkpoint = self.parse_call(
+            ray.train.report.call_args_list[1]
+        )
+        # We should always include metrics, even if it isn't during one of the events
+        # specified in `report_metrics_on`.
+        assert second_metric == {"loss": 1}
+        assert second_checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_report_delete_tempdir(self, mock_report, model):
+        # This tests `ReportCheckpointCallback`. The test simulates the end of an epoch,
+        # and asserts that the temporary checkpoint directory is deleted afterwards.
+        callback = ReportCheckpointCallback()
+        callback.model = model
+
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 1
+        metrics, checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert metrics == {"loss": 0}
+        assert checkpoint is not None
+        assert checkpoint.path is not None
+        assert not os.path.exists(checkpoint.path)
+
+    def parse_call(self, call) -> Tuple[Dict, train.Checkpoint]:
+        (metrics,), kwargs = call
+        checkpoint = kwargs["checkpoint"]
+        return metrics, checkpoint
+
+
+def get_dataset(a=5, b=10, size=1000):
+    items = [i / size for i in range(size)]
+    dataset = ray.data.from_items([{"x": x, "y": a * x + b} for x in items])
+    return dataset
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=()),
+            # Add feature dimension, expanding (batch_size,) to (batch_size, 1).
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_model()
+        multi_worker_model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
+            loss=tf.keras.losses.mean_squared_error,
+            metrics=[tf.keras.metrics.mean_squared_error],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    for _ in range(config.get("epoch", 3)):
+        tf_dataset = dataset.to_tf("x", "y", batch_size=32)
+        multi_worker_model.fit(tf_dataset, callbacks=[ReportCheckpointCallback()])
+
+
+def test_keras_callback_e2e():
+    epochs = 3
+    config = {
+        "epochs": epochs,
+    }
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=2),
+        datasets={TRAIN_DATASET_KEY: get_dataset()},
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_new_dataset_config.py

@@ -0,0 +1,394 @@
+import random
+from typing import Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+from ray import train
+from ray.data import DataIterator
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.tests.conftest import *  # noqa
+from ray.train import DataConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+class TestBasic(DataParallelTrainer):
+    def __init__(
+        self, num_workers: int, expect_ds: bool, expect_sizes: Optional[dict], **kwargs
+    ):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            for k, v in expect_sizes.items():
+                shard = train.get_dataset_shard(k)
+                if v == -1:
+                    assert shard is None, shard
+                else:
+                    count = 0
+                    for batch in shard.iter_batches():
+                        for arr in batch.values():
+                            count += arr.size
+                    assert count == v, shard
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_basic(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Single worker basic case.
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+    )
+    test.fit()
+
+    # Single worker, no test ds.
+    test = TestBasic(1, True, {"train": 10, "test": -1}, datasets={"train": ds})
+    test.fit()
+
+    # Two workers, train and test split.
+    test = TestBasic(
+        2, True, {"train": 5, "test": 5}, datasets={"train": ds, "test": ds}
+    )
+    test.fit()
+
+    # Two workers, both split.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        dataset_config=DataConfig(datasets_to_split=["train", "test"]),
+        datasets={"train": ds, "test": ds},
+    )
+    # Test get config.
+    assert isinstance(test.get_dataset_config(), DataConfig)
+    test.fit()
+
+
+def test_split(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Split all by default
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5, "val": 5},
+        datasets={"train": ds, "test": ds, "val": ds},
+    )
+    test.fit()
+
+    # Test flag "all"
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split="all"),
+    )
+
+    # Test split train only.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=["train"]),
+    )
+    test.fit()
+
+    # Test invalid arguments
+    for datasets_to_split in ["train", ("train"), {}]:
+        with pytest.raises(TypeError, match="`datasets_to_split` should be.*"):
+            test = TestBasic(
+                2,
+                True,
+                {"train": 5, "test": 10},
+                datasets={"train": ds, "test": ds},
+                dataset_config=DataConfig(datasets_to_split=datasets_to_split),
+            )
+
+    # Test empty `datasets_to_split` list
+    test = TestBasic(
+        2,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def test_configure_execution_options_carryover_context(ray_start_4_cpus):
+    """Tests that execution options in DataContext are carried over to DatConfig
+    automatically."""
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.preserve_order = True
+    ctx.execution_options.verbose_progress = True
+
+    data_config = DataConfig()
+
+    ingest_options = data_config.default_ingest_options()
+    assert ingest_options.preserve_order is True
+    assert ingest_options.verbose_progress is True
+
+
+@pytest.mark.parametrize("enable_locality", [True, False])
+def test_configure_locality(enable_locality):
+    data_config = DataConfig(enable_shard_locality=enable_locality)
+
+    mock_ds = MagicMock()
+    mock_ds.streaming_split = MagicMock()
+    mock_ds.copy = MagicMock(return_value=mock_ds)
+    world_size = 2
+    worker_handles = [MagicMock() for _ in range(world_size)]
+    worker_node_ids = ["node" + str(i) for i in range(world_size)]
+    data_config.configure(
+        datasets={"train": mock_ds},
+        world_size=world_size,
+        worker_handles=worker_handles,
+        worker_node_ids=worker_node_ids,
+    )
+    mock_ds.streaming_split.assert_called_once()
+    mock_ds.streaming_split.assert_called_with(
+        world_size,
+        equal=True,
+        locality_hints=worker_node_ids if enable_locality else None,
+    )
+
+
+class CustomConfig(DataConfig):
+    def __init__(self):
+        pass
+
+    def configure(self, *args, **kwargs):
+        ds = ray.data.range(10)
+        return [
+            {"train": ds.iterator()},
+            {"train": ds.iterator()},
+        ]
+
+
+def test_custom_config_subclass(ray_start_4_cpus):
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10},
+        dataset_config=CustomConfig(),
+    )
+    test.fit()
+
+
+class TestRandom(DataParallelTrainer):
+    def __init__(self, num_workers: int, expect_random: bool, **kwargs):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            epoch1 = list(data_shard.iter_rows())
+            epoch2 = list(data_shard.iter_rows())
+            print("Epochs", epoch1, "\n", epoch2)
+            if expect_random:
+                assert epoch1 != epoch2
+            else:
+                assert epoch1 == epoch2
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_per_epoch_preprocessing(ray_start_4_cpus):
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+
+def test_materialized_preprocessing(ray_start_4_cpus):
+    # TODO(ekl) we should test all these configs with splitting enabled, but this
+    # requires implementing deterministic streaming split.
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def _run_data_config_resource_test(data_config):
+    cluster_cpus, cluster_gpus = 20, 10
+    num_workers = 2
+    # Resources used by training workers.
+    cpus_per_worker, gpus_per_worker = 2, 1
+
+    original_execution_options = data_config._get_execution_options("train")
+
+    ray.init(num_cpus=cluster_cpus, num_gpus=cluster_gpus)
+
+    class MyTrainer(DataParallelTrainer):
+        def __init__(self, **kwargs):
+            def train_loop_fn():
+                train_ds = train.get_dataset_shard("train")
+                new_execution_options = train_ds.get_context().execution_options
+                if original_execution_options.is_resource_limits_default():
+                    # If the original resource limits are default, the new resource
+                    # limits should be the default as well.
+                    assert new_execution_options.is_resource_limits_default()
+                    exclude_resources = new_execution_options.exclude_resources
+                    assert (
+                        exclude_resources.cpu
+                        == original_execution_options.exclude_resources.cpu
+                        + cpus_per_worker * num_workers
+                        + 1  # trainer coordinator
+                    )
+                    assert (
+                        exclude_resources.gpu
+                        == original_execution_options.exclude_resources.gpu
+                        + gpus_per_worker * num_workers
+                    )
+                else:
+                    # If the original resource limits are not default, the new resource
+                    # limits should be the same as the original ones.
+                    # And the new exclude_resources should be zero.
+                    assert (
+                        new_execution_options.resource_limits
+                        == original_execution_options.resource_limits
+                    )
+                    assert (
+                        new_execution_options.exclude_resources
+                        == ExecutionResources.zero()
+                    )
+
+            kwargs.pop("scaling_config", None)
+
+            super().__init__(
+                train_loop_per_worker=train_loop_fn,
+                scaling_config=ScalingConfig(
+                    num_workers=num_workers,
+                    use_gpu=True,
+                    resources_per_worker={
+                        "CPU": cpus_per_worker,
+                        "GPU": gpus_per_worker,
+                    },
+                ),
+                datasets={"train": ray.data.range(10)},
+                dataset_config=data_config,
+                **kwargs,
+            )
+
+    trainer = MyTrainer()
+    trainer.fit()
+
+
+def test_data_config_default_resource_limits(shutdown_only):
+    """Test that DataConfig preserves user-configured exclude_resources."""
+    execution_options = ExecutionOptions()
+    execution_options.exclude_resources = execution_options.exclude_resources.copy(
+        cpu=2, gpu=1
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_data_config_manual_resource_limits(shutdown_only):
+    """Test manually setting resource limits in DataConfig."""
+    execution_options = ExecutionOptions()
+    execution_options.resource_limits = execution_options.resource_limits.copy(
+        cpu=10, gpu=5
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_v1_train_with_v2_data_autoscaler_sets_exclude_resources(
+    shutdown_only, monkeypatch
+):
+    """Regression test for the Train V1 + V2 cluster autoscaler combination."""
+    monkeypatch.setenv("RAY_DATA_CLUSTER_AUTOSCALER", "V2")
+
+    ray.init(num_cpus=10, num_gpus=2)
+
+    num_train_cpus, num_train_gpus = 4.0, 2.0
+    data_config = DataConfig()
+    data_config.set_train_total_resources(
+        num_train_cpus=num_train_cpus, num_train_gpus=num_train_gpus
+    )
+
+    iterators = data_config.configure(
+        datasets={"train": ray.data.range(10)},
+        world_size=2,
+        worker_handles=None,
+        worker_node_ids=None,
+    )
+
+    exclude_resources = (
+        iterators[0]["train"].get_context().execution_options.exclude_resources
+    )
+    assert exclude_resources.cpu == num_train_cpus
+    assert exclude_resources.gpu == num_train_gpus
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_remote_storage_hdfs.py

@@ -0,0 +1,50 @@
+"""Test remote_storage in a ci environment with real hdfs setup."""
+
+import os
+
+import pytest
+
+from ray.train.v2._internal.execution.storage import (
+    _list_at_fs_path,
+    _upload_to_fs_path,
+    get_fs_and_path,
+)
+
+
+@pytest.fixture
+def setup_hdfs():
+    """Set env vars required by pyarrow to talk to hdfs correctly.
+
+    Returns hostname and port needed for the hdfs uri."""
+
+    # the following file is written in `install-hdfs.sh`.
+    with open("/tmp/hdfs_env", "r") as f:
+        for line in f.readlines():
+            line = line.rstrip("\n")
+            tokens = line.split("=", maxsplit=1)
+            os.environ[tokens[0]] = tokens[1]
+    import sys
+
+    sys.path.insert(0, os.path.join(os.environ["HADOOP_HOME"], "bin"))
+    hostname = os.getenv("CONTAINER_ID")
+    port = os.getenv("HDFS_PORT")
+    yield hostname, port
+
+
+def test_hdfs(tmp_path, setup_hdfs):
+    pytest.skip("TODO: Fix this test")
+
+    hostname, port = setup_hdfs
+    hdfs_uri = f"hdfs://{hostname}:{port}/test/"
+    fs, path = get_fs_and_path(hdfs_uri)
+
+    dummy_file = tmp_path.joinpath("dummy.txt")
+    dummy_file.write_text("dummy")
+    _upload_to_fs_path(dummy_file, fs, path)
+    assert _list_at_fs_path(fs, path) == ["dummy.txt"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_tracebacks.py

@@ -0,0 +1,102 @@
+import pytest
+from tblib import pickling_support
+
+import ray
+from ray import cloudpickle
+from ray.air._internal.util import StartTraceback, exception_cause, skip_exceptions
+from ray.tune import Tuner
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def _failing_recursive(levels: int = 0, start_traceback: int = -1):
+    if levels > 0:
+        if start_traceback == 0:
+            try:
+                _failing_recursive(
+                    levels=levels - 1, start_traceback=start_traceback - 1
+                )
+            except Exception as e:
+                raise StartTraceback from e
+        else:
+            _failing_recursive(levels=levels - 1, start_traceback=start_traceback - 1)
+    else:
+        raise RuntimeError("Failing")
+
+
+@pytest.mark.parametrize("levels", [4, 5, 6, 7, 8, 9, 10])
+def test_short_traceback(levels):
+    start_traceback = 3
+    with pytest.raises(StartTraceback) as exc_info:
+        _failing_recursive(levels=levels, start_traceback=start_traceback)
+
+    exc = skip_exceptions(exc_info.value)
+    tb = exc.__traceback__
+    i = 0
+    while tb:
+        i += 1
+        tb = tb.tb_next
+
+    assert i == levels - start_traceback + 1
+
+
+def test_recursion():
+    """Test that the skipped exception does not point to the original exception."""
+    root_exception = None
+
+    with pytest.raises(StartTraceback) as exc_info:
+        try:
+            raise Exception("Root Exception")
+        except Exception as e:
+            root_exception = e
+            raise StartTraceback from root_exception
+
+    assert root_exception, "Root exception was not captured."
+
+    start_traceback = exc_info.value
+    skipped_exception = skip_exceptions(start_traceback)
+
+    assert (
+        root_exception != skipped_exception
+    ), "Skipped exception points to the original exception."
+
+
+def test_tblib():
+    """Test that tblib does not cause a maximum recursion error."""
+
+    with pytest.raises(Exception) as exc_info:
+        try:
+            try:
+                raise Exception("Root Exception")
+            except Exception as root_exception:
+                raise StartTraceback from root_exception
+        except Exception as start_traceback:
+            raise skip_exceptions(start_traceback) from exception_cause(start_traceback)
+
+    pickling_support.install()
+    reraised_exception = exc_info.value
+    # This should not raise a RecursionError/PicklingError.
+    cloudpickle.dumps(reraised_exception)
+
+
+def test_traceback_tuner(ray_start_2_cpus):
+    """Ensure that the Tuner's stack trace is not too long."""
+
+    def failing(config):
+        raise RuntimeError("Error")
+
+    tuner = Tuner(failing)
+    results = tuner.fit()
+    assert len(str(results[0].error).split("\n")) <= 20
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_utils.py

@@ -0,0 +1,34 @@
+"""Test AIR internal utilities (under ray.air._internal)."""
+
+import os
+
+import pytest
+
+import ray
+from ray.air._internal.filelock import RAY_LOCKFILE_DIR, TempFileLock
+
+
+def test_temp_file_lock(tmp_path, monkeypatch):
+    """Test that the directory where temp file locks are saved can be configured
+    via the env variable that configures the global Ray temp dir."""
+    monkeypatch.setenv("RAY_TMPDIR", str(tmp_path))
+    assert str(tmp_path) in ray._common.utils.get_default_system_temp_dir()
+    with TempFileLock(path="abc.txt"):
+        assert RAY_LOCKFILE_DIR in os.listdir(tmp_path)
+        assert os.listdir(tmp_path / RAY_LOCKFILE_DIR)
+
+
+def test_multiple_file_locks(tmp_path, monkeypatch):
+    """Test that a new file lock is created for unique paths."""
+    monkeypatch.setenv("RAY_TMPDIR", str(tmp_path))
+    with TempFileLock(path="abc.txt"):
+        with TempFileLock(path="subdir/abc.txt"):
+            assert RAY_LOCKFILE_DIR in os.listdir(tmp_path)
+            # We should have 2 locks, one for abc.txt and one for subdir/abc.txt
+            assert len(os.listdir(tmp_path / RAY_LOCKFILE_DIR)) == 2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))