diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 632ea2b..357804f 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -52,7 +52,7 @@ jobs: activate-environment: true - name: Install dependencies - run: uv sync --group dev + run: uv sync --group dev --extra gateway - name: Report the sandbox this runner can build # Printed before the suite so a failure downstream can be read diff --git a/README.md b/README.md index 6483fcf..c82e03a 100644 --- a/README.md +++ b/README.md @@ -44,11 +44,43 @@ drafted with any AI coding assistant. - **Multiverse analysis** — declare methodological decisions with multiple defensible options; `lc` materializes your analysis across every universe you define - **Provenance by construction** — every output is committed to git together with a content-addressed manifest and a re-runnable run record; git-annex carries the bytes, so results travel with the repository - **Locked, isolated execution** — a project's environment is `pyproject.toml` + `uv.lock`; recipes run in it under a sandbox (Landlock on Linux, Seatbelt on macOS) that keeps undeclared files out and stray writes contained -- **Containers and HPC** — declare `[tool.lightcone.image]` and recipes run in a content-addressed image archived in the repository itself; a SLURM allocation is detected and used automatically, every node included +- **Containers and HPC** — declare `[tool.lightcone.image]` and recipes run in a content-addressed image archived in the repository itself; use every node of a SLURM allocation or attach to a cluster from JupyterLab's Lightcone sidebar - **Publication view** — declare a license and `lc materialize` maintains an [RO-Crate](https://www.researchobject.org/ro-crate/) of the project and its provenance, ready to archive or deposit → [Full documentation](https://docs.lightconeresearch.org) +## Compute clusters + +Start a cluster in JupyterLab's **Lightcone sidebar › Compute**, then run +`lc materialize` as usual. The CLI chooses an execution target in this order: + +1. An existing SLURM allocation, when `SLURM_JOB_ID` is set. +2. The one compatible live cluster registered by the sidebar. +3. This machine, subject to the existing HPC login-node guard. + +The run names its target, and `--json` includes a `venue` object. If several +clusters can serve the project, stop all but one. A selected cluster that is +queued, unreachable, or incompatible with the installed engine produces an +actionable refusal; the run does not silently switch targets. Workers need +the same Lightcone, Dask Distributed and Python versions as the CLI, and must +see the project at the same filesystem path. + +Clusters remain available after each run. The CLI closes its connection; +cluster creation, scaling and shutdown remain with the sidebar and its +backend. Interrupting an attached run leaves unfinished outputs uncommitted: +stop the cluster and inspect them before restoring or rerunning. Direct-mode +recipes receive the invoking shell's environment while retaining the workers' +own host and job settings; containers retain their existing environment policy. + +Local and SLURM clusters also support containerized projects. Using several +worker hosts requires a shared image store such as `podman-hpc`. Dask Gateway +supports direct-mode projects in the same JupyterHub image; install its +optional client and use the hub's Gateway configuration: + +```bash +uv tool install 'lightcone-cli[gateway]' +``` + ## License BSD 3-Clause — see [LICENSE](LICENSE) for details. diff --git a/pyproject.toml b/pyproject.toml index 41ada83..5d4ad87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,9 +29,13 @@ dependencies = [ "rich>=13.0", "git-annex>=10.2026", "distributed>=2026.7", + "psutil>=5.8", "rocrate>=0.15", ] +[project.optional-dependencies] +gateway = ["dask-gateway>=2025.4"] + [dependency-groups] dev = [ "pytest>=8.0", @@ -92,6 +96,10 @@ namespace_packages = true explicit_package_bases = true mypy_path = "src" +[[tool.mypy.overrides]] +module = ["dask_gateway.*"] +ignore_missing_imports = true + [[tool.mypy.overrides]] module = ["astra.*", "rocrate.*"] ignore_missing_imports = true diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 8fbf4ca..f7fde13 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -292,27 +292,37 @@ def materialize( records the environment and the commit that produced it. Pass --refresh to remake those too. """ + from rich.markup import escape + from lightcone.engine import container as engine_container from lightcone.engine import materialize as engine from lightcone.engine.project import current_project root = current_project() - if not check_only and not as_json: - # The engine never prints, and the build it may be about to run - # can take minutes — so the one place that owns the console says - # so before handing over. Conditional mood, deliberately: the - # engine's own refusals (a dirty tree, an invalid spec) come - # first and cost no build, so this must promise nothing. + + def announce(selected: dict[str, str | int]) -> None: + if selected["kind"] == "cluster": + target = f"{selected['label']} ({selected['backend']}, {selected['id']})" + elif selected["kind"] == "allocation": + target = f"SLURM allocation ({selected['nodes']} nodes)" + else: + target = "this host" + _console().print(f"Running on {escape(target)}") + # Announce the build before it starts; engine refusals still come + # first, so this promises nothing until preparation succeeds. state, tag, _ = engine_container.image_state(root) if state == "absent": _console().print( f"image absent — the run rebuilds [bold]{tag}[/bold] first " "(this can take minutes)" ) + if check_only: report = engine.check(root, targets, refresh=refresh) else: - report = engine.materialize(root, targets, refresh=refresh) + report = engine.materialize( + root, targets, refresh=refresh, on_venue=None if as_json else announce + ) if as_json: click.echo(json.dumps(report.as_dict(), indent=2)) diff --git a/src/lightcone/engine/clusters.py b/src/lightcone/engine/clusters.py new file mode 100644 index 0000000..7bc01ab --- /dev/null +++ b/src/lightcone/engine/clusters.py @@ -0,0 +1,406 @@ +"""Read and attach to the sidebar's clusters; their lifecycle belongs to the sidebar. + +The registry is the interface between the two packages. Backend state, rather +than a stale scheduler file, decides which records can serve a run. +""" + +from __future__ import annotations + +import asyncio +import importlib.metadata +import json +import os +import platform +import re +import socket +import subprocess +import time +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import dask.config +import psutil # type: ignore[import-untyped] + +from . import project +from .project import ProjectError + +FORMAT = "lightcone.cluster/1" +_ID = re.compile(r"[0-9]{8}-[0-9]{6}-[a-z0-9]{4}") +_WAIT = 120 +_CONNECT_TIMEOUT = 30 +_COMPUTE = "Lightcone sidebar › Compute" + + +@dataclass(frozen=True) +class Record: + """A validated registry record and its last observed backend state.""" + + directory: Path + data: dict[str, Any] + state: str = "unknown" + start_estimate: str | None = None + + @property + def id(self) -> str: + return str(self.data["id"]) + + @property + def backend(self) -> str: + return str(self.data["backend"]) + + @property + def label(self) -> str: + return str(self.data["label"]) + + def section(self, name: str) -> dict[str, Any]: + """Return a validated section of the record.""" + section = self.data.get(name) + return section if isinstance(section, dict) else {} + + +def registry_root() -> Path: + """The per-user registry, shared with the JupyterLab extension.""" + return Path.home() / ".lightcone" / "clusters" + + +def _records() -> Iterator[Record]: + try: + directories = sorted(registry_root().iterdir()) + except OSError: + return + for directory in directories: + if not _ID.fullmatch(directory.name): + continue + try: + data = json.loads((directory / "cluster.json").read_text()) + except (OSError, ValueError): + continue + if not isinstance(data, dict) or data.get("format") != FORMAT: + continue + backend = data.get("backend") + if ( + data.get("id") != directory.name + or backend not in ("local", "slurm", "gateway") + or not isinstance(data.get("label"), str) + or not isinstance(data.get(backend), dict) + or not isinstance(data.get("workers"), dict) + ): + continue + record = Record(directory, data) + if not all( + isinstance(data["workers"].get(key), str) + for key in ("lightcone", "distributed", "python") + ): + continue + section = record.section(backend) + if backend == "slurm" and not isinstance(section.get("job"), str): + continue + if backend == "gateway": + if not all( + isinstance(section.get(key), str) and section[key] for key in ("name", "address") + ): + continue + else: + tls = record.section("tls") + if not all( + isinstance(tls.get(key), str) + and tls[key] + and not Path(tls[key]).is_absolute() + and ".." not in Path(tls[key]).parts + for key in ("ca", "cert", "key") + ): + continue + yield record + + +def _local_process(record: Record, key: str) -> bool: + """Reject reused PIDs, including records predating creation timestamps.""" + local = record.section("local") + pid = local.get(key) + if isinstance(pid, bool) or not isinstance(pid, int) or pid <= 0: + return False + try: + process = psutil.Process(pid) + if process.status() == psutil.STATUS_ZOMBIE or not process.is_running(): + return False + started = local.get(f"{key}_started") + if started is not None: + if process.create_time() != started: + return False + else: + command = process.cmdline() + module = "dask_scheduler" if key == "pid" else "dask_worker" + if command[1:3] != ["-m", f"distributed.cli.{module}"]: + return False + try: + path = command[command.index("--scheduler-file") + 1] + except (ValueError, IndexError): + return False + if Path(path) != record.directory / "scheduler.json": + return False + return Path(process.cwd()) == record.directory.resolve() + except psutil.NoSuchProcess: + return False + + +def _scheduler_address(record: Record) -> str | None: + try: + data = json.loads((record.directory / "scheduler.json").read_text()) + except (OSError, ValueError): + return None + address = data.get("address") if isinstance(data, dict) else None + return address if isinstance(address, str) and address else None + + +def _slurm_states() -> dict[str, tuple[str, str | None]]: + """One queue snapshot for all the user's Slurm records.""" + result = subprocess.run( + ["squeue", "--me", "--noheader", "--format=%i|%T|%S|%L"], + capture_output=True, + text=True, + check=True, + timeout=_CONNECT_TIMEOUT, + ) + states = {} + for line in result.stdout.splitlines(): + fields = line.strip().split("|") + if len(fields) == 4: + job, state, start, _remaining = fields + states[job] = (state, None if start in ("N/A", "Unknown", "") else start) + return states + + +def _gateway_address() -> str | None: + address = dask.config.get("gateway.address", None) + if not isinstance(address, str) or not address: + return None + try: + return address.format(**os.environ).rstrip("/") or None + except (KeyError, ValueError): + return None + + +def _gateway_type() -> Any: + try: + from dask_gateway import Gateway + except ImportError as error: + raise ProjectError( + "This cluster needs dask-gateway. Install lightcone-cli with its gateway extra: " + "uv tool install 'lightcone-cli[gateway]'." + ) from error + return Gateway + + +def _gateway_states() -> dict[str, str]: + gateway_type = _gateway_type() + + async def query() -> dict[str, str]: + async with gateway_type(asynchronous=True) as gateway: + reports = await asyncio.wait_for( + gateway.list_clusters(status=["pending", "running", "stopping"]), + timeout=_CONNECT_TIMEOUT, + ) + return {report.name: report.status.name for report in reports} + + return asyncio.run(query()) + + +def _gateway_connection(record: Record) -> tuple[str, Any]: + """Fetch native credentials without creating a cluster lifecycle handle. + + GatewayCluster.get_client() cannot accept a connection timeout. Its address + and security come from the same public report, and suffice for a Client. + """ + gateway_type = _gateway_type() + + async def connect() -> tuple[str, Any]: + async with gateway_type(asynchronous=True) as gateway: + report = await asyncio.wait_for( + gateway.get_cluster(record.section("gateway")["name"]), timeout=_CONNECT_TIMEOUT + ) + if report.status.name != "RUNNING": + raise ProjectError( + f"Cluster {record.label} is {report.status.name.lower()}. " + f"Check it in {_COMPUTE}." + ) + return report.scheduler_address, report.security + + return asyncio.run(connect()) + + +def _states(records: list[Record]) -> list[Record]: + """Unknown backends are not candidates; discovery never mutates the registry.""" + slurm: dict[str, tuple[str, str | None]] | None = None + gateway: dict[str, str] | None = None + if any(record.backend == "slurm" for record in records): + try: + slurm = _slurm_states() + except (OSError, subprocess.SubprocessError): + pass + if any(record.backend == "gateway" for record in records): + try: + gateway = _gateway_states() + except ProjectError: + raise + except Exception: + pass + observed = [] + for record in records: + state, estimate = "unknown", None + if record.backend == "local": + try: + if _local_process(record, "pid"): + state = "running" if _scheduler_address(record) else "starting" + else: + state = "stopping" if _local_process(record, "worker") else "gone" + except psutil.AccessDenied: + pass + elif record.backend == "slurm" and slurm is not None: + status, estimate = slurm.get(record.section("slurm")["job"], ("GONE", None)) + if status == "PENDING": + state = "queued" + elif status in ("RUNNING", "CONFIGURING"): + state = "running" if _scheduler_address(record) else "starting" + else: + state = "gone" if status == "GONE" else "stopping" + elif record.backend == "gateway" and gateway is not None: + status = gateway.get(record.section("gateway")["name"], "GONE") + state = {"PENDING": "starting", "RUNNING": "running", "GONE": "gone"}.get( + status, "stopping" + ) + observed.append(replace(record, state=state, start_estimate=estimate)) + return observed + + +def attached_cluster(root: Path) -> Record | None: + """Select the one compatible live cluster, refusing ambiguous choices.""" + records = [] + for record in _records(): + image = record.section("workers").get("image") + if record.backend == "gateway": + server_image = os.environ.get("JUPYTER_IMAGE_SPEC") or os.environ.get("JUPYTER_IMAGE") + if ( + not server_image + or image != server_image + or record.section("gateway").get("address") != _gateway_address() + or project.mode(root) != "direct" + ): + continue + elif image is not None: + continue + elif ( + record.backend == "local" + and record.section("local").get("host") != socket.gethostname() + ): + continue + records.append(record) + candidates = [ + record for record in _states(records) if record.state in ("queued", "starting", "running") + ] + if len(candidates) > 1: + names = ", ".join(f"{record.label} ({record.id})" for record in candidates) + raise ProjectError( + f"Several clusters could run this project: {names}. Stop all but one in {_COMPUTE}." + ) + return candidates[0] if candidates else None + + +def _probe(root: str) -> dict[str, str | bool]: + """Run on each worker before any recipe is submitted.""" + versions: dict[str, str | bool] = {"python": platform.python_version()} + for name, distribution in (("lightcone", "lightcone-cli"), ("distributed", "distributed")): + try: + versions[name] = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + versions[name] = "not installed" + versions["project"] = Path(root).is_dir() and os.access(root, os.R_OK | os.X_OK) + return versions + + +def _verify(client: Any, record: Record, root: Path) -> None: + expected = _probe(str(root)) + reports = client.run(_probe, str(root)) + if not reports: + raise ProjectError(f"Cluster {record.label} has no workers. Check it in {_COMPUTE}.") + for address, actual in reports.items(): + for key in ("lightcone", "distributed", "python"): + if actual[key] != expected[key]: + raise ProjectError( + f"Cluster {record.label}, worker {address}, runs {key} {actual[key]}; " + f"this process runs {expected[key]}. Replace the cluster in {_COMPUTE}." + ) + if not actual["project"]: + raise ProjectError( + f"Cluster {record.label}, worker {address}, cannot read project {root}. " + "Move the project to a filesystem the workers see, such as $SCRATCH or home." + ) + + +@contextmanager +def client(record: Record, root: Path) -> Iterator[Any]: + """Attach, verify and close only this run's client, never shared workers.""" + from distributed import Client, Security + + expected = _probe(str(root)) + for key in ("lightcone", "distributed", "python"): + actual = record.section("workers").get(key) + if actual != expected[key]: + raise ProjectError( + f"Cluster {record.label} records {key} {actual}; " + f"this process runs {expected[key]}. " + f"Replace the cluster in {_COMPUTE}." + ) + deadline = time.monotonic() + _WAIT + while record.state != "running": + if record.state == "queued": + estimate = ( + f" (estimated start {record.start_estimate})" if record.start_estimate else "" + ) + raise ProjectError( + f"Cluster {record.label} is queued{estimate}. Run again once it starts." + ) + if record.state != "starting": + raise ProjectError(f"Cluster {record.label} is {record.state}. Check it in {_COMPUTE}.") + if time.monotonic() >= deadline: + raise ProjectError( + f"Cluster {record.label} did not start in {_WAIT} seconds. " + f"Check its logs in {record.directory} and {_COMPUTE}." + ) + time.sleep(0.2) + record = _states([record])[0] + + with ExitStack() as stack: + try: + if record.backend == "gateway": + address, security = _gateway_connection(record) + else: + # Read the address ourselves: Client(scheduler_file=...) can wait + # forever if the scheduler removes the file between check and open. + scheduler_address = _scheduler_address(record) + if not scheduler_address: + raise OSError("the scheduler file is missing or incomplete") + address = scheduler_address + tls = record.section("tls") + security = Security( # type: ignore[no-untyped-call] + tls_ca_file=str(record.directory / tls["ca"]), + tls_client_cert=str(record.directory / tls["cert"]), + tls_client_key=str(record.directory / tls["key"]), + require_encryption=True, + ) + connected = Client( # type: ignore[no-untyped-call] + address, security=security, timeout=_CONNECT_TIMEOUT, set_as_default=False + ) + stack.callback(connected.close) + connected.wait_for_workers(1, timeout=_WAIT) + _verify(connected, record, root) + except ProjectError: + raise + except Exception as error: + raise ProjectError( + f"Cluster {record.label} could not be reached or verified: {error}. " + f"Check or stop it in {_COMPUTE}." + ) from error + yield connected diff --git a/src/lightcone/engine/materialize.py b/src/lightcone/engine/materialize.py index 9d441b4..75b2bc8 100644 --- a/src/lightcone/engine/materialize.py +++ b/src/lightcone/engine/materialize.py @@ -36,13 +36,24 @@ import json import os import re -from collections.abc import Iterator, Sequence +import uuid +from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Protocol -from lightcone.engine import assets, container, dataset, identity, plan, project, venue, worker +from lightcone.engine import ( + assets, + clusters, + container, + dataset, + identity, + plan, + project, + venue, + worker, +) from lightcone.engine.plan import Graph, Key, Task from lightcone.engine.project import ProjectError @@ -74,6 +85,8 @@ class MaterializeReport: #: width breaks the one thing a denial message is for. The caller #: prints these unwrapped, exactly as ``lc run`` does. notes: list[str] = field(default_factory=list) + #: Execution target; check mode never selects or connects to a venue. + venue: dict[str, str | int] | None = None @property def ok(self) -> bool: @@ -472,7 +485,11 @@ def _sandbox_line(mode: str) -> str: def materialize( - root: Path, targets: Sequence[str], *, refresh: bool = False + root: Path, + targets: Sequence[str], + *, + refresh: bool = False, + on_venue: Callable[[dict[str, str | int]], None] | None = None, ) -> MaterializeReport: """Make everything *targets* names, committing each output as it lands. @@ -482,6 +499,8 @@ def materialize( output asks for what it is made of. refresh: Also remake outputs that are merely behind — still what the spec asks for, but made under an earlier environment. + on_venue: Announce the selected execution target before preparing + or running work. The same selection is recorded in the report. Returns: What was made, what was current or behind, what failed or was @@ -495,12 +514,29 @@ def materialize( # First, because its remedy is the one with queue latency: the user # can submit the allocation and fix anything the later refusals name # while waiting for it. - venue.require_compute_node() + attached = venue.attached_cluster(root) + if attached is None or attached.backend == "local": + venue.require_compute_node() + nodes = venue.allocation_nodes() + selected: dict[str, str | int] + if attached is not None: + selected = { + "kind": "cluster", + "backend": attached.backend, + "id": attached.id, + "label": attached.label, + } + elif nodes: + selected = {"kind": "allocation", "nodes": nodes} + else: + selected = {"kind": "local"} + if on_venue is not None: + on_venue(selected) project.require_uv() project.require_git() project.require_git_annex() dataset.require_committer(root) - report = MaterializeReport() + report = MaterializeReport(venue=selected) if warning := project.uv_scrub_warning(): report.warnings.append(warning) # The dirty check comes before anything that writes: the image @@ -523,61 +559,52 @@ def materialize( return report _fetch_inputs(root, graph, report) # Before the runtime resolves, because the refusal must not cost an - # image build: a containerized graph can span an allocation only if + # image build: a containerized graph can span several hosts only if # every node can see the image — the hint suffices, since which # stores span nodes is `container._SHARED_STORE_RUNTIMES`'s fact and # a wholly missing runtime gets `runtime_for_run`'s own refusal. Off # the driver's node a task would otherwise fail to find an image # `--pull=never` forbids it to fetch. - if ( - (nodes := venue.allocation_nodes()) > 1 - and project.mode(root) == "containerized" - and (name := container.runtime_hint()) - and name not in container._SHARED_STORE_RUNTIMES - ): - raise ProjectError( - f"this allocation spans {nodes} nodes and `{name}`'s image store is " - "node-local, so recipes scheduled on the other nodes would not find " - "the image. Use a single-node allocation, or a system whose runtime " - "shares images across nodes (NERSC's podman-hpc)." - ) - # Materialize is one of the two verbs allowed to build the image (the - # other is `lc build`); the probe and the rerun entry point only find - # one. Resolved once, then handed to every task — the HEAD discipline. - runtime = container.runtime_for_run(root, build=True) - # Converge the environment: workers pass `--no-sync`, so this is the - # only place on a run's path where it is made to match the lock. (A - # rerun does not come through here; its entry point converges too.) - report.warnings.extend(f"uv: {w}" for w in container.converge(runtime)) - - # The run's driver-resolved facts, each read once: HEAD because the - # driver commits as outputs land and a per-task read would stamp - # later manifests with a commit this run created; the uv probe - # because attestation is a fact about the run (and empty is an - # answer, not a failure); one content-hash memo because a declared - # input shared by several outputs is the same bytes every time. - context = worker.RunContext( - env_version=env_version, - head=dataset.head(root), - versions=assets.Versions(), - runtime=runtime, - uv_version=project.uv_version(root), - ) - # The history question is the driver's to answer — workers have no - # git, by design — so each task is told up front whether its - # directory was last written by something other than its own run - # record. A foreign write contradicts the manifest, and a worker that - # trusted the recorded digest would skip the output forever. Guarded - # on the manifest's presence, as `_classified` is: without one the - # answer is dead — the output is remade regardless — and each ask is - # a git process. - foreign = { - key: _foreign_write(root, task) if task.manifest_path.is_file() else None - for key, task in graph.tasks.items() - } - outstanding: dict[Key, Task] = dict(graph.tasks) + _require_shared_image_store(root, nodes) + outstanding: dict[Key, Task] = {} try: - with cluster_for_run() as scheduler: + with cluster_for_run(root, attached) as scheduler: + # Materialize is one of the two verbs allowed to build the image (the + # other is `lc build`); the probe and the rerun entry point only find + # one. Resolved once, then handed to every task — the HEAD discipline. + runtime = container.runtime_for_run(root, build=True) + # Converge the environment: workers pass `--no-sync`, so this is the + # only place on a run's path where it is made to match the lock. (A + # rerun does not come through here; its entry point converges too.) + report.warnings.extend(f"uv: {w}" for w in container.converge(runtime)) + + # The run's driver-resolved facts, each read once: HEAD because the + # driver commits as outputs land and a per-task read would stamp + # later manifests with a commit this run created; the uv probe + # because attestation is a fact about the run (and empty is an + # answer, not a failure); one content-hash memo because a declared + # input shared by several outputs is the same bytes every time. + context = worker.RunContext( + env_version=env_version, + head=dataset.head(root), + versions=assets.Versions(), + runtime=runtime, + uv_version=project.uv_version(root), + ) + # The history question is the driver's to answer — workers have no + # git, by design — so each task is told up front whether its + # directory was last written by something other than its own run + # record. A foreign write contradicts the manifest, and a worker that + # trusted the recorded digest would skip the output forever. Guarded + # on the manifest's presence, as `_classified` is: without one the + # answer is dead — the output is remade regardless — and each ask is + # a git process. + foreign = { + key: _foreign_write(root, task) if task.manifest_path.is_file() else None + for key, task in graph.tasks.items() + } + outstanding = dict(graph.tasks) + prefix = f"lc/{project.project_name(root)}/{uuid.uuid4().hex}" pending: dict[Key, Any] = {} # Submitted in dependency order so a task's upstream futures # exist to be passed to it. Dask still derives the *execution* @@ -592,18 +619,29 @@ def materialize( refresh, foreign[key], *[pending[dep] for dep in task.depends_on], - key=_name(key), + key=f"{prefix}/{_name(key)}", ) for result in scheduler.completed(list(pending.values())): _consume(root, graph.tasks[result.key], result, dsid, runtime, report) outstanding.pop(result.key, None) + except BaseException as error: + if attached is not None and outstanding: + # Disconnecting does not interrupt a recipe already running + # in a shared worker. Restoring its output now would race it. + detail = str(error) or type(error).__name__ + raise ProjectError( + f"Run on {attached.label} interrupted: {detail}. Running recipes may " + "still finish; their uncommitted outputs have been left in place. " + "Stop the cluster in Lightcone sidebar › Compute, then inspect " + "those outputs before restoring them or rerunning." + ) from error + raise finally: - # Whatever never reported — an interrupt, a dead cluster — left a - # reset output directory behind. Scoped to this run's outputs and - # never to the whole tree, so edits made while the graph ran - # survive. - for task in outstanding.values(): - dataset.restore(root, _owned(root, task)) + # Owned clusters have stopped their workers by now. Restore only + # this run's outputs, so unrelated edits made during it survive. + if attached is None: + for task in outstanding.values(): + dataset.restore(root, _owned(root, task)) # The tree was clean at the start-of-run refusal and save/restore # keeps `results/` clean, so anything dirty *now* was edited while # the graph ran — and every manifest records the starting commit, @@ -656,11 +694,9 @@ def _consume( class Scheduler(Protocol): """How the driver talks to whatever is running the graph. - Two methods, because that is all the driver needs and all a venue has - to supply: hand over a task with its upstream handles, and iterate the - results as they land. Keeping it this narrow is what lets the suite - run the graph inline — and what will let a venue larger than a laptop - land behind :func:`cluster_for_run` without the driver noticing. + Hand over a task with its upstream handles, then iterate results as + they land. All venues share this interface; the suite can also run + the graph inline. """ def submit(self, fn: Any, *args: Any, key: str) -> Any: @@ -669,7 +705,7 @@ def submit(self, fn: Any, *args: Any, key: str) -> Any: Args: fn: The function to run. *args: Its arguments, upstream handles included. - key: A display name for the task. + key: A unique run-scoped scheduler key. Returns: A handle to pass to dependents. @@ -704,19 +740,19 @@ def completed(self, handles: list[Any]) -> Iterator[worker.TaskResult]: # annotated rather than the module exempted. from distributed import as_completed - for _, result in as_completed(handles, with_results=True): # type: ignore[no-untyped-call] + for _, result in as_completed( # type: ignore[no-untyped-call] + handles, with_results=True, loop=self.client.loop + ): yield result @contextmanager -def cluster_for_run() -> Iterator[Scheduler]: +def cluster_for_run(root: Path, attached: clusters.Record | None) -> Iterator[Scheduler]: """Open a scheduler for one run — the venue ladder, and nothing else. - Every core, with no knob to say otherwise: how much of a machine a run - may use, and which machine, is one question, and the venue answers it — - a SLURM allocation spans every node it was granted, and the local - machine is the whole of itself. Detected, never configured, and only - here: nothing outside this function asks where a run executes. + Selection happens once, before the login guard and venue announcement. + An attached client borrows its cluster; only clusters created here are + closed with the run. Threads rather than processes on the local branch — every task's real work happens in a subprocess behind the exec boundary, so a worker @@ -731,6 +767,17 @@ def cluster_for_run() -> Iterator[Scheduler]: with venue.slurm_client() as client: yield _Dask(client) return + if attached is not None: + with clusters.client(attached, root) as client: + hosts = {worker["host"] for worker in client.scheduler_info()["workers"].values()} + # Slurm workers can still be joining after the first one is + # ready. Its declared size prevents a transient single-host + # view from admitting a node-local image store. + requested = attached.section("slurm").get("nodes", 1) + nodes = max(len(hosts), requested if isinstance(requested, int) else 1) + _require_shared_image_store(root, nodes) + yield _Dask(client) + return from distributed import Client, LocalCluster with LocalCluster( # type: ignore[no-untyped-call] @@ -743,6 +790,22 @@ def cluster_for_run() -> Iterator[Scheduler]: yield _Dask(client) +def _require_shared_image_store(root: Path, nodes: int) -> None: + """Refuse a multi-host container run before paying for an image build.""" + if ( + nodes > 1 + and project.mode(root) == "containerized" + and (name := container.runtime_hint()) + and name not in container._SHARED_STORE_RUNTIMES + ): + raise ProjectError( + f"this run spans {nodes} nodes and `{name}`'s image store is " + "node-local, so recipes scheduled on the other nodes would not find " + "the image. Use a single-node cluster or allocation, or a system whose " + "runtime shares images across nodes (NERSC's podman-hpc)." + ) + + def _fetch_inputs(root: Path, graph: Graph, report: MaterializeReport) -> None: """Bring declared inputs' bytes into this clone before anything hashes. diff --git a/src/lightcone/engine/project.py b/src/lightcone/engine/project.py index c907216..fb6a4c6 100644 --- a/src/lightcone/engine/project.py +++ b/src/lightcone/engine/project.py @@ -8,7 +8,7 @@ import shutil import subprocess import uuid -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import asdict, dataclass, field from functools import partial from pathlib import Path @@ -752,22 +752,26 @@ def _uv_scrubbed(name: str) -> bool: ) -def child_env() -> dict[str, str]: +def child_env(environment: Mapping[str, str] | None = None) -> dict[str, str]: """Build the environment external tools run in. - Ours, minus ``VIRTUAL_ENV`` and minus every ``UV_*`` variable outside - the :data:`_UV_KEPT` plumbing allowlist. Every uv invocation names its - project explicitly, so an activated environment elsewhere is never + Ours (or the supplied mapping), minus ``VIRTUAL_ENV`` and every + ``UV_*`` variable outside the :data:`_UV_KEPT` plumbing allowlist. + Every uv invocation names its project explicitly, so an activated + environment elsewhere is never what we mean — and an ambient install setting would change what a sync installs without moving ``env_version``, which is the identity hole the scrub closes. + Args: + environment: An explicit environment, or the current process's. + Returns: - The current environment without ``VIRTUAL_ENV`` or scrubbed ``UV_*``. + The environment without ``VIRTUAL_ENV`` or scrubbed ``UV_*``. """ return { k: v - for k, v in os.environ.items() + for k, v in (os.environ if environment is None else environment).items() if k != "VIRTUAL_ENV" and not _uv_scrubbed(k) } diff --git a/src/lightcone/engine/venue.py b/src/lightcone/engine/venue.py index 9acdcbd..be02c18 100644 --- a/src/lightcone/engine/venue.py +++ b/src/lightcone/engine/venue.py @@ -1,20 +1,18 @@ -"""Where a run executes: the venue a materialization finds itself on. +"""Where a run executes: an allocation, a managed cluster, or this host. -A venue is host state, never project state — nothing here reads the -project or enters any identity. The one venue beyond the local machine is -a SLURM allocation, detected rather than configured: the user already -declared every resource question to SLURM (`salloc -N4 …`), so the -allocation *is* the declaration, and lc's job is to span it — one Dask -worker per allocated node, launched with a single `srun`, all connected -to a scheduler living in the driver process. +An active SLURM allocation takes precedence over the user's managed +clusters. Otherwise :func:`attached_cluster` selects a compatible cluster +from the registry; absent one, the login-node guard protects this host. +Managed clusters outlive runs: only the allocation workers launched here +belong to the engine and are retired when its client closes. -Workers run the driver's own interpreter (`sys.executable -m`), which on +Allocation workers run the driver's interpreter (`sys.executable -m`), which on an HPC system is the lc tool environment on the shared filesystem — so driver and workers are the identical installation, which is all a worker process needs: `lightcone.engine` importable at the driver's version. Workers need no git and no git-annex; the driver owns git alone. -If the driver dies uncleanly, workers exit on their own (death timeout) +If the allocation's driver dies, its workers exit on their own (death timeout) and the allocation's walltime is the backstop; whatever the interrupted run left behind meets the next run's dirty-tree refusal, which names the `results/` paths to discard — that is the designed recovery, not a @@ -32,10 +30,15 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass +from pathlib import Path from typing import Any +from lightcone.engine import clusters from lightcone.engine.project import ProjectError +#: The complete attachment contract understood by this engine. +CLUSTER_RECORD_FORMAT = clusters.FORMAT + #: How long the allocation's workers get to connect before the run #: refuses. Generous because the first import of `distributed` from a #: cold parallel filesystem is seconds, not milliseconds. @@ -81,6 +84,18 @@ class _Site: ) +def attached_cluster(root: Path) -> clusters.Record | None: + """Select this project's managed cluster, unless already in an allocation. + + Raises ProjectError when several compatible live clusters exist. A + selected cluster's connection or compatibility failure never falls + back to this host. + """ + if "SLURM_JOB_ID" in os.environ: + return None + return clusters.attached_cluster(root) + + def require_compute_node(command: str = "lc materialize") -> None: """Refuse to execute recipes on a known HPC center's login node. @@ -102,6 +117,11 @@ def require_compute_node(command: str = "lc materialize") -> None: site = next((s for s in _SITES if s.marker in os.environ), None) if site is None: return + cluster_remedy = ( + "Or start a cluster in Lightcone sidebar › Compute, then run again.\n\n" + if command == "lc materialize" + else "" + ) raise ProjectError( f"{command} executes recipes on compute nodes, and this is a " f"{site.name} login node ({site.marker} is set with no SLURM " @@ -117,6 +137,7 @@ def require_compute_node(command: str = "lc materialize") -> None: f" {site.sbatch} \\\n" f" --wrap '{command}'\n" "\n" + f"{cluster_remedy}" "lc materialize --check, lc status and lc run work anywhere." ) diff --git a/src/lightcone/engine/worker.py b/src/lightcone/engine/worker.py index ffffce7..a1365f6 100644 --- a/src/lightcone/engine/worker.py +++ b/src/lightcone/engine/worker.py @@ -28,9 +28,10 @@ from __future__ import annotations import functools +import os import sys from collections.abc import Mapping -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from datetime import UTC, datetime from pathlib import Path from typing import Literal @@ -56,6 +57,18 @@ _SHELL = "bash" +def _worker_local(name: str) -> bool: + """Identify variables whose meaning belongs to the worker's host or job.""" + return name in {"HOSTNAME", "TMPDIR", "XDG_RUNTIME_DIR", "DISPLAY"} or name.startswith( + ("SLURM_", "SLURMD_", "JUPYTER_", "JPY_", "SSH_") + ) + + +def _driver_environment() -> dict[str, str]: + """Snapshot the driver's recipe environment without its host/job settings.""" + return child_env({name: value for name, value in os.environ.items() if not _worker_local(name)}) + + @dataclass(frozen=True) class TaskResult: """What one task did. Returned, never raised, and handed to dependents.""" @@ -108,6 +121,9 @@ class RunContext: runtime: container.Runtime #: The uv that converges environments this run. Attestation only. uv_version: str + #: Recipe settings captured once by the driver, never installed in the + #: shared worker's environment or recorded in provenance (may contain credentials). + environment: dict[str, str] = field(default_factory=_driver_environment, repr=False) # ============================================================================= @@ -250,7 +266,12 @@ def execute( [_SHELL, "-c", task.recipe], cwd=root, prefix=uv_prefix(root, sync=False), - env=child_env(), + env=child_env( + { + **context.environment, + **{name: value for name, value in os.environ.items() if _worker_local(name)}, + } + ), ) finished_at = _now() diff --git a/tests/conftest.py b/tests/conftest.py index 04a0188..bdd50af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,7 +23,7 @@ def runner() -> CliRunner: @pytest.fixture(autouse=True) -def venue_env(monkeypatch: pytest.MonkeyPatch) -> None: +def venue_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: """Strip the host's venue out of the suite's environment. On a known center's login node every materialize test would otherwise @@ -32,7 +32,9 @@ def venue_env(monkeypatch: pytest.MonkeyPatch) -> None: own table, so a center added there is scrubbed here for free; the venue tests set these back deliberately. """ - from lightcone.engine import venue + from lightcone.engine import clusters, venue + + monkeypatch.setattr(clusters, "registry_root", lambda: tmp_path / "clusters") for name in ( *(site.marker for site in venue._SITES), @@ -41,6 +43,8 @@ def venue_env(monkeypatch: pytest.MonkeyPatch) -> None: "SLURM_JOB_NUM_NODES", "SLURM_NNODES", "SLURM_CPUS_ON_NODE", + "JUPYTER_IMAGE_SPEC", + "JUPYTER_IMAGE", ): monkeypatch.delenv(name, raising=False) @@ -183,7 +187,7 @@ def inline(monkeypatch: pytest.MonkeyPatch) -> None: from lightcone.engine import materialize @contextmanager - def fake() -> Iterator[_Inline]: + def fake(root: Path, attached: object) -> Iterator[_Inline]: yield _Inline() monkeypatch.setattr(materialize, "cluster_for_run", fake) diff --git a/tests/test_cli.py b/tests/test_cli.py index 1f5bd9e..055fcd5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -327,8 +327,12 @@ def _stub(monkeypatch: pytest.MonkeyPatch, **outcomes: object) -> list[tuple[str def record(name: str) -> object: def call(root: Path, targets: object, **kwargs: object) -> object: + announce = kwargs.pop("on_venue", None) seen.append((name, (list(targets), kwargs))) - return outcomes.get(name, engine.MaterializeReport()) + report = outcomes.get(name, engine.MaterializeReport()) + if announce is not None: + announce(report.venue or {"kind": "local"}) + return report return call @@ -469,9 +473,50 @@ def test_the_json_report_is_machine_readable( "planned": {}, "warnings": [], "notes": [], + "venue": None, } +@pytest.mark.parametrize( + ("venue", "description"), + [ + ({"kind": "local"}, "this host"), + ({"kind": "allocation", "nodes": 4}, "SLURM allocation (4 nodes)"), + ( + {"kind": "cluster", "backend": "slurm", "id": "cluster-id", "label": "[debug]"}, + "[debug] (slurm, cluster-id)", + ), + ], +) +def test_materialize_announces_its_venue_and_includes_it_in_json( + runner: CliRunner, + project: Path, + monkeypatch: pytest.MonkeyPatch, + venue: dict[str, str | int], + description: str, +) -> None: + from lightcone.engine.materialize import MaterializeReport + + _stub(monkeypatch, materialize=MaterializeReport(venue=venue)) + + result = runner.invoke(main, ["materialize"]) + assert result.exit_code == 0, result.output + assert result.output.startswith(f"Running on {description}\n") + + result = runner.invoke(main, ["materialize", "--json"]) + assert result.exit_code == 0, result.output + assert json.loads(result.output)["venue"] == venue + + +def test_check_never_announces_an_execution_venue( + runner: CliRunner, project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _stub(monkeypatch) + result = runner.invoke(main, ["materialize", "--check"]) + assert result.exit_code == 0, result.output + assert "Running on" not in result.output + + def test_an_engine_refusal_is_a_clean_error( runner: CliRunner, project: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/test_clusters.py b/tests/test_clusters.py new file mode 100644 index 0000000..bf927ef --- /dev/null +++ b/tests/test_clusters.py @@ -0,0 +1,427 @@ +"""The registry is read-only, and attached clusters outlive their clients.""" + +from __future__ import annotations + +import asyncio +import json +import os +import socket +import subprocess +import sys +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import dask.config +import psutil +import pytest + +from lightcone.engine import clusters +from lightcone.engine.project import ProjectError + + +@pytest.fixture +def registry(tmp_path, monkeypatch): + path = tmp_path / "clusters" + path.mkdir() + monkeypatch.setattr(clusters, "registry_root", lambda: path) + return path + + +def record(registry, backend="slurm", suffix="aaaa", **overrides): + directory = registry / f"20260925-120000-{suffix}" + directory.mkdir(exist_ok=True) + data = { + "format": clusters.FORMAT, + "id": directory.name, + "backend": backend, + "label": f"Example {suffix}", + "workers": {**clusters._probe(str(registry)), "image": None}, + "tls": {"ca": "tls/cert.pem", "cert": "tls/cert.pem", "key": "tls/key.pem"}, + "local": {"host": socket.gethostname(), "pid": 123, "worker": 456}, + "slurm": {"job": "12345"}, + "gateway": {"name": "alice.cluster", "address": "https://gateway.example"}, + **overrides, + } + (directory / "cluster.json").write_text(json.dumps(data)) + return clusters.Record(directory, data) + + +def scheduler_file(record): + (record.directory / "scheduler.json").write_text(json.dumps({"address": "tls://127.0.0.1:1"})) + + +def test_registry_skips_foreign_malformed_and_invalid_directories(registry, monkeypatch): + expected = record(registry) + record(registry, suffix="bbbb", format="future") + record(registry, suffix="cccc", slurm={"job": []}) + record(registry, suffix="dddd", tls={"ca": "/outside", "cert": "a", "key": "b"}) + bad = registry / "20260925-120000-eeee" + bad.mkdir() + (bad / "cluster.json").write_text("{") + unrelated = registry / "not-a-cluster" + unrelated.mkdir() + (unrelated / "cluster.json").write_text(json.dumps(expected.data)) + assert [item.id for item in clusters._records()] == [expected.id] + assert len(list(registry.iterdir())) == 6 + + +@pytest.mark.parametrize( + ("status", "has_file", "expected"), + [ + ("PENDING", False, "queued"), + ("RUNNING", False, "starting"), + ("CONFIGURING", True, "running"), + ("COMPLETING", True, None), + (None, True, None), + ], +) +def test_slurm_backend_decides_liveness(registry, monkeypatch, status, has_file, expected): + item = record(registry) + if has_file: + scheduler_file(item) + monkeypatch.setattr( + clusters, "_slurm_states", lambda: {"12345": (status, "soon")} if status else {} + ) + selected = clusters.attached_cluster(registry) + assert (selected.state if selected else None) == expected + + +def test_queue_snapshot_is_batched_and_ambiguity_refuses(registry, monkeypatch): + record(registry) + record(registry, suffix="bbbb") + command = MagicMock( + return_value=SimpleNamespace(stdout="12345|PENDING|2026-09-25T15:00|1:00\n") + ) + monkeypatch.setattr(clusters.subprocess, "run", command) + with pytest.raises(ProjectError, match="Several clusters.*aaaa.*bbbb.*Stop all but one"): + clusters.attached_cluster(registry) + command.assert_called_once() + assert command.call_args.args[0] == ["squeue", "--me", "--noheader", "--format=%i|%T|%S|%L"] + + +def test_unavailable_backend_does_not_make_a_stale_file_live(registry, monkeypatch): + scheduler_file(record(registry)) + monkeypatch.setattr(clusters, "_slurm_states", MagicMock(side_effect=FileNotFoundError)) + assert clusters.attached_cluster(registry) is None + + +@pytest.mark.parametrize("field", ["host", "pid_started", "cwd", "zombie"]) +def test_local_selection_requires_process_identity(registry, monkeypatch, field): + item = record(registry, backend="local") + item.data["local"]["pid_started"] = 42 + process = MagicMock() + process.create_time.return_value = 43 if field == "pid_started" else 42 + process.cwd.return_value = "/elsewhere" if field == "cwd" else str(item.directory) + process.status.return_value = ( + psutil.STATUS_ZOMBIE if field == "zombie" else psutil.STATUS_RUNNING + ) + monkeypatch.setattr(clusters.psutil, "Process", lambda pid: process) + if field == "host": + item.data["local"]["host"] = "another-machine" + (item.directory / "cluster.json").write_text(json.dumps(item.data)) + scheduler_file(item) + assert clusters.attached_cluster(registry) is None + + +def test_legacy_process_requires_exact_module_and_scheduler_file(registry, monkeypatch): + item = record(registry, backend="local") + process = MagicMock() + process.cwd.return_value = str(item.directory) + process.cmdline.return_value = [ + sys.executable, + "-m", + "distributed.cli.dask_scheduler", + "--scheduler-file", + str(item.directory / "scheduler.json"), + ] + monkeypatch.setattr(clusters.psutil, "Process", lambda pid: process) + assert clusters._local_process(item, "pid") + process.cmdline.return_value[-1] = "/other/scheduler.json" + assert not clusters._local_process(item, "pid") + + +@pytest.mark.parametrize("incompatible", ["address", "image", "containerized", "missing-image"]) +def test_gateway_compatibility_filters_before_connect(registry, monkeypatch, incompatible): + item = record(registry, backend="gateway") + item.data["workers"]["image"] = "hub:image" + if incompatible == "address": + item.data["gateway"]["address"] = "https://elsewhere" + if incompatible == "image": + item.data["workers"]["image"] = "hub:other" + if incompatible != "missing-image": + monkeypatch.setenv("JUPYTER_IMAGE_SPEC", "hub:image") + else: + monkeypatch.delenv("JUPYTER_IMAGE_SPEC", raising=False) + monkeypatch.delenv("JUPYTER_IMAGE", raising=False) + monkeypatch.setattr( + clusters.project, + "mode", + lambda root: "containerized" if incompatible == "containerized" else "direct", + ) + (item.directory / "cluster.json").write_text(json.dumps(item.data)) + monkeypatch.setattr( + clusters, "_gateway_states", MagicMock(side_effect=AssertionError("not queried")) + ) + with dask.config.set({"gateway.address": "https://gateway.example"}): + assert clusters.attached_cluster(registry) is None + + +@pytest.mark.parametrize( + "status,expected", + [("PENDING", "starting"), ("RUNNING", "running"), ("STOPPING", None), (None, None)], +) +def test_gateway_liveness(registry, monkeypatch, status, expected): + item = record(registry, backend="gateway") + item.data["workers"]["image"] = "hub:image" + (item.directory / "cluster.json").write_text(json.dumps(item.data)) + monkeypatch.setenv("JUPYTER_IMAGE_SPEC", "hub:image") + monkeypatch.setattr( + clusters, "_gateway_states", lambda: {"alice.cluster": status} if status else {} + ) + with dask.config.set({"gateway.address": "https://gateway.example/"}): + selected = clusters.attached_cluster(registry) + assert (selected.state if selected else None) == expected + + +def test_queued_refuses_with_estimate_and_never_connects(registry, monkeypatch): + item = replace(record(registry), state="queued", start_estimate="tomorrow") + connected = MagicMock(side_effect=AssertionError("must not connect")) + monkeypatch.setattr("distributed.Client", connected) + with pytest.raises(ProjectError, match="queued.*tomorrow.*Run again"): + with clusters.client(item, registry): + pytest.fail("yielded queued cluster") + connected.assert_not_called() + + +def test_startup_wait_is_bounded(registry, monkeypatch): + item = replace(record(registry), state="starting") + monkeypatch.setattr(clusters, "_WAIT", 0) + with pytest.raises(ProjectError, match="did not start.*Check its logs"): + with clusters.client(item, registry): + pytest.fail("yielded starting cluster") + + +def test_scheduler_disappearing_does_not_wait_forever(registry, monkeypatch): + item = replace(record(registry), state="running") + monkeypatch.setattr( + "distributed.Client", MagicMock(side_effect=AssertionError("must not connect")) + ) + with pytest.raises(ProjectError, match="scheduler file is missing"): + with clusters.client(item, registry): + pytest.fail("yielded missing scheduler") + + +@pytest.mark.parametrize("mismatch", ["lightcone", "distributed", "python", "project"]) +def test_probe_checks_every_worker(registry, mismatch): + item = record(registry) + report = clusters._probe(str(registry)) + other = {**report, mismatch: False if mismatch == "project" else "different"} + connected = MagicMock() + connected.run.return_value = {"worker-1": report, "worker-2": other} + with pytest.raises(ProjectError, match="worker-2.*(Replace|Move)"): + clusters._verify(connected, item, registry) + + +@pytest.fixture +def gateway(monkeypatch): + gateway = MagicMock() + gateway.__aenter__.return_value = gateway + gateway.get_cluster = AsyncMock( + return_value=SimpleNamespace( + scheduler_address="gateway://example/alice.cluster", + security=object(), + status=SimpleNamespace(name="RUNNING"), + ) + ) + monkeypatch.setattr(clusters, "_gateway_type", lambda: lambda **kwargs: gateway) + return gateway + + +def test_gateway_uses_native_credentials_and_only_closes_client(registry, monkeypatch, gateway): + item = replace(record(registry, backend="gateway"), state="running") + report = gateway.get_cluster.return_value + connected = MagicMock() + constructor = MagicMock(return_value=connected) + monkeypatch.setattr("distributed.Client", constructor) + connected.run.return_value = {"worker-1": clusters._probe(str(registry))} + with pytest.raises(ValueError, match="recipe failed"): + with clusters.client(item, registry): + raise ValueError("recipe failed") + gateway.get_cluster.assert_awaited_once_with("alice.cluster") + constructor.assert_called_once_with( + report.scheduler_address, security=report.security, timeout=30, set_as_default=False + ) + connected.wait_for_workers.assert_called_once_with(1, timeout=120) + connected.close.assert_called_once() + gateway.__aexit__.assert_awaited_once() + gateway.connect.assert_not_called() + gateway.stop_cluster.assert_not_called() + connected.retire_workers.assert_not_called() + + +def test_gateway_timeout_closes_connection_without_shutdown(registry, monkeypatch, gateway): + item = replace(record(registry, backend="gateway"), state="running") + monkeypatch.setattr(clusters, "_CONNECT_TIMEOUT", 0.01) + + async def pending(name): + await asyncio.sleep(1) + + gateway.get_cluster.side_effect = pending + with pytest.raises(ProjectError, match="could not be reached"): + with clusters.client(item, registry): + pytest.fail("yielded pending cluster") + gateway.__aexit__.assert_awaited_once() + gateway.connect.assert_not_called() + gateway.stop_cluster.assert_not_called() + + +def test_gateway_states_use_native_async_lifecycle(gateway): + gateway.list_clusters = AsyncMock( + return_value=[SimpleNamespace(name="alice.cluster", status=SimpleNamespace(name="RUNNING"))] + ) + assert clusters._gateway_states() == {"alice.cluster": "RUNNING"} + gateway.list_clusters.assert_awaited_once_with(status=["pending", "running", "stopping"]) + gateway.__aexit__.assert_awaited_once() + + +def test_missing_gateway_dependency_is_actionable(monkeypatch): + monkeypatch.setitem(sys.modules, "dask_gateway", None) + with pytest.raises(ProjectError, match=r"lightcone-cli\[gateway\]"): + clusters._gateway_type() + + +def test_gateway_report_credentials_survive_api_cleanup(registry, gateway): + """Pin the optional client's public API, including its real TLS security.""" + api = pytest.importorskip("dask_gateway.client") + from distributed import Security + + generated = Security.temporary() + gateway.get_cluster.return_value = api.ClusterReport( + name="alice.cluster", + options={}, + status=api.ClusterStatus.RUNNING, + scheduler_address="gateway://example/alice.cluster", + dashboard_link=None, + start_time=None, + stop_time=None, + tls_cert=generated.tls_ca_file, + tls_key=generated.tls_client_key, + ) + address, security = clusters._gateway_connection(record(registry, backend="gateway")) + gateway.__aexit__.assert_awaited_once() + assert address == "gateway://example/alice.cluster" + assert security.get_connection_args("client")["require_encryption"] + + +def test_gateway_outage_is_unknown(registry, monkeypatch): + item = record(registry, backend="gateway") + monkeypatch.setattr(clusters, "_gateway_states", MagicMock(side_effect=ConnectionError)) + assert clusters._states([item])[0].state == "unknown" + + +@pytest.mark.parametrize("key", ["lightcone", "distributed", "python"]) +def test_record_version_mismatch_refuses_before_connect(registry, gateway, key): + item = replace(record(registry, backend="gateway"), state="running") + item.data["workers"][key] = "obsolete" + with pytest.raises(ProjectError, match=f"{key} obsolete.*Replace"): + with clusters.client(item, registry): + pytest.fail("yielded obsolete cluster") + gateway.get_cluster.assert_not_called() + + +def test_worker_wait_failure_closes_client(registry, monkeypatch, gateway): + item = replace(record(registry, backend="gateway"), state="running") + connected = MagicMock() + connected.wait_for_workers.side_effect = TimeoutError("no workers arrived") + monkeypatch.setattr("distributed.Client", MagicMock(return_value=connected)) + with pytest.raises(ProjectError, match="no workers arrived"): + with clusters.client(item, registry): + pytest.fail("yielded empty cluster") + connected.close.assert_called_once() + connected.retire_workers.assert_not_called() + + +def test_real_tls_cluster_survives_two_runs(registry): + """The writer's actual record and command lines work across processes.""" + from distributed import Security + + item = record(registry, backend="local") + security = Security.temporary() + tls = item.directory / "tls" + tls.mkdir() + (tls / "cert.pem").write_text(security.tls_ca_file) + (tls / "key.pem").write_text(security.tls_client_key) + common = [ + "--scheduler-file", + str(item.directory / "scheduler.json"), + "--protocol", + "tls", + "--tls-ca-file", + str(tls / "cert.pem"), + "--tls-cert", + str(tls / "cert.pem"), + "--tls-key", + str(tls / "key.pem"), + ] + argv = { + "pid": [ + "distributed.cli.dask_scheduler", + *common, + "--host", + "127.0.0.1", + "--port", + "0", + "--no-dashboard", + ], + "worker": [ + "distributed.cli.dask_worker", + *common, + "--nthreads", + "1", + "--nworkers", + "1", + "--no-nanny", + "--no-dashboard", + "--memory-limit", + "0", + "--death-timeout", + "20", + ], + } + processes = [] + with (item.directory / "test.log").open("w") as log: + try: + for key, args in argv.items(): + process = subprocess.Popen( + [sys.executable, "-m", *args], + cwd=item.directory, + stdout=log, + stderr=log, + env=os.environ.copy(), + ) + processes.append(process) + item.data["local"][key] = process.pid + item.data["local"][f"{key}_started"] = psutil.Process(process.pid).create_time() + path = item.directory / "cluster.json" + path.write_text(json.dumps(item.data)) + initial = path.read_bytes() + for number in (1, 2): + selected = clusters.attached_cluster(registry) + assert selected is not None + with clusters.client(selected, registry) as connected: + assert ( + connected.submit(sum, [number, 2], key=f"run-{number}").result() + == number + 2 + ) + assert all(process.poll() is None for process in processes) + assert path.read_bytes() == initial + finally: + for process in processes: + process.terminate() + for process in processes: + try: + process.wait(timeout=15) + except subprocess.TimeoutExpired: + process.kill() + process.wait() diff --git a/tests/test_materialize.py b/tests/test_materialize.py index 6a99c36..0f3000d 100644 --- a/tests/test_materialize.py +++ b/tests/test_materialize.py @@ -77,7 +77,7 @@ def _cluster(monkeypatch: pytest.MonkeyPatch, scheduler: _Inline) -> None: """Point the run at a custom scheduler — the one monkeypatch point.""" @contextmanager - def fake() -> Iterator[_Inline]: + def fake(root: Path, attached: object) -> Iterator[_Inline]: yield scheduler monkeypatch.setattr(engine, "cluster_for_run", fake) @@ -967,7 +967,7 @@ def test_a_processes_cluster_fits_through_the_seam( results travel back whole.""" @contextmanager - def processes() -> Iterator[engine._Dask]: + def processes(root: Path, attached: object) -> Iterator[engine._Dask]: from distributed import Client, LocalCluster with LocalCluster( # type: ignore[no-untyped-call] diff --git a/tests/test_materialize_clusters.py b/tests/test_materialize_clusters.py new file mode 100644 index 0000000..9eb69b5 --- /dev/null +++ b/tests/test_materialize_clusters.py @@ -0,0 +1,287 @@ +"""Managed venues share the existing execution path without sharing run identity.""" + +from __future__ import annotations + +import json +import threading +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from conftest import _Inline +from test_materialize import _SPEC, _UNIVERSE, _clone + +from lightcone.engine import clusters, container, dataset, project +from lightcone.engine import materialize as engine +from lightcone.engine.project import ProjectError + + +@pytest.fixture +def root(analysis: Callable[..., Path]) -> Path: + return analysis(_SPEC, universes={"baseline": _UNIVERSE}) + + +def _record(directory: Path, backend: str = "slurm") -> clusters.Record: + return clusters.Record( + directory, + {"id": "20260925-120000-test", "backend": backend, "label": "Shared compute"}, + state="running", + ) + + +def test_selection_is_shared_by_execution_announcement_and_report( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + selected = _record(root) + selections = [] + connections = [] + announcements = [] + + def choose(path: Path) -> clusters.Record: + selections.append(path) + return selected + + @contextmanager + def connect(path: Path, attached: clusters.Record | None) -> Iterator[_Inline]: + connections.append((path, attached)) + yield _Inline() + + monkeypatch.setattr(clusters, "attached_cluster", choose) + monkeypatch.setattr(engine, "cluster_for_run", connect) + + report = engine.materialize(root, ["first"], on_venue=announcements.append) + + expected = {"kind": "cluster", "backend": "slurm", "id": selected.id, "label": selected.label} + assert selections == [root] + assert connections == [(root, selected)] + assert announcements == [expected] + assert json.loads(json.dumps(report.as_dict()))["venue"] == expected + assert report.made == ["baseline/first"] + + +def test_allocation_precedence_ignores_the_registry( + root: Path, inline: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("SLURM_JOB_ID", "31415926") + monkeypatch.setenv("SLURM_JOB_NUM_NODES", "4") + + def forbidden(path: Path) -> None: + pytest.fail("an active allocation must not inspect managed clusters") + + monkeypatch.setattr(clusters, "attached_cluster", forbidden) + + report = engine.materialize(root, ["first"]) + + assert report.venue == {"kind": "allocation", "nodes": 4} + + +@pytest.mark.parametrize("backend", ["slurm", "gateway"]) +def test_remote_cluster_allows_the_driver_on_a_login_node( + root: Path, inline: None, monkeypatch: pytest.MonkeyPatch, backend: str +) -> None: + monkeypatch.setenv("NERSC_HOST", "perlmutter") + monkeypatch.setattr(clusters, "attached_cluster", lambda path: _record(root, backend)) + + report = engine.materialize(root, ["first"]) + + assert report.made == ["baseline/first"] + assert report.venue is not None and report.venue["backend"] == backend + + +@pytest.mark.parametrize("backend", [None, "local"]) +def test_local_execution_still_refuses_a_login_node( + root: Path, monkeypatch: pytest.MonkeyPatch, backend: str | None +) -> None: + monkeypatch.setenv("NERSC_HOST", "perlmutter") + selected = None if backend is None else _record(root, backend) + monkeypatch.setattr(clusters, "attached_cluster", lambda path: selected) + + with pytest.raises(ProjectError, match="login node"): + engine.materialize(root, []) + + +def test_announcement_precedes_connection_and_image_preparation( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + events = [] + monkeypatch.setattr(clusters, "attached_cluster", lambda path: _record(root)) + + @contextmanager + def connect(path: Path, attached: clusters.Record | None) -> Iterator[_Inline]: + events.append("connect") + yield _Inline() + + def prepare(path: Path, *, build: bool) -> None: + events.append("prepare") + raise ProjectError("stop before image preparation") + + monkeypatch.setattr(engine, "cluster_for_run", connect) + monkeypatch.setattr(container, "runtime_for_run", prepare) + + with pytest.raises(ProjectError, match="stop before image preparation"): + engine.materialize(root, [], on_venue=lambda selected: events.append("announce")) + + assert events == ["announce", "connect", "prepare"] + + +def test_failed_attachment_never_starts_local_work( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(clusters, "attached_cluster", lambda path: _record(root)) + + def unreachable(record: clusters.Record, path: Path) -> None: + raise ProjectError("Shared compute could not be reached") + + def forbidden(path: Path, *, build: bool) -> None: + pytest.fail("failed attachment must refuse before image preparation or local work") + + monkeypatch.setattr(clusters, "client", unreachable) + monkeypatch.setattr(container, "runtime_for_run", forbidden) + + with pytest.raises(ProjectError, match="Shared compute could not be reached"): + engine.materialize(root, []) + + assert not (root / "results/baseline/first.txt").exists() + + +@pytest.mark.parametrize( + ("hosts", "nodes", "message"), + [ + (("node-a", "node-b"), 2, "node-local"), + (("node-a", "node-a"), 1, "reached image preparation"), + (("node-a", "node-a"), 2, "node-local"), + ], +) +def test_attached_container_guard_counts_hosts_before_building( + root: Path, monkeypatch: pytest.MonkeyPatch, hosts: tuple[str, str], nodes: int, message: str +) -> None: + selected = _record(root) + selected.data["slurm"] = {"nodes": nodes} + monkeypatch.setattr(clusters, "attached_cluster", lambda path: selected) + monkeypatch.setattr(project, "mode", lambda path: "containerized") + monkeypatch.setattr(container, "runtime_hint", lambda: "podman") + + @contextmanager + def connect(record: clusters.Record, path: Path) -> Iterator[SimpleNamespace]: + workers = {str(index): {"host": host} for index, host in enumerate(hosts)} + yield SimpleNamespace(scheduler_info=lambda: {"workers": workers}) + + def prepare(path: Path, *, build: bool) -> None: + raise ProjectError("reached image preparation") + + monkeypatch.setattr(clusters, "client", connect) + monkeypatch.setattr(container, "runtime_for_run", prepare) + + with pytest.raises(ProjectError, match=message): + engine.materialize(root, []) + + +def test_concurrent_clones_do_not_share_tasks_or_close_the_cluster( + root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Identical project and output names must still execute in both clones.""" + from distributed import Client, LocalCluster + + other = tmp_path / "other" + other.mkdir() + clone = _clone(root, other).rename(other / root.name) + selected = _record(tmp_path) + monkeypatch.setattr(clusters, "attached_cluster", lambda path: selected) + entered = threading.Barrier(2, timeout=30) + finished = threading.Barrier(2, timeout=30) + keys: list[str] = [] + submit = engine._Dask.submit + + def capture(self: engine._Dask, fn: Any, *args: Any, key: str) -> Any: + keys.append(key) + return submit(self, fn, *args, key=key) + + monkeypatch.setattr(engine._Dask, "submit", capture) + + with LocalCluster( + n_workers=2, threads_per_worker=1, processes=False, dashboard_address=None + ) as cluster: + + @contextmanager + def connect(record: clusters.Record, path: Path) -> Iterator[Client]: + with Client(cluster.scheduler_address, set_as_default=False) as client: + entered.wait() + try: + yield client + finally: + # Keep both clients alive until both runs finish, so a + # reused scheduler key cannot escape by being forgotten. + finished.wait() + + monkeypatch.setattr(clusters, "client", connect) + with Client(cluster.scheduler_address, set_as_default=False) as observer: + unrelated = observer.submit(str, "another project", key="unrelated") + with ThreadPoolExecutor(max_workers=2) as pool: + runs = [pool.submit(engine.materialize, path, []) for path in (root, clone)] + reports = [run.result(timeout=90) for run in runs] + + assert unrelated.result() == "another project" + after = observer.submit(str, "still running", key="after-runs") + assert after.result() == "still running" + assert len(observer.scheduler_info()["workers"]) == 2 + + assert len(keys) == len(set(keys)) == 4 + prefixes = {key.rsplit("/", 2)[0] for key in keys} + assert len(prefixes) == 2 and all(prefix.startswith("lc/analysis/") for prefix in prefixes) + for path, report in zip((root, clone), reports, strict=True): + assert report.made == ["baseline/first", "baseline/second"] + assert (path / "results/baseline/second.txt").read_text() == "alpha\n" + assert not dataset.status(path) + + +@pytest.mark.parametrize("attached", [False, True]) +def test_interruption_restores_only_after_owned_workers_stop( + root: Path, monkeypatch: pytest.MonkeyPatch, attached: bool +) -> None: + """A shared recipe can outlive disconnect; its output must not be restored underneath it.""" + selected = _record(root) if attached else None + monkeypatch.setattr(clusters, "attached_cluster", lambda path: selected) + events = [] + output = root / "results/baseline/first.txt" + restore = dataset.restore + + class Interrupted(_Inline): + def submit(self, fn: Callable[..., object], *args: object, key: str) -> object: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("recipe still writing\n") + return object() + + def completed(self, handles: list[object]) -> Iterator[object]: + raise RuntimeError("scheduler disconnected") + + @contextmanager + def connect(path: Path, record: clusters.Record | None) -> Iterator[Interrupted]: + try: + yield Interrupted() + finally: + events.append("disconnected") + output.write_text("last write before disconnect\n") + + def record_restore(path: Path, paths: Any) -> None: + events.append("restore") + restore(path, paths) + + monkeypatch.setattr(engine, "cluster_for_run", connect) + monkeypatch.setattr(dataset, "restore", record_restore) + + error = ProjectError if attached else RuntimeError + with pytest.raises(error, match="scheduler disconnected") as raised: + engine.materialize(root, []) + + if attached: + assert events == ["disconnected"] + assert output.read_text() == "last write before disconnect\n" + assert "Stop the cluster" in str(raised.value) + else: + assert events == ["disconnected", "restore", "restore"] + assert not output.exists() + assert not dataset.status(root) diff --git a/tests/test_venue.py b/tests/test_venue.py index 1884170..93124c2 100644 --- a/tests/test_venue.py +++ b/tests/test_venue.py @@ -111,6 +111,7 @@ def test_a_login_node_refuses_with_both_commands(monkeypatch: pytest.MonkeyPatch assert "salloc" in message assert "sbatch" in message assert "--wrap 'lc materialize'" in message + assert "Lightcone sidebar › Compute" in message def test_the_guard_fires_before_anything_else( @@ -321,7 +322,7 @@ def test_an_unresolvable_node_name_is_a_refusal_not_a_traceback( def test_a_multi_node_allocation_refuses_a_node_local_image_store( - root: Path, monkeypatch: pytest.MonkeyPatch + root: Path, monkeypatch: pytest.MonkeyPatch, inline: None ) -> None: """podman's and docker's stores are node-local; only podman-hpc's migrate makes an image visible to the allocation's other nodes. @@ -359,3 +360,4 @@ def test_the_rerun_entry_point_is_guarded_like_materialize( assert worker.main(["baseline/first"]) == 2 err = capsys.readouterr().err assert "login node" in err and "salloc" in err + assert "Compute" not in err # this standalone recipe does not attach diff --git a/tests/test_worker_environment.py b/tests/test_worker_environment.py new file mode 100644 index 0000000..2adacc0 --- /dev/null +++ b/tests/test_worker_environment.py @@ -0,0 +1,163 @@ +"""Per-run recipe settings on a worker shared by several projects.""" + +from __future__ import annotations + +import os +import pickle +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier +from typing import Literal + +import pytest + +from lightcone.engine import assets, container, plan, project, sandbox, worker + +_LOCAL_SETTINGS = ( + "SLURM_JOB_ID", + "SLURMD_NODENAME", + "HOSTNAME", + "TMPDIR", + "XDG_RUNTIME_DIR", + "JUPYTER_IMAGE_SPEC", + "JPY_PARENT_PID", + "DISPLAY", + "SSH_AUTH_SOCK", +) + + +def _context( + root: Path, mode: Literal["direct", "containerized"] = "direct" +) -> worker.RunContext: + return worker.RunContext( + env_version="test", + head=("commit", "origin"), + versions=assets.Versions(), + runtime=container.Runtime(root, mode, root / ".venv"), + uv_version="test", + ) + + +def _task(root: Path, name: str = "result") -> plan.Task: + return plan.Task("baseline", name, root / f"{name}.txt", name, {}, {}, {}, "test") + + +def test_context_snapshots_the_driver_without_host_or_uv_install_settings( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + for name in _LOCAL_SETTINGS: + monkeypatch.setenv(name, "driver") + monkeypatch.setenv("LC_RECIPE_SETTING", "captured") + monkeypatch.setenv("VIRTUAL_ENV", "/driver/venv") + monkeypatch.setenv("UV_PYTHON", "3.11") + monkeypatch.setenv("UV_CACHE_DIR", "/shared/uv") + monkeypatch.setenv("UV_INDEX_PRIVATE_PASSWORD", "driver-secret") + + context = _context(tmp_path) + monkeypatch.setenv("LC_RECIPE_SETTING", "later") + + assert context.environment["LC_RECIPE_SETTING"] == "captured" + assert context.environment["UV_CACHE_DIR"] == "/shared/uv" + assert context.environment["UV_INDEX_PRIVATE_PASSWORD"] == "driver-secret" + assert not set(_LOCAL_SETTINGS) & context.environment.keys() + assert "VIRTUAL_ENV" not in context.environment + assert "UV_PYTHON" not in context.environment + assert pickle.loads(pickle.dumps(context)).environment == context.environment + assert "driver-secret" not in repr(context) + + +@pytest.mark.parametrize("mode", ["direct", "containerized"]) +def test_recipe_uses_driver_settings_and_only_host_local_worker_settings( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: Literal["direct", "containerized"] +) -> None: + monkeypatch.setenv("LC_RECIPE_SETTING", "driver") + monkeypatch.delenv("LC_OLD_CREDENTIAL", raising=False) + monkeypatch.delenv("HTTPS_PROXY", raising=False) + for name in _LOCAL_SETTINGS: + monkeypatch.setenv(name, "driver") + context = _context(tmp_path, mode) + + monkeypatch.setenv("LC_RECIPE_SETTING", "worker") + monkeypatch.setenv("LC_OLD_CREDENTIAL", "stale-secret") + monkeypatch.setenv("HTTPS_PROXY", "https://stale.invalid") + for name in _LOCAL_SETTINGS: + monkeypatch.setenv(name, "worker") + monkeypatch.delenv("SSH_AUTH_SOCK") + before = dict(os.environ) + observed: dict[str, str] = {} + + def capture( + backend: sandbox.Backend, + policy: sandbox.Policy, + argv: Sequence[str], + *, + cwd: Path, + env: dict[str, str], + prefix: Sequence[str], + ) -> sandbox.Outcome: + observed.update(env) + return sandbox.Outcome(1, sandbox.Attestation("none", "open")) + + monkeypatch.setattr(worker, "_gate", lambda *_: "") + monkeypatch.setattr(container, "backend", lambda _: sandbox.Unavailable()) + monkeypatch.setattr(sandbox, "run", capture) + worker.execute(tmp_path, _task(tmp_path), {}, context) + + assert observed["LC_RECIPE_SETTING"] == "driver" + assert "LC_OLD_CREDENTIAL" not in observed + assert "HTTPS_PROXY" not in observed + assert "SSH_AUTH_SOCK" not in observed + assert all(observed[name] == "worker" for name in _LOCAL_SETTINGS if name != "SSH_AUTH_SOCK") + assert dict(os.environ) == before + + +def test_concurrent_runs_do_not_change_the_worker_or_each_others_environment( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LC_RECIPE_SETTING", "first") + first = _context(tmp_path) + monkeypatch.setenv("LC_RECIPE_SETTING", "second") + second = _context(tmp_path) + monkeypatch.setenv("LC_RECIPE_SETTING", "worker") + before = dict(os.environ) + together = Barrier(2) + observed: dict[str, str] = {} + + def capture( + backend: sandbox.Backend, + policy: sandbox.Policy, + argv: Sequence[str], + *, + cwd: Path, + env: dict[str, str], + prefix: Sequence[str], + ) -> sandbox.Outcome: + together.wait(timeout=5) + observed[argv[-1]] = env["LC_RECIPE_SETTING"] + assert dict(os.environ) == before + return sandbox.Outcome(1, sandbox.Attestation("none", "open")) + + monkeypatch.setattr(worker, "_gate", lambda *_: "") + monkeypatch.setattr(container, "backend", lambda _: sandbox.Unavailable()) + monkeypatch.setattr(sandbox, "run", capture) + with ThreadPoolExecutor(max_workers=2) as pool: + calls = [ + pool.submit(worker.execute, tmp_path, _task(tmp_path, "one"), {}, first), + pool.submit(worker.execute, tmp_path, _task(tmp_path, "two"), {}, second), + ] + for call in calls: + call.result(timeout=10) + + assert observed == {"one": "first", "two": "second"} + assert dict(os.environ) == before + + +def test_explicit_child_environment_does_not_add_worker_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LC_STALE_VALUE", "worker") + assert project.child_env({}) == {} + assert project.child_env({"VIRTUAL_ENV": "/other", "UV_PYTHON": "3.11", "KEY": "value"}) == { + "KEY": "value" + }