diff --git a/README.md b/README.md index 151f74b19..d04d2ca8f 100644 --- a/README.md +++ b/README.md @@ -158,7 +158,25 @@ disables workers until the study is terminated and relaunched. If you are runnin or with a small number of workers, this could halt your entire study until you manually kill and relaunch it. In the Ray parallel backend, we've implemented a system to automatically terminate jobs exceeding a specified timeout. This feature is particularly useful when task hanging limits your -experiments. +experiments. + +Each Ray experiment uses its own `env_args.max_steps * avg_step_timeout` budget, +capped by `episode_timeout` when set. Without `max_steps`, the default is ten +hours. Budgets apply to running attempts, excluding time queued for workers or +dependencies. A timeout first requests cooperative cancellation; after 60 seconds +without completion, the worker is force killed. Runtime measurements use Ray's +task-state service; when state information is unavailable, no deadline is inferred +from submission time. Sequential and joblib execution do not enforce these budgets. + +Dependency graphs are validated before agent servers start or experiment +directories are prepared. Duplicate IDs, missing dependencies, and cycles raise +`ValueError`. Ray and sequential execution both honor dependencies regardless of +input order; joblib rejects nonempty dependency graphs, so use Ray or sequential +for these studies. Failed Ray tasks are returned as exceptions by experiment ID; +dependents of a raised task do not run, while independent branches finish. The +launcher prepares and closes each shared agent argument object once, and releases +acquired agents even if preparation or execution fails. Cleanup attempts continue +when one resource fails to close, preserving the original execution error. ### Debugging diff --git a/src/agentlab/experiments/graph_execution_ray.py b/src/agentlab/experiments/graph_execution_ray.py index f7aad7808..1d9842f77 100644 --- a/src/agentlab/experiments/graph_execution_ray.py +++ b/src/agentlab/experiments/graph_execution_ray.py @@ -1,9 +1,15 @@ +"""Dependency-aware Ray execution and runtime timeout monitoring.""" + import logging +import math import time +from collections.abc import Mapping +from numbers import Real import bgym import ray from agentlab.experiments.exp_utils import _episode_timeout, run_exp +from agentlab.experiments.task_graph import TaskGraph from ray.util import state logger = logging.getLogger(__name__) @@ -11,97 +17,188 @@ run_exp = ray.remote(run_exp) -def execute_task_graph(exp_args_list: list[bgym.ExpArgs], avg_step_timeout=60): - """Execute a task graph in parallel while respecting dependencies using Ray.""" +def _validate_timeout(value, name, allow_zero=False): + if ( + isinstance(value, bool) + or not isinstance(value, Real) + or not math.isfinite(value) + or value < 0 + or (value == 0 and not allow_zero) + ): + qualifier = "nonnegative" if allow_zero else "positive" + raise ValueError(f"{name} must be a finite {qualifier} number, got {value!r}") + return float(value) - exp_args_map = {exp_args.exp_id: exp_args for exp_args in exp_args_list} - task_map = {} - def get_task(exp_arg: bgym.ExpArgs): - if exp_arg.exp_id not in task_map: - # Get all dependency tasks first - dependency_tasks = [get_task(exp_args_map[dep_key]) for dep_key in exp_arg.depends_on] +def _experiment_timeouts(experiments, avg_step_timeout): + _validate_timeout(avg_step_timeout, "avg_step_timeout") + timeouts = {} + for exp in experiments: + if hasattr(exp, "episode_timeout"): + _validate_timeout(exp.episode_timeout, f"episode_timeout for {exp.exp_id!r}") + timeouts[exp.exp_id] = _validate_timeout( + _episode_timeout(exp, avg_step_timeout), f"timeout for {exp.exp_id!r}" + ) + return timeouts - # Create new task that depends on the dependency results - task_map[exp_arg.exp_id] = run_exp.options(name=f"{exp_arg.exp_name}").remote( - exp_arg, *dependency_tasks, avg_step_timeout=avg_step_timeout - ) - return task_map[exp_arg.exp_id] - # Build task graph - for exp_arg in exp_args_list: - get_task(exp_arg) - - max_timeout = max([_episode_timeout(exp_args, avg_step_timeout) for exp_args in exp_args_list]) +def execute_task_graph(exp_args_list: list[bgym.ExpArgs], avg_step_timeout=60): + """Execute a validated experiment DAG, returning results by experiment ID. - return poll_for_timeout(task_map, max_timeout, poll_interval=max_timeout * 0.1) + The complete graph and all deadlines are checked before submitting any work. + Dependencies are top-level Ray object references, so Ray does not run a child + until its parents finish. A raised parent exception prevents its descendants + from running; independent branches are still collected. Exceptions are values + in the returned mapping, as in ``poll_for_timeout``. + Args: + exp_args_list: Experiments with assigned IDs and their dependency IDs. + avg_step_timeout: Runtime budget per environment step, in seconds. -def poll_for_timeout(tasks: dict[str, ray.ObjectRef], timeout: float, poll_interval: float = 1.0): - """Cancel tasks that exceeds the timeout + Returns: + A mapping from experiment ID to its result or exception. - I tried various different methods for killing a job that hangs. so far it's - the only one that seems to work reliably (hopefully) + Raises: + BaseException: If graph validation, submission, or monitoring fails. + """ + graph = TaskGraph(exp_args_list) + timeouts = _experiment_timeouts(graph.experiments, avg_step_timeout) + tasks = {} + try: + for exp_arg in graph.ordered_experiments: + dependencies = [tasks[exp_id] for exp_id in graph.dependencies[exp_arg.exp_id]] + tasks[exp_arg.exp_id] = run_exp.options(name=f"{exp_arg.exp_name}").remote( + exp_arg, *dependencies, avg_step_timeout=avg_step_timeout + ) + except BaseException: + _cancel_tasks(tasks.values()) + raise + + # Keep the result mapping in caller order, independently of submission order. + tasks = {exp.exp_id: tasks[exp.exp_id] for exp in graph.experiments} + poll_interval = min(1.0, min(timeouts.values()) * 0.1) if timeouts else 1.0 + return poll_for_timeout(tasks, timeouts, poll_interval=poll_interval) + + +def _cancel_tasks(tasks): + """Best-effort cleanup that never hides an execution or interruption error.""" + for task in tasks: + try: + ray.cancel(task, force=True, recursive=False) + except Exception: + logger.exception("Could not cancel outstanding experiment task %r", task) + + +def poll_for_timeout( + tasks: dict[str, ray.ObjectRef], + timeout: float | Mapping[str, float], + poll_interval: float = 1.0, + *, + cancel_grace_period: float = 60.0, +): + """Collect results and cancel tasks that exceed their own running time limit. + + ``timeout`` can be one positive number for every task, or a mapping covering + exactly the experiment IDs in ``tasks``. Time waiting for dependencies or an + available worker is excluded. Ray task-state information is best effort; when + unavailable, no runtime deadline is inferred from submission time. + + A timed-out task receives one cooperative cancellation, then one force kill + if it is still pending after ``cancel_grace_period`` seconds. The grace period + uses a monotonic clock and continues even when task-state queries fail. Result + collection waits for Ray to acknowledge completion/cancellation so dependent + references reach a terminal state. Ordinary task errors are returned alongside + successful results. Driver errors and interrupts cancel outstanding work and + propagate to the caller. Args: - tasks: dict[str, ray.ObjectRef] - Dictionary of task_id: task_ref - timeout: float - Timeout in seconds - poll_interval: float - Polling interval in seconds + tasks: Mapping of experiment IDs to unique Ray object references. + timeout: Shared runtime limit or a per-experiment mapping, in seconds. + poll_interval: Maximum time to wait for each batch of ready references. + cancel_grace_period: Seconds between cooperative cancellation and force kill. Returns: - dict[str, Any]: Dictionary of task_id: result + A mapping in input order containing each task result or exception. + + Raises: + ValueError: If deadlines, intervals, or object references are invalid. + BaseException: If the driver is interrupted or Ray monitoring fails. """ - task_list = list(tasks.values()) - task_ids = list(tasks.keys()) - - logger.warning(f"Any task exceeding {timeout} seconds will be cancelled.") - - while True: - ready, not_ready = ray.wait(task_list, num_returns=len(task_list), timeout=poll_interval) - for task in not_ready: - elapsed_time = get_elapsed_time(task) - # print(f"Task {task.task_id().hex()} elapsed time: {elapsed_time}") - if elapsed_time is not None and elapsed_time > timeout: - msg = f"Task {task.task_id().hex()} hase been running for {elapsed_time}s, more than the timeout: {timeout}s." - if elapsed_time < timeout + 60 + poll_interval: - logger.warning(msg + " Cancelling task.") - ray.cancel(task, force=False, recursive=False) - else: - logger.warning(msg + " Force killing.") - ray.cancel(task, force=True, recursive=False) - if len(ready) == len(task_list): - results = [] + _validate_timeout(poll_interval, "poll_interval") + _validate_timeout(cancel_grace_period, "cancel_grace_period", allow_zero=True) + if isinstance(timeout, Mapping): + if set(timeout) != set(tasks): + raise ValueError("timeout mapping must cover exactly the submitted experiment IDs") + timeouts = { + task_id: _validate_timeout(value, f"timeout for {task_id!r}") + for task_id, value in timeout.items() + } + else: + limit = _validate_timeout(timeout, "timeout") + timeouts = dict.fromkeys(tasks, limit) + + pending = {ref: task_id for task_id, ref in tasks.items()} + if len(pending) != len(tasks): + raise ValueError("Each experiment must have a distinct Ray object reference") + results = {} + cancelled_at = {} + force_cancelled = set() + try: + while pending: + ready, not_ready = ray.wait( + list(pending), num_returns=len(pending), timeout=poll_interval + ) for task in ready: + task_id = pending.pop(task) try: - result = ray.get(task) - except Exception as e: - logger.exception(f"Task failed: {e}") - result = e - results.append(result) - - return {task_id: result for task_id, result in zip(task_ids, results)} + results[task_id] = ray.get(task) + except Exception as error: + logger.exception("Experiment %s failed", task_id) + results[task_id] = error + + for task in not_ready: + task_id = pending[task] + if task in cancelled_at: + if ( + task not in force_cancelled + and time.monotonic() - cancelled_at[task] >= cancel_grace_period + ): + logger.warning( + "Force killing experiment %s after cancellation grace", task_id + ) + ray.cancel(task, force=True, recursive=False) + force_cancelled.add(task) + continue + elapsed_time = get_elapsed_time(task) + if elapsed_time is not None and elapsed_time > timeouts[task_id]: + logger.warning( + "Experiment %s exceeded its timeout (%ss > %ss); cancelling", + task_id, + elapsed_time, + timeouts[task_id], + ) + ray.cancel(task, force=False, recursive=False) + cancelled_at[task] = time.monotonic() + except BaseException: + _cancel_tasks(pending) + raise + return {task_id: results[task_id] for task_id in tasks} def get_elapsed_time(task_ref: ray.ObjectRef): + """Return seconds since the latest Ray task attempt started, or None.""" + task_id = "" try: task_id = task_ref.task_id().hex() task_info = state.get_task(task_id, address="auto") if not task_info: return None - if not isinstance(task_info, list): - task_info = [task_info] - - start_times_ms = [getattr(t, "start_time_ms", None) for t in task_info] - start_time_s = max([t / 1000.0 if t is not None else -1 for t in start_times_ms]) - if start_time_s < 0: - return None # Task has not started yet - - current_time_s = time.time() - elapsed_time = current_time_s - start_time_s - return elapsed_time - except Exception as e: - logger.warning(f"Could not get elapsed time for task {task_id}: {e}") + attempts = task_info if isinstance(task_info, list) else [task_info] + latest_attempt = max(attempts, key=lambda task: getattr(task, "attempt_number", 0)) + start_time_ms = getattr(latest_attempt, "start_time_ms", None) + if start_time_ms is None or start_time_ms <= 0: + return None + return max(0, time.time() - start_time_ms / 1000.0) + except Exception as error: + logger.warning("Could not get elapsed time for task %s: %s", task_id, error) return None diff --git a/src/agentlab/experiments/launch_exp.py b/src/agentlab/experiments/launch_exp.py index 199454c58..097e7746e 100644 --- a/src/agentlab/experiments/launch_exp.py +++ b/src/agentlab/experiments/launch_exp.py @@ -7,6 +7,7 @@ from agentlab.experiments.exp_utils import run_exp from agentlab.experiments.loop import ExpArgs, yield_all_exp_results +from agentlab.experiments.task_graph import TaskGraph RAY_PUBLIC_DASHBOARD = os.environ.get("RAY_PUBLIC_DASHBOARD", "false") == "true" RAY_DASHBOARD_PORT = os.environ.get("RAY_DASHBOARD_PORT") @@ -19,96 +20,125 @@ def run_experiments( parallel_backend="ray", avg_step_timeout=60, ): - """Run a list of ExpArgs in parallel. + """Run experiments with validated dependencies and managed agent resources. - To ensure optimal parallelism, make sure ExpArgs.depend_on is set correctly - and the backend is set to dask. + Ray and sequential execution support ``ExpArgs.depends_on`` regardless of + input order. Joblib supports independent experiments only; a dependency graph + is rejected before preparing agents or creating experiment directories. + Shared agent argument objects are prepared and closed once per invocation. Args: - n_jobs: int - Number of parallel jobs. - exp_args_list: list[ExpArgs] - List of ExpArgs objects. - study_dir: Path - Directory where the experiments will be saved. - parallel_backend: str - Parallel backend to use. Either "joblib", "ray" or "sequential". - The only backend that supports webarena graph dependencies correctly is ray or sequential. - avg_step_timeout: int - Will raise a TimeoutError if the episode is not finished after env_args.max_steps * avg_step_timeout seconds. + n_jobs: Number of parallel workers (ignored for sequential execution). + exp_args_list: Experiments to prepare and run. + study_dir: Directory where experiments will be saved. + parallel_backend: One of "joblib", "ray", or "sequential". + avg_step_timeout: Ray runtime budget per environment step. Each episode + uses its own max_steps and optional episode_timeout cap. Other + backends do not enforce episode timeouts. Raises: - ValueError: If the parallel_backend is not recognized. + ValueError: If the backend, dependency graph, or Ray deadlines are invalid. + BaseException: If preparation, execution, or resource cleanup fails. """ - - if len(exp_args_list) == 0: + if not exp_args_list: logging.warning("No experiments to run.") return + if parallel_backend not in ("joblib", "ray", "sequential"): + raise ValueError(f"Unknown parallel_backend: {parallel_backend}") - study_dir = Path(study_dir) - study_dir.mkdir(parents=True, exist_ok=True) + # Assign identities without preparing directories or starting model servers. + for exp_args in exp_args_list: + if exp_args.exp_id is None: + exp_args.make_id() + graph = TaskGraph(exp_args_list) + if parallel_backend == "joblib" and graph.has_dependencies: + raise ValueError("joblib does not support experiment dependencies; use ray or sequential") + if parallel_backend == "ray": + from agentlab.experiments.graph_execution_ray import _experiment_timeouts - # if n_jobs == 1 and parallel_backend != "sequential": - # logging.warning("Only 1 job, switching to sequential backend.") - # parallel_backend = "sequential" + _experiment_timeouts(graph.experiments, avg_step_timeout) + study_dir = Path(study_dir) + study_dir.mkdir(parents=True, exist_ok=True) logging.info(f"Saving experiments to {study_dir}") - for exp_args in exp_args_list: - exp_args.agent_args.prepare() - exp_args.prepare(exp_root=study_dir) + agents = [] + agent_ids = set() + primary_error = None try: + for exp_args in graph.experiments: + agent = exp_args.agent_args + if id(agent) not in agent_ids: + agent_ids.add(id(agent)) + # Register before preparation: a failing prepare may already + # have started a server and still needs its matching close. + agents.append(agent) + agent.prepare() + exp_args.prepare(exp_root=study_dir) + if parallel_backend == "joblib": from joblib import Parallel, delayed - # split sequential (should be no longer needed with dependencies) - sequential_exp_args, exp_args_list = _split_sequential_exp(exp_args_list) - - logging.info( - f"Running {len(sequential_exp_args)} in sequential first. The remaining {len(exp_args_list)} will be run in parallel." - ) - for exp_args in sequential_exp_args: + sequential, parallel = _split_sequential_exp(graph.experiments) + for exp_args in sequential: run_exp(exp_args, avg_step_timeout=avg_step_timeout) - Parallel(n_jobs=n_jobs, prefer="processes")( delayed(run_exp)(exp_args, avg_step_timeout=avg_step_timeout) - for exp_args in exp_args_list + for exp_args in parallel ) - - # dask will be deprecated, as there was issues. use ray instead - # elif parallel_backend == "dask": - # from agentlab.experiments.graph_execution_dask import ( - # execute_task_graph, - # make_dask_client, - # ) - - # with make_dask_client(n_worker=n_jobs): - # execute_task_graph(exp_args_list) elif parallel_backend == "ray": - from agentlab.experiments.graph_execution_ray import execute_task_graph, ray - - ray.init( - num_cpus=n_jobs, - dashboard_host="0.0.0.0" if RAY_PUBLIC_DASHBOARD else "127.0.0.1", - dashboard_port=None if RAY_DASHBOARD_PORT is None else int(RAY_DASHBOARD_PORT), - ) - try: - execute_task_graph(exp_args_list, avg_step_timeout=avg_step_timeout) - finally: - ray.shutdown() - elif parallel_backend == "sequential": - for exp_args in exp_args_list: - run_exp(exp_args, avg_step_timeout=avg_step_timeout) + _run_ray_experiments(n_jobs, graph.experiments, avg_step_timeout) else: - raise ValueError(f"Unknown parallel_backend: {parallel_backend}") + for exp_args in graph.ordered_experiments: + run_exp(exp_args, avg_step_timeout=avg_step_timeout) + except BaseException as error: + primary_error = error + raise finally: - # will close servers even if there is an exception or ctrl+c - # servers won't be closed if the script is killed with kill -9 or segfaults. - logging.info("All jobs are finished. Calling agent_args.close() on all agents...") - for exp_args in exp_args_list: - exp_args.agent_args.close() + logging.info("Closing prepared agent resources...") + _cleanup_resources([agent.close for agent in reversed(agents)], primary_error) logging.info(f"Experiment finished and saved in {study_dir}.") +def _cleanup_resources(callbacks, primary_error): + """Attempt all releases, preserving the original failure when there is one.""" + errors = [] + for close in callbacks: + try: + close() + except BaseException as error: + logging.exception("Experiment resource cleanup failed") + errors.append(error) + if primary_error is not None: + for error in errors: + primary_error.add_note(f"Resource cleanup also failed: {error!r}") + elif len(errors) == 1: + raise errors[0] + elif errors: + raise BaseExceptionGroup("Experiment resource cleanup failed", errors) + + +def _run_ray_experiments(n_jobs, experiments, avg_step_timeout): + from agentlab.experiments.graph_execution_ray import execute_task_graph, ray + + # shutdown() also cleans partially initialized local processes. Never take + # ownership of a runtime that existed before this launcher was called. + owns_runtime = not ray.is_initialized() + primary_error = None + try: + ray.init( + num_cpus=n_jobs, + dashboard_host="0.0.0.0" if RAY_PUBLIC_DASHBOARD else "127.0.0.1", + dashboard_port=None if RAY_DASHBOARD_PORT is None else int(RAY_DASHBOARD_PORT), + ) + execute_task_graph(experiments, avg_step_timeout=avg_step_timeout) + except BaseException as error: + primary_error = error + raise + finally: + if owns_runtime: + _cleanup_resources([ray.shutdown], primary_error) + + def find_incomplete(study_dir: str | Path, include_errors=True): """Find all incomplete experiments for relaunching. diff --git a/src/agentlab/experiments/task_graph.py b/src/agentlab/experiments/task_graph.py new file mode 100644 index 000000000..bd2b0be45 --- /dev/null +++ b/src/agentlab/experiments/task_graph.py @@ -0,0 +1,68 @@ +"""Validate experiment dependencies before launching workers or model servers.""" + +from heapq import heappop, heappush + + +class TaskGraph: + """A validated snapshot of experiments with a stable topological order. + + IDs must already have been assigned (``ExpArgs.make_id()`` does this without + preparing an experiment directory). Original experiment objects are retained; + neither the caller's list nor its dependency collections are modified. Among + runnable experiments, the earliest input position is chosen first. + """ + + def __init__(self, experiments): + self.experiments = tuple(experiments) + self.by_id = {} + self.dependencies = {} + for index, experiment in enumerate(self.experiments): + exp_id = experiment.exp_id + if not isinstance(exp_id, str) or not exp_id: + raise ValueError(f"Invalid experiment ID at position {index}: {exp_id!r}") + if exp_id in self.by_id: + raise ValueError(f"Duplicate experiment ID: {exp_id!r}") + self.by_id[exp_id] = experiment + + for experiment in self.experiments: + exp_id = experiment.exp_id + dependencies = experiment.depends_on + if not isinstance(dependencies, (tuple, list)): + raise ValueError(f"depends_on for {exp_id!r} must be a list or tuple of IDs") + unique_dependencies = [] + seen = set() + for dependency in dependencies: + if not isinstance(dependency, str) or not dependency: + raise ValueError(f"Invalid dependency ID for {exp_id!r}: {dependency!r}") + if dependency not in self.by_id: + raise ValueError(f"Experiment {exp_id!r} depends on missing ID {dependency!r}") + if dependency not in seen: + seen.add(dependency) + unique_dependencies.append(dependency) + self.dependencies[exp_id] = tuple(unique_dependencies) + + positions = {exp_id: index for index, exp_id in enumerate(self.by_id)} + remaining = {exp_id: len(deps) for exp_id, deps in self.dependencies.items()} + dependents = {exp_id: [] for exp_id in self.by_id} + for exp_id, dependencies in self.dependencies.items(): + for dependency in dependencies: + dependents[dependency].append(exp_id) + + ready = [] + for exp_id, count in remaining.items(): + if count == 0: + heappush(ready, positions[exp_id]) + ordered = [] + while ready: + experiment = self.experiments[heappop(ready)] + ordered.append(experiment) + for dependent in dependents[experiment.exp_id]: + remaining[dependent] -= 1 + if remaining[dependent] == 0: + heappush(ready, positions[dependent]) + + if len(ordered) != len(self.experiments): + unresolved = [exp_id for exp_id, count in remaining.items() if count] + raise ValueError(f"Dependency cycle involving experiments: {', '.join(unresolved)}") + self.ordered_experiments = tuple(ordered) + self.has_dependencies = any(self.dependencies.values()) diff --git a/tests/experiments/test_execution_integration.py b/tests/experiments/test_execution_integration.py new file mode 100644 index 000000000..a048b9698 --- /dev/null +++ b/tests/experiments/test_execution_integration.py @@ -0,0 +1,188 @@ +"""Exercise native experiment persistence and real local Ray without a browser.""" + +import json +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +import pytest +import ray + +from agentlab.experiments.graph_execution_ray import execute_task_graph +from agentlab.experiments.launch_exp import find_incomplete, run_experiments +from agentlab.experiments.loop import AbstractAgentArgs, EnvArgs, ExpArgs + + +@dataclass +class LocalAgentArgs(AbstractAgentArgs): + resource_log: str = None + + def prepare(self): + with open(self.resource_log, "a") as stream: + stream.write("prepare\n") + + def close(self): + with open(self.resource_log, "a") as stream: + stream.write("close\n") + + def make_agent(self): + raise AssertionError("The scheduler integration test never creates a browser agent") + + +@dataclass +class FileExperiment(ExpArgs): + """Persist scheduler effects using the real ExpArgs preparation contract.""" + + marker_dir: str = None + delay: float = 0 + fail: bool = False + episode_timeout: float = 20 + + def run(self): + markers = Path(self.marker_dir) + for dependency in self.depends_on: + assert (markers / dependency).exists(), f"dependency {dependency} has not finished" + time.sleep(self.delay) + if self.fail: + raise RuntimeError(f"experiment {self.exp_id} failed") + (markers / self.exp_id).write_text("finished") + if self.exp_dir: + (Path(self.exp_dir) / "summary_info.json").write_text( + json.dumps({"terminated": True, "truncated": False, "err_msg": None}) + ) + return self.exp_id + + +def file_experiment(exp_id, marker_dir, *dependencies, agent=None, **kwargs): + return FileExperiment( + agent_args=agent, + env_args=EnvArgs(task_name=f"offline.{exp_id}", task_seed=1, max_steps=100), + exp_id=exp_id, + depends_on=dependencies, + marker_dir=str(marker_dir), + **kwargs, + ) + + +def test_native_exp_args_persist_reordered_dependency_execution(tmp_path): + markers = tmp_path / "markers" + markers.mkdir() + agent = LocalAgentArgs(resource_log=str(tmp_path / "resources.log")) + experiments = [ + file_experiment("child", markers, "parent", agent=agent), + file_experiment("parent", markers, agent=agent), + ] + run_experiments(1, experiments, tmp_path / "study", parallel_backend="sequential") + assert agent.resource_log and Path(agent.resource_log).read_text() == "prepare\nclose\n" + assert {path.name for path in markers.iterdir()} == {"child", "parent"} + assert all((exp.exp_dir / "exp_args.pkl").exists() for exp in experiments) + restored = find_incomplete(tmp_path / "study") + assert len(restored) == 2 + assert all(exp.is_dummy for exp in restored) + assert {exp.exp_id for exp in restored} == {"child", "parent"} + + +def test_relaunch_keeps_completed_dependency_as_dummy(tmp_path): + markers = tmp_path / "markers" + markers.mkdir() + agent = LocalAgentArgs(resource_log=str(tmp_path / "resources.log")) + parent = file_experiment("parent", markers, agent=agent) + child = file_experiment("child", markers, "parent", agent=agent) + run_experiments(1, [parent, child], tmp_path / "study", parallel_backend="sequential") + (child.exp_dir / "summary_info.json").unlink() + (markers / "child").unlink() + restored = find_incomplete(tmp_path / "study") + assert {exp.exp_id: exp.is_dummy for exp in restored} == {"parent": True, "child": False} + parent_marker_time = (markers / "parent").stat().st_mtime_ns + # Reversing the persisted launch order must still leave the dummy dependency first. + restored.sort(key=lambda exp: exp.exp_id) + run_experiments(1, restored, tmp_path / "study", parallel_backend="sequential") + assert (markers / "child").exists() + assert (markers / "parent").stat().st_mtime_ns == parent_marker_time + assert len(find_incomplete(tmp_path / "study")) == 2 + + +def test_native_exp_args_get_ids_before_directory_preparation(tmp_path): + markers = tmp_path / "markers" + markers.mkdir() + agent = LocalAgentArgs(resource_log=str(tmp_path / "resources.log")) + exp = file_experiment(None, markers, agent=agent) + run_experiments(1, [exp], tmp_path / "study", parallel_backend="sequential") + assert isinstance(exp.exp_id, str) + assert (markers / exp.exp_id).exists() + + +@pytest.fixture(scope="module") +def local_ray(): + # Workers must receive the test class by value, not import pytest's module name. + ray.cloudpickle.register_pickle_by_value(sys.modules[__name__]) + ray.init(num_cpus=2, include_dashboard=True, log_to_driver=False) + try: + yield + finally: + ray.shutdown() + ray.cloudpickle.unregister_pickle_by_value(sys.modules[__name__]) + + +def test_real_ray_out_of_order_diamond_and_result_identity(tmp_path, local_ray): + experiments = [ + file_experiment("finish", tmp_path, "left", "right"), + file_experiment("right", tmp_path, "start", delay=0.05), + file_experiment("left", tmp_path, "start", delay=0.2), + file_experiment("start", tmp_path), + ] + result = execute_task_graph(experiments) + assert result == {exp.exp_id: exp.exp_id for exp in experiments} + assert {path.name for path in tmp_path.iterdir()} == {"finish", "right", "left", "start"} + + +def test_real_ray_parent_failure_blocks_child_but_not_other_branch(tmp_path, local_ray): + experiments = [ + file_experiment("child", tmp_path, "failed"), + file_experiment("failed", tmp_path, fail=True), + file_experiment("independent", tmp_path), + ] + result = execute_task_graph(experiments) + assert isinstance(result["failed"], Exception) + assert isinstance(result["child"], Exception) + assert result["independent"] == "independent" + assert not (tmp_path / "child").exists() + assert not (tmp_path / "failed").exists() + + +def test_real_ray_per_episode_timeout_does_not_inherit_long_sibling_budget(tmp_path, local_ray): + experiments = [ + file_experiment("short", tmp_path, delay=5, episode_timeout=0.5), + file_experiment("long", tmp_path, delay=1, episode_timeout=20), + ] + started = time.monotonic() + result = execute_task_graph(experiments) + assert isinstance(result["short"], Exception) + assert result["long"] == "long" + assert not (tmp_path / "short").exists() + assert time.monotonic() - started < 15 + + +def test_real_ray_dependency_wait_does_not_use_child_runtime_budget(tmp_path, local_ray): + experiments = [ + file_experiment("child", tmp_path, "parent", delay=0.01, episode_timeout=0.5), + file_experiment("parent", tmp_path, delay=1.2, episode_timeout=20), + ] + assert execute_task_graph(experiments) == {"child": "child", "parent": "parent"} + + +def test_invalid_native_graph_preserves_existing_experiment_directory(tmp_path): + markers = tmp_path / "markers" + markers.mkdir() + agent = LocalAgentArgs(resource_log=str(tmp_path / "resources.log")) + exp = file_experiment("existing", markers, agent=agent) + exp.prepare(tmp_path / "study") + saved_directory = exp.exp_dir + (saved_directory / "keep.txt").write_text("original results") + exp.depends_on = ("missing",) + with pytest.raises(ValueError, match="missing"): + run_experiments(1, [exp], tmp_path / "study", parallel_backend="sequential") + assert (saved_directory / "keep.txt").read_text() == "original results" + assert not saved_directory.with_name("_" + saved_directory.name).exists() + assert not Path(agent.resource_log).exists() diff --git a/tests/experiments/test_execution_regressions.py b/tests/experiments/test_execution_regressions.py new file mode 100644 index 000000000..e095f6ea1 --- /dev/null +++ b/tests/experiments/test_execution_regressions.py @@ -0,0 +1,362 @@ +"""Provider-free regressions for dependency execution and resource ownership.""" + +from dataclasses import dataclass, field +from types import SimpleNamespace + +import pytest + +from agentlab.experiments import graph_execution_ray as graph_ray +from agentlab.experiments import launch_exp + + +@dataclass +class AgentResources: + events: list + name: str = "agent" + prepare_error: BaseException = None + close_error: BaseException = None + + def prepare(self): + self.events.append(("prepare", self.name)) + if self.prepare_error: + raise self.prepare_error + + def close(self): + self.events.append(("close", self.name)) + if self.close_error: + raise self.close_error + + +@dataclass +class Experiment: + exp_id: str + depends_on: tuple = () + agent_args: AgentResources = None + events: list = field(default_factory=list) + run_error: BaseException = None + prepare_error: BaseException = None + env_args: object = field(default_factory=lambda: SimpleNamespace(max_steps=1)) + + @property + def exp_name(self): + return self.exp_id + + def prepare(self, exp_root): + self.events.append(("directory", self.exp_id)) + if self.prepare_error: + raise self.prepare_error + + def run(self): + self.events.append(("run", self.exp_id)) + if self.run_error: + raise self.run_error + return self.exp_id + + +def test_sequential_backend_honors_reversed_dependency_order(tmp_path): + events = [] + agent = AgentResources(events) + child = Experiment("child", ("parent",), agent, events) + parent = Experiment("parent", agent_args=agent, events=events) + launch_exp.run_experiments(1, [child, parent], tmp_path, parallel_backend="sequential") + assert [event for event in events if event[0] == "run"] == [ + ("run", "parent"), + ("run", "child"), + ] + + +def test_shared_agent_has_one_resource_lifecycle(tmp_path): + events = [] + agent = AgentResources(events) + experiments = [Experiment(str(i), agent_args=agent, events=events) for i in range(3)] + launch_exp.run_experiments(1, experiments, tmp_path, parallel_backend="sequential") + assert events.count(("prepare", "agent")) == 1 + assert events.count(("close", "agent")) == 1 + + +def test_partial_preparation_closes_every_acquired_agent(tmp_path): + events = [] + first = AgentResources(events, "first") + second = AgentResources(events, "second", prepare_error=RuntimeError("acquisition failed")) + experiments = [Experiment("a", agent_args=first), Experiment("b", agent_args=second)] + with pytest.raises(RuntimeError, match="acquisition failed"): + launch_exp.run_experiments(1, experiments, tmp_path, parallel_backend="sequential") + assert ("close", "first") in events + assert ("close", "second") in events + + +def test_cleanup_failure_does_not_skip_other_agents(tmp_path): + events = [] + agents = [ + AgentResources(events, "first", close_error=RuntimeError("close failed")), + AgentResources(events, "second"), + ] + with pytest.raises(Exception, match="close failed"): + launch_exp.run_experiments( + 1, + [Experiment(str(i), agent_args=agent) for i, agent in enumerate(agents)], + tmp_path, + parallel_backend="sequential", + ) + assert ("close", "second") in events + + +def test_ray_ready_order_does_not_change_experiment_identity(monkeypatch): + first, second = object(), object() + monkeypatch.setattr(graph_ray.ray, "wait", lambda *args, **kwargs: ([second, first], [])) + monkeypatch.setattr(graph_ray.ray, "get", lambda ref: "first" if ref is first else "second") + assert graph_ray.poll_for_timeout({"a": first, "b": second}, 10) == { + "a": "first", + "b": "second", + } + + +def test_invalid_graph_submits_nothing(monkeypatch): + submitted = [] + + class Remote: + def options(self, **kwargs): + return self + + def remote(self, exp, *args, **kwargs): + submitted.append(exp.exp_id) + return object() + + monkeypatch.setattr(graph_ray, "run_exp", Remote()) + with pytest.raises(ValueError, match="missing"): + graph_ray.execute_task_graph([Experiment("valid"), Experiment("bad", ("missing",))]) + assert submitted == [] + + +@pytest.mark.parametrize("backend", ["sequential", "ray", "joblib"]) +def test_unknown_dependency_fails_before_resource_or_directory_side_effects(tmp_path, backend): + events = [] + agent = AgentResources(events) + exp = Experiment("child", ("missing",), agent, events) + destination = tmp_path / "must-not-exist" + with pytest.raises(ValueError, match="missing"): + launch_exp.run_experiments(1, [exp], destination, parallel_backend=backend) + assert events == [] + assert not destination.exists() + + +def test_unknown_backend_does_not_prepare_agents(tmp_path): + events = [] + exp = Experiment("a", agent_args=AgentResources(events), events=events) + with pytest.raises(ValueError, match="Unknown parallel_backend"): + launch_exp.run_experiments(1, [exp], tmp_path, parallel_backend="dask") + assert events == [] + + +def test_joblib_rejects_dependencies_before_preparation(tmp_path): + events = [] + agent = AgentResources(events) + experiments = [Experiment("a", agent_args=agent), Experiment("b", ("a",), agent)] + with pytest.raises(ValueError, match="joblib.*dependencies"): + launch_exp.run_experiments(2, experiments, tmp_path, parallel_backend="joblib") + assert events == [] + + +def test_failed_experiment_preparation_closes_agents(tmp_path): + events = [] + agent = AgentResources(events) + exp = Experiment("a", agent_args=agent, prepare_error=OSError("directory unavailable")) + with pytest.raises(OSError, match="directory unavailable"): + launch_exp.run_experiments(1, [exp], tmp_path, parallel_backend="sequential") + assert events == [("prepare", "agent"), ("close", "agent")] + + +def test_agents_not_reached_during_preparation_are_not_closed(tmp_path): + events = [] + first = AgentResources(events, "first", prepare_error=RuntimeError("prepare failed")) + second = AgentResources(events, "second") + with pytest.raises(RuntimeError, match="prepare failed"): + launch_exp.run_experiments( + 1, + [Experiment("a", agent_args=first), Experiment("b", agent_args=second)], + tmp_path, + parallel_backend="sequential", + ) + assert events == [("prepare", "first"), ("close", "first")] + + +def test_primary_run_error_survives_cleanup_error(tmp_path): + events = [] + agent = AgentResources(events, close_error=RuntimeError("cleanup failed")) + failure = ValueError("execution failed") + exp = Experiment("a", agent_args=agent, run_error=failure) + with pytest.raises(ValueError, match="execution failed") as error: + launch_exp.run_experiments(1, [exp], tmp_path, parallel_backend="sequential") + assert error.value is failure + assert any("cleanup failed" in note for note in error.value.__notes__) + + +def test_all_cleanup_failures_are_reported(tmp_path): + events = [] + agents = [ + AgentResources(events, "first", close_error=RuntimeError("first close failed")), + AgentResources(events, "second", close_error=OSError("second close failed")), + ] + with pytest.raises(ExceptionGroup) as error: + launch_exp.run_experiments( + 1, + [Experiment(str(i), agent_args=agent) for i, agent in enumerate(agents)], + tmp_path, + parallel_backend="sequential", + ) + assert {str(exc) for exc in error.value.exceptions} == { + "first close failed", + "second close failed", + } + assert len([event for event in events if event[0] == "close"]) == 2 + + +def test_keyboard_interrupt_during_prepare_closes_partial_agent(tmp_path): + events = [] + agent = AgentResources(events, prepare_error=KeyboardInterrupt()) + with pytest.raises(KeyboardInterrupt): + launch_exp.run_experiments( + 1, [Experiment("a", agent_args=agent)], tmp_path, parallel_backend="sequential" + ) + assert ("close", "agent") in events + + +def test_keyboard_interrupt_during_run_is_not_masked_by_close(tmp_path): + events = [] + agent = AgentResources(events, close_error=RuntimeError("cleanup failed")) + with pytest.raises(KeyboardInterrupt): + launch_exp.run_experiments( + 1, + [Experiment("a", agent_args=agent, run_error=KeyboardInterrupt())], + tmp_path, + parallel_backend="sequential", + ) + assert ("close", "agent") in events + + +def test_empty_study_has_no_directory_side_effect(tmp_path): + destination = tmp_path / "empty" + assert launch_exp.run_experiments(1, [], destination) is None + assert not destination.exists() + + +def test_distinct_equal_agent_objects_have_separate_lifecycles(tmp_path): + events = [] + first, second = AgentResources(events), AgentResources(events) + assert first == second + launch_exp.run_experiments( + 1, + [Experiment("a", agent_args=first), Experiment("b", agent_args=second)], + tmp_path, + parallel_backend="sequential", + ) + assert events.count(("prepare", "agent")) == 2 + assert events.count(("close", "agent")) == 2 + + +def test_joblib_sequential_partition_does_not_lose_cleanup_ownership(tmp_path, monkeypatch): + import joblib + + events = [] + first, second = AgentResources(events, "first"), AgentResources(events, "second") + sequential = Experiment("a", agent_args=first) + sequential.sequential = True + parallel = Experiment("b", agent_args=second) + + class InlineParallel: + def __init__(self, **kwargs): + pass + + def __call__(self, jobs): + return [fn(*args, **kwargs) for fn, args, kwargs in jobs] + + monkeypatch.setattr(joblib, "Parallel", InlineParallel) + launch_exp.run_experiments(2, [sequential, parallel], tmp_path, parallel_backend="joblib") + assert ("close", "first") in events + assert ("close", "second") in events + + +@pytest.fixture +def ray_launcher(monkeypatch): + from unittest.mock import Mock + + init, shutdown, execute = Mock(), Mock(), Mock() + initialized = Mock(return_value=False) + monkeypatch.setattr(graph_ray.ray, "is_initialized", initialized) + monkeypatch.setattr(graph_ray.ray, "init", init) + monkeypatch.setattr(graph_ray.ray, "shutdown", shutdown) + monkeypatch.setattr(graph_ray, "execute_task_graph", execute) + return SimpleNamespace(init=init, shutdown=shutdown, execute=execute, initialized=initialized) + + +def test_ray_startup_failure_still_closes_agents(tmp_path, ray_launcher): + events = [] + ray_launcher.init.side_effect = RuntimeError("startup failed") + with pytest.raises(RuntimeError, match="startup failed"): + launch_exp.run_experiments( + 2, [Experiment("a", agent_args=AgentResources(events))], tmp_path + ) + assert ("close", "agent") in events + ray_launcher.execute.assert_not_called() + ray_launcher.shutdown.assert_called_once() + + +def test_ray_run_error_survives_shutdown_and_agent_cleanup_failures(tmp_path, ray_launcher): + events = [] + failure = RuntimeError("graph failed") + ray_launcher.execute.side_effect = failure + ray_launcher.shutdown.side_effect = OSError("shutdown failed") + agent = AgentResources(events, close_error=RuntimeError("agent cleanup failed")) + with pytest.raises(RuntimeError, match="graph failed") as error: + launch_exp.run_experiments(2, [Experiment("a", agent_args=agent)], tmp_path) + assert error.value is failure + assert any("shutdown failed" in note for note in failure.__notes__) + assert any("agent cleanup failed" in note for note in failure.__notes__) + assert ("close", "agent") in events + ray_launcher.shutdown.assert_called_once() + + +def test_ray_shutdown_failure_is_reported_after_successful_run(tmp_path, ray_launcher): + events = [] + ray_launcher.shutdown.side_effect = OSError("shutdown failed") + with pytest.raises(OSError, match="shutdown failed"): + launch_exp.run_experiments( + 2, [Experiment("a", agent_args=AgentResources(events))], tmp_path + ) + assert ("close", "agent") in events + + +def test_invalid_ray_deadline_does_not_start_agents_or_cluster(tmp_path, ray_launcher): + events = [] + exp = Experiment("a", agent_args=AgentResources(events)) + with pytest.raises(ValueError, match="avg_step_timeout"): + launch_exp.run_experiments(2, [exp], tmp_path, avg_step_timeout=0) + assert events == [] + ray_launcher.init.assert_not_called() + + +def test_preexisting_ray_runtime_is_not_closed_on_reinitialization_failure(tmp_path, ray_launcher): + events = [] + ray_launcher.initialized.return_value = True + ray_launcher.init.side_effect = RuntimeError("already initialized") + with pytest.raises(RuntimeError, match="already initialized"): + launch_exp.run_experiments( + 2, [Experiment("a", agent_args=AgentResources(events))], tmp_path + ) + ray_launcher.shutdown.assert_not_called() + assert ("close", "agent") in events + + +def test_partial_ray_startup_error_survives_shutdown_failure(tmp_path, ray_launcher): + events = [] + failure = RuntimeError("partial startup failed") + ray_launcher.init.side_effect = failure + ray_launcher.shutdown.side_effect = OSError("partial shutdown failed") + with pytest.raises(RuntimeError, match="partial startup failed") as error: + launch_exp.run_experiments( + 2, [Experiment("a", agent_args=AgentResources(events))], tmp_path + ) + assert error.value is failure + ray_launcher.shutdown.assert_called_once() + assert any("partial shutdown failed" in note for note in failure.__notes__) + assert ("close", "agent") in events diff --git a/tests/experiments/test_ray_monitor.py b/tests/experiments/test_ray_monitor.py new file mode 100644 index 000000000..5645574d9 --- /dev/null +++ b/tests/experiments/test_ray_monitor.py @@ -0,0 +1,201 @@ +"""Deterministic tests of Ray result identity, deadlines and cancellation.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from agentlab.experiments import graph_execution_ray as execution + + +class Ref: + def __init__(self, name): + self.name = name + + def task_id(self): + return self + + def hex(self): + return self.name + + +@pytest.fixture +def ray_api(monkeypatch): + api = SimpleNamespace(wait=Mock(), get=Mock(), cancel=Mock()) + for name, method in vars(api).items(): + monkeypatch.setattr(execution.ray, name, method) + return api + + +def test_empty_monitor_does_not_call_ray(ray_api): + assert execution.poll_for_timeout({}, 5) == {} + ray_api.wait.assert_not_called() + + +def test_results_are_keyed_by_reference_instead_of_ready_position(ray_api): + a, b, c = Ref("a"), Ref("b"), Ref("c") + ray_api.wait.side_effect = [([b], [a, c]), ([c, a], [])] + ray_api.get.side_effect = lambda ref: ref.name + result = execution.poll_for_timeout({"A": a, "B": b, "C": c}, 10) + assert result == {"A": "a", "B": "b", "C": "c"} + assert list(result) == ["A", "B", "C"] + assert ray_api.wait.call_args_list[1].args[0] == [a, c] + assert ray_api.get.call_count == 3 + + +def test_task_failure_is_returned_without_stopping_independent_results(ray_api): + a, b = Ref("a"), Ref("b") + failure = RuntimeError("worker failed") + ray_api.wait.return_value = ([b, a], []) + ray_api.get.side_effect = lambda ref: (_ for _ in ()).throw(failure) if ref is a else 12 + result = execution.poll_for_timeout({"a": a, "b": b}, 5) + assert result["a"] is failure + assert result["b"] == 12 + + +def test_each_task_uses_its_own_deadline(ray_api, monkeypatch): + short, long = Ref("short"), Ref("long") + ray_api.wait.side_effect = [([], [short, long]), ([long, short], [])] + ray_api.get.side_effect = lambda ref: ref.name + monkeypatch.setattr(execution, "get_elapsed_time", lambda ref: 15) + execution.poll_for_timeout({"s": short, "l": long}, {"s": 10, "l": 100}) + ray_api.cancel.assert_called_once_with(short, force=False, recursive=False) + + +def test_queued_tasks_do_not_consume_runtime_budget(ray_api, monkeypatch): + queued = Ref("queued") + ray_api.wait.side_effect = [([], [queued]), ([queued], [])] + monkeypatch.setattr(execution, "get_elapsed_time", lambda ref: None) + execution.poll_for_timeout({"q": queued}, 0.01) + ray_api.cancel.assert_not_called() + + +def test_timeout_boundary_is_not_cancelled_prematurely(ray_api, monkeypatch): + task = Ref("task") + ray_api.wait.side_effect = [([], [task]), ([task], [])] + monkeypatch.setattr(execution, "get_elapsed_time", lambda ref: 5) + execution.poll_for_timeout({"task": task}, 5) + ray_api.cancel.assert_not_called() + + +def test_cancel_happens_once_then_force_after_grace_even_if_state_disappears(ray_api, monkeypatch): + task = Ref("task") + ray_api.wait.side_effect = [([], [task])] * 4 + [([task], [])] + elapsed = iter([11, None, None, None]) + monkeypatch.setattr(execution, "get_elapsed_time", lambda ref: next(elapsed)) + times = iter([100, 101, 161, 162]) + monkeypatch.setattr(execution.time, "monotonic", lambda: next(times)) + execution.poll_for_timeout({"task": task}, 10, cancel_grace_period=60) + assert ray_api.cancel.call_args_list == [ + ((task,), {"force": False, "recursive": False}), + ((task,), {"force": True, "recursive": False}), + ] + + +def test_force_cancellation_can_be_immediate(ray_api, monkeypatch): + task = Ref("task") + ray_api.wait.side_effect = [([], [task]), ([], [task]), ([task], [])] + monkeypatch.setattr(execution, "get_elapsed_time", lambda ref: 11) + execution.poll_for_timeout({"task": task}, 10, cancel_grace_period=0) + assert [call.kwargs["force"] for call in ray_api.cancel.call_args_list] == [False, True] + + +def test_interruption_cancels_only_unfinished_tasks(ray_api): + done, pending = Ref("done"), Ref("pending") + interruption = KeyboardInterrupt() + ray_api.wait.side_effect = [([done], [pending]), interruption] + with pytest.raises(KeyboardInterrupt): + execution.poll_for_timeout({"done": done, "pending": pending}, 5) + ray_api.cancel.assert_called_once_with(pending, force=True, recursive=False) + + +def test_cleanup_continues_when_cancelling_one_reference_fails(ray_api): + a, b = Ref("a"), Ref("b") + ray_api.wait.side_effect = RuntimeError("cluster disconnected") + ray_api.cancel.side_effect = [RuntimeError("already gone"), None] + with pytest.raises(RuntimeError, match="cluster disconnected"): + execution.poll_for_timeout({"a": a, "b": b}, 5) + assert ray_api.cancel.call_count == 2 + + +@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), None, "5", True]) +def test_invalid_timeout_is_rejected_before_polling(ray_api, value): + with pytest.raises(ValueError, match="timeout"): + execution.poll_for_timeout({"a": Ref("a")}, value) + ray_api.wait.assert_not_called() + + +@pytest.mark.parametrize("timeouts", [{}, {"other": 5}, {"a": 5, "extra": 1}, {"a": -1}]) +def test_timeout_mapping_must_cover_exactly_the_submitted_tasks(ray_api, timeouts): + with pytest.raises(ValueError, match="timeout"): + execution.poll_for_timeout({"a": Ref("a")}, timeouts) + ray_api.wait.assert_not_called() + + +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan")]) +def test_invalid_poll_interval_fails_before_polling(ray_api, value): + with pytest.raises(ValueError, match="poll_interval"): + execution.poll_for_timeout({"a": Ref("a")}, 5, poll_interval=value) + ray_api.wait.assert_not_called() + + +def test_duplicate_references_are_rejected(ray_api): + ref = Ref("shared") + with pytest.raises(ValueError, match="reference"): + execution.poll_for_timeout({"a": ref, "b": ref}, 5) + ray_api.wait.assert_not_called() + + +@pytest.mark.parametrize("value", [-1, float("nan"), float("inf"), "60"]) +def test_invalid_cancellation_grace_fails_before_polling(ray_api, value): + with pytest.raises(ValueError, match="cancel_grace_period"): + execution.poll_for_timeout({"a": Ref("a")}, 5, cancel_grace_period=value) + ray_api.wait.assert_not_called() + + +def test_task_state_lookup_error_is_nonfatal(monkeypatch, caplog): + monkeypatch.setattr(execution.state, "get_task", Mock(side_effect=RuntimeError("unavailable"))) + assert execution.get_elapsed_time(Ref("task")) is None + assert "unavailable" in caplog.text + + +def test_task_id_lookup_error_does_not_mask_original_error(caplog): + class BrokenRef: + def task_id(self): + raise RuntimeError("invalid reference") + + assert execution.get_elapsed_time(BrokenRef()) is None + assert "invalid reference" in caplog.text + + +@pytest.mark.parametrize("task_info", [None, [], SimpleNamespace(start_time_ms=None)]) +def test_pending_or_missing_task_state_has_no_elapsed_time(monkeypatch, task_info): + monkeypatch.setattr(execution.state, "get_task", lambda *args, **kwargs: task_info) + assert execution.get_elapsed_time(Ref("task")) is None + + +def test_running_task_uses_wall_clock_start_time(monkeypatch): + monkeypatch.setattr( + execution.state, "get_task", lambda *args, **kwargs: SimpleNamespace(start_time_ms=12500) + ) + monkeypatch.setattr(execution.time, "time", lambda: 20) + assert execution.get_elapsed_time(Ref("task")) == 7.5 + + +def test_latest_retry_start_time_is_used(monkeypatch): + attempts = [ + SimpleNamespace(start_time_ms=1000, attempt_number=0), + SimpleNamespace(start_time_ms=15000, attempt_number=1), + ] + monkeypatch.setattr(execution.state, "get_task", lambda *args, **kwargs: attempts) + monkeypatch.setattr(execution.time, "time", lambda: 20) + assert execution.get_elapsed_time(Ref("task")) == 5 + + +def test_retry_waiting_to_start_does_not_inherit_previous_attempt_deadline(monkeypatch): + attempts = [ + SimpleNamespace(start_time_ms=1000, attempt_number=0), + SimpleNamespace(start_time_ms=None, attempt_number=1), + ] + monkeypatch.setattr(execution.state, "get_task", lambda *args, **kwargs: attempts) + assert execution.get_elapsed_time(Ref("task")) is None diff --git a/tests/experiments/test_ray_submission.py b/tests/experiments/test_ray_submission.py new file mode 100644 index 000000000..7e55b6b41 --- /dev/null +++ b/tests/experiments/test_ray_submission.py @@ -0,0 +1,149 @@ +"""Graph dispatch is validated atomically before it can create remote work.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from agentlab.experiments import graph_execution_ray as execution + + +def experiment(exp_id, *dependencies, steps=1, episode_timeout=1000): + return SimpleNamespace( + exp_id=exp_id, + exp_name=f"experiment-{exp_id}", + depends_on=dependencies, + env_args=SimpleNamespace(max_steps=steps), + episode_timeout=episode_timeout, + ) + + +@pytest.fixture +def dispatch(monkeypatch): + submitted = [] + refs = {} + deadlines = {} + + class Remote: + def options(self, **kwargs): + self.name = kwargs["name"] + return self + + def remote(self, exp, *dependencies, **kwargs): + submitted.append((exp.exp_id, dependencies, self.name, kwargs)) + refs[exp.exp_id] = object() + return refs[exp.exp_id] + + def collect(tasks, timeouts, **kwargs): + deadlines.update(timeouts) + return dict(tasks) + + api = SimpleNamespace(submitted=submitted, refs=refs, deadlines=deadlines) + monkeypatch.setattr(execution, "run_exp", Remote()) + monkeypatch.setattr(execution, "poll_for_timeout", collect) + monkeypatch.setattr(execution.ray, "cancel", Mock()) + return api + + +def test_out_of_order_diamond_submits_parent_references_once(dispatch): + experiments = [ + experiment("finish", "left", "right"), + experiment("right", "start"), + experiment("start"), + experiment("left", "start"), + ] + result = execution.execute_task_graph(experiments, avg_step_timeout=3) + assert [entry[0] for entry in dispatch.submitted] == ["start", "right", "left", "finish"] + assert dispatch.submitted[-1][1] == (dispatch.refs["left"], dispatch.refs["right"]) + assert dispatch.submitted[1][1] == (dispatch.refs["start"],) + assert dispatch.submitted[2][1] == (dispatch.refs["start"],) + assert list(result) == ["finish", "right", "start", "left"] + assert all(entry[2] == f"experiment-{entry[0]}" for entry in dispatch.submitted) + assert all(entry[3] == {"avg_step_timeout": 3} for entry in dispatch.submitted) + + +def test_episode_deadlines_are_not_replaced_by_study_maximum(dispatch): + execution.execute_task_graph( + [ + experiment("short", steps=2), + experiment("long", steps=20), + experiment("capped", steps=20, episode_timeout=4), + ], + avg_step_timeout=3, + ) + assert dispatch.deadlines == {"short": 6, "long": 60, "capped": 4} + + +def test_no_max_steps_retains_ten_hour_default(dispatch): + exp = experiment("default", steps=None, episode_timeout=1e9) + execution.execute_task_graph([exp]) + assert dispatch.deadlines == {"default": 36000} + + +@pytest.mark.parametrize( + "experiments, message", + [ + ([experiment("a"), experiment("a")], "Duplicate"), + ([experiment("ok"), experiment("bad", "missing")], "missing"), + ([experiment("ok"), experiment("a", "b"), experiment("b", "a")], "cycle"), + ], +) +def test_no_valid_prefix_is_submitted_when_later_component_is_invalid( + dispatch, experiments, message +): + with pytest.raises(ValueError, match=message): + execution.execute_task_graph(experiments) + assert dispatch.submitted == [] + + +@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf"), None, "1", True]) +def test_invalid_episode_deadline_submits_no_work(dispatch, timeout): + with pytest.raises(ValueError, match="timeout"): + execution.execute_task_graph([experiment("ok"), experiment("bad", episode_timeout=timeout)]) + assert dispatch.submitted == [] + + +@pytest.mark.parametrize("average", [0, -1, None, "1", True]) +def test_invalid_average_submits_no_work(dispatch, average): + with pytest.raises(ValueError, match="avg_step_timeout"): + execution.execute_task_graph([experiment("a")], avg_step_timeout=average) + assert dispatch.submitted == [] + + +def test_empty_graph_returns_empty_mapping(dispatch): + assert execution.execute_task_graph([]) == {} + assert dispatch.submitted == [] + + +@pytest.mark.parametrize("failure", [RuntimeError("submission failed"), KeyboardInterrupt()]) +def test_partial_submission_failure_cancels_created_work(monkeypatch, failure): + reference = object() + remote = Mock() + remote.options.return_value = remote + remote.remote.side_effect = [reference, failure] + cancel = Mock() + monkeypatch.setattr(execution, "run_exp", remote) + monkeypatch.setattr(execution.ray, "cancel", cancel) + with pytest.raises(type(failure)) as raised: + execution.execute_task_graph([experiment("a"), experiment("b", "a")]) + assert raised.value is failure + cancel.assert_called_once_with(reference, force=True, recursive=False) + + +def test_submission_cleanup_does_not_replace_original_failure(monkeypatch): + remote = Mock() + remote.options.return_value = remote + remote.remote.side_effect = [object(), RuntimeError("submission failed")] + monkeypatch.setattr(execution, "run_exp", remote) + monkeypatch.setattr(execution.ray, "cancel", Mock(side_effect=RuntimeError("cancel failed"))) + with pytest.raises(RuntimeError, match="submission failed"): + execution.execute_task_graph([experiment("a"), experiment("b")]) + + +def test_large_reversed_graph_submits_without_recursion(dispatch): + experiments = [experiment(str(i), str(i + 1)) for i in range(1500)] + experiments.append(experiment("1500")) + result = execution.execute_task_graph(experiments) + assert len(result) == 1501 + assert dispatch.submitted[0][0] == "1500" + assert dispatch.submitted[-1][0] == "0" diff --git a/tests/experiments/test_task_graph.py b/tests/experiments/test_task_graph.py new file mode 100644 index 000000000..91bb2f53d --- /dev/null +++ b/tests/experiments/test_task_graph.py @@ -0,0 +1,124 @@ +"""The graph contract is independent of browsers, agents, and Ray.""" + +from types import SimpleNamespace + +import pytest + +from agentlab.experiments.task_graph import TaskGraph + + +def experiment(exp_id, *dependencies): + return SimpleNamespace(exp_id=exp_id, depends_on=dependencies) + + +def test_empty_graph(): + graph = TaskGraph([]) + assert graph.experiments == () + assert graph.ordered_experiments == () + assert not graph.has_dependencies + + +def test_independent_experiments_preserve_input_order(): + experiments = [experiment("z"), experiment("a"), experiment("middle")] + graph = TaskGraph(experiments) + assert graph.experiments == tuple(experiments) + assert graph.ordered_experiments == tuple(experiments) + assert graph.by_id["z"] is experiments[0] + + +def test_stable_topological_order_for_diamond(): + experiments = [ + experiment("finish", "left", "right"), + experiment("right", "start"), + experiment("unrelated"), + experiment("left", "start"), + experiment("start"), + ] + graph = TaskGraph(experiments) + assert [exp.exp_id for exp in graph.ordered_experiments] == [ + "unrelated", + "start", + "right", + "left", + "finish", + ] + assert graph.has_dependencies + assert graph.dependencies["finish"] == ("left", "right") + + +def test_graph_does_not_reorder_caller_list_or_mutate_dependencies(): + child, parent = experiment("child", "parent", "parent"), experiment("parent") + experiments = [child, parent] + graph = TaskGraph(experiments) + assert experiments == [child, parent] + assert child.depends_on == ("parent", "parent") + assert graph.dependencies["child"] == ("parent",) + assert graph.ordered_experiments == (parent, child) + + +def test_long_chain_does_not_recurse(): + experiments = [experiment(str(i), str(i + 1)) for i in range(2500)] + experiments.append(experiment("2500")) + graph = TaskGraph(experiments) + assert [exp.exp_id for exp in graph.ordered_experiments] == [ + str(i) for i in range(2500, -1, -1) + ] + + +@pytest.mark.parametrize("bad_id", [None, "", [], {}, 7]) +def test_invalid_experiment_id_is_actionable(bad_id): + with pytest.raises(ValueError, match="experiment ID"): + TaskGraph([experiment(bad_id)]) + + +def test_duplicate_id_is_not_silently_overwritten(): + with pytest.raises(ValueError, match="Duplicate experiment ID.*same"): + TaskGraph([experiment("same"), experiment("same")]) + + +def test_repeated_experiment_object_is_also_a_duplicate(): + exp = experiment("same") + with pytest.raises(ValueError, match="Duplicate experiment ID"): + TaskGraph([exp, exp]) + + +def test_unknown_dependency_names_both_experiments(): + with pytest.raises(ValueError, match="child.*missing"): + TaskGraph([experiment("child", "missing")]) + + +@pytest.mark.parametrize("dependencies", [None, "parent", 17, {"parent": True}]) +def test_dependency_collection_must_be_a_list_or_tuple(dependencies): + exp = SimpleNamespace(exp_id="child", depends_on=dependencies) + with pytest.raises(ValueError, match="depends_on.*child"): + TaskGraph([exp, experiment("parent")]) + + +@pytest.mark.parametrize("dependency", [None, "", [], 7]) +def test_dependency_ids_are_validated_before_lookup(dependency): + with pytest.raises(ValueError, match="dependency ID.*child"): + TaskGraph([experiment("child", dependency)]) + + +def test_self_dependency_reports_cycle(): + with pytest.raises(ValueError, match="cycle.*self"): + TaskGraph([experiment("self", "self")]) + + +def test_cycle_in_disconnected_component_is_rejected(): + with pytest.raises(ValueError, match="cycle.*a.*b"): + TaskGraph([experiment("valid"), experiment("a", "b"), experiment("b", "a")]) + + +def test_cycle_with_dependent_leaf_reports_unresolved_ids(): + with pytest.raises(ValueError, match="cycle") as error: + TaskGraph([experiment("leaf", "a"), experiment("a", "b"), experiment("b", "a")]) + assert "a" in str(error.value) + assert "b" in str(error.value) + + +def test_snapshot_is_not_changed_by_mutating_input_collection(): + experiments = [experiment("a")] + graph = TaskGraph(experiments) + experiments.clear() + assert len(graph.experiments) == 1