diff --git a/CLAUDE.md b/CLAUDE.md index 961880e6..548f3c12 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1767,9 +1767,15 @@ does not independently guarantee a finite termination deadline for Slurm jobs. **Execution borrows a client and leaves the allocation alive.** Validate native identity and scheduler readiness. The driver keeps git and convergence. Use unique -invocation task keys. Interrupted unreported outputs remain in place because a +invocation task keys. `engine.execution` requires an atomic scheduler claim before +each recipe or probe runs. A repeated claim or missing invocation state refuses +execution, including recomputation after a completed result is lost. There are no +completion receipts, leases, heartbeats, or automatic recovery. `retries=0` alone +does not prevent Dask recomputation after worker loss. This is an invocation-local +replay guard, not a project lock or proof that a disconnected command stopped. +Interrupted unreported outputs remain in place because a client disconnect does not prove remote subprocess termination. Comprehensive -cancellation/fencing and simultaneous writers are deferred by explicit user decision. +cancellation and simultaneous writers are deferred by explicit user decision. Local containerized processes can outlive process-group shutdown; do not claim that `down` or walltime proves an external runtime's containers have stopped. Read-only project validation precedes cluster connection, and a run with no diff --git a/docs/api/compute.md b/docs/api/compute.md index c6fa9429..14f5b874 100644 --- a/docs/api/compute.md +++ b/docs/api/compute.md @@ -113,6 +113,20 @@ Dask chooses the workers and handles dependencies; invocation-specific keys prev unintended reuse across commands. There is no worker-selection layer, per-worker preflight orchestration, source fingerprinting, or login-node guard. Driver-side preparation and the existing task runtime/sandbox checks remain in their owners. + +`engine.execution` prevents automatic replay of side effects. The driver registers +an invocation in the existing Dask scheduler; before running a recipe or probe, +its worker must receive an atomic claim for that task. A repeated claim, missing +invocation state, or absent driver client fails before the command runs. This also +refuses recomputation when a completed result was lost: Lightcone does not store +or recover results. +`retries=0` alone cannot prevent [Dask's recomputation after worker loss](https://distributed.dask.org/en/stable/resilience.html). +There are no leases, heartbeats, command supervisors, or separate execution service. +This guard prevents automatic replay within an invocation; it does not lock the +project against another invocation or prove that a disconnected worker stopped. +The context removes its record on exit; later invocations prune records left by +disconnected clients. Removing a record cannot authorize a later claim. + `output.py` transports byte chunks through standard Dask events so detached workers' output reaches the invoking CLI. It uses the borrowed client's event topic, which the schedulers lc launches drop as soon as the client disconnects diff --git a/docs/architecture.md b/docs/architecture.md index 5e55fb62..c6ca928a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -169,6 +169,10 @@ client on exit. Both execution commands require a cluster ID. The materializatio scheduler keeps its `submit`/`completed` seam. Driver preparation and existing task runtime/sandbox checks remain unchanged. Tasks use ordinary Dask scheduling; there is no separate worker-selection or preflight layer, or site-marker guard. No execution command implicitly allocates compute. +Before a recipe or probe runs, `engine.execution` claims its task once in the +existing scheduler. Repeated claims and missing invocation state refuse execution, +so Dask cannot silently replay side effects after a worker or result is lost. +This adds no heartbeat or execution service and does not guarantee cancellation. See [compute internals](api/compute.md) and [deployment limits](user/cluster.md). ## The publication view diff --git a/docs/user/cluster.md b/docs/user/cluster.md index 0250908d..886253bd 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -332,9 +332,15 @@ directories and credential files still reject symlinks, retain ownership and ancestor-permission checks, and require modes `0700` and `0600`, respectively. The catalog's location is independent of the private connection files. +Lightcone prevents Dask from automatically rerunning a recipe or probe within the +same invocation: each task must claim permission once in the existing scheduler +before it runs. If a worker or its result is lost, an attempted replay fails +instead. Missing scheduler state also refuses execution. There is no automatic +recovery; inspect the allocation and outputs before starting another invocation. + Use one execution invocation per project at a time. Concurrent writers, -comprehensive cancellation, task fencing, and recovery after client/worker loss -are not guaranteed. A lost client does not prove its subprocesses stopped. +comprehensive cancellation, and recovery after client/worker loss are not +guaranteed. A lost client does not prove its subprocesses stopped. Unreported partial outputs are retained after interruption rather than restored while a task may still write them. End the allocation and establish that work has stopped before inspecting or repairing that project's outputs. diff --git a/src/lightcone/engine/execution.py b/src/lightcone/engine/execution.py new file mode 100644 index 00000000..7aa6d450 --- /dev/null +++ b/src/lightcone/engine/execution.py @@ -0,0 +1,91 @@ +"""Refuse repeated execution of side-effecting tasks in the existing Dask scheduler.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any +from uuid import uuid4 + +from lightcone.engine.project import ProjectError + + +def _register(invocation: str, owner: str, *, dask_scheduler: Any) -> None: + records = dask_scheduler.extensions.setdefault("lightcone-executions", {}) + if invocation in records: + raise ProjectError("execution is already registered") + # Collect abandoned invocations using Dask's own client membership, without + # a background service or a second liveness protocol. + for key, record in list(records.items()): + if record["client"] not in dask_scheduler.clients: + del records[key] + records[invocation] = {"client": owner, "tasks": set()} + + +def _claim(invocation: str, task: str, *, dask_scheduler: Any) -> None: + records = dask_scheduler.extensions.get("lightcone-executions", {}) + record = records.get(invocation) + if record is None or record["client"] not in dask_scheduler.clients: + raise ProjectError("execution is no longer registered or its client disconnected") + # This synchronous callback runs atomically on the scheduler's event loop. + # Claims survive task failure, worker loss, and Dask forgetting task results. + if task in record["tasks"]: + raise ProjectError(f"{task}: already claimed; refusing duplicate execution") + record["tasks"].add(task) + + +def _forget(invocation: str, *, dask_scheduler: Any) -> None: + dask_scheduler.extensions.get("lightcone-executions", {}).pop(invocation, None) + + +def _rpc(client: Any, function: Callable[..., None], *args: Any) -> None: + try: + client.sync(client.run_on_scheduler, function, *args, callback_timeout=5) + except ProjectError: + raise + except Exception as exc: + raise ProjectError(f"cannot contact the Dask execution guard: {exc}") from exc + + +def _call(invocation: str, task: str, function: Callable[..., Any], *args: Any) -> Any: + from distributed import get_client + + # A lost claim reply is ambiguous. Execute only after acknowledgment, and + # never recreate missing state or retry the claim on a worker's behalf. + _rpc(get_client(), _claim, invocation, task) + return function(*args) + + +@dataclass(frozen=True) +class Invocation: + """Submit tasks that may begin at most once within this invocation.""" + + client: Any + id: str + + def submit(self, function: Callable[..., Any], *args: Any, key: str) -> Any: + """Guard each task before its first side effect, including Dask recomputation.""" + return self.client.submit( + _call, self.id, key, function, *args, + key=f"lc-{self.id}-{key}", pure=False, retries=0, + ) + + +@contextmanager +def invocation(client: Any) -> Iterator[Invocation]: + """Register claims for one borrowed client; never stop its running commands. + + Missing records refuse admission, so forgetting an invocation also prevents + late tasks from starting. Cleanup failures cannot discard received results; + abandoned records are collected when another invocation registers. + """ + run = Invocation(client, uuid4().hex) + _rpc(client, _register, run.id, client.id) + try: + yield run + finally: + try: + _rpc(client, _forget, run.id) + except ProjectError: + pass diff --git a/src/lightcone/engine/materialize.py b/src/lightcone/engine/materialize.py index e9791540..24c67631 100644 --- a/src/lightcone/engine/materialize.py +++ b/src/lightcone/engine/materialize.py @@ -39,9 +39,8 @@ from dataclasses import asdict, dataclass, field, replace from pathlib import Path from typing import TYPE_CHECKING, Any, Protocol -from uuid import uuid4 -from lightcone.engine import assets, container, dataset, identity, plan, project, worker +from lightcone.engine import assets, container, dataset, execution, identity, plan, project, worker from lightcone.engine.plan import Graph, Key, Task from lightcone.engine.project import ProjectError @@ -676,16 +675,16 @@ class _Dask: """A borrowed Dask client, narrowed to what the graph driver needs.""" client: Any - invocation: str + invocation: execution.Invocation output: Forwarder def submit(self, fn: Any, *args: Any, key: str) -> Any: """Submit an ordinary Dask task with a unique key and forwarded output.""" from lightcone.engine.compute.output import call - return self.client.submit( + return self.invocation.submit( call, fn, self.output.topic, key, *args, - key=f"lc-{self.invocation}-{key}", pure=False, + key=key, ) def completed(self, handles: list[Any]) -> Iterator[worker.TaskResult]: @@ -725,8 +724,10 @@ def cluster_for_run(cluster_id: str) -> Iterator[Scheduler]: from lightcone.engine.compute.output import forwarding with compute.connect(cluster_id) as client: - invocation = uuid4().hex - with forwarding(client, stdout="stderr") as output: + with ( + forwarding(client, stdout="stderr") as output, + execution.invocation(client) as invocation, + ): yield _Dask(client, invocation, output) diff --git a/src/lightcone/engine/run.py b/src/lightcone/engine/run.py index d54b7ed2..be44b7d1 100644 --- a/src/lightcone/engine/run.py +++ b/src/lightcone/engine/run.py @@ -18,9 +18,8 @@ from dataclasses import replace from pathlib import Path from typing import Any -from uuid import uuid4 -from lightcone.engine import container, sandbox +from lightcone.engine import container, execution, sandbox from lightcone.engine.project import ( SPEC_FILENAME, ProjectError, @@ -53,11 +52,10 @@ def probe(project: Path, command: Sequence[str], *, cluster_id: str) -> sandbox. with compute.connect(cluster_id) as client: runtime = container.runtime_for_run(project, build=False) notes = [f"uv: {warning}" for warning in container.converge(runtime)] - invocation = uuid4().hex - with forwarding(client) as output: - future = client.submit( + with forwarding(client) as output, execution.invocation(client) as invocation: + future = invocation.submit( call, _probe, output.topic, "probe", runtime, paths, tuple(command), - key=f"lc-{invocation}-probe", pure=False, + key="probe", ) try: outcome: sandbox.Outcome = future.result() diff --git a/tests/test_execution.py b/tests/test_execution.py new file mode 100644 index 00000000..0092a3aa --- /dev/null +++ b/tests/test_execution.py @@ -0,0 +1,227 @@ +"""One scheduler claim prevents Dask from replaying a side-effecting task.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from uuid import uuid4 + +import pytest +from distributed import Client, LocalCluster, get_worker + +from lightcone.engine import execution +from lightcone.engine.project import ProjectError + + +@pytest.fixture +def execution_client() -> Iterator[Client]: + with LocalCluster( + n_workers=2, threads_per_worker=1, processes=False, + dashboard_address=None, memory_limit=0, + ) as cluster, Client(cluster) as client: + yield client + + +def _wait_until(condition: Callable[[], bool]) -> None: + deadline = time.monotonic() + 10 + while not condition(): + assert time.monotonic() < deadline, "Dask task did not reach the expected state" + time.sleep(0.01) + + +def _effect(path: Path, fail: bool = False) -> str: + address = get_worker().address + with path.open("a") as stream: + stream.write(address + "\n") + if fail: + raise ValueError("recipe failed after writing") + return address + + +def _running_effect(path: Path, release: Path, finished: Path) -> None: + _effect(path) + # Publish readiness after the append is closed, including on slow CI hosts. + path.with_suffix(".ready").touch() + try: + _wait_until(release.exists) + finally: + finished.touch() + + +def _remove_worker(client: Client, address: str) -> None: + worker = next(worker for worker in client.cluster.workers.values() + if worker.address == address) + # Discard its data while keeping an executing thread alive: scheduler loss + # does not itself prove that the original command has stopped writing. + client.cluster.sync(worker.close, executor_wait=False, timeout=1) + + +def _replay(client: Client, invocation: str, task: str, path: Path) -> Any: + return client.submit( + execution._call, invocation, task, _effect, path, + key=f"replay-{uuid4().hex}", pure=False, retries=0, + ) + + +def _lose_state(*, dask_scheduler: Any) -> None: + dask_scheduler.extensions.pop("lightcone-executions", None) + + +def test_completed_worker_loss_refuses_recomputation( + execution_client: Client, tmp_path: Path, +) -> None: + effects = tmp_path / "effects" + with execution.invocation(execution_client) as run: + future = run.submit(_effect, effects, key="recipe") + address = future.result(timeout=10) + _remove_worker(execution_client, address) + # A previously finished Future can still contain its old result until + # the scheduler tells this client that recomputation has failed. + _wait_until(lambda: future.status == "error") + with pytest.raises(ProjectError, match="already claimed"): + future.result(timeout=10) + assert effects.read_text().splitlines() == [address] + + +def test_running_worker_loss_refuses_replay_while_original_can_still_write( + execution_client: Client, tmp_path: Path, +) -> None: + effects, release, finished = (tmp_path / name for name in ("effects", "release", "finished")) + with execution.invocation(execution_client) as run: + future = run.submit(_running_effect, effects, release, finished, key="recipe") + try: + _wait_until(effects.with_suffix(".ready").exists) + address = effects.read_text().strip() + _remove_worker(execution_client, address) + with pytest.raises(ProjectError, match="already claimed"): + future.result(timeout=10) + assert not finished.exists() + assert effects.read_text().splitlines() == [address] + finally: + release.touch() + _wait_until(finished.exists) + + +def test_task_failure_does_not_release_its_claim( + execution_client: Client, tmp_path: Path, +) -> None: + effects = tmp_path / "effects" + with execution.invocation(execution_client) as run: + with pytest.raises(ValueError, match="recipe failed"): + run.submit(_effect, effects, True, key="recipe").result(timeout=10) + with pytest.raises(ProjectError, match="already claimed"): + _replay(execution_client, run.id, "recipe", effects).result(timeout=10) + assert len(effects.read_text().splitlines()) == 1 + + +def test_dask_forgetting_and_recreating_a_task_does_not_forget_its_claim( + execution_client: Client, tmp_path: Path, +) -> None: + effects = tmp_path / "effects" + with execution.invocation(execution_client) as run: + original = run.submit(_effect, effects, key="recipe") + original.result(timeout=10) + key = original.key + original.release() + _wait_until(lambda: execution_client.run_on_scheduler( + lambda dask_scheduler: key not in dask_scheduler.tasks, + )) + with pytest.raises(ProjectError, match="already claimed"): + run.submit(_effect, effects, key="recipe").result(timeout=10) + assert len(effects.read_text().splitlines()) == 1 + + +@pytest.mark.parametrize( + "state_lost", [False, True], ids=["context-exited", "scheduler-state-lost"], +) +def test_missing_invocation_refuses_replays_and_late_new_tasks( + execution_client: Client, tmp_path: Path, state_lost: bool, +) -> None: + effects = tmp_path / "effects" + with execution.invocation(execution_client) as run: + run.submit(_effect, effects, key="recipe").result(timeout=10) + if state_lost: + execution_client.run_on_scheduler(_lose_state) + with pytest.raises(ProjectError): + _replay(execution_client, run.id, "new-recipe", effects).result(timeout=10) + for key in ("recipe", "new-recipe"): + with pytest.raises(ProjectError): + _replay(execution_client, run.id, key, effects).result(timeout=10) + assert len(effects.read_text().splitlines()) == 1 + + +def test_lost_claim_reply_never_executes_and_cannot_be_retried( + execution_client: Client, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + effects = tmp_path / "effects" + request = execution._rpc + + def lose_reply(client: Any, function: Any, *args: Any) -> Any: + result = request(client, function, *args) + if function is execution._claim: + raise ProjectError("claim accepted but reply lost") + return result + + with execution.invocation(execution_client) as run: + with monkeypatch.context() as patch: + patch.setattr(execution, "_rpc", lose_reply) + with pytest.raises(ProjectError, match="reply lost"): + run.submit(_effect, effects, key="recipe").result(timeout=10) + with pytest.raises(ProjectError, match="already claimed"): + _replay(execution_client, run.id, "recipe", effects).result(timeout=10) + assert not effects.exists() + + +def test_claims_are_scoped_to_one_invocation( + execution_client: Client, tmp_path: Path, +) -> None: + effects = tmp_path / "effects" + for _ in range(2): + with execution.invocation(execution_client) as run: + run.submit(_effect, effects, key="recipe").result(timeout=10) + assert len(effects.read_text().splitlines()) == 2 + + +def test_disconnected_owner_cannot_admit_tasks_and_registration_prunes_stale_records() -> None: + scheduler = SimpleNamespace(extensions={}, clients={"owner": object()}) + execution._register("first", "owner", dask_scheduler=scheduler) + execution._claim("first", "recipe", dask_scheduler=scheduler) + with pytest.raises(ProjectError): + execution._register("first", "owner", dask_scheduler=scheduler) + with pytest.raises(ProjectError, match="already claimed"): + execution._claim("first", "recipe", dask_scheduler=scheduler) + scheduler.clients.clear() + with pytest.raises(ProjectError): + execution._claim("first", "late-recipe", dask_scheduler=scheduler) + scheduler.clients["new-owner"] = object() + execution._register("second", "new-owner", dask_scheduler=scheduler) + assert "first" not in scheduler.extensions["lightcone-executions"] + execution._claim("second", "recipe", dask_scheduler=scheduler) + + +def test_unavailable_scheduler_is_reported_as_a_project_error() -> None: + def offline(*args: Any, **kwargs: Any) -> None: + raise OSError("connection lost") + + client = SimpleNamespace(sync=offline, run_on_scheduler=None) + with pytest.raises(ProjectError, match="Dask execution guard"): + execution._rpc(client, execution._claim, "invocation", "recipe") + + +def test_forget_failure_does_not_hide_a_successful_result( + execution_client: Client, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + request = execution._rpc + + def fail_forget(client: Any, function: Any, *args: Any) -> Any: + if function is execution._forget: + raise ProjectError("scheduler disconnected during cleanup") + return request(client, function, *args) + + monkeypatch.setattr(execution, "_rpc", fail_forget) + with execution.invocation(execution_client) as run: + result = run.submit(_effect, tmp_path / "effects", key="recipe").result(timeout=10) + assert result