From 79c822520b203201b3ed29d8559af74cb9179baf Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 10:04:27 +0200 Subject: [PATCH 1/9] Port Dask Gateway groundwork from gateway-cluster-support + LCR-176 Part A Squashed port of the unmerged gateway-cluster-support branch (attach-only Gateway cluster branch, kubernetes container runtime, LCR-175 worker stderr forwarding, registry image resolution, design doc) plus alex/lcr-176-worker-image-deps (pinned dask stack, worker-capable Containerfile scaffold). Superseded PR #159 carried this work; this branch re-lands it as the base for the PRD run/build lifecycle. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/api/dask_cluster.md | 55 +- docs/contributing/backends.md | 17 +- docs/design/jupyterhub-dask-gateway-gcp.md | 635 ++++++++++++++++++ docs/user/cluster.md | 86 ++- pyproject.toml | 20 +- src/lightcone/cli/commands.py | 395 +++++++++-- src/lightcone/engine/container.py | 274 +++++++- src/lightcone/engine/dask_cluster.py | 290 +++++++- src/lightcone/engine/manifest.py | 6 + src/lightcone/engine/site_registry.py | 43 +- src/lightcone/engine/snakefile.py | 11 +- .../executor.py | 163 ++++- tests/conftest.py | 19 + tests/test_cli.py | 96 +++ tests/test_container.py | 305 +++++++++ tests/test_dask_cluster.py | 479 ++++++++++++- tests/test_dask_plugin.py | 43 +- tests/test_site_registry.py | 31 + 18 files changed, 2832 insertions(+), 136 deletions(-) create mode 100644 docs/design/jupyterhub-dask-gateway-gcp.md diff --git a/docs/api/dask_cluster.md b/docs/api/dask_cluster.md index 54e254a4..0784b5ea 100644 --- a/docs/api/dask_cluster.md +++ b/docs/api/dask_cluster.md @@ -1,24 +1,49 @@ # lightcone.engine.dask_cluster Cluster lifecycle for `lc run`. One context manager (`cluster_for_run`), -three branches, no service to manage. +four branches, no service to manage. Source: `src/lightcone/engine/dask_cluster.py`. -## `cluster_for_run(*, verbose=False) → Iterator[str]` - -Yields a Dask scheduler address valid for the duration of `lc run`. -Three branches in priority order: - -1. **`DASK_SCHEDULER_ADDRESS` already set** → yield as-is. We don't - own the cluster, so we don't tear it down. -2. **`SLURM_JOB_ID` set** → start an in-process scheduler bound to the +## `cluster_for_run(*, verbose=False, local_directory=None, expected_worker_image=None) → Iterator[dict[str, str]]` + +Yields the **env overlay** the child snakemake process needs to reach +the cluster — the parent and the executor plugin are separate +processes, so connection info travels via environment variables. +Four branches in priority order: + +1. **`DASK_SCHEDULER_ADDRESS` already set** → yield + `{"DASK_SCHEDULER_ADDRESS": addr}` as-is. We don't own the cluster, + so we don't tear it down. +2. **Dask Gateway detected** (`LIGHTCONE_GATEWAY_CLUSTER` or + `DASK_GATEWAY__ADDRESS` set, e.g. a JupyterHub pod) → **attach** to + the user's running Gateway cluster; yield + `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Attach-only by design: the + user creates clusters from JupyterLab (that's where the + image/cores/memory options widget and the dashboard live); the + Gateway API is user-scoped, so exactly one running cluster attaches + unambiguously, and zero or several raises with the fix spelled out + (`LIGHTCONE_GATEWAY_CLUSTER` disambiguates). Gateway scheduler + addresses use a `gateway://` comm scheme a bare `Client` cannot + dial, so the child rejoins **by name** through the authenticated + Gateway API. The cluster, and its scaling, are left untouched on + exit — same convention as branch 1. Startup fails fast if live + workers don't advertise the resource contract below (zero workers + is fine — adaptive clusters scale on demand), and warns when the + cluster's actual worker image (read from the scheduler pod's + `LIGHTCONE_WORKER_IMAGE`) differs from *expected_worker_image*. + Requires the optional dependency: + `pip install lightcone-cli[gateway]`. +3. **`SLURM_JOB_ID` set** → start an in-process scheduler bound to the driver's SLURM hostname (`SLURMD_NODENAME` or `gethostname()`), then `srun` one `dask worker` per node across the allocation. -3. **Neither** → `LocalCluster()` sized to the local machine. +4. **None of the above** → `LocalCluster()` sized to the local machine. -The scheduler is always in-process, so its lifetime equals the run's -lifetime: no orphaned schedulers if the driver crashes. +Outside the Gateway branch the scheduler is always in-process, so its +lifetime equals the run's lifetime: no orphaned schedulers if the +driver crashes. On the Gateway branch there is nothing to orphan +either — lc never creates the cluster, and idle clusters are reaped +by the deployment's `idle_timeout`. ## Resource keys @@ -73,6 +98,8 @@ and keeps everything in one process tree. ## Tests -`tests/test_dask_cluster.py` covers the three branches and the +`tests/test_dask_cluster.py` covers all four branches and the resource-advertising contract. The SLURM branch is tested with mocked -`subprocess.Popen` plus a stubbed `Client.wait_for_workers`. +`subprocess.Popen` plus a stubbed `Client.wait_for_workers`; the +Gateway branch (discovery, attach-only lifecycle, contract and image +verification) against a fake `dask_gateway` module. diff --git a/docs/contributing/backends.md b/docs/contributing/backends.md index 5c651e2d..3a8c3d6f 100644 --- a/docs/contributing/backends.md +++ b/docs/contributing/backends.md @@ -6,9 +6,15 @@ of these: ## Adding a container runtime -The supported runtimes are `docker`, `podman`, and `podman-hpc` (plus the -`none` no-op). They are listed in -`src/lightcone/engine/container.py::RUNTIMES`. To add a new one: +The supported OCI runtimes are `docker`, `podman`, and `podman-hpc`, +listed in `src/lightcone/engine/container.py::RUNTIMES`. Two non-OCI +values sit outside that tuple: `none` (no isolation, recipes run on the +host) and `kubernetes` (the Dask Gateway worker pod *is* the container; +recipes run unwrapped inside it and images resolve against the +deployment registry — see `REGISTRY_ENV` / `resolve_worker_image`). +`kubernetes` is never auto-detected: a site declares it +(`site_registry.py`, `container_runtime: "kubernetes"`) or the user +pins it in config. To add a new OCI runtime: 1. Append the binary name to `RUNTIMES` (detection priority is the tuple order). @@ -23,8 +29,9 @@ The supported runtimes are `docker`, `podman`, and `podman-hpc` (plus the ## Adding a Dask cluster shape -Today the cluster manager has three branches: existing scheduler, SLURM -allocation, local. To add a fourth (for example, a custom GPU farm): +Today the cluster manager has four branches: existing scheduler, Dask +Gateway (JupyterHub), SLURM allocation, local. To add another (for +example, a custom GPU farm): 1. Add a branch to `cluster_for_run()` in `src/lightcone/engine/dask_cluster.py`. diff --git a/docs/design/jupyterhub-dask-gateway-gcp.md b/docs/design/jupyterhub-dask-gateway-gcp.md new file mode 100644 index 00000000..3d299b06 --- /dev/null +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -0,0 +1,635 @@ +# JupyterHub + Dask Gateway + kbatch on GCP — deployment procedure and lightcone-cli integration + +*Status: design proposal — researched 2026-07-12 against `main` (post-#157).* +*Revised 2026-07-12 — see §7 for the implemented design, which supersedes +§4.2's owned-cluster branch and reframes the container story of §4.5.* + +## 1. Why this works almost out of the box + +The execution layer on `main` is already shaped for this. `lc run` always invokes Snakemake with +`--executor dask` (our first-party plugin in `src/snakemake_executor_plugin_dask/`), and cluster +lifecycle lives in `src/lightcone/engine/dask_cluster.py:cluster_for_run()`, which has three +branches: + +1. **`DASK_SCHEDULER_ADDRESS` set** → attach to an external scheduler verbatim, no cluster + creation or teardown (`dask_cluster.py:101`). +2. **`SLURM_JOB_ID` set** → in-process scheduler + `srun`-launched workers. +3. **Neither** → sized `LocalCluster`. + +Branch 1 is the Dask Gateway seam: a Gateway cluster hands us a scheduler address; exporting it +before `lc run` makes the whole pipeline execute on Kubernetes **with zero code changes**. The +executor plugin (`executor.py:104`) is a pure `Client(addr)` consumer of the same variable. + +Three real constraints follow from how the executor works: + +- **Shared filesystem.** The plugin declares `implies_no_shared_fs=False`; each task is a child + `snakemake` invocation running the rule body on a worker (`executor.py:113-128`, + `runner.py:52-107`). The project directory, `results/`, `.snakemake/` state, and the + `LIGHTCONE_OUT_LOCK` flock file must be visible at the **same path** on the driver and every + worker. On GKE that means an RWX PersistentVolume (Filestore) — *not* GCS FUSE, which does not + support POSIX `flock`. +- **Worker resource contract.** Workers must advertise Dask abstract resources + `cpus` / `memory` / `gpus` (`dask_cluster.py:33-35`) or per-rule `resources:` requests will + never schedule (`executor.py:78-90`). Gateway workers don't set these by default; we configure + them via `dask` config env vars in the worker pod spec (§4.3). +- **Container gap.** `wrap_recipe()` (`container.py:706-751`) shells out to + `docker` / `podman` / `podman-hpc`, none of which exist inside a pod. Kubernetes *is* the + container runtime: the worker pod image must *be* the project environment, and the project runs + with `runtime: none`. (Longer term: image-per-rule via Gateway cluster options, §6.4.) + +## 2. Current state of the upstream stack (verified July 2026) + +| Component | Status | Version to use | +|---|---|---| +| Zero to JupyterHub (z2jh) | Actively maintained | Helm chart **4.4.0** (Jun 2026), ships JupyterHub 5.5.0; requires k8s ≥ 1.28, Helm ≥ 3.5 | +| Dask Gateway | Actively maintained | Helm chart **2026.3.x** from `https://helm.dask.org`; requires k8s ≥ 1.30 | +| DaskHub combined chart | **Deprecated** — do not use for new deployments | Install z2jh + dask-gateway as separate charts; DaskHub's `values.yaml` remains the reference for the wiring between them | +| kbatch | **Lightly maintained.** Last stable release 0.4.2 (Sep 2023); 0.5.0 alphas/betas through late 2024; quiescent since | `kbatch-proxy` chart from `https://kbatch-dev.github.io/helm-chart` (pin `0.5.0-alpha.1`) | + +kbatch's quiescence is a real consideration. It is deliberately tiny (a thin authz proxy in front +of the k8s Jobs API, JupyterHub-authenticated), which limits bit-rot risk, and it's exactly the +right shape for our use case (submit a long-running `lc run` driver as a k8s Job that survives +notebook disconnects). But treat it as a component we may eventually vendor or replace +(alternatives: `jupyter-scheduler`, a 20-line FastAPI hub service of our own, or Argo Workflows +if we ever want DAG-level k8s orchestration — which we don't, Snakemake owns the DAG). + +## 3. GCP deployment procedure + +### 3.0 Prerequisites + +```bash +gcloud services enable container.googleapis.com file.googleapis.com \ + artifactregistry.googleapis.com +gcloud config set project +# local tools: kubectl, helm >= 3.5 +``` + +### 3.1 GKE cluster + +Use a **Standard** (not Autopilot) cluster: we need node pools with taints, scale-to-zero, and +Spot VMs for workers; Autopilot's per-pod model fights the Dask worker pattern and the Filestore +mount topology. + +```bash +ZONE=europe-west1-b # or --region for HA control plane (99.95% SLA) +CLUSTER=lightcone-hub + +# Core pool: hub, proxy, gateway api/traefik/controller, kbatch-proxy +gcloud container clusters create $CLUSTER \ + --zone $ZONE --cluster-version latest \ + --machine-type e2-standard-4 --num-nodes 1 \ + --enable-network-policy \ + --addons GcpFilestoreCsiDriver \ + --workload-pool=.svc.id.goog + +# User pool: JupyterLab singleuser pods (the lc/Claude driver environment) +gcloud container node-pools create user-pool \ + --cluster $CLUSTER --zone $ZONE \ + --machine-type e2-standard-8 \ + --num-nodes 0 --enable-autoscaling --min-nodes 0 --max-nodes 10 \ + --node-labels hub.jupyter.org/node-purpose=user \ + --node-taints hub.jupyter.org_dedicated=user:NoSchedule + +# Compute pool: Dask Gateway workers + kbatch jobs — Spot for ~60-90% discount +gcloud container node-pools create dask-pool \ + --cluster $CLUSTER --zone $ZONE \ + --machine-type c2-standard-16 --spot \ + --num-nodes 0 --enable-autoscaling --min-nodes 0 --max-nodes 50 \ + --node-labels lightcone.dev/node-purpose=compute \ + --node-taints lightcone.dev_dedicated=compute:NoSchedule + +kubectl create clusterrolebinding cluster-admin-binding \ + --clusterrole=cluster-admin --user= +``` + +Add a GPU pool later with `--accelerator type=nvidia-l4,count=1 --spot` when a project needs +`gpus_per_task`. + +### 3.2 Shared filesystem (Filestore RWX) + +This is the load-bearing piece for lightcone: driver and workers must share +`/home/jovyan/` (or a dedicated `/shared`) with POSIX `flock` support. + +```yaml +# filestore-pvc.yaml +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: lightcone-shared + namespace: hub +spec: + accessModes: [ReadWriteMany] + storageClassName: standard-rwx # GKE Filestore CSI, Basic HDD + resources: + requests: + storage: 1Ti # Filestore Basic minimum; ~$160-200/mo +``` + +Cost note: Filestore Basic has a 1 TiB floor. For an early/staging deployment, a single-pod NFS +server (`nfs-ganesha` or `nfs-server-provisioner` chart) backed by a cheap PD gets you RWX + +flock for ~$40/mo; swap to Filestore when it matters. Keep per-user home directories on z2jh's +default dynamic PD volumes; the RWX volume is for **project workspaces and results**. + +### 3.3 Images (Artifact Registry) + +```bash +gcloud artifacts repositories create lightcone --repository-format=docker --location=europe-west1 +``` + +Two images: + +1. **`lightcone-notebook`** — base `quay.io/jupyter/scipy-notebook` (or pangeo-notebook) plus + `lightcone-cli`, `astra-tools`, `dask-gateway` (client), `kbatch`, git, and Claude Code. This + is both the JupyterLab singleuser image and the kbatch job image. +2. **`lightcone-worker-`** — per-project compute environment: `lightcone-cli`, + `snakemake`, `dask distributed` (versions matching the notebook image — Dask is picky about + client/scheduler/worker version skew), plus the project's science deps. Initially this can be + built *from* the project's existing Containerfile with a `pip install lightcone-cli snakemake` + layer appended — `lc build` already computes content-addressed tags + (`compute_image_tag()`, `container.py:319`), so pushing the same tag to Artifact Registry is a + natural extension (§6.4). + +### 3.4 JupyterHub (z2jh chart 4.4.0) + +```bash +helm repo add jupyterhub https://hub.jupyter.org/helm-chart/ +helm repo update +``` + +```yaml +# jupyterhub-values.yaml +hub: + config: + # start with GitHub/Google OAuth; dummy auth only for a private staging cluster + GitHubOAuthenticator: + client_id: <...> + client_secret: <...> + oauth_callback_url: https://hub./hub/oauth_callback + allowed_organizations: [LightconeResearch] + JupyterHub: + authenticator_class: github + services: + dask-gateway: + display: false + apiToken: "" # openssl rand -hex 32; must match gateway values + kbatch: + display: false + apiToken: "" + url: http://kbatch-proxy.kbatch.svc.cluster.local + +proxy: + https: + enabled: true + hosts: [hub.] + letsencrypt: + contactEmail: fr.eiffel@gmail.com + +singleuser: + image: + name: europe-west1-docker.pkg.dev//lightcone/lightcone-notebook + tag: "" + cpu: {limit: 4, guarantee: 1} + memory: {limit: 16G, guarantee: 4G} + defaultUrl: /lab + storage: + capacity: 20Gi # per-user home (PD) + extraVolumes: + - name: lightcone-shared + persistentVolumeClaim: + claimName: lightcone-shared + extraVolumeMounts: + - name: lightcone-shared + mountPath: /shared # same path in worker pods — see gateway values + extraEnv: + LIGHTCONE_SCRATCH: /shared/scratch + # Dask Gateway client defaults, so Gateway() works with no args: + DASK_GATEWAY__ADDRESS: https://hub./services/dask-gateway/ + DASK_GATEWAY__PROXY_ADDRESS: gateway://traefik-dask-gateway-dask-gateway.dask-gateway:80 + DASK_GATEWAY__PUBLIC_ADDRESS: /services/dask-gateway/ + DASK_GATEWAY__AUTH__TYPE: jupyterhub +``` + +```bash +helm upgrade --cleanup-on-fail --install jupyterhub jupyterhub/jupyterhub \ + --namespace hub --create-namespace --version 4.4.0 \ + --values jupyterhub-values.yaml +``` + +Point DNS at the `proxy-public` LoadBalancer IP, wait for Let's Encrypt. + +### 3.5 Dask Gateway (chart 2026.3.x) + +The hub proxy routes `/services/dask-gateway/` to the gateway's traefik; the gateway +authenticates users against the hub with `TOKEN_A`. + +```yaml +# dask-gateway-values.yaml +gateway: + prefix: /services/dask-gateway # served behind the hub proxy + auth: + type: jupyterhub + jupyterhub: + apiToken: "" # same as hub.services.dask-gateway.apiToken + apiUrl: http://hub.hub.svc.cluster.local:8081/hub/api + backend: + image: + name: europe-west1-docker.pkg.dev//lightcone/lightcone-worker-default + tag: "" + scheduler: + cores: {request: 1, limit: 2} + memory: {request: 2G, limit: 4G} + worker: + extraPodConfig: + tolerations: + - key: lightcone.dev_dedicated + operator: Equal + value: compute + effect: NoSchedule + volumes: + - name: lightcone-shared + persistentVolumeClaim: + claimName: lightcone-shared # needs a PVC in this ns, or cross-ns PV binding + extraContainerConfig: + volumeMounts: + - name: lightcone-shared + mountPath: /shared # MUST equal the singleuser mountPath + # user-selectable knobs + the lightcone resource contract (see §4.3) + extraConfig: + clusteroptions: | + from dask_gateway_server.options import Options, Integer, Float, String + + def options_handler(options): + return { + "worker_cores": options.worker_cores, + "worker_memory": "%fG" % options.worker_memory, + "image": options.image, + "environment": { + # advertise lightcone's Dask resource contract on every worker + "DASK_DISTRIBUTED__WORKER__RESOURCES__CPUS": str(options.worker_cores), + "DASK_DISTRIBUTED__WORKER__RESOURCES__MEMORY": str(int(options.worker_memory * 1e9)), + "DASK_DISTRIBUTED__WORKER__RESOURCES__GPUS": "0", + }, + } + + c.Backend.cluster_options = Options( + Integer("worker_cores", default=4, min=1, max=16, label="Worker cores"), + Float("worker_memory", default=16, min=4, max=64, label="Worker memory (GB)"), + String("image", + default="europe-west1-docker.pkg.dev//lightcone/lightcone-worker-default:", + label="Worker image"), + handler=options_handler, + ) + +traefik: + service: + type: ClusterIP # not exposed; everything flows through the hub proxy +``` + +```bash +helm upgrade --install dask-gateway dask-gateway \ + --repo=https://helm.dask.org \ + --namespace dask-gateway --create-namespace \ + --values dask-gateway-values.yaml +``` + +Namespace note: the simplest topology is to deploy dask-gateway **in the `hub` namespace** so +the shared PVC binds directly and the daskhub-style service discovery just works; the split shown +above requires either a second PVC bound to the same Filestore PV or a `ReferenceGrant`-style +cross-namespace PV. Same-namespace is the recommended starting point. + +### 3.6 kbatch + +```bash +helm repo add kbatch https://kbatch-dev.github.io/helm-chart +kubectl create namespace kbatch +``` + +```yaml +# kbatch-values.yaml +app: + jupyterhub_api_token: "" # same as hub.services.kbatch.apiToken + jupyterhub_api_url: http://hub.hub.svc.cluster.local:8081/hub/api + extra_env: + KBATCH_PREFIX: /services/kbatch + # job template: run on the compute pool, mount the shared volume + job_template: + spec: + template: + spec: + tolerations: + - key: lightcone.dev_dedicated + operator: Equal + value: compute + effect: NoSchedule + containers: + - volumeMounts: + - name: lightcone-shared + mountPath: /shared + volumes: + - name: lightcone-shared + persistentVolumeClaim: + claimName: lightcone-shared +``` + +```bash +helm install kbatch kbatch/kbatch-proxy --version 0.5.0-alpha.1 \ + --namespace kbatch --values kbatch-values.yaml +``` + +The `hub.services.kbatch` entry from §3.4 makes the hub proxy route +`https://hub./services/kbatch/` to it. (Verify the exact values-key spelling against the +chart's `values.yaml` at install time — the alpha chart has shifted key names between releases.) + +### 3.7 Smoke test + +From a JupyterLab terminal on the hub: + +```bash +python -c " +from dask_gateway import Gateway +g = Gateway() # picks up DASK_GATEWAY__* env +c = g.new_cluster(worker_cores=2, worker_memory=4) +c.scale(2) +client = c.get_client() +print(client.scheduler_info()) +print(client.run(lambda: __import__('os').path.exists('/shared'))) +" +kbatch configure --kbatch-url https://hub./services/kbatch --token $JUPYTERHUB_API_TOKEN +kbatch job submit --name smoke --image busybox --command 'ls /shared' +``` + +## 4. Integration with lightcone-cli + +### 4.1 Phase 0 — zero code changes (works today) + +From a notebook or terminal in the singleuser pod, with the project checked out under `/shared`: + +```python +from dask_gateway import Gateway +gw = Gateway() +cluster = gw.new_cluster(worker_cores=8, worker_memory=32) +cluster.adapt(minimum=0, maximum=20) +print(cluster.scheduler_address) # tls://... +``` + +```bash +export DASK_SCHEDULER_ADDRESS='tls://...' +cd /shared/my-analysis && lc run +``` + +`cluster_for_run` branch 1 attaches and never tears down; `resources:` per rule schedule +correctly because §3.5 injected the `cpus/memory/gpus` worker resources. Requirements: project +lives on `/shared`; project `lightcone.yaml` (or `LIGHTCONE_SCRATCH`) points scratch at +`/shared/scratch`; container `runtime: none` (worker image *is* the environment). + +One wrinkle: Gateway schedulers speak TLS with per-cluster credentials, so a bare +`Client(addr)` in the executor may need the Gateway security context. If +`DASK_SCHEDULER_ADDRESS` alone proves insufficient, that's the first concrete argument for +Phase 1's native branch (which gets `cluster.get_client()` for free). Alternatively +`gw.new_cluster(...)` → `cluster.security` can be exported as cert/key paths — but at that point +the native branch is less code. + +### 4.2 Phase 1 — a `gateway` branch in `cluster_for_run` + a site entry + +Add a fourth branch to `dask_cluster.cluster_for_run()`, gated on config/env (not on +`site["backend"]`, which is inert today): + +``` +LIGHTCONE_DASK_GATEWAY=https://hub./services/dask-gateway/ # or config.yaml key +``` + +```python +def _gateway_cluster(...): + from dask_gateway import Gateway + gw = Gateway() # env-configured, jupyterhub auth + cluster = gw.new_cluster(**_options_from_config()) + cluster.adapt(minimum=0, maximum=cfg.max_workers) + client = cluster.get_client() # handles TLS security context + try: + yield cluster.scheduler_address + finally: + cluster.shutdown() +``` + +Plus a `SITE_DEFAULTS["lightcone-hub"]` entry (`site_registry.py:27`): + +```python +"lightcone-hub": { + "hostname_patterns": [], # detected via JUPYTERHUB_SERVICE_PREFIX env instead + "display_name": "Lightcone JupyterHub (GCP)", + "backend": "dask-gateway", + "container_runtime": "none", + "scratch_root": "/shared/scratch", +} +``` + +Detection: `JUPYTERHUB_USER` + `DASK_GATEWAY__ADDRESS` in the environment is a reliable "we're +on a hub" signal — cleaner than hostname patterns for pods. This is also where the +`design_review.md` aspiration ("laptop, SLURM allocation, Kubernetes pod — only configuration is +a `--target` choice") gets its third leg: `lc run --target hub|slurm|local`, defaulting to +auto-detect. + +The executor's `code_version`/manifest layer needs nothing: manifests are written host-side by +the worker process onto the shared FS, `lc status`/`lc verify` read manifests only, and the +flock serialization (`LIGHTCONE_OUT_LOCK`) works on Filestore. + +### 4.3 Worker resource contract (the one silent failure mode) + +If workers don't advertise `cpus`/`memory`/`gpus`, every submitted task hangs unscheduled — +no error. Belt and suspenders: + +- Gateway side: the `cluster_options` handler in §3.5 sets + `DASK_DISTRIBUTED__WORKER__RESOURCES__{CPUS,MEMORY,GPUS}` from the chosen worker shape. +- lightcone side (Phase 1): after `wait_for_workers`, assert the resources are present on at + least one worker and fail fast with a pointed message if not. Cheap, saves an afternoon. + +### 4.4 Phase 2 — kbatch as the headless driver channel + +`lc run` is a long-running driver (Snakemake + in-process Dask client). Running it in a notebook +terminal dies with the singleuser pod (culling, browser close). kbatch turns it into a k8s Job: + +```bash +kbatch job submit --name my-analysis \ + --image europe-west1-docker.pkg.dev//lightcone/lightcone-notebook: \ + --command "bash -c 'cd /shared/my-analysis && lc run'" +``` + +Natural CLI sugar: **`lc submit [outputs...]`** — generates the kbatch job spec (image = notebook +image, workdir = project dir on `/shared`, env = gateway address), submits via kbatch's Python +API, and prints `kbatch job logs` / `lc status` follow-up commands. `kbatch cronjob submit` +gives scheduled re-runs (nightly `lc run && lc verify`) for free. This is also the substrate for +agentic loops: a `ralph`-style long-running Claude Code session can itself be a kbatch job that +shells out to `lc run`, with the human reviewing results in JupyterLab against the same shared +volume. + +Given kbatch's maintenance state, keep the coupling thin: one module +(`src/lightcone/engine/kbatch.py` or similar) that builds the job spec, so swapping the +submission backend later is a one-file change. + +### 4.5 Phase 3 — per-rule images (closing the container gap properly) + +Today all rules in a Gateway run share one worker image, and recipes run unwrapped +(`runtime: none`). ASTRA's `astra.yaml` already carries per-output container specs, and +`compute_image_tag()` gives content-addressed tags. The clean endgame: + +1. `lc build --push` builds each rule's image and pushes `lc--` to Artifact Registry. +2. The gateway branch groups rules by image and requests one Gateway cluster per image (Gateway + cluster options already accept `image`), or — simpler and probably sufficient — `lc run` + errors early when `astra.yaml` declares more than one distinct container and suggests a + combined image. +3. Manifests keep recording the image tag either way, so provenance is intact. + +Defer this until a real project needs heterogeneous environments; most current projects use a +single environment. + +## 5. Cost & ops summary + +- Idle floor: 1× `e2-standard-4` core node (~$100/mo) + Filestore 1 TiB (~$160-200/mo, or ~$40/mo + with the NFS-pod stopgap) + LoadBalancer (~$20/mo). User and compute pools scale to zero. +- Spot VMs on the compute pool are safe: Snakemake retries failed jobs, manifests make partial + progress durable, and `lc verify` catches anything torn. +- Upgrades: z2jh and dask-gateway both publish frequent chart releases; pin versions in a + `helmfile`/terraform and bump deliberately. kbatch: pin and don't expect upstream movement. +- GPU: add an L4/A100 Spot pool + a `worker_gpus` cluster option mapping to + `nvidia.com/gpu` limits and `DASK_DISTRIBUTED__WORKER__RESOURCES__GPUS`. + +## 6. Open questions + +1. **ANSWERED 2026-07-12 (staging cluster, lightcone-hub repo): Phase 1 is required.** + Gateway's `cluster.scheduler_address` is not `tls://` but a custom scheme — + `gateway://traefik-dask-gateway.hub:80/` — that only the `dask_gateway` + client library's comm layer can dial (traefik multiplexing + per-cluster TLS + SNI). A bare + `distributed.Client(addr)`, which is what the executor creates, fails with + `unknown address scheme 'gateway'` regardless of TLS env configuration. Phase 0 + (pure `DASK_SCHEDULER_ADDRESS` env) cannot work against Gateway; the `gateway` branch in + `cluster_for_run` must create the cluster and hold the `GatewayCluster` object. The executor's + `Client(addr)` will additionally need the Gateway security context — simplest is for the + branch to pass the client through (or export the `gateway://` handling) rather than an address + string alone. Everything else validated on staging: workers advertise the `cpus/memory/gpus` + resource contract via the cluster-options env injection, `/shared` (NFS RWX) is visible on + workers, and `flock` works on it. +2. **PVC namespace topology** — same-namespace deploy (hub + gateway + kbatch jobs in one ns) vs. + the three-namespace layout with shared-PV plumbing. Recommend starting same-namespace. +3. **Home vs. shared workspace layout** — one `/shared//` convention vs. per-project + subpaths; interacts with `lc init` scaffolding and the session-start hooks. +4. **kbatch alternative** — if the 0.5 alphas misbehave, a minimal in-house hub service + (JupyterHub-authenticated POST → k8s Job) is ~150 lines and removes the dependency. +5. **On-hub container build path — UNRESOLVED, needs a decision (surfaced 2026-07-15, LCR-176).** + §7.2 currently defers all image builds off-hub: `lc build` in the JupyterLab pod prints the + `docker build && docker push` commands to run elsewhere, because there is no docker in-pod by + design. End-to-end testing on the staging cluster confirmed this *works* but is a real friction + point — every dependency change forces the user to leave the hub, build on a docker-equipped + machine with registry access, push, and only then start a Gateway cluster on the new image. The + open decision is which build story we commit to: + - **(a) Keep builds off-hub, just document/smooth them** (status quo, §7.2). Cheapest; the + friction stays. + - **(b) In-cluster rootless build** (kaniko/buildkit) that builds the project image and pushes + to Artifact Registry from within the pod. §7.2 deferred this "until someone actually cannot + reach a docker machine" — LCR-176 is arguably that moment. + - **(c) Hub-side build service** — a JupyterHub-authenticated POST → a Kubernetes build Job, + the same shape as the kbatch-alternative in item 4 and the `lc submit` service. + + The blocking sub-question is the one already raised as PRD open-question #3: **can an arbitrary + hub user be allowed to push to the deployment registry?** Options (b) and (c) both need an answer + on registry credentials / least privilege before they can be built. Two constraints any choice + must respect: the worker pod image *is* the recipe environment (so the build target is a real + Gateway-worker image, `FROM lightcone-worker-default`, not a slim base), and — from LCR-175 — + the driver (notebook) and worker images must stay on the **same lightcone-cli version**, because + `snakemake_executor_plugin_dask` runs split across both and version skew breaks execution; the + build path must keep them in lockstep the way we pin dask/distributed. *(The related but separate + defect — `lc init` scaffolding a `FROM python:3.12-slim` Containerfile that can never run as a + Gateway worker — is being handled as a code fix under LCR-176 and is not part of this open + question.)* + +## 7. Implemented design (2026-07-12) — attach-only + kubernetes as a first-class runtime + +What actually shipped diverges from §4.2/§4.5 in three deliberate ways, all +in the direction of less machinery: + +### 7.1 `kubernetes` is a container runtime, not a workaround + +§4.2's `container_runtime: "none"` was a lie with side effects: `load_runtime()` +never consulted the site entry, auto-detection found no OCI binary, and every +`lc run` on the hub fired the "no container runtime found / provenance will +not match" warning — for recipes that *are* containerized, by the worker pod. + +The implemented model: `container_runtime: "kubernetes"` (site-declared, +treated as **explicit** by `load_runtime()` — a deployment fact, not a +detection guess). `wrap_recipe()` is a no-op for it, same as `none`, but the +semantics differ: the pod is the container. Site-declared *OCI* runtimes +(Perlmutter → podman-hpc) keep their detection-order-hint semantics. + +### 7.2 Registry wiring by convention: `LIGHTCONE_REGISTRY` + +The deployment names its registry via singleuser env (exactly like +`LIGHTCONE_SCRATCH` and `DASK_GATEWAY__*`). A `container: Containerfile` spec +resolves to `$LIGHTCONE_REGISTRY/lc-:` — the same content hash +as the local `lc--` tag, so the environment is provably the +same artifact on every path. `code_version` keeps hashing the *local* tag on +all paths (`resolve_image_for_run` stays site-agnostic) so moving a project +laptop↔hub is not code drift; the registry ref (`resolve_worker_image`) is +used only to name the worker image and verify the attached cluster. + +`lc build` on the hub builds nothing (no docker in-pod, by design): it probes +the registry (metadata-server token, best-effort, never gating) and prints the +exact `docker build -t … && docker push ` commands to run from any +docker-equipped machine. In-cluster kaniko builds remain deliberately +deferred until someone actually cannot reach one. + +The worker image must carry dask/distributed/lightcone-cli at hub-matching +pins; the recommended convention is `FROM /lightcone-worker-default:` +in the project Containerfile, which then serves every path (locally it wraps; +on the hub it runs as the pod). + +### 7.3 Attach-only Gateway branch + +`lc run` never creates Gateway clusters (§4.2's `new_cluster` + adapt + +shutdown path is gone). The user creates one from JupyterLab — sidebar or +notebook, where the options widget (image/cores/memory) and dashboard live — +and `lc run` discovers it: `list_clusters()` is user-scoped under JupyterHub +auth, so exactly one running cluster attaches with zero configuration; zero +raises with a copy-pasteable `new_cluster(image=…)` snippet (image pre-filled +from the project's resolved ref); several raises listing names, with +`LIGHTCONE_GATEWAY_CLUSTER` as the disambiguator. Scaling is never touched; +the cluster is left running on exit — the Gateway flavor of the +`DASK_SCHEDULER_ADDRESS` convention, and the one intentional asymmetry with +the owned local/SLURM clusters. + +Two contract checks at attach time, both deployment-backed: + +- **Resources**: live workers must advertise `cpus`/`memory`/`gpus` + (injected by the cluster-options handler); zero live workers is fine + (adaptive clusters scale on demand). +- **Image**: the options handler also injects `LIGHTCONE_WORKER_IMAGE` into + scheduler+worker pods; lc reads it back via `run_on_scheduler` (a lambda, + cloudpickled by value, so the scheduler needn't import lightcone) and + warns on mismatch with the project's resolved image. Manifests record the + worker pod's actual image as `worker_image` (additive field, like + `slurm_job_id`), so provenance is ground truth even on a stale cluster. + +### 7.4 Path similarity, as implemented + +| | local | SLURM | hub | +|---|---|---|---| +| environment defined by | `Containerfile` | same | same | +| `lc build` | build into local store | build via podman-hpc | check ref exists in deployment registry | +| recipe isolation | `docker run` wrap | `podman-hpc run` wrap | worker pod **is** the image | +| cluster lifecycle | owned per-run | owned per-run | attached, user-owned | +| manifest truth | declared spec + code_version(tag) | same | same + `worker_image` ground truth | + +Deployment-side counterpart (lightcone-hub repo): `LIGHTCONE_REGISTRY` in +singleuser extraEnv, `LIGHTCONE_WORKER_IMAGE` in the cluster-options +environment, worker-default as the default cluster image, and a +`just build-worker` recipe mirroring `build-notebook`. + +## Sources + +- z2jh docs & GKE setup: https://z2jh.jupyter.org/en/stable/ · + https://z2jh.jupyter.org/en/stable/kubernetes/google/step-zero-gcp.html +- z2jh chart releases: https://hub.jupyter.org/helm-chart/ (4.4.0 / JupyterHub 5.5.0, Jun 2026) +- Dask Gateway on k8s: https://gateway.dask.org/install-kube.html (chart 2026.3.x, k8s ≥ 1.30) +- DaskHub deprecation & wiring reference: https://github.com/dask/helm-chart (daskhub/values.yaml); + 2i2c infra guide https://infrastructure.2i2c.org/topic/infrastructure/hub-helm-charts/ +- kbatch: https://kbatch.readthedocs.io/ · https://github.com/kbatch-dev/kbatch · + https://kbatch-dev.github.io/helm-chart/ +- Worked example of gateway-behind-hub-proxy: https://www.zonca.dev/posts/2022-04-04-dask-gateway-jupyterhub diff --git a/docs/user/cluster.md b/docs/user/cluster.md index ffa73f11..41da87c4 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -7,13 +7,17 @@ hardware to spread across. ## The big picture -`lc run` always dispatches through a Dask cluster. Three branches: +`lc run` always dispatches through a Dask cluster. Four branches: 1. On your laptop → a `LocalCluster` sized to the machine. 2. **Inside a SLURM allocation** → an in-process scheduler bound to the driver's hostname, with one `dask worker` per allocated node launched via `srun`. -3. With `DASK_SCHEDULER_ADDRESS` set → connect to whatever scheduler +3. **On a lightcone JupyterHub deployment** → attach to your running + Dask Gateway cluster (see the JupyterHub section below — don't use + `DASK_SCHEDULER_ADDRESS` with a `gateway://` address, it cannot be + dialled directly). +4. With `DASK_SCHEDULER_ADDRESS` set → connect to whatever scheduler you've pointed at. You don't pick — `lc run` detects which case applies. The only thing @@ -170,6 +174,84 @@ lc run `lc run` notices the env var and connects rather than starting its own scheduler. It does *not* tear the scheduler down on exit. +## JupyterHub / Dask Gateway + +On a lightcone JupyterHub deployment (where the `DASK_GATEWAY__*` env +vars are ambient in every pod), the model is: **you** create the +cluster, `lc run` attaches to it. + +1. Create a Dask Gateway cluster from JupyterLab — the Dask sidebar's + **+ NEW** button, or a notebook: + + ```python + from dask_gateway import Gateway + # shutdown_on_close=False keeps the cluster alive when this kernel + # exits — otherwise a kernel restart kills any lc run attached to it. + # Idle clusters are reaped by the deployment after 30 min. + cluster = Gateway().new_cluster(shutdown_on_close=False) + cluster.adapt(minimum=1, maximum=8) + ``` + +2. Run: + + ```bash + lc run + ``` + +The Gateway API only shows *your* clusters, so with a single cluster +running `lc run` attaches with zero configuration — no cluster name to +copy. It never changes the cluster's scaling and leaves it running on +exit, mirroring the `DASK_SCHEDULER_ADDRESS` convention, so you can +iterate `lc run` against the same warm cluster with its dashboard +panels docked. With no cluster running, `lc run` tells you how to +create one; with several, pick one: + +```bash +export LIGHTCONE_GATEWAY_CLUSTER= # e.g. hub.a1b2c3... +lc run +``` + +A Gateway scheduler's `gateway://` address cannot be used with +`DASK_SCHEDULER_ADDRESS` — attachment is always by name. + +Requires the optional gateway extra: +`pip install lightcone-cli[gateway]` (preinstalled on the hub image). + +### Containers on the hub: the pod is the runtime + +There is no docker or podman inside a pod — Kubernetes itself is the +container runtime (`container_runtime: kubernetes`, declared by the +site, no warning fired). Your recipes run *unwrapped inside the worker +pod*, so the worker image **is** the project environment: + +- A `container:` spec naming a registry image is used directly: create + your Gateway cluster with that image. +- A `container: Containerfile` spec resolves to a content-addressed + ref in the deployment registry: + `$LIGHTCONE_REGISTRY/lc-:` — the *same hash* the + local `lc build` tag carries, so the environment is provably the + same artifact on every path. `lc build` on the hub checks whether + that ref exists in the registry and, when it doesn't, prints the + publish commands: run `lc build` from a clone on any machine with + docker (this builds from the hash-attested, filtered build context — + don't substitute a raw `docker build .`), then + `docker tag lc-- ` and `docker push `. + +For the image to work as a Gateway worker it must contain `dask`, +`distributed`, `dask-gateway`, and `lightcone-cli` at versions +matching the hub — the simplest way is to base your Containerfile on +the deployment's worker image +(`FROM /lightcone-worker-default:`) and add your +science deps on top. That one Containerfile then serves every path: +built locally it wraps recipes on your laptop; pushed to the registry +it runs them as pods on the hub. + +At attach time `lc run` verifies the cluster's actual worker image +against the project's resolved image and warns on mismatch. Manifests +record the image the worker pod actually ran (`worker_image`), so +provenance stays truthful even if you knowingly run on a stale +cluster. + ## NERSC Perlmutter: site-specific notes !!! note "Setting up on Perlmutter for the first time?" diff --git a/pyproject.toml b/pyproject.toml index f7e5d297..6ae2649c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,11 +34,21 @@ dependencies = [ "snakemake>=9.0", "snakemake-interface-executor-plugins>=9.0", "snakemake-interface-common>=1.14", - "dask>=2024.1", - "distributed>=2024.1", + # Pinned to the hub's Gateway image (ghcr.io/dask/dask-gateway:2026.3.0): + # distributed warns/fails on client<->scheduler<->worker version skew, and + # because lightcone-cli travels into the worker image, the pin is what + # keeps driver and workers in lockstep. Bump in step with the hub's + # gateway chart/base image. + "dask==2026.3.0", + "distributed==2026.3.0", "rocrate>=0.11", ] +[project.optional-dependencies] +# Dask Gateway client, for running against a gateway-brokered cluster +# (e.g. the lightcone-hub JupyterHub deployment on GKE). +gateway = ["dask-gateway==2026.3.0"] + [dependency-groups] eval = [ "anthropic>=0.90", @@ -100,9 +110,13 @@ explicit_package_bases = true mypy_path = "src" [[tool.mypy.overrides]] -module = ["dask.*", "distributed.*"] +module = ["dask.*", "distributed.*", "dask_gateway.*"] ignore_missing_imports = true follow_untyped_imports = true +# dask_gateway re-exports Gateway from .client without __all__; strict +# mode would reject `from dask_gateway import Gateway` whenever the +# [gateway] extra happens to be installed. +implicit_reexport = true [[tool.mypy.overrides]] module = ["rocrate.*"] diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 259f8428..e52788f0 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -321,19 +321,22 @@ def init( _CONTAINERFILE = """\ FROM python:3.12-slim +# uv, pulled from its official image (fast, reproducible, no curl bootstrap). +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + WORKDIR /app +# On a Kubernetes / Dask Gateway deployment the pod image IS the execution +# environment, so it must contain lightcone-cli. Installing it here also pulls +# dask, distributed, and the dask-gateway client at pinned, hub-matching +# versions, so this image doubles as a valid Gateway worker. Add project +# dependencies to requirements.txt. +# +# --system installs into the image's Python (no venv in a container), which +# also puts dask-worker / dask-gateway-scheduler on PATH where the Gateway +# backend launches them by name. COPY requirements.txt . - -# Install curl and certificates, then clean up to keep image size small -RUN apt-get update && \ - apt-get install -y --no-install-recommends curl ca-certificates && \ - rm -rf /var/lib/apt/lists/* - -# Install uv -RUN curl -LsSf https://astral.sh/uv/install.sh | sh -ENV PATH="/root/.local/bin:$PATH" -RUN uv pip install -r requirements.txt +RUN uv pip install --system --no-cache "lightcone-cli[gateway]" -r requirements.txt COPY . . """ @@ -537,13 +540,19 @@ def run( """Materialize outputs declared in astra.yaml. Always dispatches through a Dask cluster: a ``LocalCluster`` on a - workstation, srun-launched workers inside a SLURM allocation, or an + workstation, srun-launched workers inside a SLURM allocation, the + user's running Dask Gateway cluster on a JupyterHub deployment + (created from JupyterLab; lc attaches, never creates), or an existing scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. """ _abort_on_perlmutter_login() - from lightcone.engine.container import load_runtime - from lightcone.engine.dask_cluster import cluster_for_run + from lightcone.engine.container import KUBERNETES_RUNTIME, load_runtime + from lightcone.engine.dask_cluster import ( + GATEWAY_CLUSTER_ENV, + cluster_for_run, + gateway_branch_active, + ) from lightcone.engine.scratch import ( RunLockBusyError, acquire_run_lock, @@ -569,6 +578,15 @@ def run( choice = load_runtime(project_path=project) _ensure_images(project, runtime=choice.runtime) + # On a kubernetes-runtime site the worker pod is the container, so + # resolve which image the Gateway cluster is expected to run — the + # gateway branch verifies the attached cluster against it and the + # no-cluster-yet error message names it. + expected_worker_image = ( + _expected_worker_image(project) + if choice.runtime == KUBERNETES_RUNTIME + else None + ) snakefile_path, cfg_path = generate( project, universes=universes, runtime=choice.runtime ) @@ -624,6 +642,7 @@ def run( targets=targets, force=force, has_outputs=bool(outputs), + gateway=gateway_branch_active(), ) # Hold a project-level flock for the duration of the run. Acquiring @@ -637,23 +656,48 @@ def run( except RunLockBusyError as e: raise click.ClickException(str(e)) - with cluster_for_run( - verbose=verbose, local_directory=str(rundirs.dask_local) - ) as scheduler_addr: - env = { - **os.environ, - "DASK_SCHEDULER_ADDRESS": scheduler_addr, - # The dask plugin's worker-side ``_run_shell`` takes this - # ``flock`` before forwarding a rule's lightcone output, so - # parallel rules' blocks never interleave at the line level - # — even across nodes. The lockfile sits under our scratch - # root specifically to avoid DVS on NERSC. - "LIGHTCONE_OUT_LOCK": str(rundirs.lock_path), - } - if verbose: - console.print(f"[dim]$ {' '.join(cmd)}[/dim]") - sys.exit(subprocess.run(cmd, env=env).returncode) - sys.exit(_run_silent(cmd, env=env, scratch_root=rundirs.root)) + try: + with cluster_for_run( + verbose=verbose, + local_directory=str(rundirs.dask_local), + expected_worker_image=expected_worker_image, + ) as cluster_env: + _warn_runtime_cluster_mismatch( + project, runtime=choice.runtime, cluster_env=cluster_env + ) + env = { + **os.environ, + # How the child snakemake's executor reaches the cluster: + # DASK_SCHEDULER_ADDRESS for address-dialled clusters, or + # LIGHTCONE_GATEWAY_CLUSTER for Dask Gateway (rejoined by + # name through the Gateway API). + **cluster_env, + # The dask plugin's worker-side ``_run_shell`` takes this + # ``flock`` before forwarding a rule's lightcone output, so + # parallel rules' blocks never interleave at the line level + # — even across nodes. The lockfile sits under our scratch + # root specifically to avoid DVS on NERSC. + "LIGHTCONE_OUT_LOCK": str(rundirs.lock_path), + } + # Exactly one of the two cluster variables may reach the + # child: a stale LIGHTCONE_GATEWAY_CLUSTER lingering in the + # user's shell (we tell them to export it) must not redirect + # the child to a Gateway cluster the parent never verified, + # and vice versa. + if "DASK_SCHEDULER_ADDRESS" in cluster_env: + env.pop(GATEWAY_CLUSTER_ENV, None) + elif GATEWAY_CLUSTER_ENV in cluster_env: + env.pop("DASK_SCHEDULER_ADDRESS", None) + if verbose: + console.print(f"[dim]$ {' '.join(cmd)}[/dim]") + sys.exit(subprocess.run(cmd, env=env).returncode) + sys.exit(_run_silent(cmd, env=env, scratch_root=rundirs.root)) + except RuntimeError as e: + # Cluster bootstrap errors are user guidance ("no Gateway cluster + # running — create one like this…", "workers lack the resource + # contract"), not stack traces. SystemExit from the run itself + # passes through untouched. + raise click.ClickException(str(e)) def _run_silent( @@ -712,12 +756,23 @@ def _build_snakemake_cmd( targets: list[str], force: bool, has_outputs: bool, + gateway: bool = False, ) -> list[str]: """Build the snakemake argv list for ``lc run``. ``--rerun-triggers`` uses ``nargs=+`` in snakemake's argparse, so without an explicit ``--`` separator it greedily consumes the first positional target path as an extra trigger value, causing an "invalid choice" error. + + *gateway* scopes ``--shared-fs-usage``: Gateway worker pods share + only the project volume with the driver, not its HOME or install + prefix. Snakemake's default (everything shared) makes the child + snakemake use driver-local paths — the driver's source-cache under + ``~/.cache`` (unwritable/absent in the worker pod) and the driver's + ``sys.executable`` (a conda path a slim worker image doesn't have). + Declaring what is actually shared keeps the child on worker-local + equivalents. Local and SLURM runs keep the default: there the + workers genuinely share the driver's environment. """ cmd: list[str] = [ "snakemake", @@ -734,6 +789,21 @@ def _build_snakemake_cmd( "--rerun-triggers", *rerun_triggers.split(","), ] + if gateway: + cmd += [ + "--shared-fs-usage", + "input-output", + "persistence", + "sources", + "storage-local-copies", + # Driver and workers see the project through NFS (the hub's + # RWX volume); the client-side attribute cache can hide a + # worker's freshly written outputs from the driver for tens + # of seconds. Snakemake's default 5s declares the rule + # failed ("output missing") even though it succeeded. + "--latency-wait", + "60", + ] if force: # ``--force`` scopes to explicit targets; ``rule all`` itself # has no recipe, so force-all is the only useful sense when no @@ -890,7 +960,10 @@ def verify(universe: str | None) -> None: @click.option( "--runtime", default=None, - help="docker | podman | podman-hpc (overrides ~/.lightcone/config.yaml)", + help=( + "docker | podman | podman-hpc | kubernetes " + "(overrides ~/.lightcone/config.yaml)" + ), ) def build(force: bool, runtime: str | None) -> None: """Build container images declared in astra.yaml. @@ -901,8 +974,18 @@ def build(force: bool, runtime: str | None) -> None: image store. Pre-built registry images (``python:3.12-slim``, ``ghcr.io/foo/bar:tag``) are skipped — the runtime pulls them at ``lc run`` time. + + On a kubernetes-runtime site (a lightcone-hub deployment) there is + no local image store and nothing to build in-pod: ``lc build`` + instead checks that each Containerfile's content-addressed image + exists in the deployment registry and prints exactly how to publish + it when missing. """ - from lightcone.engine.container import ContainerBuildError, load_runtime + from lightcone.engine.container import ( + KUBERNETES_RUNTIME, + ContainerBuildError, + load_runtime, + ) project = _project_root() try: @@ -910,6 +993,10 @@ def build(force: bool, runtime: str | None) -> None: except ContainerBuildError as e: raise click.ClickException(str(e)) + if resolved_runtime == KUBERNETES_RUNTIME: + _report_registry_images(project) + return + if resolved_runtime == "none": console.print( "[yellow]No container runtime available " @@ -923,34 +1010,22 @@ def build(force: bool, runtime: str | None) -> None: console.print("[green]Done.[/green]") -def _ensure_images(project: Path, *, runtime: str, force: bool = False) -> None: - """Build/pull every container image referenced in astra.yaml. +def _container_specs(project: Path) -> tuple[str, list[str]]: + """``(project_name, distinct container specs)`` from astra.yaml. - No-op when *runtime* is ``"none"``. Idempotent: skips images already - present in the runtime's local image store. Used by ``lc build`` (with - ``--force`` exposed) and as a pre-flight by ``lc run`` so the first - invocation after editing a Containerfile doesn't fail mid-DAG with a - missing image. + Specs are returned in tree order, first occurrence wins — the same + walk ``_ensure_images`` has always done, factored out so the + kubernetes-runtime paths (`lc build` registry report, expected + worker image) resolve the identical set. """ - if runtime == "none": - return - from astra.helpers import load_yaml, resolve_analysis_tree - from lightcone.engine.container import ( - ContainerBuildError, - build_image, - compute_image_tag, - image_exists_locally, - is_containerfile, - pull_image, - ) from lightcone.engine.tree import collect_tree_outputs spec = resolve_analysis_tree(load_yaml(project / "astra.yaml"), project) project_name = (spec.get("name") or project.name).lower().replace(" ", "-") - seen: set[str] = set() + specs: list[str] = [] for to in collect_tree_outputs(spec): recipe = to.output_def.get("recipe") or {} spec_str = ( @@ -958,9 +1033,223 @@ def _ensure_images(project: Path, *, runtime: str, force: bool = False) -> None: or to.analysis_spec.get("container") or spec.get("container") ) - if not spec_str or spec_str in seen: + if spec_str and spec_str not in specs: + specs.append(spec_str) + return project_name, specs + + +def _warn_runtime_cluster_mismatch( + project: Path, *, runtime: str, cluster_env: dict[str, str] +) -> None: + """Warn when the container runtime and the cluster branch disagree. + + The two are resolved from independent signals (site declaration / + user config vs. ambient cluster env), so misconfiguration can pair + them wrongly in both directions: + + * ``kubernetes`` runtime off-Gateway — recipes run wherever the + cluster puts them with no pod image lc can verify, while manifests + still record the declared ``container_image`` (provenance hazard, + e.g. ``runtime: kubernetes`` copied into a laptop config). + * An OCI runtime on a Gateway cluster — recipes arrive wrapped in + ``docker run``/``podman run`` inside worker pods that have no such + binary, so every containerized rule fails (e.g. + ``DASK_GATEWAY__ADDRESS`` set on a workstation that has docker). + + Warnings, not errors: the pairing can be intentional (a k8s-native + external scheduler; a gateway whose workers do ship podman). + """ + from lightcone.engine.container import KUBERNETES_RUNTIME, RUNTIMES + from lightcone.engine.dask_cluster import GATEWAY_CLUSTER_ENV + + on_gateway = GATEWAY_CLUSTER_ENV in cluster_env + if runtime == KUBERNETES_RUNTIME and not on_gateway: + console.print( + "[yellow]⚠ Container runtime is [cyan]kubernetes[/cyan] but " + "this run is not attached through Dask Gateway.[/yellow]\n" + " Recipes will execute unwrapped and lc cannot verify they " + "run in the declared container\n" + " image — manifests may record provenance that does not " + "match what executed. If this\n" + " machine is not a lightcone-hub deployment, remove " + "[cyan]container: {runtime: kubernetes}[/cyan]\n" + " from [cyan]~/.lightcone/config.yaml[/cyan]." + ) + elif runtime in RUNTIMES and on_gateway: + _, specs = _container_specs(project) + if specs: + console.print( + f"[yellow]⚠ Recipes are wrapped for [cyan]{runtime}" + f"[/cyan], but they will execute inside Dask Gateway " + "worker pods,[/yellow]\n" + f" which typically do not provide {runtime} — " + "containerized rules will fail. On a\n" + " Kubernetes-backed gateway the pod is the container: " + "set [cyan]container: {runtime: kubernetes}[/cyan]\n" + " in [cyan]~/.lightcone/config.yaml[/cyan] and give the " + "cluster the project's image instead." + ) + + +def _expected_worker_image(project: Path) -> str | None: + """The image the project's Gateway cluster is expected to run. + + Exactly one distinct declared container → its pullable realization + (deployment-registry ref for a Containerfile, the spec itself for a + registry image). Zero → ``None``. Several distinct images → ``None`` + with a warning: a Gateway cluster runs a single worker image, so + heterogeneous per-rule containers cannot be realized on this path. + A Containerfile that cannot be resolved (no ``LIGHTCONE_REGISTRY``) + also yields ``None`` with a warning — silently dropping it would + corrupt the single-vs-multiple decision and disable verification + without saying so. + """ + from lightcone.engine.container import is_containerfile, resolve_worker_image + + project_name, specs = _container_specs(project) + images: list[str] = [] + unresolved: list[str] = [] + for spec_str in specs: + image = resolve_worker_image( + spec_str, project_path=project, project_name=project_name + ) + if image is None and is_containerfile(spec_str, project): + unresolved.append(spec_str) + if image and image not in images: + images.append(image) + if unresolved: + console.print( + "[yellow]⚠ Cannot resolve " + + ", ".join(f"[cyan]{s}[/cyan]" for s in unresolved) + + " to a registry image: LIGHTCONE_REGISTRY is not set.[/yellow]\n" + " Worker-image verification is disabled for this run. Ask " + "the hub admin to inject\n" + " [cyan]LIGHTCONE_REGISTRY[/cyan] (see lightcone-hub helm " + "values)." + ) + return None + if len(images) > 1: + console.print( + "[yellow]⚠ astra.yaml declares multiple distinct container " + "images, but a Dask Gateway cluster runs a single worker " + "image:[/yellow]\n" + + "\n".join(f" [dim]•[/dim] {i}" for i in images) + + "\n All recipes will execute in whatever image the attached " + "cluster runs.\n" + " Consolidate on one Containerfile to restore per-recipe " + "provenance on this deployment." + ) + return None + return images[0] if images else None + + +def _report_registry_images(project: Path) -> None: + """``lc build`` on a kubernetes-runtime site. + + The worker pod is the container, so images must exist in the + deployment registry (``LIGHTCONE_REGISTRY``) for a Gateway cluster + to pull — there is no docker/podman in-pod to build with. Probe the + registry best-effort and print the exact publish commands when an + image is missing. Purely informational: never fails the command, + because the authoritative check happens when the user's cluster + starts (Kubernetes pulls the image or says why not). + """ + from lightcone.engine.container import ( + compute_image_tag, + deployment_registry, + is_containerfile, + registry_image_exists, + registry_image_ref, + ) + + project_name, specs = _container_specs(project) + if not specs: + console.print("No containers declared in astra.yaml — nothing to check.") + return + + console.print( + "[dim]This deployment runs recipes inside Dask Gateway worker " + "pods; images are pulled from the deployment registry, not built " + "here.[/dim]" + ) + registry = deployment_registry() + for spec_str in specs: + if not is_containerfile(spec_str, project): + console.print( + f"[green]•[/green] {spec_str} — registry image; use it as " + "the cluster's worker image.\n" + " [dim]Note: a Gateway worker image must contain dask, " + "distributed, dask-gateway, and\n" + " lightcone-cli at hub-matching versions — plain images " + "like python:*-slim will not\n" + " start as workers. Base project images on the " + "deployment's lightcone-worker-default.[/dim]" + ) + continue + if registry is None: + console.print( + f"[yellow]•[/yellow] {spec_str} — cannot resolve: " + "LIGHTCONE_REGISTRY is not set on this deployment. " + "Ask the hub admin to inject it (see lightcone-hub " + "helm values)." + ) continue - seen.add(spec_str) + tag = compute_image_tag(project_name, project / spec_str, project) + ref = registry_image_ref(tag, registry=registry) + exists = registry_image_exists(ref) if ref else None + if exists is True: + console.print(f"[green]✓[/green] {spec_str} → {ref} [dim](in registry)[/dim]") + continue + state = ( + "[red]missing from the registry[/red]" + if exists is False + else "[yellow]could not be verified[/yellow]" + ) + # Instruct `lc build` (not a raw `docker build .`): the tag hash + # attests the *staged* build context (Containerfile + deps + + # COPY sources, with results//.git/venvs excluded). A raw docker + # build from the project root would bake in whatever else lives + # there and push a different artifact under the attested name. + console.print( + f"[yellow]•[/yellow] {spec_str} → {ref} — {state}.\n" + " Publish it from a clone of this project on any machine " + "with docker and registry access:\n" + f" [cyan]lc build[/cyan] [dim]# builds {tag} from the " + "hash-attested build context[/dim]\n" + f" [cyan]docker tag {tag} {ref}[/cyan]\n" + f" [cyan]docker push {ref}[/cyan]\n" + " then create your Gateway cluster with " + f'[cyan]image="{ref}"[/cyan].' + ) + + +def _ensure_images(project: Path, *, runtime: str, force: bool = False) -> None: + """Build/pull every container image referenced in astra.yaml. + + No-op for non-wrapping runtimes: ``none`` uses no images, and + ``kubernetes`` has no local image store — images live in the + deployment registry (see ``_report_registry_images``). Idempotent: + skips images already present in the runtime's local image store. + Used by ``lc build`` (with ``--force`` exposed) and as a pre-flight + by ``lc run`` so the first invocation after editing a Containerfile + doesn't fail mid-DAG with a missing image. + """ + from lightcone.engine.container import KUBERNETES_RUNTIME + + if runtime in ("none", KUBERNETES_RUNTIME): + return + + from lightcone.engine.container import ( + ContainerBuildError, + build_image, + compute_image_tag, + image_exists_locally, + is_containerfile, + pull_image, + ) + + project_name, specs = _container_specs(project) + for spec_str in specs: if not is_containerfile(spec_str, project): # Registry image — pull so ``lc run`` can use ``--pull=never`` # without depending on the runtime's registry resolution. diff --git a/src/lightcone/engine/container.py b/src/lightcone/engine/container.py index bf114787..3534fd7d 100644 --- a/src/lightcone/engine/container.py +++ b/src/lightcone/engine/container.py @@ -21,6 +21,12 @@ * ``podman-hpc`` — NERSC-style login nodes; ``build`` migrates the image so compute-node apptainer can read it. ``run`` still uses ``podman-hpc`` directly. + * ``kubernetes`` — the pod is the container. On a lightcone-hub + deployment there is no OCI binary; recipes execute inside Dask + Gateway worker pods whose image *is* the project environment, so + ``wrap_recipe`` is a no-op and ``lc build`` resolves against the + deployment registry (:data:`REGISTRY_ENV`) instead of a local + image store. Declared by the site registry, never auto-detected. * ``none`` — no container; recipe runs on the host. Useful for development and for projects that don't need isolation. """ @@ -29,6 +35,7 @@ import hashlib import json import logging +import os import re import shlex import shutil @@ -53,6 +60,24 @@ #: probe so a down daemon doesn't silently win over a healthy podman. RUNTIMES: tuple[str, ...] = ("podman-hpc", "podman", "docker") +#: The pod-is-the-container runtime (Dask Gateway worker pods on a +#: lightcone-hub deployment). Not in :data:`RUNTIMES` — there is no binary +#: to detect and nothing to build locally; it enters via a site declaration +#: (see :func:`load_runtime`) or explicit user config. +KUBERNETES_RUNTIME = "kubernetes" + +#: Runtimes that never wrap recipes in an OCI invocation. ``kubernetes`` +#: still *means* containerized — isolation is delegated to the worker pod — +#: whereas ``none`` means the recipe runs on the host unisolated. +_NON_WRAPPING_RUNTIMES: frozenset[str] = frozenset({"none", KUBERNETES_RUNTIME}) + +#: Env var naming the deployment's container registry (e.g. +#: ``europe-west1-docker.pkg.dev//lightcone``). Injected into +#: singleuser pods by a lightcone-hub deployment; consumed by +#: :func:`registry_image_ref` so Containerfile specs resolve to a pullable +#: reference that a Dask Gateway cluster can be created with. +REGISTRY_ENV = "LIGHTCONE_REGISTRY" + #: Files whose contents contribute to the image tag hash. DEPENDENCY_FILES = ( "requirements.txt", @@ -123,12 +148,15 @@ class ContainerStatus: class RuntimeChoice: """Result of resolving the container runtime to use. - ``runtime`` is the resolved value (``docker | podman | podman-hpc | none``). - ``explicit`` is ``True`` when the user pinned this value in - ``~/.lightcone/config.yaml`` — i.e. they typed ``runtime: docker``, - ``runtime: podman``, … or ``runtime: none``. ``False`` means - ``runtime: auto`` (or no config), and the runtime is whatever - detection produced — including ``none`` as a silent fallback. + ``runtime`` is the resolved value + (``docker | podman | podman-hpc | kubernetes | none``). + ``explicit`` is ``True`` when the value was pinned rather than + detected: either the user typed it in ``~/.lightcone/config.yaml`` + (``runtime: docker``, ``runtime: none``, …) or the host site + *declares* a non-OCI runtime (``kubernetes``/``none``) — a + deployment fact, not a guess. ``False`` means ``runtime: auto`` + (or no config), and the runtime is whatever detection produced — + including ``none`` as a silent fallback. Callers use ``explicit`` to decide whether silently running without isolation is acceptable. When the user explicitly opted out, no @@ -207,11 +235,20 @@ def load_runtime(*, project_path: Path | None = None) -> RuntimeChoice: project_path is accepted for future per-project overrides but is not consulted today). Values: - * ``auto`` (default) — first available runtime in :data:`RUNTIMES`, - else falls back to ``"none"`` with ``explicit=False``. + * ``auto`` (default) — a site-declared non-OCI runtime wins first + (``kubernetes`` on a lightcone-hub deployment is a deployment + fact, so it resolves ``explicit=True``); otherwise the first + available runtime in :data:`RUNTIMES`, else ``"none"`` with + ``explicit=False``. * ``docker | podman | podman-hpc`` — explicit; binary must exist. + * ``kubernetes`` — explicit; recipes run unwrapped inside the + cluster's worker pods (no binary to check). * ``none`` — explicit opt-out; recipes run on the host. + Site-declared *OCI* runtimes (e.g. Perlmutter → podman-hpc) remain + detection-order hints, not explicit choices — missing-from-PATH + falls through to the next candidate as before. + Raises :class:`ContainerBuildError` if an explicit runtime is configured but its binary is missing on PATH, or if the configured value is unrecognised. @@ -227,13 +264,17 @@ def load_runtime(*, project_path: Path | None = None) -> RuntimeChoice: requested = "auto" if requested == "auto": + declared = detect_current_site().get("container_runtime") + if declared in _NON_WRAPPING_RUNTIMES: + return RuntimeChoice(runtime=declared, explicit=True) return RuntimeChoice(runtime=detect_runtime() or "none", explicit=False) - if requested == "none": - return RuntimeChoice(runtime="none", explicit=True) + if requested in _NON_WRAPPING_RUNTIMES: + return RuntimeChoice(runtime=requested, explicit=True) if requested not in RUNTIMES: raise ContainerBuildError( f"Unknown container.runtime {requested!r} in {cfg_path}. " - f"Expected one of: auto, none, {', '.join(RUNTIMES)}." + f"Expected one of: auto, none, {KUBERNETES_RUNTIME}, " + f"{', '.join(RUNTIMES)}." ) if shutil.which(requested) is None: raise ContainerBuildError( @@ -249,6 +290,31 @@ def load_runtime(*, project_path: Path | None = None) -> RuntimeChoice: # --------------------------------------------------------------------------- +def image_name_slug(name: str) -> str: + """Reduce *name* to a valid OCI image-name component. + + Registries (and docker/podman locally) require lowercase ASCII + ``[a-z0-9]`` with ``.``/``_``/``-`` separators, starting and ending + alphanumeric — but ASTRA analysis names are prose (e.g. + ``"Union2.1 Flat ΛCDM MAP Fit"``). Accented latin is transliterated + via NFKD; anything else non-ASCII is dropped; remaining invalid + characters become ``-``; separator runs collapse. Pure-ASCII names + that were already valid (``gwtest``, ``union2.1``) pass through + unchanged, so existing content-addressed tags keep their names. + """ + import unicodedata + + ascii_name = ( + unicodedata.normalize("NFKD", name) + .encode("ascii", "ignore") + .decode("ascii") + .lower() + ) + slug = re.sub(r"[^a-z0-9._]+", "-", ascii_name) + slug = re.sub(r"[-._]{2,}", "-", slug).strip("-._") + return slug or "project" + + def find_dependency_files(project_path: Path) -> list[Path]: """Return sorted list of dependency files found in *project_path*.""" found = [project_path / name for name in DEPENDENCY_FILES] @@ -350,8 +416,7 @@ def compute_image_tag( h.update(b"\0") digest = h.hexdigest()[:12] - safe_name = project_name.lower().replace(" ", "-") - return f"lc-{safe_name}-{digest}" + return f"lc-{image_name_slug(project_name)}-{digest}" def _safe_relpath(path: Path, root: Path) -> str: @@ -703,6 +768,168 @@ def resolve(spec: str | None) -> str | None: return resolve +# --------------------------------------------------------------------------- +# Deployment registry (kubernetes runtime) +# --------------------------------------------------------------------------- + + +def deployment_registry() -> str | None: + """The deployment's registry prefix from :data:`REGISTRY_ENV`, or ``None``. + + E.g. ``europe-west1-docker.pkg.dev//lightcone`` on a + lightcone-hub deployment. Trailing slashes are stripped so callers + can join with ``/`` unconditionally. + """ + value = (os.environ.get(REGISTRY_ENV) or "").strip().rstrip("/") + return value or None + + +def registry_image_ref(tag: str, *, registry: str | None = None) -> str | None: + """Translate a local content-addressed tag into a pullable registry ref. + + ``lc--`` becomes ``/lc-:`` — + the same content hash on both sides, so the environment is provably + the same artifact whether it was built locally (and wrapped per + recipe) or pulled by Kubernetes as a Gateway worker pod image. + + Returns ``None`` when no registry is configured (*registry* argument + or :data:`REGISTRY_ENV`): without one there is no pullable name to + give, and inventing a local-only tag would just fail at pod start. + """ + registry = registry or deployment_registry() + if registry is None: + return None + name, _, digest = tag.rpartition("-") + if not name or not digest: + return None + return f"{registry}/{name}:{digest}" + + +def resolve_worker_image( + spec: str | None, + *, + project_path: Path, + project_name: str, +) -> str | None: + """The image a Dask Gateway cluster should run workers with for *spec*. + + The kubernetes-runtime sibling of :func:`resolve_image_for_run`: + + * ``None`` / empty → ``None`` (no declared container). + * Path to a Containerfile → the deployment-registry ref for its + content-addressed tag, or ``None`` when :data:`REGISTRY_ENV` is + unset (the Containerfile cannot be realized as a pod image). + * Anything else (a registry image spec) → returned as-is. + + Note :func:`resolve_image_for_run` stays site-agnostic on purpose: + ``code_version`` hashes the *local* tag on every path, so moving a + project between a laptop and the hub does not register as code + drift. The registry ref is a deployment detail, used only to name + the worker image and to verify the attached cluster runs it. + """ + if not spec: + return None + if is_containerfile(spec, project_path): + tag = compute_image_tag(project_name, project_path / spec, project_path) + return registry_image_ref(tag) + return spec + + +#: GCE/GKE metadata-server endpoint for the node service account's OAuth +#: token. Reachable from pods on GKE (unless Workload Identity blocks it); +#: grants Artifact Registry read via the node SA's roles. +_METADATA_TOKEN_URL = ( + "http://metadata.google.internal/computeMetadata/v1/" + "instance/service-accounts/default/token" +) + +#: Manifest media types accepted when probing a registry — both Docker +#: and OCI, single-image and index, so multi-arch pushes don't 404. +_MANIFEST_ACCEPT = ", ".join( + ( + "application/vnd.docker.distribution.manifest.v2+json", + "application/vnd.docker.distribution.manifest.list.v2+json", + "application/vnd.oci.image.manifest.v1+json", + "application/vnd.oci.image.index.v1+json", + ) +) + + +def _metadata_access_token() -> str | None: + """OAuth access token from the GCE metadata server, or ``None``.""" + import urllib.error + import urllib.request + + req = urllib.request.Request( + _METADATA_TOKEN_URL, headers={"Metadata-Flavor": "Google"} + ) + try: + with urllib.request.urlopen(req, timeout=3) as resp: + payload = json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, OSError, ValueError): + return None + token = payload.get("access_token") + return token if isinstance(token, str) and token else None + + +def registry_image_exists(ref: str) -> bool | None: + """Best-effort probe: does *ref* exist in its registry? + + Returns ``True``/``False`` on a definitive registry answer and + ``None`` when we cannot tell (no metadata-server credentials, no + network, unexpected status). Callers must treat ``None`` as + "unverified", not "missing" — ``lc build`` on the hub uses this to + report status, never to gate execution. + """ + import urllib.error + import urllib.parse + import urllib.request + + if "/" not in ref: + return None + host, remainder = ref.split("/", 1) + last = remainder.rsplit("/", 1)[-1] + if ":" in last: + path, tag = remainder.rsplit(":", 1) + else: + path, tag = remainder, "latest" + if not host or not path or not tag: + return None + + token = _metadata_access_token() + if token is None: + return None + + # Percent-encode the path segments: http.client rejects any + # non-ASCII byte in the request line with UnicodeEncodeError, and a + # ref built from user input (project names are prose) must degrade + # to "unverified", never to a traceback. + quoted_path = "/".join( + urllib.parse.quote(seg, safe="") for seg in path.split("/") + ) + quoted_tag = urllib.parse.quote(tag, safe="") + req = urllib.request.Request( + f"https://{host}/v2/{quoted_path}/manifests/{quoted_tag}", + method="HEAD", + headers={ + "Authorization": f"Bearer {token}", + "Accept": _MANIFEST_ACCEPT, + }, + ) + try: + with urllib.request.urlopen(req, timeout=5) as resp: + return bool(200 <= resp.status < 300) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return False + return None + except (urllib.error.URLError, OSError, ValueError, UnicodeError): + # URLError/OSError: network; ValueError/UnicodeError: a ref the + # URL machinery refuses (bad host chars, non-ASCII that slipped + # through). All mean "cannot tell", not "missing". + return None + + def wrap_recipe( recipe: str, *, @@ -718,18 +945,23 @@ def wrap_recipe( No-op cases: * *image* is ``None`` → recipe returned unchanged - * *runtime* is ``"none"`` → recipe returned unchanged + * *runtime* is ``"none"`` → recipe returned unchanged (runs on + the host, unisolated) + * *runtime* is ``"kubernetes"`` → recipe returned unchanged (the + Dask Gateway worker pod *is* the container; there is no OCI + binary to invoke inside it) The recipe is shell-quoted with :func:`shlex.quote` and passed as the argument to ``bash -c`` inside the container, which keeps single quotes, dollar signs, and other shell metacharacters intact across the host bash → runtime CLI → container bash boundaries. """ - if image is None or runtime == "none": + if image is None or runtime in _NON_WRAPPING_RUNTIMES: return recipe if runtime not in RUNTIMES: raise ContainerBuildError( - f"Unsupported run runtime {runtime!r}; expected one of {RUNTIMES} or 'none'." + f"Unsupported run runtime {runtime!r}; expected one of " + f"{RUNTIMES}, '{KUBERNETES_RUNTIME}', or 'none'." ) inner = shlex.quote(recipe) # ``--pull=never`` is critical for podman, which by default does @@ -772,7 +1004,15 @@ def get_container_status( containerfile = project_path / spec tag = compute_image_tag(project_name, containerfile, project_path) - exists = image_exists_locally(tag, runtime=runtime) if runtime != "none" else None + # Non-wrapping runtimes have no local image store to consult: + # ``none`` doesn't use images at all, and ``kubernetes`` resolves + # against the deployment registry (probed by ``lc build``, not here — + # status must stay offline-cheap). + exists = ( + image_exists_locally(tag, runtime=runtime) + if runtime not in _NON_WRAPPING_RUNTIMES + else None + ) return ContainerStatus( type="build", image=tag, diff --git a/src/lightcone/engine/dask_cluster.py b/src/lightcone/engine/dask_cluster.py index 03049e5b..67bb9c89 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -1,19 +1,34 @@ # mypy: disable-error-code="no-untyped-call" """Cluster lifecycle for ``lc run``. -One context manager, three branches: +One context manager, four branches: - ``DASK_SCHEDULER_ADDRESS`` is already set → yield it as-is. We don't own the cluster, so we don't tear it down. +- A Dask Gateway environment is detected (``LIGHTCONE_GATEWAY_CLUSTER`` or + ``DASK_GATEWAY__ADDRESS`` is set, e.g. on a JupyterHub deployment) → + **attach** to the user's running Gateway cluster. ``lc run`` never + creates Gateway clusters: the user starts one from JupyterLab (Dask + sidebar or a notebook, where the options widget offers image/cores/ + memory) and lc discovers it through the user-scoped Gateway API — + exactly one running cluster attaches unambiguously; zero or several + is an error naming the fix. Gateway scheduler addresses use a custom + ``gateway://`` comm scheme that a bare ``distributed.Client`` cannot + dial, so the child snakemake process is told the *cluster name* via + ``LIGHTCONE_GATEWAY_CLUSTER`` and rejoins through the Gateway API — + the same rendezvous-by-ambient-context pattern the SLURM branch uses. - ``SLURM_JOB_ID`` is set → start an in-process scheduler via ``LocalCluster(n_workers=0)``, then ``srun`` one ``dask worker`` per node across the allocation. Workers advertise the node's full resources; per-rule ``threads`` / ``mem_mb`` / ``gpus`` map to per-task constraints. -- Neither → ``LocalCluster()`` sized to the local machine. +- None of the above → ``LocalCluster()`` sized to the local machine. -The scheduler is always in-process (driven by ``lc run`` itself) so its +Except on the Gateway branch (where the Gateway server owns scheduling), +the scheduler is always in-process (driven by ``lc run`` itself) so its lifetime equals the run's lifetime — no service to manage, no orphaned -schedulers if the driver crashes. +schedulers if the driver crashes. The Gateway branch mirrors the +``DASK_SCHEDULER_ADDRESS`` convention: we attach to a cluster someone +else owns, so we leave it running (and its scaling untouched) on exit. """ from __future__ import annotations @@ -34,6 +49,22 @@ RESOURCE_MEMORY = "memory" RESOURCE_GPUS = "gpus" +#: Env var carrying a Dask Gateway cluster name. Set by the user to pick +#: one of several running Gateway clusters (discovery attaches +#: automatically when exactly one is up); set by :func:`cluster_for_run` +#: for the child snakemake process so the executor plugin can rejoin the +#: cluster through the Gateway API (a bare ``Client`` cannot dial +#: ``gateway://``). +GATEWAY_CLUSTER_ENV = "LIGHTCONE_GATEWAY_CLUSTER" + +#: Env var carrying the image a Gateway worker pod was started with. A +#: lightcone-hub deployment's cluster-options handler injects it into +#: every scheduler/worker pod; :func:`cluster_for_run` reads it back to +#: verify the attached cluster actually runs the project's image (and +#: the manifest layer records it as ground truth — see +#: ``lightcone.engine.manifest.write_manifest``). +WORKER_IMAGE_ENV = "LIGHTCONE_WORKER_IMAGE" + @dataclass class _NodeShape: @@ -84,37 +115,278 @@ def _resources_arg(shape: _NodeShape) -> str: return " ".join(f"{k}={int(v)}" for k, v in _resource_dict(shape).items()) +def gateway_branch_active() -> bool: + """Would :func:`cluster_for_run` take the Gateway branch right now? + + Exposed so ``lc run`` can shape the parent snakemake invocation + (``--shared-fs-usage``) before entering the cluster context — the + decision is a pure function of the environment, in the same + priority order as the branches below. + """ + if os.environ.get("DASK_SCHEDULER_ADDRESS"): + return False + return ( + GATEWAY_CLUSTER_ENV in os.environ + or "DASK_GATEWAY__ADDRESS" in os.environ + ) + + @contextmanager def cluster_for_run( *, verbose: bool = False, local_directory: str | None = None, -) -> Iterator[str]: - """Yield a Dask scheduler address valid for the duration of `lc run`. + expected_worker_image: str | None = None, +) -> Iterator[dict[str, str]]: + """Yield the env overlay the child snakemake needs to reach the cluster. + + The parent (``lc run``) and the executor plugin live in different + processes, so connection info travels via environment variables. + Address-based branches yield ``{"DASK_SCHEDULER_ADDRESS": addr}``; + the Gateway branch yields ``{GATEWAY_CLUSTER_ENV: name}`` because + Gateway clusters are rejoined by name through the authenticated + Gateway API rather than dialled by address. *local_directory*, when given, is where dask workers stage their spilled task data and internal state files. ``lc run`` resolves it to a path under :mod:`lightcone.engine.scratch` so on NERSC the spill lands on Lustre instead of DVS-mounted home/CFS (where small- file I/O is slow and can pressure the gateway nodes). + + *expected_worker_image*, when given, is the image the project's + declared container resolves to on this deployment; the Gateway + branch compares it against what the attached cluster actually runs + and warns on mismatch. Ignored by the other branches (they realize + containers by wrapping recipes, not via pod images). """ if addr := os.environ.get("DASK_SCHEDULER_ADDRESS"): if verbose: print(f"→ Using existing Dask scheduler at {addr}") - yield addr + yield {"DASK_SCHEDULER_ADDRESS": addr} + return + + if GATEWAY_CLUSTER_ENV in os.environ or "DASK_GATEWAY__ADDRESS" in os.environ: + with _gateway_cluster( + verbose=verbose, expected_worker_image=expected_worker_image + ) as name: + yield {GATEWAY_CLUSTER_ENV: name} return if "SLURM_JOB_ID" in os.environ: with _slurm_backed_cluster( verbose=verbose, local_directory=local_directory ) as addr: - yield addr + yield {"DASK_SCHEDULER_ADDRESS": addr} return with _local_cluster( verbose=verbose, local_directory=local_directory ) as addr: - yield addr + yield {"DASK_SCHEDULER_ADDRESS": addr} + + +@contextmanager +def _gateway_cluster( + *, verbose: bool, expected_worker_image: str | None +) -> Iterator[str]: + """Attach to the user's running Dask Gateway cluster; yield its name. + + Attach-only by design: cluster lifecycle belongs to the user (create + one from the JupyterLab Dask sidebar or a notebook — that's where + the image/cores/memory options widget lives, and where the dashboard + is already docked). The Gateway API is user-scoped under JupyterHub + auth, so ``list_clusters()`` sees only the caller's clusters — + exactly one running cluster attaches with zero configuration; zero + or several raises with the fix spelled out. We never touch the + cluster's scaling and leave it running on exit, mirroring the + ``DASK_SCHEDULER_ADDRESS`` convention. + + The Gateway client is configured entirely by ambient dask config — + on a lightcone JupyterHub deployment the ``DASK_GATEWAY__*`` env + vars carry the API address, the JupyterHub auth mode, and the + proxy address, so ``Gateway()`` needs no arguments here. + """ + try: + from dask_gateway import Gateway + except ImportError as exc: + raise RuntimeError( + "A Dask Gateway environment was detected " + f"({GATEWAY_CLUSTER_ENV} or DASK_GATEWAY__ADDRESS is set) but " + "the dask-gateway client is not installed. Install it with " + "`pip install lightcone-cli[gateway]`." + ) from exc + + try: + gateway = Gateway() + except Exception as exc: + # Gateway() raises (ValueError) when no gateway address is + # configured — e.g. LIGHTCONE_GATEWAY_CLUSTER lingering in a + # shell off-hub, where no DASK_GATEWAY__* env exists. + raise RuntimeError( + "A Dask Gateway environment was detected but the gateway " + f"client could not be configured ({exc}). If you exported " + f"{GATEWAY_CLUSTER_ENV} on a machine without a Dask Gateway, " + "unset it." + ) from exc + + named = os.environ.get(GATEWAY_CLUSTER_ENV) + name = named or _discover_cluster_name( + gateway, expected_worker_image=expected_worker_image + ) + try: + cluster = gateway.connect(name) + except Exception as exc: + # dask-gateway raises its own error types (ValueError, + # GatewayClusterError) for a missing/stopped cluster; translate + # into guidance, because a *stale* LIGHTCONE_GATEWAY_CLUSTER is + # the likely cause — we told users to export it. + hint = ( + f"unset {GATEWAY_CLUSTER_ENV} to let lc discover your " + "running cluster, or point it at one that is running" + if named + else "it may have just stopped — check the JupyterLab Dask panel" + ) + raise RuntimeError( + f"Could not connect to Dask Gateway cluster {name!r} " + f"({exc}); {hint}." + ) from exc + + if verbose: + print( + f"→ Attached to Dask Gateway cluster {cluster.name} " + f"(dashboard: {cluster.dashboard_link})" + ) + + try: + client = cluster.get_client() + try: + # Verify the deployment contract only against what's live: + # an adaptive cluster sitting at zero workers will scale + # once the executor submits tasks, so an empty worker set + # is not an error here. + _assert_worker_resources(client) + _check_worker_image(client, expected=expected_worker_image) + finally: + client.close() + yield cluster.name + finally: + # Releases this process's connections only — the cluster (and + # its scaling) belongs to the user. + cluster.close() + + +def _discover_cluster_name( + gateway: object, expected_worker_image: str | None +) -> str: + """Name of the caller's single running Gateway cluster. + + Raises with actionable instructions when there isn't exactly one: + creation is deliberately not offered here (the user creates clusters + from JupyterLab), so the error message *is* the UX — it must say + precisely what to do next. + """ + reports = gateway.list_clusters() # type: ignore[attr-defined] + if len(reports) == 1: + return str(reports[0].name) + + if not reports: + image_arg = ( + f'image="{expected_worker_image}", ' if expected_worker_image else "" + ) + # shutdown_on_close=False matters: without it the cluster dies + # with the creating kernel/process, killing any lc run attached + # to it. The deployment's idle_timeout reaps forgotten clusters. + raise RuntimeError( + "No Dask Gateway cluster is running for your user. Create one " + "first — from the JupyterLab Dask sidebar (+ NEW), or in a " + "notebook:\n\n" + " from dask_gateway import Gateway\n" + f" cluster = Gateway().new_cluster({image_arg}" + "shutdown_on_close=False)\n" + " cluster.adapt(minimum=1, maximum=8)\n\n" + "then re-run `lc run` (it attaches to your running cluster " + "and leaves it up)." + ) + + names = ", ".join(str(r.name) for r in reports) + raise RuntimeError( + f"Multiple Dask Gateway clusters are running for your user " + f"({names}). Pick one by setting {GATEWAY_CLUSTER_ENV}, e.g.:\n\n" + f" export {GATEWAY_CLUSTER_ENV}={reports[0].name}\n\n" + "or shut the extras down from the JupyterLab Dask sidebar." + ) + + +def _check_worker_image(client: object, *, expected: str | None) -> None: + """Warn when the attached cluster doesn't run the project's image. + + The scheduler pod carries the same ``LIGHTCONE_WORKER_IMAGE`` env + the deployment injects into workers (Gateway applies cluster-config + ``environment`` to both), so this works even while an adaptive + cluster sits at zero workers. Read via a lambda — cloudpickled by + value — so the scheduler pod does not need lightcone importable. + + A warning, not an error: the user may be knowingly iterating on a + stale cluster, and the manifest layer records the *actual* image + (ground truth) either way. Silently skipped when the deployment + doesn't inject the marker or no expectation was computed. + """ + if expected is None: + return + try: + env_key = WORKER_IMAGE_ENV # captured by value into the lambda + actual = client.run_on_scheduler( # type: ignore[attr-defined] + lambda: __import__("os").environ.get(env_key) + ) + except Exception: + return # older deployment / restricted scheduler — not fatal + if actual and actual != expected: + print( + f"⚠ The attached Dask Gateway cluster runs image\n" + f" {actual}\n" + f" but this project's container resolves to\n" + f" {expected}\n" + f" Recipes will execute in the cluster's image; manifests " + f"record what actually ran.\n" + f" To use the project image, create a new cluster with " + f'image="{expected}".' + ) + + +def _assert_worker_resources(client: object) -> None: + """Fail fast when Gateway workers don't advertise the resource contract. + + Dask schedules a task only on workers advertising *every* requested + resource key, so a Gateway deployment that forgot to inject + ``cpus``/``memory``/``gpus`` (via cluster-options environment, e.g. + ``DASK_DISTRIBUTED__WORKER__RESOURCES__CPUS``) makes every rule hang + with no error. Better to refuse loudly at startup. + """ + workers = client.scheduler_info().get("workers", {}) # type: ignore[attr-defined] + if not workers: + return + # Require cpus AND memory together on at least one worker: the + # executor requests cpus for every rule (threads >= 1) and memory + # for any rule with mem_mb, so a partial injection (cpus without + # memory) still hangs mem_mb rules forever — the exact silent + # failure this check exists to catch. gpus is not required: a + # CPU-only deployment may legitimately omit it, and rules that + # request GPUs on such a cluster are a real scheduling constraint, + # not a forgotten contract. + if any( + RESOURCE_CPUS in res and RESOURCE_MEMORY in res + for w in workers.values() + if (res := w.get("resources") or {}) is not None + ): + return + raise RuntimeError( + "Dask Gateway workers do not advertise the lightcone resource " + f"contract ({RESOURCE_CPUS}+{RESOURCE_MEMORY}, plus " + f"{RESOURCE_GPUS} on GPU pools); per-rule resource requests " + "would never schedule. Fix the gateway deployment to inject " + "DASK_DISTRIBUTED__WORKER__RESOURCES__* into worker pods (see " + "the lightcone-hub dask-gateway values)." + ) @contextmanager diff --git a/src/lightcone/engine/manifest.py b/src/lightcone/engine/manifest.py index 27749fbb..7b9f2eeb 100644 --- a/src/lightcone/engine/manifest.py +++ b/src/lightcone/engine/manifest.py @@ -200,6 +200,12 @@ def write_manifest( "finished_at": time.time(), "host": socket.gethostname(), "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + # Ground truth on kubernetes-runtime sites: the image the Dask + # Gateway worker pod that executed this rule was started with + # (injected by the deployment's cluster-options handler). Unlike + # ``container_image`` (the declared spec), this records what + # actually ran. Additive, like ``slurm_job_id`` — absent off-hub. + "worker_image": os.environ.get("LIGHTCONE_WORKER_IMAGE"), } final_path = output_dir / MANIFEST_FILENAME diff --git a/src/lightcone/engine/site_registry.py b/src/lightcone/engine/site_registry.py index f4c00351..b3b7c644 100644 --- a/src/lightcone/engine/site_registry.py +++ b/src/lightcone/engine/site_registry.py @@ -14,6 +14,7 @@ """ from __future__ import annotations +import os import socket from collections.abc import Mapping from dataclasses import dataclass, field @@ -76,6 +77,25 @@ "//global/cfs/cdirs/**", ], }, + "lightcone-hub": { + # Pods have generated hostnames, so detection is by environment: + # a JupyterHub identity plus a configured Dask Gateway client + # (both injected into every singleuser pod by the lightcone-hub + # deployment) unambiguously mark the hub. + "hostname_patterns": [], + "env_markers": ["JUPYTERHUB_USER", "DASK_GATEWAY__ADDRESS"], + "display_name": "Lightcone JupyterHub", + "backend": "dask-gateway", + "connection": {}, + # Kubernetes is the container runtime: the Gateway worker pod + # image IS the recipe environment, so recipes run unwrapped but + # are still containerized. Declared (not detected) — load_runtime + # treats it as explicit, so no "no runtime found" warning fires. + "container_runtime": "kubernetes", + # Scratch comes from the deployment's LIGHTCONE_SCRATCH env (a + # path on the shared RWX volume), which outranks site defaults + # in lightcone.engine.scratch — nothing to declare here. + }, "local": { "hostname_patterns": [], "display_name": "Local", @@ -152,16 +172,31 @@ def get(self, name: str, default: Any = None) -> Any: _UNKNOWN_HOST_SITE = HostSite(key=None, defaults={}) +def _detect_site_from_env() -> str | None: + """Match a site by its declared ``env_markers`` (all must be set). + + Container/pod environments have generated hostnames, so sites like + the JupyterHub deployment declare ambient env vars instead of + hostname patterns. + """ + for site_key, site in SITE_DEFAULTS.items(): + markers = site.get("env_markers") + if markers and all(m in os.environ for m in markers): + return site_key + return None + + def detect_current_site() -> HostSite: """Return the :class:`HostSite` for the local host. Single source of truth for "which site are we on?" — everything else in the codebase should call this rather than re-deriving it from - :func:`socket.gethostname` and :func:`detect_site`. Returns a falsy - :class:`HostSite` (``key is None``) when the hostname matches no - known site. + :func:`socket.gethostname` and :func:`detect_site`. Environment + markers are checked before hostname patterns (pods have generated + hostnames). Returns a falsy :class:`HostSite` (``key is None``) + when nothing matches. """ - key = detect_site(socket.gethostname()) + key = _detect_site_from_env() or detect_site(socket.gethostname()) if key is None: return _UNKNOWN_HOST_SITE return HostSite(key=key, defaults=get_site_defaults(key) or {}) diff --git a/src/lightcone/engine/snakefile.py b/src/lightcone/engine/snakefile.py index 87b20158..f88205b6 100644 --- a/src/lightcone/engine/snakefile.py +++ b/src/lightcone/engine/snakefile.py @@ -332,10 +332,13 @@ def generate( project_path: Project root containing ``astra.yaml``. universes: Universe ids to expand rules over. runtime: Container runtime to wrap recipes with. One of - ``docker | podman | podman-hpc | none``. ``none`` runs - recipes on the host without isolation. Resolution is done - here once, not per-rule, so all rules use a consistent - runtime. See :func:`lightcone.engine.container.load_runtime`. + ``docker | podman | podman-hpc | kubernetes | none``. + ``none`` runs recipes on the host without isolation; + ``kubernetes`` also leaves recipes unwrapped, but because + the Dask Gateway worker pod *is* the container. Resolution + is done here once, not per-rule, so all rules use a + consistent runtime. See + :func:`lightcone.engine.container.load_runtime`. Returns ``(snakefile_path, config_path)``. """ diff --git a/src/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index 41bf9c1e..ae2806e1 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -17,8 +17,13 @@ from snakemake_interface_executor_plugins.jobs import ( # type: ignore[import-untyped] JobExecutorInterface, ) +from snakemake_interface_executor_plugins.utils import ( # type: ignore[import-untyped] + format_cli_arg, + join_cli_args, +) from lightcone.engine.dask_cluster import ( + GATEWAY_CLUSTER_ENV, RESOURCE_CPUS, RESOURCE_GPUS, RESOURCE_MEMORY, @@ -26,9 +31,9 @@ from lightcone.engine.runner import SENTINEL -def _run_shell(cmd: str) -> int: +def _run_shell(cmd: str) -> tuple[int, str]: """Worker-side: run the child snakemake command, forward its lightcone - output, and return its exit code. + output, and return its exit code plus the forwarded block. The command is a child snakemake invocation that loads the generated Snakefile and executes one rule's ``run:`` block. That block calls @@ -47,6 +52,11 @@ def _run_shell(cmd: str) -> int: On NERSC, ``$HOME`` and ``/global/cfs`` are mounted on compute nodes via DVS, which silently swallows ``flock``; lc run resolves the path onto Lustre via :mod:`lightcone.engine.scratch`. + + The same ``block`` (``""`` if nothing was forwarded) is also returned + to the caller so the driver can reprint it on the Dask Gateway path, + where this worker's stdout lands in the pod log rather than the + terminal running ``lc run``. """ p = subprocess.run( cmd, shell=True, capture_output=True, text=True, check=False @@ -57,6 +67,17 @@ def _run_shell(cmd: str) -> int: for line in stream.splitlines(): if line.startswith(SENTINEL): forwarded.append(line[len(SENTINEL):]) + if p.returncode != 0: + # The child failed: its own diagnostics are the only clue, and + # they exist solely in this worker process. Forward a bounded + # tail so the failure is debuggable from the driver terminal + # (exit 127 with no output is a debugging dead end). + tail = (p.stderr.strip() or p.stdout.strip()).splitlines()[-15:] + forwarded.append( + f"✗ worker-side snakemake exited {p.returncode}; last output:" + ) + forwarded.extend(f" {line}" for line in tail) + block = "" if forwarded: block = "\n".join(forwarded) + "\n" lock_path = os.environ.get("LIGHTCONE_OUT_LOCK") @@ -72,7 +93,41 @@ def _run_shell(cmd: str) -> int: sys.stdout.write(block) sys.stdout.flush() - return p.returncode + return p.returncode, block + + +def _emit_block(block: str) -> None: + """Driver-side: reprint a worker-forwarded block on our own stdout. + + Only used on the Dask Gateway path (see ``check_active_jobs``), where + the worker is a separate pod and its own stdout write in + ``_run_shell`` never reaches the terminal running ``lc run``. No + flock is needed here — ``check_active_jobs`` runs single-threaded in + the driver's event loop. + """ + if block: + sys.stdout.write(block) + sys.stdout.flush() + + +def _unpack_result(result: object) -> tuple[int, str]: + """Normalize a worker's future result to ``(exit_code, block)``. + + A worker running an older lightcone-cli returns a bare ``int`` exit + code; the current :func:`_run_shell` returns ``(exit_code, block)``. + Driver and worker images drift out of sync routinely on Dask Gateway + (the worker image lags the notebook/driver image), so accept both + rather than crash — the forwarded block is simply unavailable until + the worker image carries the tuple-returning change. + """ + if isinstance(result, tuple): + exit_code, block = result + return int(exit_code), str(block) + if isinstance(result, int): + return result, "" + raise TypeError( + f"unexpected dask task result type: {type(result).__name__}" + ) def _build_resources(job: JobExecutorInterface) -> dict[str, float]: @@ -90,9 +145,27 @@ def _build_resources(job: JobExecutorInterface) -> dict[str, float]: return res -class DaskExecutor(RemoteExecutor): # type: ignore[misc] - def __init__(self, workflow, logger): # type: ignore[no-untyped-def] - super().__init__(workflow, logger) +def _connect_client(): # type: ignore[no-untyped-def] + """Resolve the Dask client from the environment. + + Mirrors the branches of ``lightcone.engine.dask_cluster.cluster_for_run`` + from the child process side, **in the same priority order**: a plain + address in ``DASK_SCHEDULER_ADDRESS`` wins first (matching the parent, + where an explicit address outranks a Gateway environment); otherwise a + Gateway cluster is rejoined *by name* through the authenticated Gateway + API (its ``gateway://`` scheduler address cannot be dialled by a bare + ``Client``). The parent additionally strips the losing variable from + the child env (see ``lc run``), so a stale ``LIGHTCONE_GATEWAY_CLUSTER`` + lingering in a user's shell can never redirect the child to a cluster + the parent didn't verify. + + Returns ``(client, gateway_cluster_or_None)`` — the cluster handle is + kept so ``shutdown()`` can release its local connections. It is never + shut down here: on the Gateway path the cluster belongs to the *user* + (lc is attach-only); on the address paths it belongs to whoever + started the scheduler. + """ + if addr := os.environ.get("DASK_SCHEDULER_ADDRESS"): try: from dask.distributed import Client except ImportError as exc: @@ -100,15 +173,70 @@ def __init__(self, workflow, logger): # type: ignore[no-untyped-def] "dask.distributed is required for the dask executor " "(`pip install distributed`)." ) from exc + return Client(addr), None - addr = os.environ.get("DASK_SCHEDULER_ADDRESS") - if not addr: + if name := os.environ.get(GATEWAY_CLUSTER_ENV): + try: + from dask_gateway import Gateway + except ImportError as exc: raise WorkflowError( - "DASK_SCHEDULER_ADDRESS is not set. `lc run` should set this " - "before invoking snakemake; if you're calling snakemake " - "directly, point it at a running dask scheduler." - ) - self._client = Client(addr) + f"{GATEWAY_CLUSTER_ENV} is set but the dask-gateway client " + "is not installed (`pip install lightcone-cli[gateway]`)." + ) from exc + cluster = Gateway().connect(name) + return cluster.get_client(), cluster + + raise WorkflowError( + f"Neither DASK_SCHEDULER_ADDRESS nor {GATEWAY_CLUSTER_ENV} is " + "set. `lc run` should set one before invoking snakemake; if " + "you're calling snakemake directly, point it at a running dask " + "scheduler." + ) + + +class DaskExecutor(RemoteExecutor): # type: ignore[misc] + def __init__(self, workflow, logger): # type: ignore[no-untyped-def] + super().__init__(workflow, logger) + self._client, self._gateway_cluster = _connect_client() + + def get_python_executable(self) -> str: + """Python for the child snakemake command. + + The default (``sys.executable``) is right when driver and + workers share an environment — local LocalCluster workers and + srun-launched SLURM workers inherit the driver's venv, so the + absolute path exists there. On Dask Gateway the driver (e.g. a + JupyterLab pod) and the workers run *different images*: the + driver's interpreter path need not exist in the worker pod + (conda notebook vs pip-slim worker is exactly this), which + fails with exit 127 before snakemake even starts. The worker + image is the software deployment, so let the worker's own PATH + resolve the interpreter. + """ + if self._gateway_cluster is not None: + return "python3" + return super().get_python_executable() # type: ignore[no-any-return] + + def additional_general_args(self) -> str: + """Pin the child snakemake's working directory explicitly. + + Local and SLURM workers inherit the driver's cwd (LocalCluster + forks in place; srun preserves the submit directory), so the + child's relative output paths land in the project by accident + of process ancestry. Gateway worker pods start in their own + HOME — without ``--directory``, a rule "succeeds" while writing + results into the pod's ephemeral filesystem. The parent + snakemake chdir'd into the project (``lc run`` passes ``-d``), + so ``os.getcwd()`` here *is* the project directory — a path + valid on every worker because lc requires a shared project + filesystem (``implies_no_shared_fs=False``). + """ + return join_cli_args( # type: ignore[no-any-return] + [ + super().additional_general_args(), + format_cli_arg("--directory", os.getcwd()), + ] + ) def run_job(self, job: JobExecutorInterface) -> None: cmd = self.format_job_exec(job) @@ -143,7 +271,9 @@ async def check_active_jobs( ) continue - exit_code = future.result() + exit_code, block = _unpack_result(future.result()) + if self._gateway_cluster is not None: + _emit_block(block) if exit_code != 0: self.report_job_error( j, msg=f"Dask task '{j.external_jobid}' exited {exit_code}." @@ -169,5 +299,10 @@ def cancel_jobs(self, active_jobs: list[SubmittedJobInfo]) -> None: def shutdown(self) -> None: try: self._client.close() + if self._gateway_cluster is not None: + # Releases this process's connections only; the cluster + # itself belongs to the user (lc is attach-only and never + # shuts Gateway clusters down — see dask_cluster). + self._gateway_cluster.close() finally: super().shutdown() diff --git a/tests/conftest.py b/tests/conftest.py index baba0659..1459b9a6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,22 @@ """Shared test fixtures for lightcone-cli tests.""" from __future__ import annotations + +import pytest + + +@pytest.fixture(autouse=True) +def _no_ambient_hub_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Neutralize lightcone-hub env-marker site detection for every test. + + ``detect_current_site`` checks env markers *before* hostnames, so a + suite running where JUPYTERHUB_USER and DASK_GATEWAY__ADDRESS are + ambient (e.g. inside a lightcone-hub singleuser pod — exactly the + environment this codebase targets) would otherwise resolve every + site-dependent expectation (runtime detection, scratch roots, + hostname-mocked site tests) to the hub site. Tests that want the + markers set them explicitly via monkeypatch.setenv, which composes + fine with this fixture. + """ + monkeypatch.delenv("JUPYTERHUB_USER", raising=False) + monkeypatch.delenv("DASK_GATEWAY__ADDRESS", raising=False) diff --git a/tests/test_cli.py b/tests/test_cli.py index a5c03fc1..edaecd31 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -75,6 +75,28 @@ def test_init_creates_project(runner: CliRunner, tmp_path: Path) -> None: assert (project / "universes").is_dir() +def test_init_containerfile_is_gateway_worker_capable( + runner: CliRunner, tmp_path: Path +) -> None: + """The scaffold image must be able to run as a Dask Gateway worker. + + On the kubernetes runtime the pod image IS the recipe environment, so it + has to carry lightcone-cli (which pulls dask/distributed/dask-gateway at + pinned versions). A plain ``FROM python:3.12-slim`` that only installs + requirements.txt cannot start as a worker. + """ + project = tmp_path / "proj" + result = runner.invoke(main, ["init", str(project), "--no-git", "--no-venv"]) + assert result.exit_code == 0, result.output + + containerfile = (project / "Containerfile").read_text() + assert "lightcone-cli[gateway]" in containerfile + # uv refuses a non-venv install without --system (see plan/LCR-176). + assert "--system" in containerfile + # The spec points at the project Containerfile, not the slim base. + assert "container: Containerfile" in (project / "astra.yaml").read_text() + + def test_init_creates_report_template(runner: CliRunner, tmp_path: Path) -> None: project = tmp_path / "proj" result = runner.invoke(main, ["init", str(project), "--no-git", "--no-venv"]) @@ -213,6 +235,80 @@ def test_run_cmd_no_separator_when_no_targets() -> None: assert "--" not in cmd +def test_run_cmd_gateway_scopes_shared_fs_and_latency() -> None: + """Gateway workers share only the project volume with the driver, not + its HOME (source cache) or install prefix, and see it through NFS. + Validated on the lightcone-hub deployment: without --shared-fs-usage + the child snakemake mkdir's the driver's ~/.cache path (PermissionError + in the worker pod), and without --latency-wait the driver's NFS + attribute cache declares freshly written outputs missing.""" + from lightcone.cli.commands import _build_snakemake_cmd + + cmd = _build_snakemake_cmd( + snakefile_path=Path("/shared/proj/.lightcone/Snakefile"), + project=Path("/shared/proj"), + n="4", + rerun_triggers="code,input,mtime,params", + targets=[], + force=False, + has_outputs=False, + gateway=True, + ) + + fs_idx = cmd.index("--shared-fs-usage") + values = cmd[fs_idx + 1 : fs_idx + 5] + assert set(values) == { + "input-output", + "persistence", + "sources", + "storage-local-copies", + } + assert "source-cache" not in cmd, "driver ~/.cache is not visible to worker pods" + assert "software-deployment" not in cmd, "driver interpreter path differs in pods" + assert cmd[cmd.index("--latency-wait") + 1] == "60" + + +def test_run_cmd_default_keeps_snakemake_shared_fs_defaults() -> None: + """Local/SLURM workers genuinely share the driver's environment — + the gateway-only flags must not leak into those paths.""" + from lightcone.cli.commands import _build_snakemake_cmd + + cmd = _build_snakemake_cmd( + snakefile_path=Path("/proj/.lightcone/Snakefile"), + project=Path("/proj"), + n="4", + rerun_triggers="code,input,mtime,params", + targets=[], + force=False, + has_outputs=False, + ) + + assert "--shared-fs-usage" not in cmd + assert "--latency-wait" not in cmd + + +def test_gateway_branch_active_matches_cluster_for_run_priority( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """gateway_branch_active must mirror cluster_for_run's branch order: + an explicit scheduler address outranks the gateway environment.""" + from lightcone.engine.dask_cluster import gateway_branch_active + + for var in ( + "DASK_SCHEDULER_ADDRESS", + "DASK_GATEWAY__ADDRESS", + "LIGHTCONE_GATEWAY_CLUSTER", + ): + monkeypatch.delenv(var, raising=False) + assert gateway_branch_active() is False + + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + assert gateway_branch_active() is True + + monkeypatch.setenv("DASK_SCHEDULER_ADDRESS", "tcp://existing:8786") + assert gateway_branch_active() is False + + def test_run_cmd_multiple_triggers_all_before_separator() -> None: """All four trigger tokens must precede the '--' separator.""" from lightcone.cli.commands import _build_snakemake_cmd diff --git a/tests/test_container.py b/tests/test_container.py index 2fb10a74..b1f2655b 100644 --- a/tests/test_container.py +++ b/tests/test_container.py @@ -14,10 +14,13 @@ import yaml from lightcone.engine.container import ( + KUBERNETES_RUNTIME, + REGISTRY_ENV, RUNTIMES, ContainerBuildError, build_image, compute_image_tag, + deployment_registry, detect_runtime, find_dependency_files, get_container_status, @@ -26,7 +29,10 @@ is_containerfile, load_runtime, pull_image, + registry_image_exists, + registry_image_ref, resolve_image_for_run, + resolve_worker_image, wrap_recipe, ) @@ -568,6 +574,64 @@ def test_unknown_runtime_raises( with pytest.raises(ContainerBuildError, match="Unknown container.runtime"): load_runtime() + def test_site_declared_kubernetes_is_explicit( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """On a lightcone-hub deployment (site declares kubernetes) auto + must resolve to kubernetes with explicit=True — the pod IS the + container, so no 'no runtime found' warning is owed even though + docker/podman are absent.""" + from lightcone.engine.site_registry import HostSite + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.setattr( + "lightcone.engine.container.detect_current_site", + lambda: HostSite( + key="lightcone-hub", + defaults={"container_runtime": "kubernetes"}, + ), + ) + monkeypatch.setattr( + "lightcone.engine.container.detect_runtime", + lambda: pytest.fail("detection must not run on a declared site"), + ) + choice = load_runtime() + assert choice.runtime == "kubernetes" + assert choice.explicit is True + + def test_site_declared_oci_runtime_stays_a_hint( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Perlmutter-style OCI declarations keep detection-order-hint + semantics: a missing binary falls through instead of resolving + explicit (that would error or lie about availability).""" + from lightcone.engine.site_registry import HostSite + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.setattr( + "lightcone.engine.container.detect_current_site", + lambda: HostSite( + key="perlmutter", + defaults={"container_runtime": "podman-hpc"}, + ), + ) + monkeypatch.setattr( + "lightcone.engine.container.detect_runtime", lambda: "podman" + ) + choice = load_runtime() + assert choice.runtime == "podman" + assert choice.explicit is False + + def test_explicit_kubernetes_config_accepted( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """kubernetes in user config is explicit and needs no binary.""" + monkeypatch.setattr(Path, "home", lambda: tmp_path) + self._write_config(tmp_path, {"container": {"runtime": "kubernetes"}}) + choice = load_runtime() + assert choice.runtime == KUBERNETES_RUNTIME + assert choice.explicit is True + # ---- resolve_image_for_run ----------------------------------------------- @@ -595,6 +659,240 @@ def test_containerfile_resolves_to_tag(self, project: Path) -> None: assert result is not None assert result.startswith("lc-test-") + def test_resolution_ignores_deployment_registry( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """resolve_image_for_run stays site-agnostic on purpose: + code_version hashes the local tag on every path, so moving a + project between a laptop and the hub must not register as code + drift. The registry ref is resolve_worker_image's job.""" + monkeypatch.setenv(REGISTRY_ENV, "eu-docker.pkg.dev/proj/lightcone") + result = resolve_image_for_run( + "Containerfile", project_path=project, project_name="test" + ) + assert result is not None + assert result.startswith("lc-test-") + assert "pkg.dev" not in result + + +# ---- image_name_slug -------------------------------------------------------- + + +class TestImageNameSlug: + def test_valid_ascii_names_pass_through(self) -> None: + """Existing content-addressed tags must keep their names.""" + from lightcone.engine.container import image_name_slug + + for name in ("gwtest", "union2.1", "my_proj", "a-b-c"): + assert image_name_slug(name) == name + + def test_prose_astra_name_with_greek(self, project: Path) -> None: + """Regression: 'Union2.1 Flat ΛCDM MAP Fit' produced an invalid + OCI name (λ survived .lower()), which docker/podman reject and + which crashed the registry probe with UnicodeEncodeError.""" + from lightcone.engine.container import image_name_slug + + slug = image_name_slug("Union2.1 Flat ΛCDM MAP Fit") + assert slug == "union2.1-flat-cdm-map-fit" + tag = compute_image_tag( + "Union2.1 Flat ΛCDM MAP Fit", project / "Containerfile", project + ) + assert tag.startswith("lc-union2.1-flat-cdm-map-fit-") + assert tag.isascii() + + def test_accented_latin_transliterates(self) -> None: + from lightcone.engine.container import image_name_slug + + assert image_name_slug("Étude Numérique") == "etude-numerique" + + def test_degenerate_name_falls_back(self) -> None: + from lightcone.engine.container import image_name_slug + + assert image_name_slug("ΛΩΔ") == "project" + assert image_name_slug("...") == "project" + + +# ---- deployment registry (kubernetes runtime) ------------------------------ + + +class TestDeploymentRegistry: + def test_unset_env_returns_none(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(REGISTRY_ENV, raising=False) + assert deployment_registry() is None + + def test_strips_trailing_slash(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(REGISTRY_ENV, "eu-docker.pkg.dev/proj/lightcone/") + assert deployment_registry() == "eu-docker.pkg.dev/proj/lightcone" + + def test_blank_env_returns_none(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(REGISTRY_ENV, " ") + assert deployment_registry() is None + + +class TestRegistryImageRef: + def test_local_tag_becomes_name_colon_hash(self) -> None: + ref = registry_image_ref( + "lc-my-analysis-abc123def456", + registry="eu-docker.pkg.dev/proj/lightcone", + ) + # Same content hash on both sides — the environment is provably + # the same artifact locally and as a worker pod image. + assert ref == "eu-docker.pkg.dev/proj/lightcone/lc-my-analysis:abc123def456" + + def test_env_registry_used_when_not_passed( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv(REGISTRY_ENV, "eu-docker.pkg.dev/proj/lightcone") + ref = registry_image_ref("lc-proj-cafe01234567") + assert ref == "eu-docker.pkg.dev/proj/lightcone/lc-proj:cafe01234567" + + def test_no_registry_returns_none( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv(REGISTRY_ENV, raising=False) + assert registry_image_ref("lc-proj-cafe01234567") is None + + +class TestResolveWorkerImage: + def test_none_spec(self, project: Path) -> None: + assert resolve_worker_image( + None, project_path=project, project_name="test" + ) is None + + def test_registry_image_passes_through( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv(REGISTRY_ENV, raising=False) + assert resolve_worker_image( + "ghcr.io/foo/bar:tag", project_path=project, project_name="test" + ) == "ghcr.io/foo/bar:tag" + + def test_containerfile_resolves_to_registry_ref( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv(REGISTRY_ENV, "eu-docker.pkg.dev/proj/lightcone") + ref = resolve_worker_image( + "Containerfile", project_path=project, project_name="test" + ) + assert ref is not None + assert ref.startswith("eu-docker.pkg.dev/proj/lightcone/lc-test:") + # Hash identical to the local tag's — split differs (:), content + # hash doesn't. + local = resolve_image_for_run( + "Containerfile", project_path=project, project_name="test" + ) + assert local is not None + assert ref.rsplit(":", 1)[1] == local.rsplit("-", 1)[1] + + def test_containerfile_without_registry_returns_none( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv(REGISTRY_ENV, raising=False) + assert resolve_worker_image( + "Containerfile", project_path=project, project_name="test" + ) is None + + +class TestRegistryImageExists: + def test_no_credentials_returns_unknown( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Off-GKE (no metadata server) the probe must answer 'can't + tell' (None), never False — lc build treats None as unverified.""" + monkeypatch.setattr( + "lightcone.engine.container._metadata_access_token", lambda: None + ) + assert registry_image_exists( + "eu-docker.pkg.dev/proj/lightcone/lc-x:abc" + ) is None + + def test_malformed_ref_returns_unknown(self) -> None: + assert registry_image_exists("not-a-ref") is None + + def test_non_ascii_ref_never_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Regression: a λ in the image name reached http.client's + ASCII-only request line and escaped as UnicodeEncodeError, + crashing `lc build`. The probe must degrade to None instead + (the ref is percent-encoded now; the registry answers 404).""" + monkeypatch.setattr( + "lightcone.engine.container._metadata_access_token", lambda: "tok" + ) + + captured: dict[str, str] = {} + + class _Resp: + status = 200 + + def __enter__(self): # noqa: ANN204 + return self + + def __exit__(self, *exc: object) -> None: + pass + + def _open(req, timeout=0): # noqa: ANN001, ANN202 + # Real http.client encodes the request line as ASCII — + # reproduce that constraint so a regression fails here. + req.full_url.encode("ascii") + captured["url"] = req.full_url + return _Resp() + + monkeypatch.setattr("urllib.request.urlopen", _open) + result = registry_image_exists( + "eu-docker.pkg.dev/proj/lightcone/lc-union2.1-flat-λcdm:abc" + ) + assert result is not None # quoted URL went through + assert "%CE%BB" in captured["url"] # λ percent-encoded + + def test_404_means_missing(self, monkeypatch: pytest.MonkeyPatch) -> None: + import io + import urllib.error + + monkeypatch.setattr( + "lightcone.engine.container._metadata_access_token", lambda: "tok" + ) + + def _raise_404(req, timeout=0): # noqa: ANN001, ANN202 + raise urllib.error.HTTPError( + req.full_url, 404, "Not Found", hdrs=None, fp=io.BytesIO(b"") + ) + + monkeypatch.setattr("urllib.request.urlopen", _raise_404) + assert registry_image_exists( + "eu-docker.pkg.dev/proj/lightcone/lc-x:abc" + ) is False + + def test_200_means_present(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "lightcone.engine.container._metadata_access_token", lambda: "tok" + ) + + class _Resp: + status = 200 + + def __enter__(self): # noqa: ANN204 + return self + + def __exit__(self, *exc: object) -> None: + pass + + captured: dict[str, object] = {} + + def _ok(req, timeout=0): # noqa: ANN001, ANN202 + captured["url"] = req.full_url + captured["method"] = req.get_method() + return _Resp() + + monkeypatch.setattr("urllib.request.urlopen", _ok) + assert registry_image_exists( + "eu-docker.pkg.dev/proj/lightcone/lc-x:abc" + ) is True + assert captured["url"] == ( + "https://eu-docker.pkg.dev/v2/proj/lightcone/lc-x/manifests/abc" + ) + assert captured["method"] == "HEAD" + # ---- wrap_recipe ---------------------------------------------------------- @@ -608,6 +906,13 @@ def test_runtime_none_passthrough(self) -> None: "echo hi", image="python:3.12-slim", runtime="none" ) == "echo hi" + def test_runtime_kubernetes_passthrough(self) -> None: + """On the hub the worker pod IS the container — there is no OCI + binary inside it, so the recipe must run unwrapped.""" + assert wrap_recipe( + "echo hi", image="python:3.12-slim", runtime="kubernetes" + ) == "echo hi" + def test_podman_wrap_basic(self) -> None: wrapped = wrap_recipe( "echo hi", image="python:3.12-slim", runtime="podman" diff --git a/tests/test_dask_cluster.py b/tests/test_dask_cluster.py index 4bb6b473..904fece2 100644 --- a/tests/test_dask_cluster.py +++ b/tests/test_dask_cluster.py @@ -28,6 +28,9 @@ def _clean_env(monkeypatch: pytest.MonkeyPatch) -> None: for var in ( "DASK_SCHEDULER_ADDRESS", + "DASK_GATEWAY__ADDRESS", + "LIGHTCONE_GATEWAY_CLUSTER", + "LIGHTCONE_WORKER_IMAGE", "SLURM_JOB_ID", "SLURM_NNODES", "SLURM_CPUS_ON_NODE", @@ -69,8 +72,8 @@ def test_existing_scheduler_address_yields_unchanged( ) -> None: monkeypatch.setenv("DASK_SCHEDULER_ADDRESS", "tcp://example:8786") - with cluster_for_run() as addr: - assert addr == "tcp://example:8786" + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://example:8786"} def test_no_env_uses_local_cluster() -> None: @@ -83,8 +86,8 @@ def _fake_local(*, verbose: bool, local_directory: str | None = None): yield "tcp://stub:9999" with patch("lightcone.engine.dask_cluster._local_cluster", _fake_local): - with cluster_for_run() as addr: - assert addr == "tcp://stub:9999" + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://stub:9999"} assert sentinel["called"] == "local" @@ -98,8 +101,8 @@ def _fake_slurm(*, verbose: bool, local_directory: str | None = None): yield "tcp://stub:9999" with patch("lightcone.engine.dask_cluster._slurm_backed_cluster", _fake_slurm): - with cluster_for_run() as addr: - assert addr == "tcp://stub:9999" + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://stub:9999"} assert sentinel["called"] == "slurm" @@ -116,8 +119,8 @@ def _should_not_run(*, verbose: bool, local_directory: str | None = None): yield # pragma: no cover with patch("lightcone.engine.dask_cluster._slurm_backed_cluster", _should_not_run): - with cluster_for_run() as addr: - assert addr == "tcp://existing:8786" + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://existing:8786"} def test_slurm_backed_cluster_binds_to_routable_host( @@ -276,6 +279,466 @@ def close(self) -> None: assert set(resources.keys()) == {RESOURCE_CPUS, RESOURCE_MEMORY, RESOURCE_GPUS} +# --------------------------------------------------------------------------- +# Dask Gateway branch +# --------------------------------------------------------------------------- + + +class _FakeGatewayClient: + def __init__( + self, + resources: dict[str, float] | None, + worker_image: str | None = None, + ) -> None: + self._resources = resources + self._worker_image = worker_image + + def scheduler_info(self) -> dict[str, object]: + if self._resources is None: + return {"workers": {}} + return {"workers": {"w0": {"resources": self._resources}}} + + def run_on_scheduler(self, fn): # noqa: ANN001, ANN202 + # The real client executes fn inside the scheduler pod, where the + # deployment injects LIGHTCONE_WORKER_IMAGE. Simulate that env. + import os + from unittest.mock import patch as _patch + + env = dict(os.environ) + if self._worker_image is not None: + env["LIGHTCONE_WORKER_IMAGE"] = self._worker_image + else: + env.pop("LIGHTCONE_WORKER_IMAGE", None) + with _patch.dict("os.environ", env, clear=True): + return fn() + + def close(self) -> None: + pass + + +class _FakeGatewayCluster: + def __init__( + self, + name: str, + log: list[str], + resources: dict[str, float] | None, + worker_image: str | None = None, + ) -> None: + self.name = name + self.dashboard_link = f"http://hub/services/dask-gateway/{name}/status" + self._log = log + self._resources = resources + self._worker_image = worker_image + + def adapt(self, minimum: int, maximum: int) -> None: + self._log.append(f"adapt({minimum},{maximum})") + + def get_client(self) -> _FakeGatewayClient: + return _FakeGatewayClient(self._resources, self._worker_image) + + def shutdown(self) -> None: + self._log.append("shutdown") + + def close(self) -> None: + self._log.append("close") + + +class _FakeClusterReport: + def __init__(self, name: str) -> None: + self.name = name + + +#: A worker resource set satisfying the deployment contract: cpus AND +#: memory must be advertised together (either alone hangs some rules). +_GOOD_RESOURCES = {RESOURCE_CPUS: 2.0, RESOURCE_MEMORY: 4e9} + + +def _install_fake_gateway( + monkeypatch: pytest.MonkeyPatch, + log: list[str], + resources: dict[str, float] | None = None, + running: list[str] | None = None, + worker_image: str | None = None, + connect_error: Exception | None = None, + init_error: Exception | None = None, +): + """Inject a fake ``dask_gateway`` module; return it for inspection. + + *running* is the list of cluster names ``list_clusters()`` reports + (the user's running clusters). ``new_cluster`` is deliberately + absent from the fake: lc must never create Gateway clusters, and a + call would fail loudly with AttributeError. *connect_error* / + *init_error* simulate the real client's failure modes (stale + cluster name; no gateway address configured), which raise + non-RuntimeError types. + """ + import sys + import types + + class _FakeGateway: + def __init__(self) -> None: + if init_error is not None: + raise init_error + + def list_clusters(self) -> list[_FakeClusterReport]: + log.append("list_clusters()") + return [_FakeClusterReport(n) for n in (running or [])] + + def connect(self, name: str) -> _FakeGatewayCluster: + log.append(f"connect({name})") + if connect_error is not None: + raise connect_error + return _FakeGatewayCluster(name, log, resources, worker_image) + + module = types.ModuleType("dask_gateway") + module.Gateway = _FakeGateway # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "dask_gateway", module) + return module + + +def test_gateway_env_attaches_to_single_running_cluster( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ambient DASK_GATEWAY__ADDRESS (a hub pod) routes to the gateway + branch, which discovers the user's one running cluster and attaches — + zero configuration, no cluster creation, no scaling changes.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, log, resources=dict(_GOOD_RESOURCES), running=["hub.run1"] + ) + + with cluster_for_run() as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.run1"} + assert "list_clusters()" in log + assert "connect(hub.run1)" in log + assert not any(call.startswith("adapt") for call in log), ( + "attached clusters keep the user's scaling" + ) + + assert "shutdown" not in log, "lc never owns Gateway clusters" + assert "close" in log + + +def test_gateway_no_running_cluster_raises_with_instructions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Zero running clusters is an error whose message IS the UX: it must + tell the user to create one from JupyterLab (lc never creates).""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, running=[]) + + with pytest.raises(RuntimeError, match="No Dask Gateway cluster") as exc: + with cluster_for_run(): + pass # pragma: no cover + assert "new_cluster" in str(exc.value) # the copy-pasteable snippet + + +def test_gateway_no_cluster_error_names_expected_image( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, running=[]) + + with pytest.raises(RuntimeError, match=r"reg\.example/lc-proj:abc") as exc: + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc"): + pass # pragma: no cover + assert 'image="reg.example/lc-proj:abc"' in str(exc.value) + + +def test_gateway_multiple_clusters_raises_with_disambiguator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, running=["hub.a", "hub.b"]) + + with pytest.raises(RuntimeError, match="LIGHTCONE_GATEWAY_CLUSTER") as exc: + with cluster_for_run(): + pass # pragma: no cover + assert "hub.a" in str(exc.value) + assert "hub.b" in str(exc.value) + + +def test_gateway_wins_over_slurm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + monkeypatch.setenv("SLURM_JOB_ID", "12345") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, log, resources=dict(_GOOD_RESOURCES), running=["hub.run1"] + ) + + @contextmanager + def _should_not_run(*, verbose: bool, local_directory: str | None = None): + raise AssertionError("slurm path should not have been taken") + yield # pragma: no cover + + with patch("lightcone.engine.dask_cluster._slurm_backed_cluster", _should_not_run): + with cluster_for_run() as env: + assert "LIGHTCONE_GATEWAY_CLUSTER" in env + + +def test_explicit_address_wins_over_gateway( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DASK_SCHEDULER_ADDRESS", "tcp://existing:8786") + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://existing:8786"} + + +def test_gateway_named_cluster_skips_discovery( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIGHTCONE_GATEWAY_CLUSTER picks a cluster directly (the + several-clusters disambiguator) — no list_clusters round-trip, no + shutdown on exit.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) + + with cluster_for_run() as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.abc999"} + assert "connect(hub.abc999)" in log + assert "list_clusters()" not in log + assert not any(call.startswith("adapt") for call in log) + + assert "shutdown" not in log + assert "close" in log + + +def test_gateway_missing_client_raises_helpfully( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import sys + + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + monkeypatch.setitem(sys.modules, "dask_gateway", None) # forces ImportError + + with pytest.raises(RuntimeError, match=r"lightcone-cli\[gateway\]"): + with cluster_for_run(): + pass # pragma: no cover + + +def test_gateway_rejects_workers_without_resource_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Workers not advertising cpus/memory/gpus would hang every task with + no error — the branch must refuse loudly at startup instead.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, log, resources={}, running=["hub.run1"] + ) # live workers, no contract + + with pytest.raises(RuntimeError, match="resource contract"): + with cluster_for_run(): + pass # pragma: no cover + + assert "close" in log, "connections must be released on failure" + assert "shutdown" not in log, "the user's cluster must survive our failure" + + +def test_gateway_rejects_partial_resource_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """cpus without memory is the sneaky variant: rules schedule until + the first mem_mb rule, which then hangs forever. The startup check + must demand cpus AND memory together.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, log, resources={RESOURCE_CPUS: 2.0}, running=["hub.run1"] + ) + + with pytest.raises(RuntimeError, match="resource contract"): + with cluster_for_run(): + pass # pragma: no cover + + +def test_gateway_stale_named_cluster_raises_with_guidance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A stale LIGHTCONE_GATEWAY_CLUSTER (cluster since stopped) makes the + real client raise its own error types (ValueError/GatewayClusterError). + lc told the user to export that variable, so lc owns the remediation: + a RuntimeError naming the env var, not a bare dask_gateway traceback.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.gone42") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, log, connect_error=ValueError("Cluster 'hub.gone42' not found") + ) + + with pytest.raises(RuntimeError, match="LIGHTCONE_GATEWAY_CLUSTER") as exc: + with cluster_for_run(): + pass # pragma: no cover + assert "hub.gone42" in str(exc.value) + + +def test_gateway_unconfigured_client_raises_with_guidance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIGHTCONE_GATEWAY_CLUSTER lingering in a shell off-hub: Gateway() + itself raises ValueError (no gateway address configured). Must become + guidance to unset the variable, not a traceback.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, + log, + init_error=ValueError("No dask-gateway address provided"), + ) + + with pytest.raises(RuntimeError, match="unset") as exc: + with cluster_for_run(): + pass # pragma: no cover + assert "LIGHTCONE_GATEWAY_CLUSTER" in str(exc.value) + + +def test_gateway_zero_workers_is_not_an_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An adaptive cluster sitting at zero workers scales up once tasks + arrive — the contract check must only judge live workers.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources=None, running=["hub.run1"]) + + with cluster_for_run() as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.run1"} + + +def test_gateway_warns_on_worker_image_mismatch( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """The attached cluster runs a different image than the project's + container resolves to → loud warning naming both (manifests record + the actual image, so provenance stays truthful either way).""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, + log, + resources=dict(_GOOD_RESOURCES), + running=["hub.run1"], + worker_image="reg.example/lc-proj:stale00", + ) + + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): + pass + + out = capsys.readouterr().out + assert "reg.example/lc-proj:stale00" in out + assert "reg.example/lc-proj:abc123" in out + + +def test_gateway_matching_worker_image_is_silent( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, + log, + resources=dict(_GOOD_RESOURCES), + running=["hub.run1"], + worker_image="reg.example/lc-proj:abc123", + ) + + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): + pass + + assert "⚠" not in capsys.readouterr().out + + +def test_gateway_image_check_skipped_without_deployment_marker( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """A deployment that doesn't inject LIGHTCONE_WORKER_IMAGE can't be + verified — skip silently rather than warn on every run.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, + log, + resources=dict(_GOOD_RESOURCES), + running=["hub.run1"], + worker_image=None, + ) + + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): + pass + + assert "⚠" not in capsys.readouterr().out + + +def test_executor_connects_via_gateway_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The executor side of the rendezvous: LIGHTCONE_GATEWAY_CLUSTER set → + rejoin by name through the Gateway API, never a bare Client.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") + monkeypatch.delenv("DASK_SCHEDULER_ADDRESS", raising=False) + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) + + from snakemake_executor_plugin_dask.executor import _connect_client + + client, cluster = _connect_client() + assert "connect(hub.abc999)" in log + assert cluster is not None and cluster.name == "hub.abc999" + assert isinstance(client, _FakeGatewayClient) + + +def test_executor_address_wins_over_gateway_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The child must mirror the parent's branch priority: an explicit + scheduler address outranks a (possibly stale, shell-exported) + LIGHTCONE_GATEWAY_CLUSTER. Otherwise the child silently rejoins a + Gateway cluster the parent never verified while the scheduler the + parent reported sits idle.""" + monkeypatch.setenv("DASK_SCHEDULER_ADDRESS", "tcp://existing:8786") + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.stale99") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log) + + captured: dict[str, str] = {} + + class _FakeClient: + def __init__(self, addr: str) -> None: + captured["addr"] = addr + + monkeypatch.setattr("dask.distributed.Client", _FakeClient) + + from snakemake_executor_plugin_dask.executor import _connect_client + + client, cluster = _connect_client() + assert captured["addr"] == "tcp://existing:8786" + assert cluster is None + assert not any(c.startswith("connect") for c in log) + + +def test_executor_requires_some_connection_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DASK_SCHEDULER_ADDRESS", raising=False) + monkeypatch.delenv("LIGHTCONE_GATEWAY_CLUSTER", raising=False) + + from snakemake_interface_common.exceptions import WorkflowError + + from snakemake_executor_plugin_dask.executor import _connect_client + + with pytest.raises(WorkflowError, match="LIGHTCONE_GATEWAY_CLUSTER"): + _connect_client() + + @pytest.mark.slow def test_local_cluster_smoke() -> None: """End-to-end: a real LocalCluster spins up, accepts a task, tears down.""" diff --git a/tests/test_dask_plugin.py b/tests/test_dask_plugin.py index ba1ad686..172f8c64 100644 --- a/tests/test_dask_plugin.py +++ b/tests/test_dask_plugin.py @@ -12,7 +12,9 @@ from snakemake_executor_plugin_dask.executor import ( _build_resources, + _emit_block, _run_shell, + _unpack_result, ) @@ -21,13 +23,48 @@ def _job(threads: int = 1, **resources: float) -> SimpleNamespace: def test_run_shell_propagates_exit_code() -> None: - assert _run_shell("true") == 0 - assert _run_shell("false") != 0 + rc, _ = _run_shell("true") + assert rc == 0 + rc, _ = _run_shell("false") + assert rc != 0 def test_run_shell_runs_under_shell() -> None: """We rely on shell=True so recipes can use pipes and env expansion.""" - assert _run_shell("echo hi | grep hi >/dev/null") == 0 + rc, _ = _run_shell("echo hi | grep hi >/dev/null") + assert rc == 0 + + +def test_run_shell_returns_failure_tail() -> None: + rc, block = _run_shell("echo boom >&2; exit 2") + assert rc == 2 + assert "✗ worker-side snakemake exited 2" in block + assert "boom" in block + + +def test_run_shell_forwards_sentinel_lines() -> None: + from lightcone.engine.runner import SENTINEL + + rc, block = _run_shell(f"echo '{SENTINEL}hello'; echo noise") + assert rc == 0 + assert "hello" in block + assert "noise" not in block + + +def test_emit_block_writes_to_stdout(capsys) -> None: # type: ignore[no-untyped-def] + _emit_block("hi\n") + assert capsys.readouterr().out == "hi\n" + + _emit_block("") + assert capsys.readouterr().out == "" + + +def test_unpack_result_accepts_tuple_and_legacy_int() -> None: + # Current worker: (exit_code, block). + assert _unpack_result((2, "boom\n")) == (2, "boom\n") + # Legacy worker (older image) returns a bare int — must not crash. + assert _unpack_result(0) == (0, "") + assert _unpack_result(1) == (1, "") def test_build_resources_default_uses_threads() -> None: diff --git a/tests/test_site_registry.py b/tests/test_site_registry.py index 61396e1a..2d339059 100644 --- a/tests/test_site_registry.py +++ b/tests/test_site_registry.py @@ -95,3 +95,34 @@ def test_unknown_host_get_returns_default( # returning an empty HostSite rather than None. fake_hostname("generic-laptop") assert detect_current_site().get("scratch_root", "/tmp") == "/tmp" + + def test_env_markers_detect_hub_despite_generated_hostname( + self, + fake_hostname: Callable[[str], None], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # JupyterHub pods have generated hostnames (jupyter-), so + # the hub site is matched by ambient env instead. + fake_hostname("jupyter-eiffl") + monkeypatch.setenv("JUPYTERHUB_USER", "eiffl") + monkeypatch.setenv( + "DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/" + ) + site = detect_current_site() + assert site.key == "lightcone-hub" + # Kubernetes (the worker pod) is the container runtime on the + # hub — recipes run unwrapped but are still containerized. + assert site.get("container_runtime") == "kubernetes" + assert site.get("backend") == "dask-gateway" + + def test_partial_env_markers_do_not_match( + self, + fake_hostname: Callable[[str], None], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # JUPYTERHUB_USER alone (any JupyterHub, no gateway) must not + # claim the lightcone-hub site. + fake_hostname("jupyter-somebody") + monkeypatch.delenv("DASK_GATEWAY__ADDRESS", raising=False) + monkeypatch.setenv("JUPYTERHUB_USER", "somebody") + assert not detect_current_site() From 56d5f8fb9d22d9d5da13cc5251f57e01fd0340c8 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 10:17:06 +0200 Subject: [PATCH 2/9] Gateway lifecycle per PRD: create run-scoped clusters, cull on exit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gateway branch now mirrors the local/SLURM lifecycle instead of being attach-only: `lc run` creates a Gateway cluster with the project's worker image, scales it adaptively (1..--jobs), waits for the first worker (fail fast on unpullable images instead of hanging at zero workers), and shuts the cluster down when the run finishes — resolved decision #1 of the JupyterHub/Kubernetes PRD, and a hard requirement of the native-k8s model where picking up a freshly built image needs a fresh cluster. LIGHTCONE_GATEWAY_CLUSTER= remains as the attach escape hatch for long-lived clusters: connect, never rescale, leave running on exit; the image-mismatch warning now only applies there (created clusters run exactly the requested image). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- src/lightcone/engine/dask_cluster.py | 280 ++++++++++++++++++--------- tests/test_dask_cluster.py | 206 ++++++++++++-------- 2 files changed, 320 insertions(+), 166 deletions(-) diff --git a/src/lightcone/engine/dask_cluster.py b/src/lightcone/engine/dask_cluster.py index 67bb9c89..26bf94fa 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -7,16 +7,18 @@ the cluster, so we don't tear it down. - A Dask Gateway environment is detected (``LIGHTCONE_GATEWAY_CLUSTER`` or ``DASK_GATEWAY__ADDRESS`` is set, e.g. on a JupyterHub deployment) → - **attach** to the user's running Gateway cluster. ``lc run`` never - creates Gateway clusters: the user starts one from JupyterLab (Dask - sidebar or a notebook, where the options widget offers image/cores/ - memory) and lc discovers it through the user-scoped Gateway API — - exactly one running cluster attaches unambiguously; zero or several - is an error naming the fix. Gateway scheduler addresses use a custom - ``gateway://`` comm scheme that a bare ``distributed.Client`` cannot - dial, so the child snakemake process is told the *cluster name* via - ``LIGHTCONE_GATEWAY_CLUSTER`` and rejoins through the Gateway API — - the same rendezvous-by-ambient-context pattern the SLURM branch uses. + **create** a run-scoped Gateway cluster with the project's worker image + and shut it down when the run finishes — the same create/cull lifecycle + the local and SLURM branches have always had, which is also what lets + every run pick up a freshly built image (a Gateway cluster's image is + fixed at creation). Setting ``LIGHTCONE_GATEWAY_CLUSTER=`` + instead **attaches** to that existing cluster and leaves it running on + exit — the escape hatch for iterating against a long-lived cluster. + Gateway scheduler addresses use a custom ``gateway://`` comm scheme + that a bare ``distributed.Client`` cannot dial, so the child snakemake + process is told the *cluster name* via ``LIGHTCONE_GATEWAY_CLUSTER`` + and rejoins through the Gateway API — the same rendezvous-by-ambient- + context pattern the SLURM branch uses. - ``SLURM_JOB_ID`` is set → start an in-process scheduler via ``LocalCluster(n_workers=0)``, then ``srun`` one ``dask worker`` per node across the allocation. Workers advertise the node's full resources; @@ -26,9 +28,11 @@ Except on the Gateway branch (where the Gateway server owns scheduling), the scheduler is always in-process (driven by ``lc run`` itself) so its lifetime equals the run's lifetime — no service to manage, no orphaned -schedulers if the driver crashes. The Gateway branch mirrors the -``DASK_SCHEDULER_ADDRESS`` convention: we attach to a cluster someone -else owns, so we leave it running (and its scaling untouched) on exit. +schedulers if the driver crashes. On the Gateway branch the same +lifetime contract is enforced server-side: clusters lc creates are shut +down on exit (and reaped by the deployment's idle timeout if lc dies +uncleanly); clusters the user named are left untouched, mirroring the +``DASK_SCHEDULER_ADDRESS`` convention. """ from __future__ import annotations @@ -49,14 +53,21 @@ RESOURCE_MEMORY = "memory" RESOURCE_GPUS = "gpus" -#: Env var carrying a Dask Gateway cluster name. Set by the user to pick -#: one of several running Gateway clusters (discovery attaches -#: automatically when exactly one is up); set by :func:`cluster_for_run` -#: for the child snakemake process so the executor plugin can rejoin the +#: Env var carrying a Dask Gateway cluster name. Set by the user to +#: attach to an existing long-lived cluster instead of the default +#: create-per-run lifecycle; set by :func:`cluster_for_run` for the +#: child snakemake process so the executor plugin can rejoin the #: cluster through the Gateway API (a bare ``Client`` cannot dial #: ``gateway://``). GATEWAY_CLUSTER_ENV = "LIGHTCONE_GATEWAY_CLUSTER" +#: Env var bounding how long the Gateway branch waits for the first +#: worker of a cluster it created (seconds; default 600). The wait +#: exists to fail fast with a real error when the requested image +#: cannot be pulled or the pool cannot schedule a worker — without it +#: the run would sit silently at zero workers forever. +GATEWAY_WORKER_TIMEOUT_ENV = "LIGHTCONE_GATEWAY_WORKER_TIMEOUT" + #: Env var carrying the image a Gateway worker pod was started with. A #: lightcone-hub deployment's cluster-options handler injects it into #: every scheduler/worker pod; :func:`cluster_for_run` reads it back to @@ -137,6 +148,7 @@ def cluster_for_run( verbose: bool = False, local_directory: str | None = None, expected_worker_image: str | None = None, + max_workers: int | None = None, ) -> Iterator[dict[str, str]]: """Yield the env overlay the child snakemake needs to reach the cluster. @@ -154,10 +166,16 @@ def cluster_for_run( file I/O is slow and can pressure the gateway nodes). *expected_worker_image*, when given, is the image the project's - declared container resolves to on this deployment; the Gateway - branch compares it against what the attached cluster actually runs - and warns on mismatch. Ignored by the other branches (they realize - containers by wrapping recipes, not via pod images). + declared container resolves to on this deployment. A Gateway cluster + lc creates is started with exactly this image; an attached cluster + is compared against it with a warning on mismatch. Ignored by the + other branches (they realize containers by wrapping recipes, not + via pod images). + + *max_workers* bounds the adaptive scaling of a Gateway cluster lc + creates (``lc run`` passes its ``--jobs`` value — there is never a + reason to hold more workers than dispatchable rules). Ignored + everywhere else. """ if addr := os.environ.get("DASK_SCHEDULER_ADDRESS"): if verbose: @@ -167,7 +185,9 @@ def cluster_for_run( if GATEWAY_CLUSTER_ENV in os.environ or "DASK_GATEWAY__ADDRESS" in os.environ: with _gateway_cluster( - verbose=verbose, expected_worker_image=expected_worker_image + verbose=verbose, + expected_worker_image=expected_worker_image, + max_workers=max_workers, ) as name: yield {GATEWAY_CLUSTER_ENV: name} return @@ -187,19 +207,28 @@ def cluster_for_run( @contextmanager def _gateway_cluster( - *, verbose: bool, expected_worker_image: str | None + *, + verbose: bool, + expected_worker_image: str | None, + max_workers: int | None, ) -> Iterator[str]: - """Attach to the user's running Dask Gateway cluster; yield its name. - - Attach-only by design: cluster lifecycle belongs to the user (create - one from the JupyterLab Dask sidebar or a notebook — that's where - the image/cores/memory options widget lives, and where the dashboard - is already docked). The Gateway API is user-scoped under JupyterHub - auth, so ``list_clusters()`` sees only the caller's clusters — - exactly one running cluster attaches with zero configuration; zero - or several raises with the fix spelled out. We never touch the - cluster's scaling and leave it running on exit, mirroring the - ``DASK_SCHEDULER_ADDRESS`` convention. + """Create (or attach to) a Dask Gateway cluster; yield its name. + + Default lifecycle — create and cull, per run: a new cluster is + created with the project's worker image, scaled adaptively between + one worker and *max_workers*, and shut down when the run finishes. + This mirrors the local/SLURM branches (compute lives exactly as + long as the run) and is what makes image updates seamless: a + Gateway cluster's image is fixed at creation, so picking up a + freshly built project image *requires* a fresh cluster. + + Attach — when ``LIGHTCONE_GATEWAY_CLUSTER`` names a cluster: connect + to it, leave its scaling alone, and leave it running on exit + (mirroring the ``DASK_SCHEDULER_ADDRESS`` convention). This is the + escape hatch for iterating repeatedly against one long-lived + cluster (create it from the JupyterLab Dask sidebar or a notebook + with ``shutdown_on_close=False``); the cost is that lc can only + warn — not fix — when its image is stale. The Gateway client is configured entirely by ambient dask config — on a lightcone JupyterHub deployment the ``DASK_GATEWAY__*`` env @@ -230,25 +259,138 @@ def _gateway_cluster( ) from exc named = os.environ.get(GATEWAY_CLUSTER_ENV) - name = named or _discover_cluster_name( - gateway, expected_worker_image=expected_worker_image - ) + if named: + with _attached_gateway_cluster( + gateway, + name=named, + verbose=verbose, + expected_worker_image=expected_worker_image, + ) as name: + yield name + return + + with _created_gateway_cluster( + gateway, + verbose=verbose, + worker_image=expected_worker_image, + max_workers=max_workers, + ) as name: + yield name + + +@contextmanager +def _created_gateway_cluster( + gateway: object, + *, + verbose: bool, + worker_image: str | None, + max_workers: int | None, +) -> Iterator[str]: + """Create a run-scoped Gateway cluster; shut it down on exit.""" + options: dict[str, object] = {} + if worker_image: + # `image` is a server-side cluster option (declared by the + # deployment's option handler). Passing it as a kwarg merges it + # into the cluster options; deployments that don't declare it + # reject the request — surfaced below with guidance. + options["image"] = worker_image try: - cluster = gateway.connect(name) + cluster = gateway.new_cluster( # type: ignore[attr-defined] + shutdown_on_close=True, **options + ) + except Exception as exc: + detail = ( + f" (requested image={worker_image!r} — if the deployment " + "does not expose an `image` cluster option, ask the hub " + "admin to add it to the gateway's cluster-options handler)" + if worker_image + else "" + ) + raise RuntimeError( + f"Could not create a Dask Gateway cluster ({exc}){detail}." + ) from exc + + bound = max(1, max_workers or 1) + if verbose: + image_note = f" with image {worker_image}" if worker_image else "" + print( + f"→ Created Dask Gateway cluster {cluster.name}{image_note}; " + f"scaling adaptively up to {bound} worker(s) " + f"(dashboard: {cluster.dashboard_link})" + ) + try: + cluster.adapt(minimum=1, maximum=bound) + client = cluster.get_client() + try: + _wait_first_worker(client, image=worker_image) + _assert_worker_resources(client) + finally: + client.close() + yield str(cluster.name) + finally: + # We created it, we cull it — the PRD lifecycle. shutdown() + # stops the cluster server-side; if lc dies before reaching + # this, shutdown_on_close and the deployment's idle timeout + # are the backstops. + try: + cluster.shutdown() + except Exception: + cluster.close() + if verbose: + print(f"→ Shut down Dask Gateway cluster {cluster.name}") + + +def _wait_first_worker(client: object, *, image: str | None) -> None: + """Block until the created cluster has one live worker. + + An unpullable image or an unschedulable pool otherwise leaves the + run sitting at zero workers with no error at all — the classic + silent-hang failure mode. Bounded by + :data:`GATEWAY_WORKER_TIMEOUT_ENV` (default 600 s: a first-time + image pull on a fresh node is minutes, not seconds). + """ + try: + timeout = int(os.environ.get(GATEWAY_WORKER_TIMEOUT_ENV) or 600) + except ValueError: + timeout = 600 + try: + client.wait_for_workers(n_workers=1, timeout=timeout) # type: ignore[attr-defined] + except Exception as exc: + image_hint = ( + f"the worker image ({image}) cannot be pulled, " + if image + else "the worker image cannot be pulled, " + ) + raise RuntimeError( + f"No Dask Gateway worker became ready within {timeout}s " + f"({exc}). Likely causes: {image_hint}the node pool cannot " + "schedule a worker (capacity/quota), or the cluster was " + "culled. Check the JupyterLab Dask panel for the cluster's " + f"state, or raise {GATEWAY_WORKER_TIMEOUT_ENV}." + ) from exc + + +@contextmanager +def _attached_gateway_cluster( + gateway: object, + *, + name: str, + verbose: bool, + expected_worker_image: str | None, +) -> Iterator[str]: + """Attach to the named running Gateway cluster; leave it running.""" + try: + cluster = gateway.connect(name) # type: ignore[attr-defined] except Exception as exc: # dask-gateway raises its own error types (ValueError, # GatewayClusterError) for a missing/stopped cluster; translate # into guidance, because a *stale* LIGHTCONE_GATEWAY_CLUSTER is # the likely cause — we told users to export it. - hint = ( - f"unset {GATEWAY_CLUSTER_ENV} to let lc discover your " - "running cluster, or point it at one that is running" - if named - else "it may have just stopped — check the JupyterLab Dask panel" - ) raise RuntimeError( f"Could not connect to Dask Gateway cluster {name!r} " - f"({exc}); {hint}." + f"({exc}); unset {GATEWAY_CLUSTER_ENV} to let lc create a " + "run-scoped cluster, or point it at a cluster that is " + "running." ) from exc if verbose: @@ -268,55 +410,13 @@ def _gateway_cluster( _check_worker_image(client, expected=expected_worker_image) finally: client.close() - yield cluster.name + yield str(cluster.name) finally: # Releases this process's connections only — the cluster (and # its scaling) belongs to the user. cluster.close() -def _discover_cluster_name( - gateway: object, expected_worker_image: str | None -) -> str: - """Name of the caller's single running Gateway cluster. - - Raises with actionable instructions when there isn't exactly one: - creation is deliberately not offered here (the user creates clusters - from JupyterLab), so the error message *is* the UX — it must say - precisely what to do next. - """ - reports = gateway.list_clusters() # type: ignore[attr-defined] - if len(reports) == 1: - return str(reports[0].name) - - if not reports: - image_arg = ( - f'image="{expected_worker_image}", ' if expected_worker_image else "" - ) - # shutdown_on_close=False matters: without it the cluster dies - # with the creating kernel/process, killing any lc run attached - # to it. The deployment's idle_timeout reaps forgotten clusters. - raise RuntimeError( - "No Dask Gateway cluster is running for your user. Create one " - "first — from the JupyterLab Dask sidebar (+ NEW), or in a " - "notebook:\n\n" - " from dask_gateway import Gateway\n" - f" cluster = Gateway().new_cluster({image_arg}" - "shutdown_on_close=False)\n" - " cluster.adapt(minimum=1, maximum=8)\n\n" - "then re-run `lc run` (it attaches to your running cluster " - "and leaves it up)." - ) - - names = ", ".join(str(r.name) for r in reports) - raise RuntimeError( - f"Multiple Dask Gateway clusters are running for your user " - f"({names}). Pick one by setting {GATEWAY_CLUSTER_ENV}, e.g.:\n\n" - f" export {GATEWAY_CLUSTER_ENV}={reports[0].name}\n\n" - "or shut the extras down from the JupyterLab Dask sidebar." - ) - - def _check_worker_image(client: object, *, expected: str | None) -> None: """Warn when the attached cluster doesn't run the project's image. @@ -348,7 +448,9 @@ def _check_worker_image(client: object, *, expected: str | None) -> None: f" {expected}\n" f" Recipes will execute in the cluster's image; manifests " f"record what actually ran.\n" - f" To use the project image, create a new cluster with " + f" To run on the project image, unset {GATEWAY_CLUSTER_ENV} " + f"and let lc create a run-scoped\n" + f" cluster, or recreate your cluster with " f'image="{expected}".' ) diff --git a/tests/test_dask_cluster.py b/tests/test_dask_cluster.py index 904fece2..d4a6eaa8 100644 --- a/tests/test_dask_cluster.py +++ b/tests/test_dask_cluster.py @@ -289,15 +289,21 @@ def __init__( self, resources: dict[str, float] | None, worker_image: str | None = None, + wait_error: Exception | None = None, ) -> None: self._resources = resources self._worker_image = worker_image + self._wait_error = wait_error def scheduler_info(self) -> dict[str, object]: if self._resources is None: return {"workers": {}} return {"workers": {"w0": {"resources": self._resources}}} + def wait_for_workers(self, n_workers: int, timeout: int) -> None: + if self._wait_error is not None: + raise self._wait_error + def run_on_scheduler(self, fn): # noqa: ANN001, ANN202 # The real client executes fn inside the scheduler pod, where the # deployment injects LIGHTCONE_WORKER_IMAGE. Simulate that env. @@ -323,18 +329,22 @@ def __init__( log: list[str], resources: dict[str, float] | None, worker_image: str | None = None, + wait_error: Exception | None = None, ) -> None: self.name = name self.dashboard_link = f"http://hub/services/dask-gateway/{name}/status" self._log = log self._resources = resources self._worker_image = worker_image + self._wait_error = wait_error def adapt(self, minimum: int, maximum: int) -> None: self._log.append(f"adapt({minimum},{maximum})") def get_client(self) -> _FakeGatewayClient: - return _FakeGatewayClient(self._resources, self._worker_image) + return _FakeGatewayClient( + self._resources, self._worker_image, self._wait_error + ) def shutdown(self) -> None: self._log.append("shutdown") @@ -343,11 +353,6 @@ def close(self) -> None: self._log.append("close") -class _FakeClusterReport: - def __init__(self, name: str) -> None: - self.name = name - - #: A worker resource set satisfying the deployment contract: cpus AND #: memory must be advertised together (either alone hangs some rules). _GOOD_RESOURCES = {RESOURCE_CPUS: 2.0, RESOURCE_MEMORY: 4e9} @@ -357,20 +362,23 @@ def _install_fake_gateway( monkeypatch: pytest.MonkeyPatch, log: list[str], resources: dict[str, float] | None = None, - running: list[str] | None = None, worker_image: str | None = None, connect_error: Exception | None = None, init_error: Exception | None = None, + new_cluster_error: Exception | None = None, + wait_error: Exception | None = None, ): """Inject a fake ``dask_gateway`` module; return it for inspection. - *running* is the list of cluster names ``list_clusters()`` reports - (the user's running clusters). ``new_cluster`` is deliberately - absent from the fake: lc must never create Gateway clusters, and a - call would fail loudly with AttributeError. *connect_error* / - *init_error* simulate the real client's failure modes (stale - cluster name; no gateway address configured), which raise - non-RuntimeError types. + ``new_cluster`` logs the options it was called with and returns a + fresh cluster named ``hub.new1`` — the create/cull default path. + ``connect`` returns a cluster with the given name — the attach + path. *connect_error* / *init_error* / *new_cluster_error* simulate + the real client's failure modes (stale cluster name; no gateway + address configured; rejected cluster options), which raise + non-RuntimeError types. *wait_error* makes the first + ``wait_for_workers`` call fail (unpullable image / unschedulable + pool). """ import sys import types @@ -380,9 +388,19 @@ def __init__(self) -> None: if init_error is not None: raise init_error - def list_clusters(self) -> list[_FakeClusterReport]: - log.append("list_clusters()") - return [_FakeClusterReport(n) for n in (running or [])] + def new_cluster( + self, shutdown_on_close: bool = False, **options: object + ) -> _FakeGatewayCluster: + log.append( + "new_cluster(" + + ",".join(f"{k}={v}" for k, v in sorted(options.items())) + + f";shutdown_on_close={shutdown_on_close})" + ) + if new_cluster_error is not None: + raise new_cluster_error + return _FakeGatewayCluster( + "hub.new1", log, resources, worker_image, wait_error + ) def connect(self, name: str) -> _FakeGatewayCluster: log.append(f"connect({name})") @@ -396,79 +414,103 @@ def connect(self, name: str) -> _FakeGatewayCluster: return module -def test_gateway_env_attaches_to_single_running_cluster( +def test_gateway_env_creates_run_scoped_cluster( monkeypatch: pytest.MonkeyPatch, ) -> None: """Ambient DASK_GATEWAY__ADDRESS (a hub pod) routes to the gateway - branch, which discovers the user's one running cluster and attaches — - zero configuration, no cluster creation, no scaling changes.""" + branch, which creates a run-scoped cluster, scales it adaptively up + to max_workers, and shuts it down when the run finishes — the PRD + create/cull lifecycle.""" monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway( - monkeypatch, log, resources=dict(_GOOD_RESOURCES), running=["hub.run1"] - ) + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) - with cluster_for_run() as env: - assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.run1"} - assert "list_clusters()" in log - assert "connect(hub.run1)" in log - assert not any(call.startswith("adapt") for call in log), ( - "attached clusters keep the user's scaling" - ) + with cluster_for_run(max_workers=8) as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.new1"} + assert "new_cluster(;shutdown_on_close=True)" in log + assert "adapt(1,8)" in log + assert "shutdown" not in log - assert "shutdown" not in log, "lc never owns Gateway clusters" - assert "close" in log + assert "shutdown" in log, "run-scoped clusters are culled on exit" -def test_gateway_no_running_cluster_raises_with_instructions( +def test_gateway_creates_with_project_image( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Zero running clusters is an error whose message IS the UX: it must - tell the user to create one from JupyterLab (lc never creates).""" + """The project's resolved worker image is passed as the `image` + cluster option — the whole point of create-per-run: every run picks + up the freshly built image.""" monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway(monkeypatch, log, running=[]) + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) - with pytest.raises(RuntimeError, match="No Dask Gateway cluster") as exc: - with cluster_for_run(): - pass # pragma: no cover - assert "new_cluster" in str(exc.value) # the copy-pasteable snippet + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): + pass + + assert ( + "new_cluster(image=reg.example/lc-proj:abc123;shutdown_on_close=True)" + in log + ) -def test_gateway_no_cluster_error_names_expected_image( +def test_gateway_creation_failure_names_the_image_option( monkeypatch: pytest.MonkeyPatch, ) -> None: + """A deployment without an `image` cluster option rejects creation + with a dask-gateway error type; the message must say what to fix.""" monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway(monkeypatch, log, running=[]) + _install_fake_gateway( + monkeypatch, + log, + new_cluster_error=ValueError("Unknown fields ['image']"), + ) - with pytest.raises(RuntimeError, match=r"reg\.example/lc-proj:abc") as exc: + with pytest.raises(RuntimeError, match="image") as exc: with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc"): pass # pragma: no cover - assert 'image="reg.example/lc-proj:abc"' in str(exc.value) + assert "cluster option" in str(exc.value) -def test_gateway_multiple_clusters_raises_with_disambiguator( +def test_gateway_default_bound_is_one_worker( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway(monkeypatch, log, running=["hub.a", "hub.b"]) + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) - with pytest.raises(RuntimeError, match="LIGHTCONE_GATEWAY_CLUSTER") as exc: - with cluster_for_run(): + with cluster_for_run(): + pass + + assert "adapt(1,1)" in log + + +def test_gateway_worker_wait_timeout_raises_and_culls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No worker within the deadline (unpullable image, no capacity) must + fail with guidance — and still shut the created cluster down.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway( + monkeypatch, + log, + resources=dict(_GOOD_RESOURCES), + wait_error=TimeoutError("timed out"), + ) + + with pytest.raises(RuntimeError, match="worker became ready") as exc: + with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc"): pass # pragma: no cover - assert "hub.a" in str(exc.value) - assert "hub.b" in str(exc.value) + assert "reg.example/lc-proj:abc" in str(exc.value) + assert "shutdown" in log, "a failed creation must not leak the cluster" def test_gateway_wins_over_slurm(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") monkeypatch.setenv("SLURM_JOB_ID", "12345") log: list[str] = [] - _install_fake_gateway( - monkeypatch, log, resources=dict(_GOOD_RESOURCES), running=["hub.run1"] - ) + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) @contextmanager def _should_not_run(*, verbose: bool, local_directory: str | None = None): @@ -490,12 +532,12 @@ def test_explicit_address_wins_over_gateway( assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://existing:8786"} -def test_gateway_named_cluster_skips_discovery( +def test_gateway_named_cluster_attaches_and_leaves_running( monkeypatch: pytest.MonkeyPatch, ) -> None: - """LIGHTCONE_GATEWAY_CLUSTER picks a cluster directly (the - several-clusters disambiguator) — no list_clusters round-trip, no - shutdown on exit.""" + """LIGHTCONE_GATEWAY_CLUSTER= is the escape hatch for a + long-lived cluster: attach to it, never create, never rescale, + leave it running on exit.""" monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") log: list[str] = [] _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) @@ -503,10 +545,10 @@ def test_gateway_named_cluster_skips_discovery( with cluster_for_run() as env: assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.abc999"} assert "connect(hub.abc999)" in log - assert "list_clusters()" not in log + assert not any(call.startswith("new_cluster") for call in log) assert not any(call.startswith("adapt") for call in log) - assert "shutdown" not in log + assert "shutdown" not in log, "attached clusters belong to the user" assert "close" in log @@ -527,12 +569,27 @@ def test_gateway_rejects_workers_without_resource_contract( monkeypatch: pytest.MonkeyPatch, ) -> None: """Workers not advertising cpus/memory/gpus would hang every task with - no error — the branch must refuse loudly at startup instead.""" + no error — the branch must refuse loudly at startup instead (and + still cull the cluster it just created).""" monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway( - monkeypatch, log, resources={}, running=["hub.run1"] - ) # live workers, no contract + _install_fake_gateway(monkeypatch, log, resources={}) # workers, no contract + + with pytest.raises(RuntimeError, match="resource contract"): + with cluster_for_run(): + pass # pragma: no cover + + assert "shutdown" in log, "a failed creation must not leak the cluster" + + +def test_gateway_attach_contract_failure_leaves_cluster_running( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """On the attach path the same contract failure must NOT shut the + user's cluster down — it isn't ours.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources={}) with pytest.raises(RuntimeError, match="resource contract"): with cluster_for_run(): @@ -550,9 +607,7 @@ def test_gateway_rejects_partial_resource_contract( must demand cpus AND memory together.""" monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") log: list[str] = [] - _install_fake_gateway( - monkeypatch, log, resources={RESOURCE_CPUS: 2.0}, running=["hub.run1"] - ) + _install_fake_gateway(monkeypatch, log, resources={RESOURCE_CPUS: 2.0}) with pytest.raises(RuntimeError, match="resource contract"): with cluster_for_run(): @@ -598,14 +653,14 @@ def test_gateway_unconfigured_client_raises_with_guidance( assert "LIGHTCONE_GATEWAY_CLUSTER" in str(exc.value) -def test_gateway_zero_workers_is_not_an_error( +def test_gateway_attached_zero_workers_is_not_an_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - """An adaptive cluster sitting at zero workers scales up once tasks - arrive — the contract check must only judge live workers.""" - monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + """An attached adaptive cluster sitting at zero workers scales up once + tasks arrive — the contract check must only judge live workers.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") log: list[str] = [] - _install_fake_gateway(monkeypatch, log, resources=None, running=["hub.run1"]) + _install_fake_gateway(monkeypatch, log, resources=None) with cluster_for_run() as env: assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.run1"} @@ -618,13 +673,12 @@ def test_gateway_warns_on_worker_image_mismatch( """The attached cluster runs a different image than the project's container resolves to → loud warning naming both (manifests record the actual image, so provenance stays truthful either way).""" - monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") log: list[str] = [] _install_fake_gateway( monkeypatch, log, resources=dict(_GOOD_RESOURCES), - running=["hub.run1"], worker_image="reg.example/lc-proj:stale00", ) @@ -640,13 +694,12 @@ def test_gateway_matching_worker_image_is_silent( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") log: list[str] = [] _install_fake_gateway( monkeypatch, log, resources=dict(_GOOD_RESOURCES), - running=["hub.run1"], worker_image="reg.example/lc-proj:abc123", ) @@ -662,13 +715,12 @@ def test_gateway_image_check_skipped_without_deployment_marker( ) -> None: """A deployment that doesn't inject LIGHTCONE_WORKER_IMAGE can't be verified — skip silently rather than warn on every run.""" - monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") log: list[str] = [] _install_fake_gateway( monkeypatch, log, resources=dict(_GOOD_RESOURCES), - running=["hub.run1"], worker_image=None, ) From cd64f9468d3ffc033e9af83bc18a3bc920d500ac Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 10:17:06 +0200 Subject: [PATCH 3/9] lc build on the hub: real image builds through the BinderHub service MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LCR-176 Part B. New engine/binder.py drives the deployment's binderhub-service (API-only mode) from inside a user pod: - Environment ref: the last commit touching the env-defining files (Containerfile + dependency files + named COPY sources; a whole-tree COPY . is excluded — code reaches workers via the shared home, the image is the environment). Code-only commits reuse the image. - Dirty env files are committed (scoped to those paths; --no-commit to refuse instead) and pushed — build pods clone from the remote. - A root 'Dockerfile -> Containerfile' symlink is created and committed when missing (repo2docker only recognizes Dockerfile). - The build endpoint is streamed (SSE) until ready/failed; BinderHub checks its registry first, so an already-built ref is a single round-trip — cheap enough that lc run calls it every time. lc run on a kubernetes-runtime site now ensures the worker image through the same path before creating its Gateway cluster, completing the PRD user story: lc run → image up to date (rebuilt if necessary) → cluster started with it → pipeline runs → cluster culled. Off-hub kubernetes falls back to the passive registry report/verification as before. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- src/lightcone/cli/commands.py | 169 ++++++++++-- src/lightcone/engine/binder.py | 417 ++++++++++++++++++++++++++++++ src/lightcone/engine/container.py | 30 +++ tests/conftest.py | 6 + tests/test_binder.py | 363 ++++++++++++++++++++++++++ tests/test_cli.py | 150 +++++++++++ tests/test_container.py | 38 +++ 7 files changed, 1158 insertions(+), 15 deletions(-) create mode 100644 src/lightcone/engine/binder.py create mode 100644 tests/test_binder.py diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index e52788f0..3a523317 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -25,6 +25,7 @@ import shutil import subprocess import sys +from collections.abc import Callable from pathlib import Path import click @@ -540,10 +541,13 @@ def run( """Materialize outputs declared in astra.yaml. Always dispatches through a Dask cluster: a ``LocalCluster`` on a - workstation, srun-launched workers inside a SLURM allocation, the - user's running Dask Gateway cluster on a JupyterHub deployment - (created from JupyterLab; lc attaches, never creates), or an - existing scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. + workstation, srun-launched workers inside a SLURM allocation, a + run-scoped Dask Gateway cluster on a JupyterHub deployment (created + with the project's worker image — kept up to date through the hub's + build service — and culled when the run finishes), or an existing + scheduler if ``DASK_SCHEDULER_ADDRESS`` is set + (``LIGHTCONE_GATEWAY_CLUSTER=`` similarly attaches to an + existing Gateway cluster and leaves it running). """ _abort_on_perlmutter_login() @@ -578,12 +582,14 @@ def run( choice = load_runtime(project_path=project) _ensure_images(project, runtime=choice.runtime) - # On a kubernetes-runtime site the worker pod is the container, so - # resolve which image the Gateway cluster is expected to run — the - # gateway branch verifies the attached cluster against it and the - # no-cluster-yet error message names it. + # On a kubernetes-runtime site the worker pod is the container: + # resolve which image the Gateway cluster must run. On a hub this + # ensures the image is up to date first (committing/pushing env + # changes and driving a BinderHub build when needed — a fast no-op + # round-trip when the registry already has it); the run-scoped + # cluster is then created with exactly that image. expected_worker_image = ( - _expected_worker_image(project) + _worker_image_for_run(project, verbose=verbose) if choice.runtime == KUBERNETES_RUNTIME else None ) @@ -661,6 +667,7 @@ def run( verbose=verbose, local_directory=str(rundirs.dask_local), expected_worker_image=expected_worker_image, + max_workers=int(n), ) as cluster_env: _warn_runtime_cluster_mismatch( project, runtime=choice.runtime, cluster_env=cluster_env @@ -965,7 +972,15 @@ def verify(universe: str | None) -> None: "(overrides ~/.lightcone/config.yaml)" ), ) -def build(force: bool, runtime: str | None) -> None: +@click.option( + "--no-commit", + is_flag=True, + help=( + "On a hub: fail instead of auto-committing environment-file " + "changes before the image build" + ), +) +def build(force: bool, runtime: str | None, no_commit: bool) -> None: """Build container images declared in astra.yaml. Containerfile syntax is Dockerfile syntax — we use ``docker``, @@ -976,10 +991,12 @@ def build(force: bool, runtime: str | None) -> None: ``lc run`` time. On a kubernetes-runtime site (a lightcone-hub deployment) there is - no local image store and nothing to build in-pod: ``lc build`` - instead checks that each Containerfile's content-addressed image - exists in the deployment registry and prints exactly how to publish - it when missing. + no docker in-pod: ``lc build`` instead commits any environment-file + changes, pushes them, and drives an image build through the hub's + BinderHub service into the deployment registry — the image a + Gateway cluster will run. Without a reachable build service it + falls back to probing the registry and printing how to publish the + image off-hub. """ from lightcone.engine.container import ( KUBERNETES_RUNTIME, @@ -994,7 +1011,12 @@ def build(force: bool, runtime: str | None) -> None: raise click.ClickException(str(e)) if resolved_runtime == KUBERNETES_RUNTIME: - _report_registry_images(project) + from lightcone.engine.binder import binder_available + + if binder_available(): + _build_via_hub(project, commit=not no_commit) + else: + _report_registry_images(project) return if resolved_runtime == "none": @@ -1143,6 +1165,123 @@ def _expected_worker_image(project: Path) -> str | None: return images[0] if images else None +def _binder_progress(verbose: bool) -> Callable[[str, str], None]: + """Progress callback for hub image builds. + + Git steps (commit/push) always print — they mutate the user's repo + and must never be silent. Build phases print on transition; the + per-line repo2docker build log only with ``--verbose``. + """ + state = {"phase": ""} + + def cb(phase: str, message: str) -> None: + if phase in ("commit", "push"): + console.print(f" [dim]{message}[/dim]") + return + if phase and phase != state["phase"]: + state["phase"] = phase + console.print(f" [dim]binder: {phase}[/dim]") + if verbose and message: + console.print(f" [dim]{message}[/dim]") + + return cb + + +def _worker_image_for_run(project: Path, *, verbose: bool) -> str | None: + """The worker image a kubernetes-runtime run should use. + + On a hub (BinderHub service reachable) a declared Containerfile is + *ensured*: environment changes are committed and pushed, and the + image is built through the hub's build service when the registry + doesn't have it yet — so ``lc run`` always starts its cluster on an + up-to-date image. Registry-image specs pass through unchanged. + Elsewhere, falls back to the passive registry resolution of + :func:`_expected_worker_image` (verification-only). + """ + from lightcone.engine.binder import ( + BinderBuildError, + binder_available, + ensure_worker_image, + ) + from lightcone.engine.container import is_containerfile + + if not binder_available(): + return _expected_worker_image(project) + + _, specs = _container_specs(project) + if not specs: + # No container declared: the cluster runs the deployment's + # default worker image, which ships the lightcone stack. + return None + if len(specs) > 1: + console.print( + "[yellow]⚠ astra.yaml declares multiple distinct containers, " + "but a Dask Gateway cluster runs a single worker image:" + "[/yellow]\n" + + "\n".join(f" [dim]•[/dim] {s}" for s in specs) + + "\n Falling back to the deployment's default worker image. " + "Consolidate on one\n" + " Containerfile to run recipes in the project environment." + ) + return None + spec = specs[0] + if not is_containerfile(spec, project): + return spec + console.print( + f"[cyan]Ensuring worker image[/cyan] for [cyan]{spec}[/cyan] " + "[dim](hub build service)[/dim]" + ) + try: + image = ensure_worker_image( + project, spec, on_progress=_binder_progress(verbose) + ) + except BinderBuildError as e: + raise click.ClickException(str(e)) + console.print(f"[green]✓[/green] Worker image: [cyan]{image}[/cyan]") + return image + + +def _build_via_hub(project: Path, *, commit: bool) -> None: + """``lc build`` on a hub: publish every declared Containerfile. + + Drives the BinderHub service (build + push into the deployment + registry) for each declared Containerfile — the explicit, + verbose-by-default form of the ensure step ``lc run`` performs + implicitly. Registry-image specs need no build and are reported + as-is. + """ + from lightcone.engine.binder import BinderBuildError, ensure_worker_image + from lightcone.engine.container import is_containerfile + + _, specs = _container_specs(project) + if not specs: + console.print("No containers declared in astra.yaml — nothing to build.") + return + for spec in specs: + if not is_containerfile(spec, project): + console.print( + f"[green]•[/green] {spec} — registry image; the Gateway " + "cluster pulls it directly (must ship the lightcone " + "worker stack)." + ) + continue + console.print(f"[cyan]Building[/cyan] {spec} [dim](hub build service)[/dim]") + try: + image = ensure_worker_image( + project, + spec, + commit=commit, + on_progress=_binder_progress(verbose=True), + ) + except BinderBuildError as e: + raise click.ClickException(str(e)) + console.print( + f"[green]✓[/green] {spec} → [cyan]{image}[/cyan] " + "[dim](in the deployment registry; `lc run` will start its " + "cluster with it)[/dim]" + ) + + def _report_registry_images(project: Path) -> None: """``lc build`` on a kubernetes-runtime site. diff --git a/src/lightcone/engine/binder.py b/src/lightcone/engine/binder.py new file mode 100644 index 00000000..a7fd4ed2 --- /dev/null +++ b/src/lightcone/engine/binder.py @@ -0,0 +1,417 @@ +"""On-hub worker-image builds through the BinderHub service. + +A lightcone-hub deployment runs 2i2c's *binderhub-service* in API-only +mode: an authenticated HTTP endpoint that clones a git ref, builds it +with repo2docker, and pushes the image to the deployment's container +registry (LCR-176 Part B). There is no docker daemon in a user pod, so +this is *the* way `lc build` publishes images from the hub. + +The flow implemented by :func:`ensure_worker_image`: + +1. Identify the **environment ref** — the last commit that touched any + environment-defining file (Containerfile, dependency files, COPY'd + sources; see :func:`lightcone.engine.container.env_context_paths`). + Code-only commits reuse the previous image. +2. If any of those files have uncommitted changes, commit them (scoped + to exactly those paths) — the PRD's "lc build commits the current + version of the code" step, kept minimal and predictable. +3. Push the ref when the remote doesn't have it yet (BinderHub build + pods clone from the remote; they cannot see the hub filesystem). +4. Ask the BinderHub service to build that ref (an SSE stream). The + service resolves the ref, checks the registry, and returns the image + name immediately when it is already built — so this is cheap to call + on every ``lc run``, which is exactly what makes "the image is + always up to date" transparent. + +Auth rides on the ambient JupyterHub identity: every singleuser pod +carries ``JUPYTERHUB_API_TOKEN``, which the binderhub-service accepts as +a JupyterHub-authenticated caller. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import time +import urllib.error +import urllib.parse +import urllib.request +from collections import deque +from collections.abc import Callable +from pathlib import Path + +#: Override for the BinderHub service URL. Defaults to the in-cluster +#: address every lightcone-hub singleuser pod can reach. +BINDER_URL_ENV = "LIGHTCONE_BINDER_URL" + +#: In-cluster address of the binderhub-service behind the JupyterHub +#: proxy on a lightcone-hub deployment (``proxy-public`` resolves inside +#: the hub namespace; the service is mounted at ``/services/binder``). +DEFAULT_BINDER_URL = "http://proxy-public/services/binder" + +#: JupyterHub-issued API token, present in every singleuser pod. Used +#: both as the auth credential for the binder service and as the "are we +#: on a hub?" marker for :func:`binder_available`. +API_TOKEN_ENV = "JUPYTERHUB_API_TOKEN" + +#: Hard ceiling on one build, seconds. repo2docker builds of a slim +#: python image finish in single-digit minutes; an hour means something +#: is wedged, not slow. +_BUILD_DEADLINE_S = 3600 + +#: Per-read socket timeout, seconds. BinderHub streams an SSE event at +#: least every few seconds while building (log lines) and keepalives +#: while waiting; minutes of silence means the stream is dead. +_STREAM_READ_TIMEOUT_S = 300 + +#: Type of the progress callback: ``(phase, message)``. *message* may be +#: empty; *phase* is BinderHub's event phase (waiting/fetching/building/ +#: pushing/built/ready/failed) or ``""`` for log-only events. +ProgressFn = Callable[[str, str], None] + + +class BinderBuildError(RuntimeError): + """A worker image could not be produced through the binder service.""" + + +def binder_service_url() -> str | None: + """The BinderHub service URL to use, or ``None`` when unavailable. + + An explicit :data:`BINDER_URL_ENV` always wins (also the seam tests + and off-hub experiments use). Otherwise the in-cluster default + applies only where it can work: on a hub pod, marked by the ambient + :data:`API_TOKEN_ENV` credential the service authenticates with. + """ + explicit = (os.environ.get(BINDER_URL_ENV) or "").strip().rstrip("/") + if explicit: + return explicit + if os.environ.get(API_TOKEN_ENV): + return DEFAULT_BINDER_URL + return None + + +def binder_available() -> bool: + """Can this process reach an authenticated BinderHub service?""" + return binder_service_url() is not None and bool( + os.environ.get(API_TOKEN_ENV) + ) + + +# --------------------------------------------------------------------------- +# Git plumbing +# --------------------------------------------------------------------------- + + +def _git(project: Path, *args: str) -> str: + """Run git in *project*, returning stripped stdout; raise on failure.""" + result = subprocess.run( + ["git", *args], + cwd=project, + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() + raise BinderBuildError( + f"`git {' '.join(args)}` failed in {project}: {detail}" + ) + return result.stdout.strip() + + +def _github_org_repo(remote_url: str) -> tuple[str, str] | None: + """``(org, repo)`` when *remote_url* points at GitHub, else ``None``. + + Handles the three shapes git produces: ``git@github.com:org/repo.git``, + ``https://github.com/org/repo(.git)``, ``ssh://git@github.com/org/repo``. + """ + url = remote_url.strip() + path = None + if url.startswith("git@github.com:"): + path = url[len("git@github.com:"):] + else: + parsed = urllib.parse.urlparse(url) + if parsed.hostname == "github.com": + path = parsed.path.lstrip("/") + if not path: + return None + if path.endswith(".git"): + path = path[: -len(".git")] + parts = [p for p in path.split("/") if p] + if len(parts) != 2: + return None + return parts[0], parts[1] + + +def repo_provider_spec(remote_url: str, ref: str) -> str: + """BinderHub ``/build/`` path for *remote_url* at *ref*. + + GitHub remotes use the ``gh`` provider (``gh///``); + anything else falls back to the generic ``git`` provider, whose spec + is the URL-escaped clone URL. Both accept a full commit sha as ref, + which is what we always pass — the resolved ref *is* the image tag, + so a moving branch name would break image identity. + """ + gh = _github_org_repo(remote_url) + if gh is not None: + org, repo = gh + return f"gh/{urllib.parse.quote(org)}/{urllib.parse.quote(repo)}/{ref}" + return f"git/{urllib.parse.quote(remote_url, safe='')}/{ref}" + + +def _ensure_dockerfile( + project: Path, containerfile_spec: str, pathspecs: list[str] +) -> None: + """Make the repo buildable by repo2docker (it looks for ``Dockerfile``). + + lightcone projects declare a ``Containerfile``; repo2docker only + recognizes ``Dockerfile`` (repo root or ``binder/``/``.binder/``). + A relative symlink at the repo root bridges the two — committed like + any other file, it survives the clone the build pod makes. + """ + if containerfile_spec == "Dockerfile": + return + dockerfile = project / "Dockerfile" + if dockerfile.is_symlink(): + if os.readlink(dockerfile) == containerfile_spec: + if "Dockerfile" not in pathspecs: + pathspecs.append("Dockerfile") + return + raise BinderBuildError( + f"Dockerfile is a symlink to {os.readlink(dockerfile)!r}, not " + f"to the declared container {containerfile_spec!r}. Fix the " + "symlink (repo2docker builds the Dockerfile, so it must point " + "at the declared Containerfile)." + ) + if dockerfile.exists(): + raise BinderBuildError( + "The project has both a Dockerfile and a declared container " + f"{containerfile_spec!r}. repo2docker will build the " + "Dockerfile, so either declare `container: Dockerfile` in " + "astra.yaml or replace the Dockerfile with a symlink to " + f"{containerfile_spec}." + ) + if "/" in containerfile_spec: + raise BinderBuildError( + f"The declared container {containerfile_spec!r} is not at the " + "project root; repo2docker needs a root-level Dockerfile. Add " + "one (e.g. a symlink to the Containerfile) and commit it." + ) + dockerfile.symlink_to(containerfile_spec) + pathspecs.append("Dockerfile") + + +def _resolve_env_ref( + project: Path, + pathspecs: list[str], + *, + commit: bool, + on_progress: ProgressFn | None, +) -> str: + """Commit env-file changes when needed; return the environment sha.""" + dirty = _git(project, "status", "--porcelain", "--", *pathspecs) + if dirty: + if not commit: + raise BinderBuildError( + "Environment-defining files have uncommitted changes:\n" + + dirty + + "\nCommit them (or run `lc build` without --no-commit) " + "so the image can be built from a git ref." + ) + if on_progress: + on_progress( + "commit", + "committing environment changes (" + + ", ".join( + sorted({line[3:].strip() for line in dirty.splitlines()}) + ) + + ")", + ) + _git(project, "add", "--", *pathspecs) + _git( + project, + "commit", + "--quiet", + "-m", + "lc build: update worker environment", + "--", + *pathspecs, + ) + + sha = _git(project, "log", "-1", "--format=%H", "--", *pathspecs) + if not sha: + raise BinderBuildError( + "No commit touches the environment-defining files " + f"({', '.join(pathspecs)}); commit them first." + ) + return sha + + +def _remote(project: Path) -> tuple[str, str]: + """``(name, clone_url)`` of the remote builds are fetched from. + + ``origin`` by convention, else the first configured remote. + """ + remotes = _git(project, "remote").splitlines() + if not remotes: + raise BinderBuildError( + "The project has no git remote. The hub's BinderHub service " + "builds images from a git hosting service (the build pod " + "clones the repo — it cannot see the hub filesystem), so " + "push the project to GitHub first:\n" + " gh repo create --source . --public --push\n" + "(note: this deployment's binder service can only clone " + "public repos until an access token is configured)." + ) + name = "origin" if "origin" in remotes else remotes[0] + return name, _git(project, "remote", "get-url", name) + + +def _ensure_pushed( + project: Path, remote: str, sha: str, on_progress: ProgressFn | None +) -> None: + """Make sure *sha* is reachable from *remote*; push if not.""" + if _git(project, "branch", "-r", "--contains", sha): + return + if on_progress: + on_progress("push", f"pushing {sha[:12]} so the build pod can clone it") + try: + _git(project, "push", "--quiet", remote, "HEAD") + except BinderBuildError as exc: + raise BinderBuildError( + f"Could not push the environment commit to the remote ({exc}). " + "The BinderHub build pod clones from the remote, so the " + "commit must be pushed. Fix your push access (or push " + "manually) and re-run." + ) from exc + + +# --------------------------------------------------------------------------- +# BinderHub build API (SSE) +# --------------------------------------------------------------------------- + + +def build_via_binder( + provider_spec: str, + *, + on_progress: ProgressFn | None = None, +) -> str: + """Build *provider_spec* through the binder service; return the image ref. + + Streams the service's SSE events until a terminal one arrives. The + service checks its registry first, so an already-built ref returns + in one round-trip — callers can treat this as an idempotent + "ensure built" primitive. + """ + base = binder_service_url() + token = os.environ.get(API_TOKEN_ENV) + if base is None or not token: + raise BinderBuildError( + "No BinderHub service is reachable from here (set " + f"{BINDER_URL_ENV}, or run on a JupyterHub pod where " + f"{API_TOKEN_ENV} is provided)." + ) + + url = f"{base}/build/{provider_spec}?build_only=true" + req = urllib.request.Request( + url, headers={"Authorization": f"token {token}"} + ) + deadline = time.monotonic() + _BUILD_DEADLINE_S + tail: deque[str] = deque(maxlen=30) + image: str | None = None + failed = False + try: + with urllib.request.urlopen(req, timeout=_STREAM_READ_TIMEOUT_S) as resp: + for raw in resp: + if time.monotonic() > deadline: + raise BinderBuildError( + f"Image build exceeded {_BUILD_DEADLINE_S}s; last " + "output:\n" + "".join(tail) + ) + line = raw.decode("utf-8", errors="replace").strip() + if not line.startswith("data:"): + continue + try: + event = json.loads(line[len("data:"):].strip()) + except json.JSONDecodeError: + continue + phase = str(event.get("phase") or "") + message = str(event.get("message") or "").rstrip("\n") + if message: + tail.append(message + "\n") + if on_progress: + on_progress(phase, message) + if isinstance(event.get("imageName"), str): + image = event["imageName"] + if phase in ("failed", "failure"): + failed = True + break + if phase in ("ready", "built") and image: + break + except urllib.error.HTTPError as exc: + body = "" + try: + body = exc.read().decode("utf-8", errors="replace")[:500] + except OSError: + pass + raise BinderBuildError( + f"BinderHub build request failed: HTTP {exc.code} for " + f"{url}\n{body}" + ) from exc + except (urllib.error.URLError, OSError) as exc: + raise BinderBuildError( + f"Could not reach the BinderHub service at {base} ({exc})." + ) from exc + + if failed or image is None: + raise BinderBuildError( + "BinderHub could not build the image" + + (f" for {provider_spec}" if image is None else "") + + ". Last build output:\n" + + ("".join(tail) or " (no output received)") + ) + return image + + +# --------------------------------------------------------------------------- +# High-level entry point +# --------------------------------------------------------------------------- + + +def ensure_worker_image( + project: Path, + containerfile_spec: str, + *, + commit: bool = True, + on_progress: ProgressFn | None = None, +) -> str: + """Make sure the project's worker image is built; return its registry ref. + + Implements the PRD's "makes sure the container image is up to date, + rebuilds it if necessary": commit env-file changes (when *commit*), + push the environment ref, and drive the BinderHub build for it — + a no-op round-trip when the registry already holds the image. + """ + from lightcone.engine.container import env_context_paths + + _git(project, "rev-parse", "--git-dir") # fail early off-git + containerfile = project / containerfile_spec + if not containerfile.is_file(): + raise BinderBuildError( + f"Declared container {containerfile_spec!r} not found in " + f"{project}." + ) + + pathspecs = env_context_paths(containerfile, project) + if containerfile_spec not in pathspecs: + pathspecs.insert(0, containerfile_spec) + _ensure_dockerfile(project, containerfile_spec, pathspecs) + + sha = _resolve_env_ref( + project, pathspecs, commit=commit, on_progress=on_progress + ) + remote_name, remote_url = _remote(project) + _ensure_pushed(project, remote_name, sha, on_progress) + spec = repo_provider_spec(remote_url, sha) + return build_via_binder(spec, on_progress=on_progress) diff --git a/src/lightcone/engine/container.py b/src/lightcone/engine/container.py index 3534fd7d..2d1d3984 100644 --- a/src/lightcone/engine/container.py +++ b/src/lightcone/engine/container.py @@ -529,6 +529,36 @@ def _expand_copy_source(src: str, project_path: Path) -> list[Path]: return [] +def env_context_paths(containerfile: Path, project_path: Path) -> list[str]: + """Project-relative paths that define the image's *environment*. + + The build-context set of :func:`_iter_build_context_entries` minus + any ``COPY .`` of the whole project root. On a Gateway deployment + the worker-pod image is the execution *environment* — recipes run + in the project directory on the shared filesystem (the executor pins + ``--directory``), so a whole-tree ``COPY . .`` (kept for image + portability) does not affect what executes. Excluding it means + code-only commits don't demand an image rebuild; edits to the + Containerfile, dependency files, or specifically COPY'd sources do. + + Used as a git pathspec to decide image freshness (the last commit + touching any of these paths identifies the environment). + """ + root = project_path.resolve() + paths: list[str] = [] + for _, path in _iter_build_context_entries(containerfile, project_path): + resolved = path.resolve() + if resolved == root: + continue + try: + rel = resolved.relative_to(root).as_posix() + except ValueError: + continue + if rel not in paths: + paths.append(rel) + return paths + + def is_containerfile(spec: str, project_path: Path) -> bool: """Return ``True`` if *spec* refers to an existing file (Containerfile).""" return (project_path / spec).is_file() diff --git a/tests/conftest.py b/tests/conftest.py index 1459b9a6..9467232e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,3 +20,9 @@ def _no_ambient_hub_env(monkeypatch: pytest.MonkeyPatch) -> None: """ monkeypatch.delenv("JUPYTERHUB_USER", raising=False) monkeypatch.delenv("DASK_GATEWAY__ADDRESS", raising=False) + # Same reasoning for the BinderHub build seam: an ambient JupyterHub + # API token (or an explicit binder URL) would make binder_available() + # true and route `lc build`/`lc run` image resolution through a real + # HTTP endpoint mid-suite. + monkeypatch.delenv("JUPYTERHUB_API_TOKEN", raising=False) + monkeypatch.delenv("LIGHTCONE_BINDER_URL", raising=False) diff --git a/tests/test_binder.py b/tests/test_binder.py new file mode 100644 index 00000000..e31dfa7d --- /dev/null +++ b/tests/test_binder.py @@ -0,0 +1,363 @@ +"""Unit tests for the BinderHub build seam (`lightcone.engine.binder`). + +The git plumbing runs against real throwaway repos (with a local bare +"remote" — pushes are exercised for real); the BinderHub HTTP side is a +fake SSE stream injected at the ``urllib`` boundary. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest + +from lightcone.engine.binder import ( + BINDER_URL_ENV, + BinderBuildError, + binder_available, + binder_service_url, + build_via_binder, + ensure_worker_image, + repo_provider_spec, +) + +# --------------------------------------------------------------------------- +# Service discovery +# --------------------------------------------------------------------------- + + +def test_binder_url_absent_off_hub() -> None: + assert binder_service_url() is None + assert not binder_available() + + +def test_binder_url_defaults_on_hub(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + assert binder_service_url() == "http://proxy-public/services/binder" + assert binder_available() + + +def test_binder_url_env_override_wins(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + monkeypatch.setenv(BINDER_URL_ENV, "http://elsewhere/binder/") + assert binder_service_url() == "http://elsewhere/binder" + + +def test_binder_url_override_without_token_is_not_available( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A binder URL alone is unusable — the service authenticates via the + JupyterHub token, so availability requires both.""" + monkeypatch.setenv(BINDER_URL_ENV, "http://elsewhere/binder") + assert binder_service_url() == "http://elsewhere/binder" + assert not binder_available() + + +# --------------------------------------------------------------------------- +# Provider specs +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "url", + [ + "git@github.com:LightconeResearch/demo.git", + "https://github.com/LightconeResearch/demo", + "https://github.com/LightconeResearch/demo.git", + "ssh://git@github.com/LightconeResearch/demo.git", + ], +) +def test_github_remotes_use_gh_provider(url: str) -> None: + assert ( + repo_provider_spec(url, "abc123") + == "gh/LightconeResearch/demo/abc123" + ) + + +def test_non_github_remote_uses_git_provider() -> None: + spec = repo_provider_spec("https://gitlab.com/org/repo.git", "abc123") + assert spec == "git/https%3A%2F%2Fgitlab.com%2Forg%2Frepo.git/abc123" + + +# --------------------------------------------------------------------------- +# SSE build stream +# --------------------------------------------------------------------------- + + +class _FakeSSE: + def __init__(self, events: list[dict[str, object] | str]) -> None: + self._lines: list[bytes] = [] + for e in events: + if isinstance(e, str): + self._lines.append(e.encode()) + else: + self._lines.append(b"data: " + json.dumps(e).encode() + b"\n") + + def __enter__(self) -> _FakeSSE: + return self + + def __exit__(self, *args: object) -> bool: + return False + + def __iter__(self): # noqa: ANN204 + return iter(self._lines) + + +def _install_fake_binder( + monkeypatch: pytest.MonkeyPatch, events: list[dict[str, object] | str] +) -> dict[str, str]: + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + monkeypatch.setenv(BINDER_URL_ENV, "http://binder.test") + seen: dict[str, str] = {} + + def fake_urlopen(req, timeout=None): # noqa: ANN001, ANN202 + seen["url"] = req.full_url + seen["auth"] = req.get_header("Authorization") or "" + return _FakeSSE(events) + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + return seen + + +def test_build_returns_image_on_ready(monkeypatch: pytest.MonkeyPatch) -> None: + seen = _install_fake_binder( + monkeypatch, + [ + {"phase": "waiting", "message": "queued"}, + ": keepalive\n", + {"phase": "building", "message": "Step 1/5 ..."}, + {"phase": "ready", "imageName": "reg/binder/x:sha", "message": "done"}, + ], + ) + image = build_via_binder("gh/org/repo/sha") + assert image == "reg/binder/x:sha" + assert seen["url"] == "http://binder.test/build/gh/org/repo/sha?build_only=true" + assert seen["auth"] == "token tok" + + +def test_build_cached_image_short_circuits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BinderHub answers a single `built` event when the registry already + has the image — the everyday `lc run` fast path.""" + _install_fake_binder( + monkeypatch, + [{"phase": "built", "imageName": "reg/binder/x:sha", "message": "found"}], + ) + assert build_via_binder("gh/org/repo/sha") == "reg/binder/x:sha" + + +def test_build_failure_raises_with_log_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_fake_binder( + monkeypatch, + [ + {"phase": "building", "message": "Step 1/5 ..."}, + {"phase": "building", "message": "error: nope"}, + {"phase": "failed", "message": "Build failed"}, + ], + ) + with pytest.raises(BinderBuildError, match="could not build") as exc: + build_via_binder("gh/org/repo/sha") + assert "error: nope" in str(exc.value) + + +def test_build_stream_without_image_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_fake_binder(monkeypatch, [{"phase": "waiting", "message": "queued"}]) + with pytest.raises(BinderBuildError, match="could not build"): + build_via_binder("gh/org/repo/sha") + + +def test_build_reports_progress(monkeypatch: pytest.MonkeyPatch) -> None: + _install_fake_binder( + monkeypatch, + [ + {"phase": "building", "message": "Step 1/5 ..."}, + {"phase": "ready", "imageName": "reg/x:sha"}, + ], + ) + phases: list[str] = [] + build_via_binder("gh/org/repo/sha", on_progress=lambda p, m: phases.append(p)) + assert phases == ["building", "ready"] + + +# --------------------------------------------------------------------------- +# ensure_worker_image: the git flow against real repos +# --------------------------------------------------------------------------- + + +def _git(cwd: Path, *args: str) -> str: + return subprocess.run( + ["git", *args], cwd=cwd, capture_output=True, text=True, check=True + ).stdout.strip() + + +@pytest.fixture() +def project(tmp_path: Path) -> Path: + """A committed project with a Containerfile and a local bare remote.""" + remote = tmp_path / "remote.git" + subprocess.run( + ["git", "init", "--bare", "-q", str(remote)], check=True + ) + proj = tmp_path / "proj" + proj.mkdir() + _git(proj, "init", "-q", "-b", "main") + _git(proj, "config", "user.email", "test@example.com") + _git(proj, "config", "user.name", "Test") + (proj / "Containerfile").write_text( + "FROM python:3.12-slim\nCOPY requirements.txt .\n" + "RUN pip install -r requirements.txt\n" + ) + (proj / "requirements.txt").write_text("numpy\n") + (proj / "analysis.py").write_text("print('hi')\n") + _git(proj, "add", "-A") + _git(proj, "commit", "-q", "-m", "initial") + _git(proj, "remote", "add", "origin", str(remote)) + _git(proj, "push", "-q", "-u", "origin", "main") + return proj + + +def _capture_build(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: + built: dict[str, str] = {} + + def fake_build(spec: str, *, on_progress=None): # noqa: ANN001, ANN202 + built["spec"] = spec + return "reg/binder/proj:" + spec.rsplit("/", 1)[-1] + + monkeypatch.setattr( + "lightcone.engine.binder.build_via_binder", fake_build + ) + return built + + +def test_ensure_builds_env_ref_and_links_dockerfile( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + built = _capture_build(monkeypatch) + image = ensure_worker_image(project, "Containerfile") + + # The Dockerfile symlink was created (repo2docker only reads + # `Dockerfile`), committed, and pushed. + assert (project / "Dockerfile").is_symlink() + assert _git(project, "status", "--porcelain") == "" + sha = _git(project, "log", "-1", "--format=%H") + remote_sha = _git(project, "ls-remote", "origin", "main").split()[0] + assert remote_sha == sha + + # Non-GitHub remote → generic git provider, ref = the env sha. + assert built["spec"].startswith("git/") + assert built["spec"].endswith(f"/{sha}") + assert image.endswith(sha) + + +def test_ensure_reuses_env_ref_across_code_commits( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + """Code-only commits must not change the environment ref — no new + image build for edits that don't touch env-defining files.""" + built = _capture_build(monkeypatch) + ensure_worker_image(project, "Containerfile") + env_sha = built["spec"].rsplit("/", 1)[-1] + + (project / "analysis.py").write_text("print('changed')\n") + _git(project, "add", "-A") + _git(project, "commit", "-q", "-m", "code only") + _git(project, "push", "-q", "origin", "main") + + ensure_worker_image(project, "Containerfile") + assert built["spec"].rsplit("/", 1)[-1] == env_sha + assert env_sha != _git(project, "log", "-1", "--format=%H") + + +def test_ensure_dirty_env_files_are_committed_and_pushed( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + built = _capture_build(monkeypatch) + (project / "requirements.txt").write_text("numpy\nscipy\n") + + ensure_worker_image(project, "Containerfile") + + assert _git(project, "status", "--porcelain", "--", "requirements.txt") == "" + head = _git(project, "log", "-1", "--format=%H") + assert built["spec"].endswith(f"/{head}") + assert _git(project, "ls-remote", "origin", "main").split()[0] == head + msg = _git(project, "log", "-1", "--format=%s") + assert "environment" in msg + + +def test_ensure_dirty_code_files_are_left_alone( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + """Only env-defining files are auto-committed — a dirty analysis + script is none of lc build's business.""" + _capture_build(monkeypatch) + (project / "analysis.py").write_text("print('wip')\n") + + ensure_worker_image(project, "Containerfile") + + dirty = _git(project, "status", "--porcelain") + assert "analysis.py" in dirty + + +def test_ensure_no_commit_flag_refuses_dirty_env( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + _capture_build(monkeypatch) + (project / "requirements.txt").write_text("numpy\nscipy\n") + + with pytest.raises(BinderBuildError, match="uncommitted"): + ensure_worker_image(project, "Containerfile", commit=False) + + +def test_ensure_without_remote_raises_with_guidance( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + _capture_build(monkeypatch) + _git(project, "remote", "remove", "origin") + + with pytest.raises(BinderBuildError, match="no git remote"): + ensure_worker_image(project, "Containerfile") + + +def test_ensure_conflicting_dockerfile_raises( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + _capture_build(monkeypatch) + (project / "Dockerfile").write_text("FROM something:else\n") + + with pytest.raises(BinderBuildError, match="Dockerfile"): + ensure_worker_image(project, "Containerfile") + + +def test_ensure_github_remote_uses_gh_provider( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + built = _capture_build(monkeypatch) + # First ensure pushes the Dockerfile-symlink commit to the real bare + # remote; the remote-tracking ref that records it survives the URL + # swap below, so the second ensure needs no push (which would fail + # against the fake GitHub URL). + ensure_worker_image(project, "Containerfile") + _git( + project, + "remote", + "set-url", + "origin", + "git@github.com:LightconeResearch/proj.git", + ) + image = ensure_worker_image(project, "Containerfile") + assert built["spec"].startswith("gh/LightconeResearch/proj/") + assert image + + +def test_ensure_declared_container_missing( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + _capture_build(monkeypatch) + with pytest.raises(BinderBuildError, match="not found"): + ensure_worker_image(project, "envs/Containerfile") diff --git a/tests/test_cli.py b/tests/test_cli.py index edaecd31..81fdf090 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -328,3 +328,153 @@ def test_run_cmd_multiple_triggers_all_before_separator() -> None: for trigger in ("code", "input", "mtime", "params"): assert trigger in cmd, f"trigger '{trigger}' missing from cmd" assert cmd.index(trigger) < sep_idx, f"trigger '{trigger}' must come before '--'" + + +# ---- lc build / lc run worker image on a hub (kubernetes + BinderHub) ------ + + +@pytest.fixture +def hub_project( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> Path: + """A minimal project declaring a project-level Containerfile.""" + project = tmp_path / "proj" + project.mkdir() + (project / "astra.yaml").write_text( + "name: proj\n" + "container: Containerfile\n" + "outputs:\n - id: foo\n recipe:\n command: echo\n" + ) + (project / "Containerfile").write_text("FROM python:3.12-slim\n") + monkeypatch.chdir(project) + return project + + +def test_build_kubernetes_builds_via_hub_service( + runner: CliRunner, hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """On a hub (BinderHub reachable) `lc build` drives a real image build + instead of printing off-hub publish instructions.""" + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + captured: dict[str, object] = {} + + def fake_ensure(project, spec, *, commit=True, on_progress=None): # noqa: ANN001, ANN202 + captured["spec"] = spec + captured["commit"] = commit + return "reg/binder/proj:abc123" + + monkeypatch.setattr( + "lightcone.engine.binder.ensure_worker_image", fake_ensure + ) + result = runner.invoke(main, ["build", "--runtime", "kubernetes"]) + assert result.exit_code == 0, result.output + assert "reg/binder/proj:abc123" in result.output + assert captured == {"spec": "Containerfile", "commit": True} + + +def test_build_kubernetes_no_commit_flag( + runner: CliRunner, hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + captured: dict[str, object] = {} + + def fake_ensure(project, spec, *, commit=True, on_progress=None): # noqa: ANN001, ANN202 + captured["commit"] = commit + return "reg/binder/proj:abc123" + + monkeypatch.setattr( + "lightcone.engine.binder.ensure_worker_image", fake_ensure + ) + result = runner.invoke( + main, ["build", "--runtime", "kubernetes", "--no-commit"] + ) + assert result.exit_code == 0, result.output + assert captured["commit"] is False + + +def test_build_kubernetes_build_error_is_user_facing( + runner: CliRunner, hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + + def fake_ensure(project, spec, *, commit=True, on_progress=None): # noqa: ANN001, ANN202 + from lightcone.engine.binder import BinderBuildError + + raise BinderBuildError("the build broke") + + monkeypatch.setattr( + "lightcone.engine.binder.ensure_worker_image", fake_ensure + ) + result = runner.invoke(main, ["build", "--runtime", "kubernetes"]) + assert result.exit_code != 0 + assert "the build broke" in result.output + assert "Traceback" not in result.output + + +def test_build_kubernetes_off_hub_falls_back_to_registry_report( + runner: CliRunner, hub_project: Path +) -> None: + """Without a reachable BinderHub service (e.g. kubernetes runtime + configured off-hub) the old passive registry report still runs.""" + result = runner.invoke(main, ["build", "--runtime", "kubernetes"]) + assert result.exit_code == 0, result.output + assert "not built" in result.output + + +def test_worker_image_for_run_ensures_via_binder( + hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.cli.commands import _worker_image_for_run + + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + monkeypatch.setattr( + "lightcone.engine.binder.ensure_worker_image", + lambda p, s, *, commit=True, on_progress=None: "reg/binder/proj:sha1", + ) + assert ( + _worker_image_for_run(hub_project, verbose=False) + == "reg/binder/proj:sha1" + ) + + +def test_worker_image_for_run_registry_spec_passes_through( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A declared registry image needs no build — used as the worker image + verbatim, even on a hub.""" + from lightcone.cli.commands import _worker_image_for_run + + project = tmp_path / "proj" + project.mkdir() + (project / "astra.yaml").write_text( + "name: proj\n" + "container: ghcr.io/org/worker:1.0\n" + "outputs:\n - id: foo\n recipe:\n command: echo\n" + ) + monkeypatch.chdir(project) + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + + def _boom(*args: object, **kwargs: object) -> str: + raise AssertionError("registry specs must not trigger a build") + + monkeypatch.setattr("lightcone.engine.binder.ensure_worker_image", _boom) + assert ( + _worker_image_for_run(project, verbose=False) == "ghcr.io/org/worker:1.0" + ) + + +def test_worker_image_for_run_no_container_uses_default( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.cli.commands import _worker_image_for_run + + project = tmp_path / "proj" + project.mkdir() + (project / "astra.yaml").write_text( + "name: proj\noutputs:\n - id: foo\n recipe:\n command: echo\n" + ) + monkeypatch.chdir(project) + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + assert _worker_image_for_run(project, verbose=False) is None diff --git a/tests/test_container.py b/tests/test_container.py index b1f2655b..281b41ac 100644 --- a/tests/test_container.py +++ b/tests/test_container.py @@ -1013,3 +1013,41 @@ def test_existing_file(self, project: Path) -> None: def test_missing_file(self, project: Path) -> None: assert is_containerfile("python:3.12-slim", project) is False + + +# ---- env_context_paths ---------------------------------------------------- + + +class TestEnvContextPaths: + def test_excludes_whole_root_copy(self, tmp_path: Path) -> None: + """`COPY . .` (kept for image portability) must not make every + commit an environment change — on a Gateway deployment code + reaches workers via the shared filesystem, not the image.""" + from lightcone.engine.container import env_context_paths + + cf = tmp_path / "Containerfile" + cf.write_text( + "FROM python:3.12-slim\n" + "COPY requirements.txt .\n" + "COPY . .\n" + ) + (tmp_path / "requirements.txt").write_text("numpy\n") + (tmp_path / "analysis.py").write_text("print('hi')\n") + + paths = env_context_paths(cf, tmp_path) + assert "Containerfile" in paths + assert "requirements.txt" in paths + assert "analysis.py" not in paths + assert "." not in paths and "" not in paths + + def test_includes_named_copy_sources(self, tmp_path: Path) -> None: + from lightcone.engine.container import env_context_paths + + cf = tmp_path / "Containerfile" + cf.write_text("FROM python:3.12-slim\nCOPY envs/ /envs/\n") + envs = tmp_path / "envs" + envs.mkdir() + (envs / "lock.txt").write_text("pinned\n") + + paths = env_context_paths(cf, tmp_path) + assert "envs" in paths From 64301b1c1751305576d9d25e3dd118d21acb04f3 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 10:19:46 +0200 Subject: [PATCH 4/9] =?UTF-8?q?Docs:=20create/cull=20lifecycle=20+=20Binde?= =?UTF-8?q?rHub=20build=20path=20(design=20doc=20=C2=A78)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User cluster guide, lc build reference, and the api/dask_cluster page now describe the PRD lifecycle; the design doc gains §8 recording the implemented decisions (run-scoped clusters; binderhub-service resolves the §6.5 open build-path question) and marks the superseded §7.3/§7.4. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/api/dask_cluster.md | 51 +++++---- docs/cli/build.md | 24 +++- docs/design/jupyterhub-dask-gateway-gcp.md | 72 +++++++++++- docs/user/cluster.md | 126 ++++++++++++--------- 4 files changed, 196 insertions(+), 77 deletions(-) diff --git a/docs/api/dask_cluster.md b/docs/api/dask_cluster.md index 0784b5ea..5a449b10 100644 --- a/docs/api/dask_cluster.md +++ b/docs/api/dask_cluster.md @@ -5,7 +5,7 @@ four branches, no service to manage. Source: `src/lightcone/engine/dask_cluster.py`. -## `cluster_for_run(*, verbose=False, local_directory=None, expected_worker_image=None) → Iterator[dict[str, str]]` +## `cluster_for_run(*, verbose=False, local_directory=None, expected_worker_image=None, max_workers=None) → Iterator[dict[str, str]]` Yields the **env overlay** the child snakemake process needs to reach the cluster — the parent and the executor plugin are separate @@ -16,23 +16,30 @@ Four branches in priority order: `{"DASK_SCHEDULER_ADDRESS": addr}` as-is. We don't own the cluster, so we don't tear it down. 2. **Dask Gateway detected** (`LIGHTCONE_GATEWAY_CLUSTER` or - `DASK_GATEWAY__ADDRESS` set, e.g. a JupyterHub pod) → **attach** to - the user's running Gateway cluster; yield - `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Attach-only by design: the - user creates clusters from JupyterLab (that's where the - image/cores/memory options widget and the dashboard live); the - Gateway API is user-scoped, so exactly one running cluster attaches - unambiguously, and zero or several raises with the fix spelled out - (`LIGHTCONE_GATEWAY_CLUSTER` disambiguates). Gateway scheduler + `DASK_GATEWAY__ADDRESS` set, e.g. a JupyterHub pod) → yield + `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Two sub-modes: + + - **Create/cull (default).** A run-scoped cluster is created with + *expected_worker_image* as its `image` cluster option, scaled + adaptively `1..max_workers`, and shut down on exit — the same + lifetime contract as the local/SLURM branches, and the mechanism + that makes a freshly built project image take effect (a Gateway + cluster's image is fixed at creation). Startup blocks until the + first worker is live (bounded by + `LIGHTCONE_GATEWAY_WORKER_TIMEOUT`, default 600 s) so an + unpullable image fails loudly instead of hanging at zero workers. + - **Attach** (`LIGHTCONE_GATEWAY_CLUSTER=` set by the user) → + connect to that cluster, leave its scaling and lifetime untouched + (same convention as branch 1), and warn when its actual worker + image (read from the scheduler pod's `LIGHTCONE_WORKER_IMAGE`) + differs from *expected_worker_image*. + + Either way, startup fails fast if live workers don't advertise the + resource contract below (zero live workers is fine on the attach + path — adaptive clusters scale on demand). Gateway scheduler addresses use a `gateway://` comm scheme a bare `Client` cannot dial, so the child rejoins **by name** through the authenticated - Gateway API. The cluster, and its scaling, are left untouched on - exit — same convention as branch 1. Startup fails fast if live - workers don't advertise the resource contract below (zero workers - is fine — adaptive clusters scale on demand), and warns when the - cluster's actual worker image (read from the scheduler pod's - `LIGHTCONE_WORKER_IMAGE`) differs from *expected_worker_image*. - Requires the optional dependency: + Gateway API. Requires the optional dependency: `pip install lightcone-cli[gateway]`. 3. **`SLURM_JOB_ID` set** → start an in-process scheduler bound to the driver's SLURM hostname (`SLURMD_NODENAME` or `gethostname()`), @@ -41,9 +48,10 @@ Four branches in priority order: Outside the Gateway branch the scheduler is always in-process, so its lifetime equals the run's lifetime: no orphaned schedulers if the -driver crashes. On the Gateway branch there is nothing to orphan -either — lc never creates the cluster, and idle clusters are reaped -by the deployment's `idle_timeout`. +driver crashes. On the Gateway branch the same contract is enforced +server-side: created clusters are shut down on exit +(`shutdown_on_close=True` and the deployment's `idle_timeout` backstop +a crashed driver); attached clusters belong to the user. ## Resource keys @@ -101,5 +109,6 @@ and keeps everything in one process tree. `tests/test_dask_cluster.py` covers all four branches and the resource-advertising contract. The SLURM branch is tested with mocked `subprocess.Popen` plus a stubbed `Client.wait_for_workers`; the -Gateway branch (discovery, attach-only lifecycle, contract and image -verification) against a fake `dask_gateway` module. +Gateway branch (create/cull lifecycle, image option, worker-wait +failure, attach mode, contract and image verification) against a fake +`dask_gateway` module. diff --git a/docs/cli/build.md b/docs/cli/build.md index eab38f0a..2cb35316 100644 --- a/docs/cli/build.md +++ b/docs/cli/build.md @@ -14,7 +14,8 @@ lc build [OPTIONS] | Option | Default | Effect | |--------|---------|--------| | `--force` | off | Rebuild / re-pull even if the tag already exists locally. | -| `--runtime {docker,podman,podman-hpc}` | resolved from `~/.lightcone/config.yaml` | Override the runtime for this build. | +| `--runtime {docker,podman,podman-hpc,kubernetes}` | resolved from `~/.lightcone/config.yaml` (or the detected site) | Override the runtime for this build. | +| `--no-commit` | off | On a hub: fail instead of auto-committing environment-file changes before the image build. | ## What it does @@ -33,6 +34,27 @@ If the runtime is `none` (either by config or because `auto` couldn't find one), `lc build` prints a friendly note and exits 0. There is nothing to build. +## On a JupyterHub deployment (`kubernetes` runtime) + +There is no docker in a hub pod, so `lc build` instead drives an image +build through the deployment's BinderHub service: + +1. Commits any changes to the environment-defining files (the + Containerfile, dependency files, and named COPY sources — never your + analysis code), maintaining a `Dockerfile → Containerfile` symlink + for repo2docker. `--no-commit` refuses instead. +2. Pushes, so the build pods can clone the ref (the project needs a + GitHub remote). +3. Streams the BinderHub build (repo2docker → the deployment registry) + and prints the resulting image ref — the image `lc run` will start + its Gateway cluster with. Already-built refs return immediately. + +`lc run` performs the same ensure step automatically on every run, so +running `lc build` explicitly is optional on the hub — useful to +pre-build after a dependency change or to see the build log. Without a +reachable build service, `lc build` falls back to probing the registry +and printing off-hub publish instructions. + ## Tag computation ```text diff --git a/docs/design/jupyterhub-dask-gateway-gcp.md b/docs/design/jupyterhub-dask-gateway-gcp.md index 3d299b06..c45054bd 100644 --- a/docs/design/jupyterhub-dask-gateway-gcp.md +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -514,7 +514,12 @@ single environment. subpaths; interacts with `lc init` scaffolding and the session-start hooks. 4. **kbatch alternative** — if the 0.5 alphas misbehave, a minimal in-house hub service (JupyterHub-authenticated POST → k8s Job) is ~150 lines and removes the dependency. -5. **On-hub container build path — UNRESOLVED, needs a decision (surfaced 2026-07-15, LCR-176).** +5. **On-hub container build path — RESOLVED 2026-07-21: BinderHub service (see §8).** + The deployment enables 2i2c's *binderhub-service* in API-only mode with a push credential for + the deployment registry, which answers the least-privilege sub-question: users never hold + registry credentials — the build service holds the only writer key, and users reach it through + JupyterHub auth. This is a variant of (c) that we don't have to build or operate ourselves. + Original framing kept below for the record. §7.2 currently defers all image builds off-hub: `lc build` in the JupyterLab pod prints the `docker build && docker push` commands to run elsewhere, because there is no docker in-pod by design. End-to-end testing on the staging cluster confirmed this *works* but is a real friction @@ -543,6 +548,11 @@ single environment. ## 7. Implemented design (2026-07-12) — attach-only + kubernetes as a first-class runtime +> **Partially superseded by §8 (2026-07-21):** §7.1 and §7.2 stand; +> §7.3's attach-only lifecycle and §7.4's build/lifecycle rows are +> replaced by the PRD create/cull lifecycle and the BinderHub build +> path. + What actually shipped diverges from §4.2/§4.5 in three deliberate ways, all in the direction of less machinery: @@ -633,3 +643,63 @@ environment, worker-default as the default cluster image, and a - kbatch: https://kbatch.readthedocs.io/ · https://github.com/kbatch-dev/kbatch · https://kbatch-dev.github.io/helm-chart/ - Worked example of gateway-behind-hub-proxy: https://www.zonca.dev/posts/2022-04-04-dask-gateway-jupyterhub + +## 8. Implemented design (2026-07-21) — PRD lifecycle: create/cull + BinderHub builds + +Implements the resolved decisions of the *Integration with Kubernetes and +JupyterHub* PRD (LCR-174) on top of §7's runtime model. Two changes. + +### 8.1 Run-scoped Gateway clusters (replaces §7.3) + +`lc run` on a hub now creates a Gateway cluster per run — with the project's +worker image as the `image` cluster option, adaptive `1..--jobs` — and shuts +it down when the run finishes. This is PRD decision #1, and it is what the +native-k8s model requires: a Gateway cluster's image is fixed at creation, so +"the image is up to date" is only achievable with a fresh cluster. Startup +waits for the first worker (`LIGHTCONE_GATEWAY_WORKER_TIMEOUT`, default +600 s) so an unpullable image is a loud error, not a silent zero-worker hang. + +`LIGHTCONE_GATEWAY_CLUSTER=` keeps the §7.3 behavior as an explicit +attach mode (long-lived cluster iteration): never rescaled, never shut down, +image drift warned about via the `LIGHTCONE_WORKER_IMAGE` check. + +### 8.2 On-hub image builds through binderhub-service (resolves §6.5) + +The deployment runs 2i2c's binderhub-service (API-only mode, repo2docker +build pods, pushing to the deployment registry). `lightcone.engine.binder` +drives it from the user pod: + +- **Auth**: the ambient `JUPYTERHUB_API_TOKEN`; service URL defaults to + `http://proxy-public/services/binder` (`LIGHTCONE_BINDER_URL` overrides). +- **Environment ref**: the last commit touching the env-defining files — + `env_context_paths()` = Containerfile + dependency files + named COPY + sources, *excluding* a whole-tree `COPY .` (code reaches workers via the + shared home; the image is the environment). Code-only commits therefore + reuse the previous image; env edits are auto-committed (scoped to those + paths; `lc build --no-commit` refuses instead) and pushed, since build + pods clone from the git remote. +- **repo2docker bridge**: a committed root `Dockerfile → Containerfile` + symlink (repo2docker does not read `Containerfile`). +- **Build**: `GET /build///?build_only=true` streamed + as SSE until `ready`/`failed`; the terminal event carries `imageName`. + BinderHub consults its registry first, so an already-built ref is one + round-trip — `lc run` calls this ensure step on every kubernetes-runtime + run and then creates the cluster with the returned image. + +Trade-offs accepted: the project must have a (public, until the deployment +configures a provider token) GitHub remote; the image tag is the env sha +rather than the content-addressed `lc--` scheme (§7.2's +registry probing remains as the off-hub fallback path); repo2docker builds +at the env sha, so a whole-tree `COPY .` bakes code as of that commit — a +non-issue on the gateway path where recipes run from the shared home via +`--directory`. + +### 8.3 Path similarity, as implemented (updates §7.4) + +| | local | SLURM | hub | +|---|---|---|---| +| environment defined by | `Containerfile` | same | same | +| `lc build` | build into local store | build via podman-hpc | BinderHub service → deployment registry | +| recipe isolation | `docker run` wrap | `podman-hpc run` wrap | worker pod **is** the image | +| cluster lifecycle | owned per-run | owned per-run | owned per-run (attach opt-in) | +| manifest truth | declared spec + code_version(tag) | same | same + `worker_image` ground truth | diff --git a/docs/user/cluster.md b/docs/user/cluster.md index 41da87c4..e0339810 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -13,8 +13,9 @@ hardware to spread across. 2. **Inside a SLURM allocation** → an in-process scheduler bound to the driver's hostname, with one `dask worker` per allocated node launched via `srun`. -3. **On a lightcone JupyterHub deployment** → attach to your running - Dask Gateway cluster (see the JupyterHub section below — don't use +3. **On a lightcone JupyterHub deployment** → create a run-scoped Dask + Gateway cluster with the project's worker image, and cull it when + the run finishes (see the JupyterHub section below — don't use `DASK_SCHEDULER_ADDRESS` with a `gateway://` address, it cannot be dialled directly). 4. With `DASK_SCHEDULER_ADDRESS` set → connect to whatever scheduler @@ -177,40 +178,65 @@ own scheduler. It does *not* tear the scheduler down on exit. ## JupyterHub / Dask Gateway On a lightcone JupyterHub deployment (where the `DASK_GATEWAY__*` env -vars are ambient in every pod), the model is: **you** create the -cluster, `lc run` attaches to it. - -1. Create a Dask Gateway cluster from JupyterLab — the Dask sidebar's - **+ NEW** button, or a notebook: - - ```python - from dask_gateway import Gateway - # shutdown_on_close=False keeps the cluster alive when this kernel - # exits — otherwise a kernel restart kills any lc run attached to it. - # Idle clusters are reaped by the deployment after 30 min. - cluster = Gateway().new_cluster(shutdown_on_close=False) - cluster.adapt(minimum=1, maximum=8) - ``` - -2. Run: - - ```bash - lc run - ``` - -The Gateway API only shows *your* clusters, so with a single cluster -running `lc run` attaches with zero configuration — no cluster name to -copy. It never changes the cluster's scaling and leaves it running on -exit, mirroring the `DASK_SCHEDULER_ADDRESS` convention, so you can -iterate `lc run` against the same warm cluster with its dashboard -panels docked. With no cluster running, `lc run` tells you how to -create one; with several, pick one: +vars are ambient in every pod), `lc run` manages compute exactly like +it does on your laptop — you just run it: + +```bash +cd ~/my-analysis +lc run +``` + +Under the hood, each run: + +1. **Makes sure the worker image is up to date.** The project's + `Containerfile` (plus its dependency files) defines the worker-pod + environment. If those files changed since the last build, `lc run` + commits them, pushes, and drives an image build through the hub's + BinderHub service into the deployment registry. When the registry + already holds the image — the common case — this is a single fast + round-trip. Code-only edits never trigger a rebuild: your code + reaches the workers through your shared home directory, not the + image. +2. **Creates a run-scoped Dask Gateway cluster** with that image, + scaled adaptively between one worker and `--jobs`. +3. Runs the pipeline (same executor, same per-recipe resource hints as + everywhere else; failed recipes get their output tail forwarded + back to your terminal). +4. **Shuts the cluster down.** Nothing to clean up; a crashed run's + cluster is reaped by the deployment's idle timeout. + +Because the project must be reachable by the BinderHub build pods, it +needs a git remote (a public GitHub repo today). `lc build` runs the +same ensure-image step explicitly and prints the resulting image ref; +`lc build --no-commit` refuses instead of auto-committing. + +### Attaching to a long-lived cluster + +To iterate repeatedly against one warm cluster (dashboard panels +docked, no per-run cluster startup), create it yourself from JupyterLab +and point `lc run` at it by name: + +```python +from dask_gateway import Gateway +# shutdown_on_close=False keeps the cluster alive when this kernel +# exits. Idle clusters are reaped by the deployment after 30 min. +cluster = Gateway().new_cluster(shutdown_on_close=False) +cluster.adapt(minimum=1, maximum=8) +cluster.name +``` ```bash export LIGHTCONE_GATEWAY_CLUSTER= # e.g. hub.a1b2c3... lc run ``` +Attached clusters are yours: `lc run` never rescales them and leaves +them running on exit, mirroring the `DASK_SCHEDULER_ADDRESS` +convention. The trade-off: an attached cluster's image is fixed at its +creation, so `lc run` can only *warn* when it drifts from the +project's current image (unset `LIGHTCONE_GATEWAY_CLUSTER` to get back +to the always-fresh default). + A Gateway scheduler's `gateway://` address cannot be used with `DASK_SCHEDULER_ADDRESS` — attachment is always by name. @@ -224,33 +250,25 @@ container runtime (`container_runtime: kubernetes`, declared by the site, no warning fired). Your recipes run *unwrapped inside the worker pod*, so the worker image **is** the project environment: -- A `container:` spec naming a registry image is used directly: create - your Gateway cluster with that image. -- A `container: Containerfile` spec resolves to a content-addressed - ref in the deployment registry: - `$LIGHTCONE_REGISTRY/lc-:` — the *same hash* the - local `lc build` tag carries, so the environment is provably the - same artifact on every path. `lc build` on the hub checks whether - that ref exists in the registry and, when it doesn't, prints the - publish commands: run `lc build` from a clone on any machine with - docker (this builds from the hash-attested, filtered build context — - don't substitute a raw `docker build .`), then - `docker tag lc-- ` and `docker push `. +- A `container:` spec naming a registry image is used directly as the + cluster's worker image — no build involved. +- A `container: Containerfile` spec is built *on the hub* through the + BinderHub service, as described above. repo2docker only recognizes + `Dockerfile`, so `lc build` maintains a committed + `Dockerfile → Containerfile` symlink at the project root. +- No `container:` at all → the deployment's default worker image, + which ships the lightcone stack. For the image to work as a Gateway worker it must contain `dask`, `distributed`, `dask-gateway`, and `lightcone-cli` at versions -matching the hub — the simplest way is to base your Containerfile on -the deployment's worker image -(`FROM /lightcone-worker-default:`) and add your -science deps on top. That one Containerfile then serves every path: -built locally it wraps recipes on your laptop; pushed to the registry -it runs them as pods on the hub. - -At attach time `lc run` verifies the cluster's actual worker image -against the project's resolved image and warns on mismatch. Manifests -record the image the worker pod actually ran (`worker_image`), so -provenance stays truthful even if you knowingly run on a stale -cluster. +matching the hub — the scaffold `lc init` writes installs +`lightcone-cli[gateway]` at pinned versions for exactly this reason. +That one Containerfile then serves every path: built locally it wraps +recipes on your laptop; built by the hub it runs them as worker pods. + +Manifests record the image the worker pod actually ran +(`worker_image`), so provenance stays truthful even when you knowingly +iterate on a stale attached cluster. ## NERSC Perlmutter: site-specific notes From 38a2898e0baa4a52eed94a08a95a3edaa46659c3 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 12:06:21 +0200 Subject: [PATCH 5/9] lc init: GitHub connect step (device-flow auth, create/connect repo) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A lightcone project wants a GitHub remote from day one — it backs the analysis up and it is what the hub's image builder clones — so lc init now actively offers to connect one, streamlined to two prompts: - Auth resolves silently from GH_TOKEN/GITHUB_TOKEN or an authenticated gh CLI; otherwise lc runs GitHub's OAuth device flow natively (stdlib urllib, no new dependencies) using the deployment-injected LIGHTCONE_GITHUB_CLIENT_ID, and hands the token to gh (gh auth login --with-token + gh auth setup-git) so one authorization powers git push and every gh workflow an agent may need. Credentials land in the user's home — the NFS volume on a hub, so they survive server restarts. - One free-form repository prompt (name | owner/name | URL): existing repos are connected, missing ones created — private supported, with the visibility default flipped to public on a hub site (the binder build pods clone anonymously) and a caveat printed when private is chosen there. - The scaffold is committed (repo-local git identity derived from the GitHub login when none is configured — fresh hub pods have none) and pushed with origin set upstream. Never fatal: failures and Ctrl-C degrade to a printed 'connect later' hint and init completes locally. Non-interactive runs use --github/--private/--public/--no-github; with no flags they get a one-line hint and never block. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/cli/init.md | 40 +++- docs/user/cluster.md | 9 +- src/lightcone/cli/commands.py | 181 ++++++++++++++- src/lightcone/engine/github.py | 388 +++++++++++++++++++++++++++++++++ tests/test_cli.py | 85 ++++++++ tests/test_github.py | 318 +++++++++++++++++++++++++++ 6 files changed, 1013 insertions(+), 8 deletions(-) create mode 100644 src/lightcone/engine/github.py create mode 100644 tests/test_github.py diff --git a/docs/cli/init.md b/docs/cli/init.md index 045d7a50..c7591c44 100644 --- a/docs/cli/init.md +++ b/docs/cli/init.md @@ -36,12 +36,46 @@ universes/ # placeholder; populate via `astra universe genera |--------|---------|--------| | `--no-git` | off | Skip `git init`. | | `--no-venv` | off | Skip `python -m venv .venv`. | +| `--github [NAME\|OWNER/NAME\|URL]` | prompt (TTY only) | Connect the project to this GitHub repository without prompting; created if it doesn't exist. | +| `--private` / `--public` | prompt (public on a hub, private elsewhere) | Visibility when `--github` creates a new repository. | +| `--no-github` | off | Skip the GitHub connection step. | | `--permissions {yolo,recommended,minimal}` | `recommended` | Which `.claude/settings.json` permission tier to install. | > The historical `--target`, `--existing-project`, and `--sub-analysis` -> flags have been removed; today's `lc init` only knows the three flags -> above. For migrating an existing project, run `lc init` in a fresh -> directory and use the `/lc-from-code` skill from inside Claude Code. +> flags have been removed. For migrating an existing project, run +> `lc init` in a fresh directory and use the `/lc-from-code` skill from +> inside Claude Code. + +## The GitHub step + +A repository backs the analysis up, makes it shareable, and on a +lightcone JupyterHub deployment it is what the image builder clones for +cloud runs — so `lc init` actively offers to connect one. On a TTY it +asks; scripted runs use the flags (no flags → a one-line hint, nothing +blocks). The flow: + +1. **Auth**: an ambient `GH_TOKEN`/`GITHUB_TOKEN` or an authenticated + `gh` CLI is used silently. Otherwise, where a device-flow client id + is configured (`LIGHTCONE_GITHUB_CLIENT_ID` — injected on the hub), + `lc` runs GitHub's device authorization natively: enter a one-time + code at github.com/login/device, done. The token is handed to `gh` + (`gh auth login --with-token` + `gh auth setup-git`) so the single + authorization also powers `git push` and everything an agent might + do with `gh`; it lives in your home directory, which on the hub is + the NFS volume — it survives server restarts. +2. **Repository**: one free-form prompt accepting a bare `name` (under + your account), `owner/name`, or the URL of an existing repo. If it + exists it is connected; if not it is created — private supported + (note: the hub's image builder can only clone *public* repos until + the deployment configures a clone token, so on the hub the + visibility prompt defaults to public). +3. **Push**: the scaffold is committed (a repo-local git identity is + derived from your GitHub login if none is configured) and pushed + with `origin` set upstream. + +The step is never fatal: any failure or Ctrl-C leaves a fully +scaffolded local project and prints how to connect later +(`gh repo create --source . --push`). ## Permission tiers diff --git a/docs/user/cluster.md b/docs/user/cluster.md index e0339810..4b1dbee1 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -206,9 +206,12 @@ Under the hood, each run: cluster is reaped by the deployment's idle timeout. Because the project must be reachable by the BinderHub build pods, it -needs a git remote (a public GitHub repo today). `lc build` runs the -same ensure-image step explicitly and prints the resulting image ref; -`lc build --no-commit` refuses instead of auto-committing. +needs a GitHub remote (public, today) — `lc init` offers to create and +connect one as part of scaffolding, including GitHub authentication via +a one-time device code (see the [`lc init` GitHub step](../cli/init.md)). +`lc build` runs the same ensure-image step explicitly and prints the +resulting image ref; `lc build --no-commit` refuses instead of +auto-committing. ### Attaching to a long-lived cluster diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 3a523317..af00d65c 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -144,6 +144,27 @@ def _project_root(start: Path | None = None) -> Path: @click.argument("directory", type=click.Path(path_type=Path), default=".") @click.option("--no-git", is_flag=True, help="Skip git init") @click.option("--no-venv", is_flag=True, help="Skip Python venv creation") +@click.option( + "--github", + "github_repo", + default=None, + metavar="[NAME|OWNER/NAME|URL]", + help=( + "Connect the project to this GitHub repository (created if it " + "doesn't exist) without prompting" + ), +) +@click.option( + "--private/--public", + "github_private", + default=None, + help="Visibility when --github creates a new repository", +) +@click.option( + "--no-github", + is_flag=True, + help="Skip the GitHub connection step entirely", +) @click.option( "--permissions", type=click.Choice(["yolo", "recommended", "minimal"]), @@ -165,6 +186,9 @@ def init( directory: Path, no_git: bool, no_venv: bool, + github_repo: str | None, + github_private: bool | None, + no_github: bool, permissions: str, scratch_override: str | None, ) -> None: @@ -174,8 +198,8 @@ def init( base ``.gitignore``, ``src/``) to ``astra init``, then layers on the lightcone-specific bits: ``Containerfile`` + ``requirements.txt``, ``.lightcone/`` project state, ``.claude/`` plugin bundle, ``CLAUDE.md``, - a template MyST report (``myst.yml`` + ``index.md``), and an optional - Python venv. + a template MyST report (``myst.yml`` + ``index.md``), an optional + GitHub repository connection, and an optional Python venv. """ console.print(f"[cyan]{_LIGHTCONE}[/cyan]") @@ -245,6 +269,13 @@ def init( subprocess.run(["git", "init", "-q"], cwd=directory, check=False) console.print("[green]✓[/green] Initialized git repository") + # GitHub connection: encouraged (a repo backs the analysis up and is + # what a hub deployment's image builder clones), never forced. + if not no_github and not no_git: + _connect_github( + directory, repo_input=github_repo, private=github_private + ) + # venv if not no_venv: if shutil.which("uv"): @@ -319,6 +350,152 @@ def init( ) +def _connect_github( + directory: Path, *, repo_input: str | None, private: bool | None +) -> None: + """The ``lc init`` GitHub step: authenticate, create/connect, push. + + Deliberately non-fatal: scaffolding must never be lost to a network + hiccup or a declined prompt — any failure degrades to a printed + "here's how to do it later" and init continues. Interactive prompts + appear only on a TTY; scripted/agent runs use ``--github``/ + ``--no-github`` (no flags → a one-line hint, no blocking). + """ + from lightcone.engine import github as gh_engine + from lightcone.engine.site_registry import detect_current_site + + probe = subprocess.run( + ["git", "remote", "get-url", "origin"], + cwd=directory, + capture_output=True, + text=True, + check=False, + ) + if probe.returncode == 0: + console.print( + f"[green]✓[/green] GitHub remote already configured " + f"([cyan]{probe.stdout.strip()}[/cyan])" + ) + return + + interactive = ( + repo_input is None and sys.stdin.isatty() and sys.stdout.isatty() + ) + if repo_input is None and not interactive: + console.print( + "[dim]No GitHub repository connected (non-interactive run). " + "Connect later with `lc init --github ` semantics: " + "`gh repo create --source . --push`.[/dim]" + ) + return + + on_hub = detect_current_site().key == "lightcone-hub" + try: + console.print( + "\n[bold]GitHub[/bold] [dim]— a repository backs your analysis " + "up and is what the hub's image builder clones for cloud " + "runs.[/dim]" + ) + if interactive and not click.confirm( + " Connect this project to a GitHub repository?", default=True + ): + console.print( + " [dim]Skipped. Later: `gh repo create --source . " + "--push`.[/dim]" + ) + return + + identity = gh_engine.discover_identity() + if identity is None and gh_engine.device_flow_client_id(): + + def _show_code(code: str, uri: str) -> None: + console.print( + f" To authorize, enter code [bold cyan]{code}[/bold cyan] " + f"at [cyan]{uri}[/cyan]\n" + " [dim]Waiting for authorization… (Ctrl-C to skip)[/dim]" + ) + + identity = gh_engine.device_flow(_show_code) + where = gh_engine.persist_token(identity) + console.print( + f"[green]✓[/green] Authenticated as " + f"[cyan]{identity.login}[/cyan]" + + (f" [dim]({where})[/dim]" if where else "") + ) + elif identity is None and interactive and shutil.which("gh"): + if click.confirm( + " No GitHub credential found. Run `gh auth login` now?", + default=True, + ): + subprocess.run(["gh", "auth", "login"], check=False) + identity = gh_engine.discover_identity() + if identity is None: + console.print( + " [yellow]No GitHub credential available.[/yellow] " + "[dim]Authenticate (`gh auth login`) and connect later " + "with `gh repo create --source . --push`.[/dim]" + ) + return + if identity.source != "device": + console.print( + f"[green]✓[/green] Authenticated as " + f"[cyan]{identity.login}[/cyan] [dim](via {identity.source})[/dim]" + ) + + raw = repo_input or click.prompt( + " Repository (name, owner/name, or URL of an existing repo)", + default=directory.name, + ) + target = gh_engine.resolve_repo(identity, raw) + if target.exists: + console.print( + f"[green]✓[/green] Connecting to existing repository " + f"[cyan]{target.full_name}[/cyan]" + ) + else: + if private is None: + # The hub's image builder clones anonymously, so public + # is the default that keeps `lc run` working there; + # everywhere else default to private. + private = ( + click.confirm( + " Keep the repository private?", default=not on_hub + ) + if interactive + else not on_hub + ) + gh_engine.create_repo(identity, target, private=private) + console.print( + f"[green]✓[/green] Created " + f"{'private' if private else 'public'} repository " + f"[cyan]{target.full_name}[/cyan]" + ) + if private and on_hub: + console.print( + " [yellow]⚠ The hub's image builder can only clone " + "public repositories today[/yellow] — cloud runs will " + "fail to build until the repo is public\n" + " ([cyan]gh repo edit --visibility public[/cyan]) or " + "the deployment gets a clone token." + ) + gh_engine.connect_and_push(directory, identity, target) + console.print( + f"[green]✓[/green] Pushed initial commit → " + f"[cyan]{target.url}[/cyan]" + ) + except KeyboardInterrupt: + console.print( + "\n [dim]GitHub step skipped. Later: `gh repo create " + "--source . --push`.[/dim]" + ) + except gh_engine.GitHubError as e: + console.print( + f" [yellow]GitHub step failed:[/yellow] {e}\n" + " [dim]The project is fully scaffolded locally; connect " + "later with `gh repo create --source . --push`.[/dim]" + ) + + _CONTAINERFILE = """\ FROM python:3.12-slim diff --git a/src/lightcone/engine/github.py b/src/lightcone/engine/github.py new file mode 100644 index 00000000..e38e6a8d --- /dev/null +++ b/src/lightcone/engine/github.py @@ -0,0 +1,388 @@ +"""GitHub connection for ``lc init``: device-flow auth + repo bootstrap. + +A lightcone project wants a GitHub remote from day one: it backs the +analysis up, makes it shareable, and on a JupyterHub deployment it is +what the image builder clones (see :mod:`lightcone.engine.binder`). +This module gives ``lc init`` a streamlined connect step: + +- **Credential discovery** (:func:`discover_identity`): an ambient + ``GH_TOKEN``/``GITHUB_TOKEN`` or an already-authenticated ``gh`` CLI + is used silently — users who set nothing up twice are never asked + twice. +- **Device-flow authorization** (:func:`device_flow`): when no + credential exists, the standard OAuth device flow (the same one + ``gh auth login`` uses) — print a one-time code, the user approves it + at github.com/login/device in any browser, we poll for the token. + Needs only a public OAuth *client id* (:data:`GITHUB_CLIENT_ID_ENV`; + a lightcone-hub deployment injects its hub OAuth app's id into every + user pod, so on the hub this works with zero configuration). +- **Token persistence** (:func:`persist_token`): the token is handed to + ``gh`` (``gh auth login --with-token`` + ``gh auth setup-git``) so a + single authorization also powers ``git push`` and everything an agent + may want to do with ``gh`` (PRs, issues, releases). Without ``gh``, + git's ``credential.helper store`` is configured instead. Either way + the credential lives in the user's home — on the hub that is the NFS + volume, so it survives server restarts. +- **Repo bootstrap** (:func:`resolve_repo`, :func:`create_repo`): + create a new repository (private supported) or connect to an existing + one, from a single free-form input (``name``, ``owner/name``, or a + full URL). + +Everything speaks to GitHub over plain HTTPS (stdlib ``urllib``): the +three endpoints involved don't justify an SDK dependency. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import time +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +#: Public OAuth-app client id used for the device flow. Injected into +#: user pods by a lightcone-hub deployment (the hub login app, with +#: device flow enabled); settable by hand anywhere else. +GITHUB_CLIENT_ID_ENV = "LIGHTCONE_GITHUB_CLIENT_ID" + +#: Scopes requested by the device flow: `repo` to create/push (private +#: included), `read:org` so org-owned repos can be targeted. +_DEVICE_SCOPE = "repo read:org" + +_DEVICE_CODE_URL = "https://github.com/login/device/code" +_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" # noqa: S105 +_API_URL = "https://api.github.com" + + +class GitHubError(RuntimeError): + """A GitHub connection step failed (auth, API, or git plumbing).""" + + +@dataclass(frozen=True) +class GitHubIdentity: + """An authenticated GitHub credential and who it belongs to.""" + + token: str + login: str + source: str # "env" | "gh" | "device" + + +@dataclass(frozen=True) +class RepoTarget: + """Where a project should live on GitHub.""" + + owner: str + name: str + exists: bool + + @property + def full_name(self) -> str: + return f"{self.owner}/{self.name}" + + @property + def url(self) -> str: + return f"https://github.com/{self.full_name}" + + +# --------------------------------------------------------------------------- +# HTTP plumbing +# --------------------------------------------------------------------------- + + +def _post_form(url: str, data: dict[str, str]) -> dict[str, object]: + req = urllib.request.Request( + url, + data=urllib.parse.urlencode(data).encode(), + headers={"Accept": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + payload = json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, OSError, ValueError) as exc: + raise GitHubError(f"Could not reach GitHub ({exc}).") from exc + if not isinstance(payload, dict): + raise GitHubError(f"Unexpected response from {url}.") + return payload + + +def _api( + token: str, method: str, path: str, payload: dict[str, object] | None = None +) -> tuple[int, dict[str, object]]: + """One GitHub REST call; returns ``(status, body)``, 4xx included.""" + req = urllib.request.Request( + f"{_API_URL}{path}", + data=json.dumps(payload).encode() if payload is not None else None, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + **({"Content-Type": "application/json"} if payload else {}), + }, + method=method, + ) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + body = json.loads(resp.read().decode("utf-8") or "{}") + return resp.status, body if isinstance(body, dict) else {} + except urllib.error.HTTPError as exc: + try: + body = json.loads(exc.read().decode("utf-8") or "{}") + except (OSError, ValueError): + body = {} + return exc.code, body if isinstance(body, dict) else {} + except (urllib.error.URLError, OSError, ValueError) as exc: + raise GitHubError(f"Could not reach the GitHub API ({exc}).") from exc + + +def whoami(token: str) -> str | None: + """The token's login, or ``None`` when the token doesn't authenticate.""" + status, body = _api(token, "GET", "/user") + login = body.get("login") + return login if status == 200 and isinstance(login, str) else None + + +# --------------------------------------------------------------------------- +# Credential discovery + device flow +# --------------------------------------------------------------------------- + + +def discover_identity() -> GitHubIdentity | None: + """An already-available GitHub credential, validated, or ``None``. + + Precedence: explicit env (``GH_TOKEN``/``GITHUB_TOKEN``) → the gh + CLI's stored credential. Tokens that no longer authenticate are + skipped rather than surfaced — a revoked token should route the + user to re-authorization, not to an API error later. + """ + for env in ("GH_TOKEN", "GITHUB_TOKEN"): + token = (os.environ.get(env) or "").strip() + if token and (login := whoami(token)): + return GitHubIdentity(token=token, login=login, source="env") + + if shutil.which("gh"): + result = subprocess.run( + ["gh", "auth", "token"], capture_output=True, text=True, check=False + ) + token = result.stdout.strip() + if result.returncode == 0 and token and (login := whoami(token)): + return GitHubIdentity(token=token, login=login, source="gh") + + return None + + +def device_flow_client_id() -> str | None: + """The OAuth client id for the device flow, or ``None`` off-hub.""" + return (os.environ.get(GITHUB_CLIENT_ID_ENV) or "").strip() or None + + +def device_flow( + on_code: Callable[[str, str], None], + *, + timeout: float = 900.0, +) -> GitHubIdentity: + """Run the OAuth device flow; return the authorized identity. + + *on_code* is called once with ``(user_code, verification_uri)`` so + the caller can render the "enter this code there" prompt however it + likes; this function then blocks, polling at GitHub's requested + interval, until the user approves (or *timeout*). + """ + client_id = device_flow_client_id() + if client_id is None: + raise GitHubError( + f"No GitHub OAuth client id configured ({GITHUB_CLIENT_ID_ENV} " + "is unset), so lc cannot start a device authorization. " + "Run `gh auth login` instead, then re-run." + ) + + grant = _post_form( + _DEVICE_CODE_URL, {"client_id": client_id, "scope": _DEVICE_SCOPE} + ) + if "device_code" not in grant: + raise GitHubError( + "GitHub refused to start a device authorization " + f"({grant.get('error_description') or grant.get('error') or grant}). " + "Is device flow enabled on the OAuth app?" + ) + on_code(str(grant["user_code"]), str(grant["verification_uri"])) + + raw_interval = grant.get("interval") + interval = ( + float(raw_interval) if isinstance(raw_interval, (int, float)) else 5.0 + ) + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + time.sleep(interval) + poll = _post_form( + _ACCESS_TOKEN_URL, + { + "client_id": client_id, + "device_code": str(grant["device_code"]), + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + }, + ) + token = poll.get("access_token") + if isinstance(token, str) and token: + login = whoami(token) + if login is None: + raise GitHubError("GitHub issued a token that does not work.") + return GitHubIdentity(token=token, login=login, source="device") + error = poll.get("error") + if error == "authorization_pending": + continue + if error == "slow_down": + interval += 5 + continue + raise GitHubError( + "GitHub device authorization failed " + f"({poll.get('error_description') or error})." + ) + raise GitHubError("Timed out waiting for the device authorization.") + + +def persist_token(identity: GitHubIdentity) -> str: + """Store a freshly obtained token where git and gh will find it. + + Preferred home is the gh CLI (one credential powering both ``gh`` + and, via ``gh auth setup-git``, https ``git push``); fallback is + git's plain ``store`` helper. Returns a short human-readable + description of where the credential went. No-op for credentials + that already came from storage (``env``/``gh`` sources). + """ + if identity.source != "device": + return "" + if shutil.which("gh"): + login = subprocess.run( + ["gh", "auth", "login", "--with-token"], + input=identity.token, + capture_output=True, + text=True, + check=False, + ) + if login.returncode == 0: + subprocess.run( + ["gh", "auth", "setup-git"], + capture_output=True, + text=True, + check=False, + ) + return "stored via gh (git push + gh CLI ready)" + cred_file = Path.home() / ".git-credentials" + with cred_file.open("a") as f: + f.write(f"https://x-access-token:{identity.token}@github.com\n") + cred_file.chmod(0o600) + subprocess.run( + ["git", "config", "--global", "credential.helper", "store"], + capture_output=True, + check=False, + ) + return "stored in ~/.git-credentials (git push ready)" + + +# --------------------------------------------------------------------------- +# Repo bootstrap +# --------------------------------------------------------------------------- + + +def resolve_repo(identity: GitHubIdentity, raw: str) -> RepoTarget: + """Interpret free-form repo input against GitHub. + + *raw* may be a bare ``name`` (owner defaults to the authenticated + user), ``owner/name``, or a full ``https://github.com/owner/name`` + / ``git@github.com:owner/name`` URL. The returned target says + whether the repository already exists (→ connect) or not (→ create). + """ + raw = raw.strip().removesuffix(".git") + if raw.startswith("git@github.com:"): + raw = raw[len("git@github.com:"):] + elif "://" in raw: + parsed = urllib.parse.urlparse(raw) + if parsed.hostname != "github.com": + raise GitHubError( + f"Only github.com repositories are supported (got {raw!r})." + ) + raw = parsed.path.strip("/") + parts = [p for p in raw.split("/") if p] + if len(parts) == 1: + owner, name = identity.login, parts[0] + elif len(parts) == 2: + owner, name = parts + else: + raise GitHubError(f"Cannot interpret {raw!r} as a GitHub repository.") + + status, _ = _api(identity.token, "GET", f"/repos/{owner}/{name}") + return RepoTarget(owner=owner, name=name, exists=status == 200) + + +def create_repo( + identity: GitHubIdentity, target: RepoTarget, *, private: bool +) -> None: + """Create *target* under the user or an organization.""" + if target.owner == identity.login: + path = "/user/repos" + else: + path = f"/orgs/{target.owner}/repos" + status, body = _api( + identity.token, + "POST", + path, + {"name": target.name, "private": private}, + ) + if status not in (200, 201): + raise GitHubError( + f"Could not create {target.full_name} " + f"({body.get('message') or f'HTTP {status}'})." + ) + + +def connect_and_push( + project: Path, identity: GitHubIdentity, target: RepoTarget +) -> None: + """Point ``origin`` at *target*, commit the scaffold, push upstream. + + Fresh environments (a new hub pod) usually have no git identity — + a repo-local one is derived from the GitHub login so the initial + commit never fails on ``user.email`` being unset. + """ + + def git(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["git", *args], cwd=project, capture_output=True, text=True, + check=False, + ) + + if git("config", "user.email").returncode != 0: + git("config", "user.name", identity.login) + git( + "config", + "user.email", + f"{identity.login}@users.noreply.github.com", + ) + + if git("rev-parse", "HEAD").returncode != 0: + git("add", "-A") + commit = git("commit", "-q", "-m", "Initial lightcone project scaffold") + if commit.returncode != 0: + raise GitHubError( + f"Could not create the initial commit: " + f"{commit.stderr.strip() or commit.stdout.strip()}" + ) + + if git("remote", "get-url", "origin").returncode == 0: + git("remote", "set-url", "origin", f"{target.url}.git") + else: + git("remote", "add", "origin", f"{target.url}.git") + + push = git("push", "-q", "-u", "origin", "HEAD") + if push.returncode != 0: + raise GitHubError( + f"Connected origin to {target.url} but the initial push failed:\n" + f"{push.stderr.strip()}\n" + "Fix and push manually (`git push -u origin HEAD`)." + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 81fdf090..1a2b2f27 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -478,3 +478,88 @@ def test_worker_image_for_run_no_container_uses_default( monkeypatch.chdir(project) monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") assert _worker_image_for_run(project, verbose=False) is None + + +# ---- lc init GitHub step --------------------------------------------------- + + +def test_init_non_interactive_skips_github_with_hint( + runner: CliRunner, tmp_path: Path +) -> None: + project = tmp_path / "proj" + result = runner.invoke(main, ["init", str(project), "--no-venv"]) + assert result.exit_code == 0, result.output + assert "No GitHub repository connected" in result.output + + +def test_init_no_github_flag_skips_silently( + runner: CliRunner, tmp_path: Path +) -> None: + project = tmp_path / "proj" + result = runner.invoke( + main, ["init", str(project), "--no-venv", "--no-github"] + ) + assert result.exit_code == 0, result.output + assert "GitHub" not in result.output.replace( + "Sign in with GitHub", "" + ) or "No GitHub repository connected" not in result.output + + +def test_init_github_flag_connects_non_interactively( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.engine.github import GitHubIdentity, RepoTarget + + created: dict[str, object] = {} + identity = GitHubIdentity(token="tok", login="eiffl", source="gh") + + monkeypatch.setattr( + "lightcone.engine.github.discover_identity", lambda: identity + ) + monkeypatch.setattr( + "lightcone.engine.github.resolve_repo", + lambda ident, raw: RepoTarget("eiffl", raw, exists=False), + ) + + def fake_create(ident, target, *, private): # noqa: ANN001, ANN202 + created["target"] = target.full_name + created["private"] = private + + def fake_push(directory, ident, target): # noqa: ANN001, ANN202 + created["pushed"] = True + + monkeypatch.setattr("lightcone.engine.github.create_repo", fake_create) + monkeypatch.setattr( + "lightcone.engine.github.connect_and_push", fake_push + ) + + project = tmp_path / "proj" + result = runner.invoke( + main, + ["init", str(project), "--no-venv", "--github", "myproj", "--private"], + ) + assert result.exit_code == 0, result.output + assert created == { + "target": "eiffl/myproj", + "private": True, + "pushed": True, + } + assert "Created private repository" in result.output + + +def test_init_github_failure_does_not_fail_init( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.engine.github import GitHubError + + def boom() -> None: + raise GitHubError("network down") + + monkeypatch.setattr("lightcone.engine.github.discover_identity", boom) + project = tmp_path / "proj" + result = runner.invoke( + main, ["init", str(project), "--no-venv", "--github", "myproj"] + ) + assert result.exit_code == 0, result.output + assert "GitHub step failed" in result.output + assert (project / "astra.yaml").exists() diff --git a/tests/test_github.py b/tests/test_github.py new file mode 100644 index 00000000..aa1538a3 --- /dev/null +++ b/tests/test_github.py @@ -0,0 +1,318 @@ +"""Unit tests for the `lc init` GitHub connection (`lightcone.engine.github`). + +All GitHub HTTP traffic is faked at the urllib boundary; the git side of +`connect_and_push` runs against real throwaway repos with subprocess +intercepted only where a network push would happen. +""" + +from __future__ import annotations + +import io +import json +import urllib.request +from typing import Any + +import pytest + +from lightcone.engine.github import ( + GITHUB_CLIENT_ID_ENV, + GitHubError, + GitHubIdentity, + RepoTarget, + create_repo, + device_flow, + device_flow_client_id, + discover_identity, + resolve_repo, +) + +_ID = GitHubIdentity(token="tok", login="eiffl", source="env") + + +class _FakeResponse: + def __init__(self, payload: object, status: int = 200) -> None: + self._body = json.dumps(payload).encode() + self.status = status + + def read(self) -> bytes: + return self._body + + def __enter__(self) -> _FakeResponse: + return self + + def __exit__(self, *args: object) -> bool: + return False + + +def _install_fake_http( + monkeypatch: pytest.MonkeyPatch, + routes: dict[tuple[str, str], list[object]], + calls: list[tuple[str, str, object]], +) -> None: + """Route (method, url) → successive responses; record request bodies. + + A response entry that is an int becomes an HTTPError with that code + (body `{}`); anything else is served as JSON. + """ + + def fake_urlopen(req: urllib.request.Request, timeout: float = 0): # noqa: ANN202 + url = req.full_url.split("?")[0] + method = req.get_method() + body = req.data.decode() if req.data else "" + calls.append((method, url, body)) + queue = routes[(method, url)] + entry = queue.pop(0) if len(queue) > 1 else queue[0] + if isinstance(entry, int): + raise urllib.error.HTTPError( + url, entry, "err", hdrs=None, fp=io.BytesIO(b"{}") # type: ignore[arg-type] + ) + return _FakeResponse(entry) + + import urllib.error + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + + +# --------------------------------------------------------------------------- +# Device flow +# --------------------------------------------------------------------------- + + +def test_device_flow_needs_client_id(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(GITHUB_CLIENT_ID_ENV, raising=False) + assert device_flow_client_id() is None + with pytest.raises(GitHubError, match="gh auth login"): + device_flow(lambda c, u: None) + + +def test_device_flow_polls_until_authorized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(GITHUB_CLIENT_ID_ENV, "cid123") + monkeypatch.setattr("time.sleep", lambda s: None) + calls: list[tuple[str, str, object]] = [] + _install_fake_http( + monkeypatch, + { + ("POST", "https://github.com/login/device/code"): [ + { + "device_code": "dev1", + "user_code": "ABCD-1234", + "verification_uri": "https://github.com/login/device", + "interval": 0, + } + ], + ("POST", "https://github.com/login/oauth/access_token"): [ + {"error": "authorization_pending"}, + {"access_token": "tok99", "token_type": "bearer"}, + ], + ("GET", "https://api.github.com/user"): [{"login": "eiffl"}], + }, + calls, + ) + + shown: list[tuple[str, str]] = [] + identity = device_flow(lambda code, uri: shown.append((code, uri))) + assert identity == GitHubIdentity( + token="tok99", login="eiffl", source="device" + ) + assert shown == [("ABCD-1234", "https://github.com/login/device")] + # The device-code request carried our client id and the repo scope. + first_body = str(calls[0][2]) + assert "cid123" in first_body and "repo" in first_body + + +def test_device_flow_denied_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(GITHUB_CLIENT_ID_ENV, "cid123") + monkeypatch.setattr("time.sleep", lambda s: None) + _install_fake_http( + monkeypatch, + { + ("POST", "https://github.com/login/device/code"): [ + { + "device_code": "dev1", + "user_code": "ABCD-1234", + "verification_uri": "https://github.com/login/device", + "interval": 0, + } + ], + ("POST", "https://github.com/login/oauth/access_token"): [ + {"error": "access_denied", "error_description": "user said no"} + ], + }, + [], + ) + with pytest.raises(GitHubError, match="user said no"): + device_flow(lambda c, u: None) + + +def test_device_flow_disabled_app_raises_helpfully( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(GITHUB_CLIENT_ID_ENV, "cid123") + _install_fake_http( + monkeypatch, + { + ("POST", "https://github.com/login/device/code"): [ + {"error": "unauthorized_client"} + ] + }, + [], + ) + with pytest.raises(GitHubError, match="device flow enabled"): + device_flow(lambda c, u: None) + + +# --------------------------------------------------------------------------- +# Credential discovery +# --------------------------------------------------------------------------- + + +def test_discover_prefers_env_token(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GH_TOKEN", "envtok") + _install_fake_http( + monkeypatch, + {("GET", "https://api.github.com/user"): [{"login": "eiffl"}]}, + [], + ) + identity = discover_identity() + assert identity is not None + assert (identity.token, identity.source) == ("envtok", "env") + + +def test_discover_skips_dead_env_token_and_uses_gh( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GH_TOKEN", "revoked") + responses = {("GET", "https://api.github.com/user"): [401]} + + def fake_urlopen(req: urllib.request.Request, timeout: float = 0): # noqa: ANN202 + import urllib.error + + auth = req.get_header("Authorization") or "" + if "goodtok" in auth: + return _FakeResponse({"login": "eiffl"}) + raise urllib.error.HTTPError( + req.full_url, 401, "bad", hdrs=None, fp=io.BytesIO(b"{}") # type: ignore[arg-type] + ) + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/gh") + + class _P: + returncode = 0 + stdout = "goodtok\n" + stderr = "" + + monkeypatch.setattr( + "lightcone.engine.github.subprocess.run", lambda *a, **k: _P() + ) + identity = discover_identity() + assert identity is not None + assert (identity.token, identity.source) == ("goodtok", "gh") + del responses # (documented intent: env token was rejected first) + + +def test_discover_none_when_nothing_works( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("GH_TOKEN", raising=False) + monkeypatch.delenv("GITHUB_TOKEN", raising=False) + monkeypatch.setattr("shutil.which", lambda name: None) + assert discover_identity() is None + + +# --------------------------------------------------------------------------- +# Repo resolution + creation +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("raw", "owner", "name"), + [ + ("my-analysis", "eiffl", "my-analysis"), + ("LightconeResearch/demo", "LightconeResearch", "demo"), + ("https://github.com/Org/repo", "Org", "repo"), + ("https://github.com/Org/repo.git", "Org", "repo"), + ("git@github.com:Org/repo.git", "Org", "repo"), + ], +) +def test_resolve_repo_forms( + monkeypatch: pytest.MonkeyPatch, raw: str, owner: str, name: str +) -> None: + _install_fake_http( + monkeypatch, + {("GET", f"https://api.github.com/repos/{owner}/{name}"): [404]}, + [], + ) + target = resolve_repo(_ID, raw) + assert (target.owner, target.name, target.exists) == (owner, name, False) + + +def test_resolve_repo_detects_existing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_fake_http( + monkeypatch, + { + ("GET", "https://api.github.com/repos/eiffl/exists"): [ + {"full_name": "eiffl/exists"} + ] + }, + [], + ) + assert resolve_repo(_ID, "exists").exists is True + + +def test_resolve_repo_rejects_non_github( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(GitHubError, match="github.com"): + resolve_repo(_ID, "https://gitlab.com/org/repo") + + +def test_create_repo_user_vs_org_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, str, Any]] = [] + _install_fake_http( + monkeypatch, + { + ("POST", "https://api.github.com/user/repos"): [{"ok": True}], + ("POST", "https://api.github.com/orgs/Lightcone/repos"): [ + {"ok": True} + ], + }, + calls, + ) + create_repo( + _ID, RepoTarget("eiffl", "mine", exists=False), private=True + ) + create_repo( + _ID, RepoTarget("Lightcone", "ours", exists=False), private=False + ) + assert calls[0][1].endswith("/user/repos") + assert json.loads(str(calls[0][2])) == {"name": "mine", "private": True} + assert calls[1][1].endswith("/orgs/Lightcone/repos") + assert json.loads(str(calls[1][2])) == {"name": "ours", "private": False} + + +def test_create_repo_surfaces_api_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fake_urlopen(req: urllib.request.Request, timeout: float = 0): # noqa: ANN202 + import urllib.error + + raise urllib.error.HTTPError( + req.full_url, + 422, + "unprocessable", + hdrs=None, # type: ignore[arg-type] + fp=io.BytesIO(b'{"message": "name already exists"}'), + ) + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + with pytest.raises(GitHubError, match="name already exists"): + create_repo( + _ID, RepoTarget("eiffl", "dupe", exists=False), private=True + ) From 3e21d1fbe51e0385509338a5aaa0a50049d665d7 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 14:27:01 +0200 Subject: [PATCH 6/9] lc init/run: kill the venv-shadowing failure class; recipe-contract docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From the first agent-driven project retro on the hub, the single largest time sink was PATH shadowing: the project venv shipped its own (older) lightcone-cli + snakemake, silently shadowing the deployment's install and producing errors three layers from the cause. Two structural fixes: - lc init's venv now installs the *project's* dependencies (requirements.txt — same list the Containerfile uses), and no longer installs lightcone-cli at all: lc/snakemake simply never exist inside the venv, so activation cannot shadow them. Bonus: the science code now actually runs locally out of the box (the venv previously had lightcone-cli but not the project's own numpy/pandas). - lc run invokes snakemake as sys.executable -m snakemake instead of a PATH lookup, so the driver snakemake is structurally the one matching the running lc. Also documents the two recipe-contract papercuts the retro hit: {output} is a directory, and input source: paths are passed verbatim (stage remote data locally, never a URL). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/cli/run.md | 20 ++++++++++++++++++ src/lightcone/cli/commands.py | 40 +++++++++++++++++++++++++++++------ tests/test_cli.py | 34 +++++++++++++++++++++++++++-- 3 files changed, 86 insertions(+), 8 deletions(-) diff --git a/docs/cli/run.md b/docs/cli/run.md index 3e311be9..82719e9f 100644 --- a/docs/cli/run.md +++ b/docs/cli/run.md @@ -43,6 +43,26 @@ everything (Snakemake's `rule all`). chatter so the output reads as lightcone's, not Snakemake's. Real error content always passes through. +## Recipe contract: outputs are directories, inputs are local paths + +Two things every recipe author (human or agent) needs to know at the +point of use: + +- **`{output}` expands to a directory, not a file.** Create it and + write your artifact(s) inside (e.g. + `os.makedirs(out, exist_ok=True)` then write `out/result.json`); + opening `{output}` itself fails with `IsADirectoryError`. The + manifest (`.lightcone-manifest.json`) is written next to your files + by the runner. +- **An input's `source:` is passed to the recipe verbatim** — it is a + literal filesystem path, never fetched. Stage remote data as a local + file first (a download script can itself be an output that later + recipes take as input) and point `source:` at the staged path. A URL + in `source:` would silently arrive at your script as a nonexistent + "path". Project-relative paths work on every backend because the + project directory is shared with (or mounted into) the workers at the + same location. + ## Output qualification When the same `output_id` appears in multiple sub-analyses, you must diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index af00d65c..a351e599 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -276,7 +276,15 @@ def init( directory, repo_input=github_repo, private=github_private ) - # venv + # venv: the project's *science* environment — requirements.txt, the + # same dependency list the Containerfile installs — so the code has a + # fast local dev loop before any container/cluster is involved. + # Deliberately NOT lightcone-cli: installing a second lc (and its + # snakemake) into the venv would shadow the deployment's install the + # moment the venv is activated, and a version-skewed lc produces + # errors that point everywhere but at the shadowing (observed on the + # hub: a stale venv lc predates the site's build service and cluster + # backend). `lc` stays a tool of the ambient environment. if not no_venv: if shutil.which("uv"): with console.status("[dim]Creating virtual environment…[/dim]"): @@ -286,7 +294,7 @@ def init( check=False, capture_output=True, ) - with console.status("[dim]Installing lightcone-cli…[/dim]"): + with console.status("[dim]Installing project dependencies…[/dim]"): subprocess.run( [ "uv", @@ -294,7 +302,8 @@ def init( "install", "--python", ".venv/bin/python", - "lightcone-cli", + "-r", + "requirements.txt", ], cwd=directory, check=False, @@ -308,15 +317,26 @@ def init( check=False, capture_output=True, ) - with console.status("[dim]Installing lightcone-cli…[/dim]"): + with console.status("[dim]Installing project dependencies…[/dim]"): subprocess.run( - [".venv/bin/python", "-m", "pip", "install", "-q", "lightcone-cli"], + [ + ".venv/bin/python", + "-m", + "pip", + "install", + "-q", + "-r", + "requirements.txt", + ], cwd=directory, check=False, capture_output=True, ) console.print( - f"[green]✓[/green] Virtual environment created in [cyan]{directory}/.venv[/cyan]" + f"[green]✓[/green] Virtual environment created in " + f"[cyan]{directory}/.venv[/cyan] " + f"[dim](project dependencies; `lc` itself stays on the ambient " + "PATH)[/dim]" ) console.print(f"\n[green]Project initialized at[/green] {directory}") @@ -959,6 +979,14 @@ def _build_snakemake_cmd( workers genuinely share the driver's environment. """ cmd: list[str] = [ + # lc's own interpreter, not a PATH lookup: an activated project + # venv (or any stray env) earlier on PATH would otherwise supply + # a different snakemake than the lightcone-cli driving it — + # version skew that surfaces as executor errors three layers + # from the cause. The worker side is unaffected (the executor + # resolves the child interpreter per-branch). + sys.executable, + "-m", "snakemake", "-s", str(snakefile_path), diff --git a/tests/test_cli.py b/tests/test_cli.py index 1a2b2f27..b1dd1fb3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -143,7 +143,13 @@ def _fake_run(cmd: list[str], **kwargs: object) -> MagicMock: assert result.exit_code == 0, result.output assert ["uv", "venv", "--python", "3.12", ".venv"] in calls - assert ["uv", "pip", "install", "--python", ".venv/bin/python", "lightcone-cli"] in calls + # Project deps only — never a second lightcone-cli that would shadow + # the ambient install once the venv is activated. + assert [ + "uv", "pip", "install", "--python", ".venv/bin/python", + "-r", "requirements.txt", + ] in calls + assert not any("lightcone-cli" in c for c in calls) def test_init_venv_falls_back_to_python_when_uv_missing( @@ -163,7 +169,11 @@ def _fake_run(cmd: list[str], **kwargs: object) -> MagicMock: assert result.exit_code == 0, result.output assert ["python", "-m", "venv", ".venv"] in calls - assert [".venv/bin/python", "-m", "pip", "install", "-q", "lightcone-cli"] in calls + assert [ + ".venv/bin/python", "-m", "pip", "install", "-q", + "-r", "requirements.txt", + ] in calls + assert not any("lightcone-cli" in c for c in calls) # ---- lc verify ------------------------------------------------------------ @@ -563,3 +573,23 @@ def boom() -> None: assert result.exit_code == 0, result.output assert "GitHub step failed" in result.output assert (project / "astra.yaml").exists() + + +def test_run_cmd_uses_own_interpreter_for_snakemake() -> None: + """The driver snakemake must be lc's own environment's, not PATH's — + an activated project venv earlier on PATH must not be able to swap + in a version-skewed snakemake under the executor.""" + import sys + + from lightcone.cli.commands import _build_snakemake_cmd + + cmd = _build_snakemake_cmd( + snakefile_path=Path("/proj/.lightcone/Snakefile"), + project=Path("/proj"), + n="1", + rerun_triggers="code", + targets=[], + force=False, + has_outputs=False, + ) + assert cmd[:3] == [sys.executable, "-m", "snakemake"] From d42921588e9c6c59d889431f6b7835da1495a5f0 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 14:38:47 +0200 Subject: [PATCH 7/9] Scaffold Containerfile: passwd entry for uid 1000 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cluster deployments run project images as uid 1000; python:slim has no passwd entry for that uid and snakemake's getpass.getuser() call crashes at startup on every rule (found live: 'getpwuid(): uid not found: 1000' inside Gateway workers — invisible from the driver on pre-release worker CLIs). The hub also now injects USER/LOGNAME as a belt (lightcone-hub), but images should stand on their own. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- src/lightcone/cli/commands.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index a351e599..63b649c6 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -522,6 +522,11 @@ def _show_code(code: str, uri: str) -> None: # uv, pulled from its official image (fast, reproducible, no curl bootstrap). COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ +# Cluster deployments run this image as uid 1000; slim bases have no +# passwd entry for it, and tools that call getpass.getuser() (snakemake, +# at startup) crash on a passwd-less uid. Give it a name. +RUN useradd --uid 1000 --create-home lightcone + WORKDIR /app # On a Kubernetes / Dask Gateway deployment the pod image IS the execution From fda67b30784643c8c1ad36c916695a07909cbf06 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 15:50:22 +0200 Subject: [PATCH 8/9] Simplify: gateway is create/cull only; hub scratch defaults to $HOME MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two deletions that shrink both the CLI and the deployment contract: - The attach escape hatch (LIGHTCONE_GATEWAY_CLUSTER as a user knob) is gone: nobody iterates against sidebar-created clusters, and removing it deletes the attached-cluster path, the worker-image mismatch warning machinery, and the stale-env-var failure modes wholesale. The env var remains as the internal parent→child rendezvous only; ambient values are ignored. The gateway branch triggers on DASK_GATEWAY__ADDRESS alone. - The hub site's scratch root now defaults to $HOME (the NFS home honors flock and is mounted into worker pods at the same path), so a deployment no longer needs a shared /shared volume or the LIGHTCONE_SCRATCH env — both existed only to serve scratch. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/api/dask_cluster.md | 49 ++--- docs/user/cluster.md | 41 +--- src/lightcone/cli/commands.py | 8 +- src/lightcone/engine/dask_cluster.py | 195 ++++-------------- src/lightcone/engine/site_registry.py | 8 +- .../executor.py | 11 +- tests/test_dask_cluster.py | 190 ++++------------- 7 files changed, 114 insertions(+), 388 deletions(-) diff --git a/docs/api/dask_cluster.md b/docs/api/dask_cluster.md index 5a449b10..88d71ed1 100644 --- a/docs/api/dask_cluster.md +++ b/docs/api/dask_cluster.md @@ -15,31 +15,22 @@ Four branches in priority order: 1. **`DASK_SCHEDULER_ADDRESS` already set** → yield `{"DASK_SCHEDULER_ADDRESS": addr}` as-is. We don't own the cluster, so we don't tear it down. -2. **Dask Gateway detected** (`LIGHTCONE_GATEWAY_CLUSTER` or - `DASK_GATEWAY__ADDRESS` set, e.g. a JupyterHub pod) → yield - `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Two sub-modes: - - - **Create/cull (default).** A run-scoped cluster is created with - *expected_worker_image* as its `image` cluster option, scaled - adaptively `1..max_workers`, and shut down on exit — the same - lifetime contract as the local/SLURM branches, and the mechanism - that makes a freshly built project image take effect (a Gateway - cluster's image is fixed at creation). Startup blocks until the - first worker is live (bounded by - `LIGHTCONE_GATEWAY_WORKER_TIMEOUT`, default 600 s) so an - unpullable image fails loudly instead of hanging at zero workers. - - **Attach** (`LIGHTCONE_GATEWAY_CLUSTER=` set by the user) → - connect to that cluster, leave its scaling and lifetime untouched - (same convention as branch 1), and warn when its actual worker - image (read from the scheduler pod's `LIGHTCONE_WORKER_IMAGE`) - differs from *expected_worker_image*. - - Either way, startup fails fast if live workers don't advertise the - resource contract below (zero live workers is fine on the attach - path — adaptive clusters scale on demand). Gateway scheduler - addresses use a `gateway://` comm scheme a bare `Client` cannot - dial, so the child rejoins **by name** through the authenticated - Gateway API. Requires the optional dependency: +2. **Dask Gateway detected** (`DASK_GATEWAY__ADDRESS` set, e.g. a + JupyterHub pod) → **create a run-scoped cluster** with + *expected_worker_image* as its `image` cluster option, scaled + adaptively `1..max_workers`, and shut down on exit — the same + lifetime contract as the local/SLURM branches, and the mechanism + that makes a freshly built project image take effect (a Gateway + cluster's image is fixed at creation). Startup blocks until the + first worker is live (bounded by `LIGHTCONE_GATEWAY_WORKER_TIMEOUT`, + default 600 s) so an unpullable image fails loudly instead of + hanging at zero workers, then fails fast if workers don't advertise + the resource contract below. Yields + `{"LIGHTCONE_GATEWAY_CLUSTER": name}` — that env var is the + parent→child rendezvous only (Gateway scheduler addresses use a + `gateway://` comm scheme a bare `Client` cannot dial, so the child + rejoins **by name** through the authenticated Gateway API); an + ambient value is ignored. Requires the optional dependency: `pip install lightcone-cli[gateway]`. 3. **`SLURM_JOB_ID` set** → start an in-process scheduler bound to the driver's SLURM hostname (`SLURMD_NODENAME` or `gethostname()`), @@ -49,9 +40,9 @@ Four branches in priority order: Outside the Gateway branch the scheduler is always in-process, so its lifetime equals the run's lifetime: no orphaned schedulers if the driver crashes. On the Gateway branch the same contract is enforced -server-side: created clusters are shut down on exit +server-side: the created cluster is shut down on exit (`shutdown_on_close=True` and the deployment's `idle_timeout` backstop -a crashed driver); attached clusters belong to the user. +a crashed driver). ## Resource keys @@ -110,5 +101,5 @@ and keeps everything in one process tree. resource-advertising contract. The SLURM branch is tested with mocked `subprocess.Popen` plus a stubbed `Client.wait_for_workers`; the Gateway branch (create/cull lifecycle, image option, worker-wait -failure, attach mode, contract and image verification) against a fake -`dask_gateway` module. +failure, stale-env-var immunity, resource-contract verification) +against a fake `dask_gateway` module. diff --git a/docs/user/cluster.md b/docs/user/cluster.md index 4b1dbee1..d6f09341 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -205,6 +205,11 @@ Under the hood, each run: 4. **Shuts the cluster down.** Nothing to clean up; a crashed run's cluster is reaped by the deployment's idle timeout. +A Gateway scheduler's `gateway://` address cannot be used with +`DASK_SCHEDULER_ADDRESS` — on a hub, `lc run` always manages its own +run-scoped cluster. Requires the optional gateway extra: +`pip install lightcone-cli[gateway]` (preinstalled on the hub image). + Because the project must be reachable by the BinderHub build pods, it needs a GitHub remote (public, today) — `lc init` offers to create and connect one as part of scaffolding, including GitHub authentication via @@ -213,39 +218,6 @@ a one-time device code (see the [`lc init` GitHub step](../cli/init.md)). resulting image ref; `lc build --no-commit` refuses instead of auto-committing. -### Attaching to a long-lived cluster - -To iterate repeatedly against one warm cluster (dashboard panels -docked, no per-run cluster startup), create it yourself from JupyterLab -and point `lc run` at it by name: - -```python -from dask_gateway import Gateway -# shutdown_on_close=False keeps the cluster alive when this kernel -# exits. Idle clusters are reaped by the deployment after 30 min. -cluster = Gateway().new_cluster(shutdown_on_close=False) -cluster.adapt(minimum=1, maximum=8) -cluster.name -``` - -```bash -export LIGHTCONE_GATEWAY_CLUSTER= # e.g. hub.a1b2c3... -lc run -``` - -Attached clusters are yours: `lc run` never rescales them and leaves -them running on exit, mirroring the `DASK_SCHEDULER_ADDRESS` -convention. The trade-off: an attached cluster's image is fixed at its -creation, so `lc run` can only *warn* when it drifts from the -project's current image (unset `LIGHTCONE_GATEWAY_CLUSTER` to get back -to the always-fresh default). - -A Gateway scheduler's `gateway://` address cannot be used with -`DASK_SCHEDULER_ADDRESS` — attachment is always by name. - -Requires the optional gateway extra: -`pip install lightcone-cli[gateway]` (preinstalled on the hub image). - ### Containers on the hub: the pod is the runtime There is no docker or podman inside a pod — Kubernetes itself is the @@ -270,8 +242,7 @@ That one Containerfile then serves every path: built locally it wraps recipes on your laptop; built by the hub it runs them as worker pods. Manifests record the image the worker pod actually ran -(`worker_image`), so provenance stays truthful even when you knowingly -iterate on a stale attached cluster. +(`worker_image`), so provenance is ground truth, not inference. ## NERSC Perlmutter: site-specific notes diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 63b649c6..5224f255 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -747,9 +747,7 @@ def run( run-scoped Dask Gateway cluster on a JupyterHub deployment (created with the project's worker image — kept up to date through the hub's build service — and culled when the run finishes), or an existing - scheduler if ``DASK_SCHEDULER_ADDRESS`` is set - (``LIGHTCONE_GATEWAY_CLUSTER=`` similarly attaches to an - existing Gateway cluster and leaves it running). + scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. """ _abort_on_perlmutter_login() @@ -1298,7 +1296,7 @@ def _warn_runtime_cluster_mismatch( if runtime == KUBERNETES_RUNTIME and not on_gateway: console.print( "[yellow]⚠ Container runtime is [cyan]kubernetes[/cyan] but " - "this run is not attached through Dask Gateway.[/yellow]\n" + "this run is not going through Dask Gateway.[/yellow]\n" " Recipes will execute unwrapped and lc cannot verify they " "run in the declared container\n" " image — manifests may record provenance that does not " @@ -1366,7 +1364,7 @@ def _expected_worker_image(project: Path) -> str | None: "images, but a Dask Gateway cluster runs a single worker " "image:[/yellow]\n" + "\n".join(f" [dim]•[/dim] {i}" for i in images) - + "\n All recipes will execute in whatever image the attached " + + "\n All recipes will execute in whatever image the " "cluster runs.\n" " Consolidate on one Containerfile to restore per-recipe " "provenance on this deployment." diff --git a/src/lightcone/engine/dask_cluster.py b/src/lightcone/engine/dask_cluster.py index 26bf94fa..b467e21e 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -5,15 +5,12 @@ - ``DASK_SCHEDULER_ADDRESS`` is already set → yield it as-is. We don't own the cluster, so we don't tear it down. -- A Dask Gateway environment is detected (``LIGHTCONE_GATEWAY_CLUSTER`` or - ``DASK_GATEWAY__ADDRESS`` is set, e.g. on a JupyterHub deployment) → - **create** a run-scoped Gateway cluster with the project's worker image - and shut it down when the run finishes — the same create/cull lifecycle - the local and SLURM branches have always had, which is also what lets - every run pick up a freshly built image (a Gateway cluster's image is - fixed at creation). Setting ``LIGHTCONE_GATEWAY_CLUSTER=`` - instead **attaches** to that existing cluster and leaves it running on - exit — the escape hatch for iterating against a long-lived cluster. +- A Dask Gateway environment is detected (``DASK_GATEWAY__ADDRESS`` is + set, e.g. on a JupyterHub deployment) → **create** a run-scoped + Gateway cluster with the project's worker image and shut it down when + the run finishes — the same create/cull lifecycle the local and SLURM + branches have always had, which is also what lets every run pick up a + freshly built image (a Gateway cluster's image is fixed at creation). Gateway scheduler addresses use a custom ``gateway://`` comm scheme that a bare ``distributed.Client`` cannot dial, so the child snakemake process is told the *cluster name* via ``LIGHTCONE_GATEWAY_CLUSTER`` @@ -29,10 +26,9 @@ the scheduler is always in-process (driven by ``lc run`` itself) so its lifetime equals the run's lifetime — no service to manage, no orphaned schedulers if the driver crashes. On the Gateway branch the same -lifetime contract is enforced server-side: clusters lc creates are shut -down on exit (and reaped by the deployment's idle timeout if lc dies -uncleanly); clusters the user named are left untouched, mirroring the -``DASK_SCHEDULER_ADDRESS`` convention. +lifetime contract is enforced server-side: the cluster lc creates is +shut down on exit, and reaped by the deployment's idle timeout if lc +dies uncleanly. """ from __future__ import annotations @@ -53,12 +49,12 @@ RESOURCE_MEMORY = "memory" RESOURCE_GPUS = "gpus" -#: Env var carrying a Dask Gateway cluster name. Set by the user to -#: attach to an existing long-lived cluster instead of the default -#: create-per-run lifecycle; set by :func:`cluster_for_run` for the -#: child snakemake process so the executor plugin can rejoin the -#: cluster through the Gateway API (a bare ``Client`` cannot dial -#: ``gateway://``). +#: Env var carrying a Dask Gateway cluster name — the parent→child +#: rendezvous: :func:`cluster_for_run` sets it for the child snakemake +#: process so the executor plugin can rejoin the run's cluster through +#: the Gateway API (a bare ``Client`` cannot dial ``gateway://``). It is +#: an internal contract, not a user knob; a value lingering in the +#: ambient environment is ignored. GATEWAY_CLUSTER_ENV = "LIGHTCONE_GATEWAY_CLUSTER" #: Env var bounding how long the Gateway branch waits for the first @@ -70,10 +66,8 @@ #: Env var carrying the image a Gateway worker pod was started with. A #: lightcone-hub deployment's cluster-options handler injects it into -#: every scheduler/worker pod; :func:`cluster_for_run` reads it back to -#: verify the attached cluster actually runs the project's image (and -#: the manifest layer records it as ground truth — see -#: ``lightcone.engine.manifest.write_manifest``). +#: every scheduler/worker pod; the manifest layer records it as ground +#: truth — see ``lightcone.engine.manifest.write_manifest``. WORKER_IMAGE_ENV = "LIGHTCONE_WORKER_IMAGE" @@ -136,10 +130,7 @@ def gateway_branch_active() -> bool: """ if os.environ.get("DASK_SCHEDULER_ADDRESS"): return False - return ( - GATEWAY_CLUSTER_ENV in os.environ - or "DASK_GATEWAY__ADDRESS" in os.environ - ) + return "DASK_GATEWAY__ADDRESS" in os.environ @contextmanager @@ -166,11 +157,10 @@ def cluster_for_run( file I/O is slow and can pressure the gateway nodes). *expected_worker_image*, when given, is the image the project's - declared container resolves to on this deployment. A Gateway cluster - lc creates is started with exactly this image; an attached cluster - is compared against it with a warning on mismatch. Ignored by the - other branches (they realize containers by wrapping recipes, not - via pod images). + declared container resolves to on this deployment; the Gateway + cluster is created with exactly this image. Ignored by the other + branches (they realize containers by wrapping recipes, not via pod + images). *max_workers* bounds the adaptive scaling of a Gateway cluster lc creates (``lc run`` passes its ``--jobs`` value — there is never a @@ -183,10 +173,10 @@ def cluster_for_run( yield {"DASK_SCHEDULER_ADDRESS": addr} return - if GATEWAY_CLUSTER_ENV in os.environ or "DASK_GATEWAY__ADDRESS" in os.environ: + if "DASK_GATEWAY__ADDRESS" in os.environ: with _gateway_cluster( verbose=verbose, - expected_worker_image=expected_worker_image, + worker_image=expected_worker_image, max_workers=max_workers, ) as name: yield {GATEWAY_CLUSTER_ENV: name} @@ -209,26 +199,18 @@ def cluster_for_run( def _gateway_cluster( *, verbose: bool, - expected_worker_image: str | None, + worker_image: str | None, max_workers: int | None, ) -> Iterator[str]: - """Create (or attach to) a Dask Gateway cluster; yield its name. - - Default lifecycle — create and cull, per run: a new cluster is - created with the project's worker image, scaled adaptively between - one worker and *max_workers*, and shut down when the run finishes. - This mirrors the local/SLURM branches (compute lives exactly as - long as the run) and is what makes image updates seamless: a - Gateway cluster's image is fixed at creation, so picking up a - freshly built project image *requires* a fresh cluster. - - Attach — when ``LIGHTCONE_GATEWAY_CLUSTER`` names a cluster: connect - to it, leave its scaling alone, and leave it running on exit - (mirroring the ``DASK_SCHEDULER_ADDRESS`` convention). This is the - escape hatch for iterating repeatedly against one long-lived - cluster (create it from the JupyterLab Dask sidebar or a notebook - with ``shutdown_on_close=False``); the cost is that lc can only - warn — not fix — when its image is stale. + """Create a run-scoped Dask Gateway cluster; yield its name. + + Create and cull, per run: a new cluster is created with the + project's worker image, scaled adaptively between one worker and + *max_workers*, and shut down when the run finishes. This mirrors + the local/SLURM branches (compute lives exactly as long as the + run) and is what makes image updates seamless: a Gateway cluster's + image is fixed at creation, so picking up a freshly built project + image *requires* a fresh cluster. The Gateway client is configured entirely by ambient dask config — on a lightcone JupyterHub deployment the ``DASK_GATEWAY__*`` env @@ -240,39 +222,23 @@ def _gateway_cluster( except ImportError as exc: raise RuntimeError( "A Dask Gateway environment was detected " - f"({GATEWAY_CLUSTER_ENV} or DASK_GATEWAY__ADDRESS is set) but " - "the dask-gateway client is not installed. Install it with " + "(DASK_GATEWAY__ADDRESS is set) but the dask-gateway client " + "is not installed. Install it with " "`pip install lightcone-cli[gateway]`." ) from exc try: gateway = Gateway() except Exception as exc: - # Gateway() raises (ValueError) when no gateway address is - # configured — e.g. LIGHTCONE_GATEWAY_CLUSTER lingering in a - # shell off-hub, where no DASK_GATEWAY__* env exists. raise RuntimeError( "A Dask Gateway environment was detected but the gateway " - f"client could not be configured ({exc}). If you exported " - f"{GATEWAY_CLUSTER_ENV} on a machine without a Dask Gateway, " - "unset it." + f"client could not be configured ({exc})." ) from exc - named = os.environ.get(GATEWAY_CLUSTER_ENV) - if named: - with _attached_gateway_cluster( - gateway, - name=named, - verbose=verbose, - expected_worker_image=expected_worker_image, - ) as name: - yield name - return - with _created_gateway_cluster( gateway, verbose=verbose, - worker_image=expected_worker_image, + worker_image=worker_image, max_workers=max_workers, ) as name: yield name @@ -370,91 +336,6 @@ def _wait_first_worker(client: object, *, image: str | None) -> None: ) from exc -@contextmanager -def _attached_gateway_cluster( - gateway: object, - *, - name: str, - verbose: bool, - expected_worker_image: str | None, -) -> Iterator[str]: - """Attach to the named running Gateway cluster; leave it running.""" - try: - cluster = gateway.connect(name) # type: ignore[attr-defined] - except Exception as exc: - # dask-gateway raises its own error types (ValueError, - # GatewayClusterError) for a missing/stopped cluster; translate - # into guidance, because a *stale* LIGHTCONE_GATEWAY_CLUSTER is - # the likely cause — we told users to export it. - raise RuntimeError( - f"Could not connect to Dask Gateway cluster {name!r} " - f"({exc}); unset {GATEWAY_CLUSTER_ENV} to let lc create a " - "run-scoped cluster, or point it at a cluster that is " - "running." - ) from exc - - if verbose: - print( - f"→ Attached to Dask Gateway cluster {cluster.name} " - f"(dashboard: {cluster.dashboard_link})" - ) - - try: - client = cluster.get_client() - try: - # Verify the deployment contract only against what's live: - # an adaptive cluster sitting at zero workers will scale - # once the executor submits tasks, so an empty worker set - # is not an error here. - _assert_worker_resources(client) - _check_worker_image(client, expected=expected_worker_image) - finally: - client.close() - yield str(cluster.name) - finally: - # Releases this process's connections only — the cluster (and - # its scaling) belongs to the user. - cluster.close() - - -def _check_worker_image(client: object, *, expected: str | None) -> None: - """Warn when the attached cluster doesn't run the project's image. - - The scheduler pod carries the same ``LIGHTCONE_WORKER_IMAGE`` env - the deployment injects into workers (Gateway applies cluster-config - ``environment`` to both), so this works even while an adaptive - cluster sits at zero workers. Read via a lambda — cloudpickled by - value — so the scheduler pod does not need lightcone importable. - - A warning, not an error: the user may be knowingly iterating on a - stale cluster, and the manifest layer records the *actual* image - (ground truth) either way. Silently skipped when the deployment - doesn't inject the marker or no expectation was computed. - """ - if expected is None: - return - try: - env_key = WORKER_IMAGE_ENV # captured by value into the lambda - actual = client.run_on_scheduler( # type: ignore[attr-defined] - lambda: __import__("os").environ.get(env_key) - ) - except Exception: - return # older deployment / restricted scheduler — not fatal - if actual and actual != expected: - print( - f"⚠ The attached Dask Gateway cluster runs image\n" - f" {actual}\n" - f" but this project's container resolves to\n" - f" {expected}\n" - f" Recipes will execute in the cluster's image; manifests " - f"record what actually ran.\n" - f" To run on the project image, unset {GATEWAY_CLUSTER_ENV} " - f"and let lc create a run-scoped\n" - f" cluster, or recreate your cluster with " - f'image="{expected}".' - ) - - def _assert_worker_resources(client: object) -> None: """Fail fast when Gateway workers don't advertise the resource contract. diff --git a/src/lightcone/engine/site_registry.py b/src/lightcone/engine/site_registry.py index b3b7c644..273f201f 100644 --- a/src/lightcone/engine/site_registry.py +++ b/src/lightcone/engine/site_registry.py @@ -92,9 +92,11 @@ # are still containerized. Declared (not detected) — load_runtime # treats it as explicit, so no "no runtime found" warning fires. "container_runtime": "kubernetes", - # Scratch comes from the deployment's LIGHTCONE_SCRATCH env (a - # path on the shared RWX volume), which outranks site defaults - # in lightcone.engine.scratch — nothing to declare here. + # The NFS home honors flock and is mounted into Gateway worker + # pods at the same path, so lightcone state can live right next + # to the user (unlike NERSC, where DVS forces it onto Lustre). + # A deployment can still override via LIGHTCONE_SCRATCH. + "scratch_root": "$HOME", }, "local": { "hostname_patterns": [], diff --git a/src/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index ae2806e1..e7bb5566 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -161,9 +161,9 @@ def _connect_client(): # type: ignore[no-untyped-def] Returns ``(client, gateway_cluster_or_None)`` — the cluster handle is kept so ``shutdown()`` can release its local connections. It is never - shut down here: on the Gateway path the cluster belongs to the *user* - (lc is attach-only); on the address paths it belongs to whoever - started the scheduler. + shut down here: the parent ``lc run`` owns the Gateway cluster's + lifecycle; on the address paths it belongs to whoever started the + scheduler. """ if addr := os.environ.get("DASK_SCHEDULER_ADDRESS"): try: @@ -300,9 +300,8 @@ def shutdown(self) -> None: try: self._client.close() if self._gateway_cluster is not None: - # Releases this process's connections only; the cluster - # itself belongs to the user (lc is attach-only and never - # shuts Gateway clusters down — see dask_cluster). + # Releases this process's connections only; the parent + # lc run owns the cluster lifecycle (see dask_cluster). self._gateway_cluster.close() finally: super().shutdown() diff --git a/tests/test_dask_cluster.py b/tests/test_dask_cluster.py index d4a6eaa8..8366a7bb 100644 --- a/tests/test_dask_cluster.py +++ b/tests/test_dask_cluster.py @@ -506,6 +506,43 @@ def test_gateway_worker_wait_timeout_raises_and_culls( assert "shutdown" in log, "a failed creation must not leak the cluster" +def test_gateway_ignores_stale_cluster_env_var( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIGHTCONE_GATEWAY_CLUSTER is the parent→child rendezvous, not a + user knob: a value lingering in the ambient environment must not + redirect lc to some old cluster — on a hub it still creates a fresh + run-scoped cluster, and off-hub it doesn't trigger the gateway + branch at all.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.stale99") + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) + + with cluster_for_run() as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.new1"} + assert not any(c.startswith("connect") for c in log) + + assert "shutdown" in log + + +def test_gateway_stale_env_var_off_hub_takes_local_branch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.stale99") + sentinel: dict[str, str] = {} + + @contextmanager + def _fake_local(*, verbose: bool, local_directory: str | None = None): + sentinel["called"] = "local" + yield "tcp://stub:9999" + + with patch("lightcone.engine.dask_cluster._local_cluster", _fake_local): + with cluster_for_run() as env: + assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://stub:9999"} + assert sentinel["called"] == "local" + + def test_gateway_wins_over_slurm(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") monkeypatch.setenv("SLURM_JOB_ID", "12345") @@ -532,26 +569,6 @@ def test_explicit_address_wins_over_gateway( assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://existing:8786"} -def test_gateway_named_cluster_attaches_and_leaves_running( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """LIGHTCONE_GATEWAY_CLUSTER= is the escape hatch for a - long-lived cluster: attach to it, never create, never rescale, - leave it running on exit.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") - log: list[str] = [] - _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) - - with cluster_for_run() as env: - assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.abc999"} - assert "connect(hub.abc999)" in log - assert not any(call.startswith("new_cluster") for call in log) - assert not any(call.startswith("adapt") for call in log) - - assert "shutdown" not in log, "attached clusters belong to the user" - assert "close" in log - - def test_gateway_missing_client_raises_helpfully( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -582,23 +599,6 @@ def test_gateway_rejects_workers_without_resource_contract( assert "shutdown" in log, "a failed creation must not leak the cluster" -def test_gateway_attach_contract_failure_leaves_cluster_running( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """On the attach path the same contract failure must NOT shut the - user's cluster down — it isn't ours.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") - log: list[str] = [] - _install_fake_gateway(monkeypatch, log, resources={}) - - with pytest.raises(RuntimeError, match="resource contract"): - with cluster_for_run(): - pass # pragma: no cover - - assert "close" in log, "connections must be released on failure" - assert "shutdown" not in log, "the user's cluster must survive our failure" - - def test_gateway_rejects_partial_resource_contract( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -614,122 +614,6 @@ def test_gateway_rejects_partial_resource_contract( pass # pragma: no cover -def test_gateway_stale_named_cluster_raises_with_guidance( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A stale LIGHTCONE_GATEWAY_CLUSTER (cluster since stopped) makes the - real client raise its own error types (ValueError/GatewayClusterError). - lc told the user to export that variable, so lc owns the remediation: - a RuntimeError naming the env var, not a bare dask_gateway traceback.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.gone42") - log: list[str] = [] - _install_fake_gateway( - monkeypatch, log, connect_error=ValueError("Cluster 'hub.gone42' not found") - ) - - with pytest.raises(RuntimeError, match="LIGHTCONE_GATEWAY_CLUSTER") as exc: - with cluster_for_run(): - pass # pragma: no cover - assert "hub.gone42" in str(exc.value) - - -def test_gateway_unconfigured_client_raises_with_guidance( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """LIGHTCONE_GATEWAY_CLUSTER lingering in a shell off-hub: Gateway() - itself raises ValueError (no gateway address configured). Must become - guidance to unset the variable, not a traceback.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") - log: list[str] = [] - _install_fake_gateway( - monkeypatch, - log, - init_error=ValueError("No dask-gateway address provided"), - ) - - with pytest.raises(RuntimeError, match="unset") as exc: - with cluster_for_run(): - pass # pragma: no cover - assert "LIGHTCONE_GATEWAY_CLUSTER" in str(exc.value) - - -def test_gateway_attached_zero_workers_is_not_an_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """An attached adaptive cluster sitting at zero workers scales up once - tasks arrive — the contract check must only judge live workers.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") - log: list[str] = [] - _install_fake_gateway(monkeypatch, log, resources=None) - - with cluster_for_run() as env: - assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.run1"} - - -def test_gateway_warns_on_worker_image_mismatch( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - """The attached cluster runs a different image than the project's - container resolves to → loud warning naming both (manifests record - the actual image, so provenance stays truthful either way).""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") - log: list[str] = [] - _install_fake_gateway( - monkeypatch, - log, - resources=dict(_GOOD_RESOURCES), - worker_image="reg.example/lc-proj:stale00", - ) - - with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): - pass - - out = capsys.readouterr().out - assert "reg.example/lc-proj:stale00" in out - assert "reg.example/lc-proj:abc123" in out - - -def test_gateway_matching_worker_image_is_silent( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") - log: list[str] = [] - _install_fake_gateway( - monkeypatch, - log, - resources=dict(_GOOD_RESOURCES), - worker_image="reg.example/lc-proj:abc123", - ) - - with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): - pass - - assert "⚠" not in capsys.readouterr().out - - -def test_gateway_image_check_skipped_without_deployment_marker( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - """A deployment that doesn't inject LIGHTCONE_WORKER_IMAGE can't be - verified — skip silently rather than warn on every run.""" - monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.run1") - log: list[str] = [] - _install_fake_gateway( - monkeypatch, - log, - resources=dict(_GOOD_RESOURCES), - worker_image=None, - ) - - with cluster_for_run(expected_worker_image="reg.example/lc-proj:abc123"): - pass - - assert "⚠" not in capsys.readouterr().out - - def test_executor_connects_via_gateway_name( monkeypatch: pytest.MonkeyPatch, ) -> None: From e1062e2c81a4b6f0743d5597757a1ef30eb7dc71 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Tue, 21 Jul 2026 20:13:11 +0200 Subject: [PATCH 9/9] Add GCP Cloud Build as the preferred on-hub build backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expert review flagged binderhub-service as heavier than the job needs on GCP (privileged dockerd DaemonSet next to user pods, dev-versioned chart, registry key in secrets, and a structural public-git-ref requirement). New lightcone.engine.cloudbuild: - Source is a tarball of the *staged build context* — the exact file set compute_image_tag hashes — uploaded to a GCS bucket. Completely git-free: no remote, no commit, no push; private projects build like public ones. The Dockerfile symlink and env-commit machinery don't apply on this path. - Image identity is the content-addressed tag (lc-:) pushed to $LIGHTCONE_REGISTRY — the same identity local builds use; freshness is one registry HEAD. - Auth is the pod's Workload Identity (metadata token): zero stored credentials. Optional dedicated build SA via LIGHTCONE_BUILD_SERVICE_ACCOUNT for least privilege. - Builds run off-cluster in Cloud Build; failures fetch the log tail from the logs bucket. Backend selection (_hub_build_backend): cloudbuild when the deployment injects LIGHTCONE_BUILD_BUCKET, else the (portable) binder backend, else the passive registry report. Docs + design doc §8.5 updated. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CBPsorm6U3Co7urs59FVtM --- docs/cli/build.md | 39 ++- docs/design/jupyterhub-dask-gateway-gcp.md | 26 ++ docs/user/cluster.md | 28 +- src/lightcone/cli/commands.py | 117 +++++-- src/lightcone/engine/cloudbuild.py | 361 +++++++++++++++++++++ tests/conftest.py | 2 + tests/test_cli.py | 35 ++ tests/test_cloudbuild.py | 205 ++++++++++++ 8 files changed, 751 insertions(+), 62 deletions(-) create mode 100644 src/lightcone/engine/cloudbuild.py create mode 100644 tests/test_cloudbuild.py diff --git a/docs/cli/build.md b/docs/cli/build.md index 2cb35316..e201c898 100644 --- a/docs/cli/build.md +++ b/docs/cli/build.md @@ -36,24 +36,33 @@ nothing to build. ## On a JupyterHub deployment (`kubernetes` runtime) -There is no docker in a hub pod, so `lc build` instead drives an image -build through the deployment's BinderHub service: - -1. Commits any changes to the environment-defining files (the - Containerfile, dependency files, and named COPY sources — never your - analysis code), maintaining a `Dockerfile → Containerfile` symlink - for repo2docker. `--no-commit` refuses instead. -2. Pushes, so the build pods can clone the ref (the project needs a - GitHub remote). -3. Streams the BinderHub build (repo2docker → the deployment registry) - and prints the resulting image ref — the image `lc run` will start - its Gateway cluster with. Already-built refs return immediately. +There is no docker in a hub pod, so `lc build` drives the deployment's +build backend instead. Two backends exist; the deployment picks one by +what it injects into user pods: + +**Cloud Build** (GCP deployments; selected by `LIGHTCONE_BUILD_BUCKET` +being set, preferred where available). Fully git-free: the *staged +build context* — the same file set the content-addressed tag hashes — +is tarred and uploaded to the build bucket, built off-cluster by GCP +Cloud Build, and pushed to `$LIGHTCONE_REGISTRY/lc-:`. +No GitHub remote, commit, or push is involved; private projects build +exactly like public ones, and an unchanged environment is a single +registry HEAD (no build). Auth is the pod's Workload Identity — no +credentials anywhere. Failures surface the build-log tail. + +**BinderHub service** (selected by an ambient `JUPYTERHUB_API_TOKEN` / +`LIGHTCONE_BINDER_URL`; the portable, non-GCP path). Builds from a git +ref: `lc build` commits environment-file changes (`--no-commit` +refuses instead), pushes (the project needs a clonable GitHub remote), +maintains a `Dockerfile → Containerfile` symlink for repo2docker, and +streams the build via SSE. Images are tagged by the environment +commit. `lc run` performs the same ensure step automatically on every run, so running `lc build` explicitly is optional on the hub — useful to -pre-build after a dependency change or to see the build log. Without a -reachable build service, `lc build` falls back to probing the registry -and printing off-hub publish instructions. +pre-build after a dependency change or to see the build log. With no +backend configured, `lc build` falls back to probing the registry and +printing off-hub publish instructions. ## Tag computation diff --git a/docs/design/jupyterhub-dask-gateway-gcp.md b/docs/design/jupyterhub-dask-gateway-gcp.md index c45054bd..82308041 100644 --- a/docs/design/jupyterhub-dask-gateway-gcp.md +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -703,3 +703,29 @@ non-issue on the gateway path where recipes run from the shared home via | recipe isolation | `docker run` wrap | `podman-hpc run` wrap | worker pod **is** the image | | cluster lifecycle | owned per-run | owned per-run | owned per-run (attach opt-in) | | manifest truth | declared spec + code_version(tag) | same | same + `worker_image` ground truth | + +### 8.5 Build backend v2 (2026-07-21): GCP Cloud Build + +On expert review, the BinderHub build path (§8.2) is heavier than the +job requires on GCP: a dev-versioned chart, a privileged dockerd +DaemonSet next to user pods, and a registry key in the deployment's +secrets — and it structurally requires a public, pushed git ref per +build. `lightcone.engine.cloudbuild` adds GCP Cloud Build as the +preferred backend where available (selected by the deployment +injecting `LIGHTCONE_BUILD_BUCKET`): + +- **Source = the staged build context** (the exact file set + `compute_image_tag` hashes), tarred and uploaded to GCS — no git + involvement at all; private projects build like public ones. +- **Image identity = the content-addressed tag** (`lc-:`) + pushed to `$LIGHTCONE_REGISTRY` — the same identity local builds use, + restoring §7.2's scheme; freshness is a registry HEAD. +- **Auth = Workload Identity** (metadata-server token + + `cloudbuild.builds.editor` + `iam.serviceAccountUser` on a dedicated + build SA that holds only registry-writer): zero stored credentials. +- Builds run in Google-managed VMs off-cluster; failures fetch the log + tail from the logs bucket. + +The binder backend remains in the CLI as the portable (non-GCP) +alternative; `_hub_build_backend()` picks per deployment. Trade-off +accepted: the on-hub build path on GCP is now GCP-specific. diff --git a/docs/user/cluster.md b/docs/user/cluster.md index d6f09341..22b60fd6 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -191,12 +191,14 @@ Under the hood, each run: 1. **Makes sure the worker image is up to date.** The project's `Containerfile` (plus its dependency files) defines the worker-pod environment. If those files changed since the last build, `lc run` - commits them, pushes, and drives an image build through the hub's - BinderHub service into the deployment registry. When the registry - already holds the image — the common case — this is a single fast - round-trip. Code-only edits never trigger a rebuild: your code - reaches the workers through your shared home directory, not the - image. + drives an image build through the hub's build backend into the + deployment registry — on GCP hubs that's Cloud Build (git-free: it + builds your working tree directly, no commit or push involved), on + others the BinderHub service (which commits env changes and pushes + so build pods can clone the ref). When the registry already holds + the image — the common case — this is a single fast round-trip. + Code-only edits never trigger a rebuild: your code reaches the + workers through your shared home directory, not the image. 2. **Creates a run-scoped Dask Gateway cluster** with that image, scaled adaptively between one worker and `--jobs`. 3. Runs the pipeline (same executor, same per-recipe resource hints as @@ -210,13 +212,13 @@ A Gateway scheduler's `gateway://` address cannot be used with run-scoped cluster. Requires the optional gateway extra: `pip install lightcone-cli[gateway]` (preinstalled on the hub image). -Because the project must be reachable by the BinderHub build pods, it -needs a GitHub remote (public, today) — `lc init` offers to create and -connect one as part of scaffolding, including GitHub authentication via -a one-time device code (see the [`lc init` GitHub step](../cli/init.md)). -`lc build` runs the same ensure-image step explicitly and prints the -resulting image ref; `lc build --no-commit` refuses instead of -auto-committing. +On BinderHub-backed hubs the project must be reachable by the build +pods, so it needs a (public, today) GitHub remote — `lc init` offers to +create and connect one, including GitHub auth via a one-time device +code (see the [`lc init` GitHub step](../cli/init.md)). Cloud +Build-backed hubs have no such requirement (GitHub remains recommended +for backup/collaboration). `lc build` runs the same ensure-image step +explicitly; see [`lc build`](../cli/build.md) for backend details. ### Containers on the hub: the pod is the runtime diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 5224f255..2139e865 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -1184,8 +1184,9 @@ def verify(universe: str | None) -> None: "--no-commit", is_flag=True, help=( - "On a hub: fail instead of auto-committing environment-file " - "changes before the image build" + "BinderHub backend only: fail instead of auto-committing " + "environment-file changes before the image build (the Cloud " + "Build backend builds the working tree directly — no commits)" ), ) def build(force: bool, runtime: str | None, no_commit: bool) -> None: @@ -1219,9 +1220,7 @@ def build(force: bool, runtime: str | None, no_commit: bool) -> None: raise click.ClickException(str(e)) if resolved_runtime == KUBERNETES_RUNTIME: - from lightcone.engine.binder import binder_available - - if binder_available(): + if _hub_build_backend() is not None: _build_via_hub(project, commit=not no_commit) else: _report_registry_images(project) @@ -1373,6 +1372,25 @@ def _expected_worker_image(project: Path) -> str | None: return images[0] if images else None +def _hub_build_backend() -> str | None: + """Which on-hub image-build backend this deployment offers. + + ``cloudbuild`` (GCP Cloud Build; selected by the deployment + injecting LIGHTCONE_BUILD_BUCKET) is preferred — off-cluster + builds, IAM-only auth, and no git remote required. ``binder`` + (BinderHub service) is the portable alternative. ``None`` → no + on-hub build path; callers fall back to passive registry checks. + """ + from lightcone.engine.binder import binder_available + from lightcone.engine.cloudbuild import cloudbuild_available + + if cloudbuild_available(): + return "cloudbuild" + if binder_available(): + return "binder" + return None + + def _binder_progress(verbose: bool) -> Callable[[str, str], None]: """Progress callback for hub image builds. @@ -1406,17 +1424,13 @@ def _worker_image_for_run(project: Path, *, verbose: bool) -> str | None: Elsewhere, falls back to the passive registry resolution of :func:`_expected_worker_image` (verification-only). """ - from lightcone.engine.binder import ( - BinderBuildError, - binder_available, - ensure_worker_image, - ) from lightcone.engine.container import is_containerfile - if not binder_available(): + backend = _hub_build_backend() + if backend is None: return _expected_worker_image(project) - _, specs = _container_specs(project) + project_name, specs = _container_specs(project) if not specs: # No container declared: the cluster runs the deployment's # default worker image, which ships the lightcone stack. @@ -1437,31 +1451,67 @@ def _worker_image_for_run(project: Path, *, verbose: bool) -> str | None: return spec console.print( f"[cyan]Ensuring worker image[/cyan] for [cyan]{spec}[/cyan] " - "[dim](hub build service)[/dim]" + f"[dim]({backend})[/dim]" ) try: - image = ensure_worker_image( - project, spec, on_progress=_binder_progress(verbose) + image = _ensure_hub_image( + backend, + project, + spec, + project_name=project_name, + on_progress=_binder_progress(verbose), ) - except BinderBuildError as e: - raise click.ClickException(str(e)) + except click.ClickException: + raise console.print(f"[green]✓[/green] Worker image: [cyan]{image}[/cyan]") return image +def _ensure_hub_image( + backend: str, + project: Path, + spec: str, + *, + project_name: str, + commit: bool = True, + on_progress: Callable[[str, str], None] | None = None, +) -> str: + """Run the selected backend's ensure step; errors become user-facing.""" + if backend == "cloudbuild": + from lightcone.engine.cloudbuild import CloudBuildError + from lightcone.engine.cloudbuild import ensure_worker_image as cb_ensure + + try: + return cb_ensure( + project, spec, project_name=project_name, on_progress=on_progress + ) + except CloudBuildError as e: + raise click.ClickException(str(e)) + from lightcone.engine.binder import BinderBuildError + from lightcone.engine.binder import ensure_worker_image as binder_ensure + + try: + return binder_ensure( + project, spec, commit=commit, on_progress=on_progress + ) + except BinderBuildError as e: + raise click.ClickException(str(e)) + + def _build_via_hub(project: Path, *, commit: bool) -> None: """``lc build`` on a hub: publish every declared Containerfile. - Drives the BinderHub service (build + push into the deployment - registry) for each declared Containerfile — the explicit, - verbose-by-default form of the ensure step ``lc run`` performs - implicitly. Registry-image specs need no build and are reported - as-is. + Drives the deployment's build backend (Cloud Build or the BinderHub + service; build + push into the deployment registry) for each + declared Containerfile — the explicit, verbose-by-default form of + the ensure step ``lc run`` performs implicitly. Registry-image + specs need no build and are reported as-is. """ - from lightcone.engine.binder import BinderBuildError, ensure_worker_image from lightcone.engine.container import is_containerfile - _, specs = _container_specs(project) + backend = _hub_build_backend() + assert backend is not None # caller gates on _hub_build_backend() + project_name, specs = _container_specs(project) if not specs: console.print("No containers declared in astra.yaml — nothing to build.") return @@ -1473,16 +1523,15 @@ def _build_via_hub(project: Path, *, commit: bool) -> None: "worker stack)." ) continue - console.print(f"[cyan]Building[/cyan] {spec} [dim](hub build service)[/dim]") - try: - image = ensure_worker_image( - project, - spec, - commit=commit, - on_progress=_binder_progress(verbose=True), - ) - except BinderBuildError as e: - raise click.ClickException(str(e)) + console.print(f"[cyan]Building[/cyan] {spec} [dim]({backend})[/dim]") + image = _ensure_hub_image( + backend, + project, + spec, + project_name=project_name, + commit=commit, + on_progress=_binder_progress(verbose=True), + ) console.print( f"[green]✓[/green] {spec} → [cyan]{image}[/cyan] " "[dim](in the deployment registry; `lc run` will start its " diff --git a/src/lightcone/engine/cloudbuild.py b/src/lightcone/engine/cloudbuild.py new file mode 100644 index 00000000..b6378fd7 --- /dev/null +++ b/src/lightcone/engine/cloudbuild.py @@ -0,0 +1,361 @@ +"""On-hub worker-image builds through GCP Cloud Build. + +The alternative to :mod:`lightcone.engine.binder` for GCP-hosted +deployments, and the preferred one where available: builds run in +Google-managed VMs entirely off-cluster (no privileged builder pods), +auth is pure IAM through the pod's Workload Identity (no stored +credentials anywhere), and — the big UX difference — the build source is +a tarball of the project's **staged build context**, not a git ref. No +GitHub remote, no commit, no push is needed to build; private projects +build exactly like public ones. + +Image identity is the content-addressed scheme the CLI already uses +everywhere else: the tag hashes the staged context +(:func:`lightcone.engine.container.compute_image_tag`), the pushed ref +is ``$LIGHTCONE_REGISTRY/lc-:``, and "is the image up to +date" is a registry HEAD on that ref — unchanged files never rebuild. + +Deployment contract (injected into user pods): + +- :data:`~lightcone.engine.container.REGISTRY_ENV` — the Artifact + Registry prefix images are pushed to (also names the GCP project and + region). +- :data:`BUCKET_ENV` — a GCS bucket for build sources and logs; its + presence is what selects this backend. +- :data:`SERVICE_ACCOUNT_ENV` (optional) — a dedicated build service + account (``projects/-/serviceAccounts/``); recommended so the + builder holds only registry-writer rights. + +The pod's Workload Identity needs ``cloudbuild.builds.editor``, +``iam.serviceAccountUser`` on the build SA, and object admin on the +bucket. Everything speaks plain REST via ``urllib`` with a +metadata-server token — no SDK dependency. +""" + +from __future__ import annotations + +import io +import json +import os +import tarfile +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Callable +from pathlib import Path + +from lightcone.engine.container import ( + _metadata_access_token, + _populate_build_context, + compute_image_tag, + deployment_registry, + registry_image_exists, + registry_image_ref, +) + +#: GCS bucket for build sources/logs. Presence selects this backend. +BUCKET_ENV = "LIGHTCONE_BUILD_BUCKET" + +#: Optional dedicated Cloud Build service account (full resource name or +#: bare email). Builds with a custom SA are the least-privilege setup. +SERVICE_ACCOUNT_ENV = "LIGHTCONE_BUILD_SERVICE_ACCOUNT" + +#: Hard ceiling on one build, seconds (also sent as the Cloud Build +#: timeout). Slim project images finish in single-digit minutes. +_BUILD_DEADLINE_S = 1800 + +_POLL_INTERVAL_S = 5.0 + +#: Progress callback ``(phase, message)`` — same shape as the binder +#: backend so the CLI renders both identically. Phases are Cloud Build +#: statuses lowercased (queued/working/success/failure...). +ProgressFn = Callable[[str, str], None] + + +class CloudBuildError(RuntimeError): + """A worker image could not be produced through Cloud Build.""" + + +def cloudbuild_available() -> bool: + """Is this deployment configured for Cloud Build image builds?""" + return bool(os.environ.get(BUCKET_ENV)) and deployment_registry() is not None + + +def _bucket() -> str: + bucket = (os.environ.get(BUCKET_ENV) or "").strip().removeprefix("gs://") + if not bucket: + raise CloudBuildError(f"{BUCKET_ENV} is not set.") + return bucket.rstrip("/") + + +def _project_from_registry(registry: str) -> str: + """GCP project id out of an Artifact Registry prefix. + + ``us-central1-docker.pkg.dev//`` → ````. + """ + parts = registry.split("/") + if len(parts) < 2 or not parts[0].endswith("-docker.pkg.dev"): + raise CloudBuildError( + f"{registry!r} is not an Artifact Registry prefix " + "(expected -docker.pkg.dev//); the " + "Cloud Build backend only targets Artifact Registry." + ) + return parts[1] + + +def _token() -> str: + token = _metadata_access_token() + if token is None: + raise CloudBuildError( + "No GCP credentials available from the metadata server. The " + "Cloud Build backend needs Workload Identity (or another " + "metadata-served identity) with cloudbuild.builds.editor." + ) + return token + + +def _request( + method: str, + url: str, + token: str, + *, + body: bytes | None = None, + content_type: str = "application/json", +) -> tuple[int, bytes]: + req = urllib.request.Request( + url, + data=body, + method=method, + headers={ + "Authorization": f"Bearer {token}", + **({"Content-Type": content_type} if body is not None else {}), + }, + ) + try: + with urllib.request.urlopen(req, timeout=120) as resp: + return resp.status, resp.read() + except urllib.error.HTTPError as exc: + return exc.code, exc.read() + except (urllib.error.URLError, OSError) as exc: + raise CloudBuildError(f"Could not reach {url.split('?')[0]} ({exc}).") from exc + + +def _json_or_error(status: int, payload: bytes, what: str) -> dict[str, object]: + if not 200 <= status < 300: + detail = payload.decode("utf-8", errors="replace")[:500] + raise CloudBuildError(f"{what} failed: HTTP {status}\n{detail}") + try: + parsed = json.loads(payload.decode("utf-8") or "{}") + except ValueError as exc: + raise CloudBuildError(f"{what} returned unparseable JSON.") from exc + return parsed if isinstance(parsed, dict) else {} + + +# --------------------------------------------------------------------------- +# Source staging + upload +# --------------------------------------------------------------------------- + + +def _staged_context_tarball(project: Path, containerfile: Path) -> bytes: + """gzip tarball of the staged build context (the hashed file set).""" + with tempfile.TemporaryDirectory(prefix="lc-cloudbuild-") as tmp: + staged = Path(tmp) + _populate_build_context(staged, containerfile, project) + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for entry in sorted(staged.rglob("*")): + tar.add(entry, arcname=str(entry.relative_to(staged))) + return buf.getvalue() + + +def _upload_source(bucket: str, object_name: str, data: bytes, token: str) -> None: + url = ( + "https://storage.googleapis.com/upload/storage/v1/b/" + f"{urllib.parse.quote(bucket, safe='')}/o?uploadType=media&name=" + f"{urllib.parse.quote(object_name, safe='')}" + ) + status, payload = _request( + "POST", url, token, body=data, content_type="application/gzip" + ) + _json_or_error(status, payload, "Source upload to the build bucket") + + +def _fetch_log_tail(bucket: str, build_id: str, token: str, lines: int = 30) -> str: + object_name = urllib.parse.quote(f"logs/log-{build_id}.txt", safe="") + url = ( + "https://storage.googleapis.com/storage/v1/b/" + f"{urllib.parse.quote(bucket, safe='')}/o/{object_name}?alt=media" + ) + try: + status, payload = _request("GET", url, token) + except CloudBuildError: + return "" + if not 200 <= status < 300: + return "" + text = payload.decode("utf-8", errors="replace") + return "\n".join(text.splitlines()[-lines:]) + + +# --------------------------------------------------------------------------- +# Build submission + polling +# --------------------------------------------------------------------------- + + +def _submit_build( + *, + project_id: str, + bucket: str, + source_object: str, + containerfile_name: str, + image_ref: str, + token: str, +) -> str: + """Create the build; return its id.""" + build: dict[str, object] = { + "source": { + "storageSource": {"bucket": bucket, "object": source_object} + }, + "steps": [ + { + "name": "gcr.io/cloud-builders/docker", + "args": [ + "build", + "-t", + image_ref, + "-f", + containerfile_name, + ".", + ], + } + ], + "images": [image_ref], + "timeout": f"{_BUILD_DEADLINE_S}s", + # GCS_ONLY logging into our own bucket: required for custom + # service accounts, and it is where the failure tail comes from. + "logsBucket": f"gs://{bucket}/logs", + "options": {"logging": "GCS_ONLY"}, + } + build_sa = (os.environ.get(SERVICE_ACCOUNT_ENV) or "").strip() + if build_sa: + if "/" not in build_sa: + build_sa = f"projects/{project_id}/serviceAccounts/{build_sa}" + build["serviceAccount"] = build_sa + + url = f"https://cloudbuild.googleapis.com/v1/projects/{project_id}/builds" + status, payload = _request( + "POST", url, token, body=json.dumps(build).encode() + ) + op = _json_or_error(status, payload, "Cloud Build submission") + meta = op.get("metadata") + build_info = meta.get("build") if isinstance(meta, dict) else None + build_id = build_info.get("id") if isinstance(build_info, dict) else None + if not isinstance(build_id, str) or not build_id: + raise CloudBuildError( + "Cloud Build submission returned no build id " + f"(response keys: {sorted(op)})." + ) + return build_id + + +def _wait_for_build( + project_id: str, + build_id: str, + token: str, + on_progress: ProgressFn | None, +) -> str: + """Poll until a terminal status; return it.""" + url = ( + f"https://cloudbuild.googleapis.com/v1/projects/{project_id}" + f"/builds/{build_id}" + ) + deadline = time.monotonic() + _BUILD_DEADLINE_S + 120 + last_status = "" + while time.monotonic() < deadline: + status_code, payload = _request("GET", url, token) + build = _json_or_error(status_code, payload, "Cloud Build status poll") + status = str(build.get("status") or "") + if status != last_status: + last_status = status + if on_progress: + on_progress(status.lower(), "") + if status in ("SUCCESS", "FAILURE", "INTERNAL_ERROR", "TIMEOUT", "CANCELLED", "EXPIRED"): + return status + time.sleep(_POLL_INTERVAL_S) + raise CloudBuildError( + f"Timed out waiting for Cloud Build {build_id} to finish." + ) + + +# --------------------------------------------------------------------------- +# High-level entry point +# --------------------------------------------------------------------------- + + +def ensure_worker_image( + project: Path, + containerfile_spec: str, + *, + project_name: str | None = None, + on_progress: ProgressFn | None = None, +) -> str: + """Make sure the project's worker image is built; return its registry ref. + + Content-addressed and git-free: the tag hashes the staged build + context, so an unchanged environment is a single registry HEAD + (no build, no upload), and any change builds from the working tree + as it is right now. + """ + containerfile = project / containerfile_spec + if not containerfile.is_file(): + raise CloudBuildError( + f"Declared container {containerfile_spec!r} not found in {project}." + ) + registry = deployment_registry() + if registry is None: + raise CloudBuildError( + "LIGHTCONE_REGISTRY is not set; the Cloud Build backend needs " + "the Artifact Registry prefix to name and probe images." + ) + # Same tag the local `lc build` computes (project name from + # astra.yaml when the caller has it) — one content-addressed identity + # across every backend. + name = (project_name or project.name).lower().replace(" ", "-") + tag = compute_image_tag(name, containerfile, project) + ref = registry_image_ref(tag, registry=registry) + if ref is None: + raise CloudBuildError(f"Could not derive a registry ref from {tag!r}.") + + if registry_image_exists(ref) is True: + if on_progress: + on_progress("cached", f"{ref} already in the registry") + return ref + + token = _token() + bucket = _bucket() + project_id = _project_from_registry(registry) + + if on_progress: + on_progress("staging", "uploading build context") + # Content-addressed object name: identical contexts collide into the + # same object, which is exactly right. + source_object = f"sources/{tag}.tar.gz" + _upload_source(bucket, source_object, _staged_context_tarball(project, containerfile), token) + + build_id = _submit_build( + project_id=project_id, + bucket=bucket, + source_object=source_object, + containerfile_name=containerfile.name, + image_ref=ref, + token=token, + ) + status = _wait_for_build(project_id, build_id, token, on_progress) + if status != "SUCCESS": + tail = _fetch_log_tail(bucket, build_id, token) + raise CloudBuildError( + f"Cloud Build {build_id} ended with status {status}." + + (f" Last build output:\n{tail}" if tail else "") + ) + return ref diff --git a/tests/conftest.py b/tests/conftest.py index 9467232e..f4e3006e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,3 +26,5 @@ def _no_ambient_hub_env(monkeypatch: pytest.MonkeyPatch) -> None: # HTTP endpoint mid-suite. monkeypatch.delenv("JUPYTERHUB_API_TOKEN", raising=False) monkeypatch.delenv("LIGHTCONE_BINDER_URL", raising=False) + monkeypatch.delenv("LIGHTCONE_BUILD_BUCKET", raising=False) + monkeypatch.delenv("LIGHTCONE_REGISTRY", raising=False) diff --git a/tests/test_cli.py b/tests/test_cli.py index b1dd1fb3..92617604 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -593,3 +593,38 @@ def test_run_cmd_uses_own_interpreter_for_snakemake() -> None: has_outputs=False, ) assert cmd[:3] == [sys.executable, "-m", "snakemake"] + + +def test_build_backend_prefers_cloudbuild( + hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.cli.commands import _hub_build_backend + + assert _hub_build_backend() is None + monkeypatch.setenv("JUPYTERHUB_API_TOKEN", "tok") + assert _hub_build_backend() == "binder" + monkeypatch.setenv("LIGHTCONE_REGISTRY", "us-central1-docker.pkg.dev/p/binder") + monkeypatch.setenv("LIGHTCONE_BUILD_BUCKET", "bkt") + assert _hub_build_backend() == "cloudbuild" + + +def test_worker_image_for_run_via_cloudbuild( + hub_project: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.cli.commands import _worker_image_for_run + + monkeypatch.setenv("LIGHTCONE_REGISTRY", "us-central1-docker.pkg.dev/p/binder") + monkeypatch.setenv("LIGHTCONE_BUILD_BUCKET", "bkt") + captured: dict[str, object] = {} + + def fake_ensure(project, spec, *, project_name=None, on_progress=None): # noqa: ANN001, ANN202 + captured["spec"] = spec + captured["project_name"] = project_name + return "us-central1-docker.pkg.dev/p/binder/lc-proj:abc123def456" + + monkeypatch.setattr( + "lightcone.engine.cloudbuild.ensure_worker_image", fake_ensure + ) + image = _worker_image_for_run(hub_project, verbose=False) + assert image == "us-central1-docker.pkg.dev/p/binder/lc-proj:abc123def456" + assert captured == {"spec": "Containerfile", "project_name": "proj"} diff --git a/tests/test_cloudbuild.py b/tests/test_cloudbuild.py new file mode 100644 index 00000000..386d966b --- /dev/null +++ b/tests/test_cloudbuild.py @@ -0,0 +1,205 @@ +"""Unit tests for the Cloud Build backend (`lightcone.engine.cloudbuild`). + +Context staging and tarring run for real against throwaway projects; +GCP (metadata, GCS, Cloud Build API) is faked at the urllib boundary. +""" + +from __future__ import annotations + +import io +import json +import tarfile +import urllib.request +from pathlib import Path + +import pytest + +from lightcone.engine.cloudbuild import ( + BUCKET_ENV, + CloudBuildError, + cloudbuild_available, + ensure_worker_image, +) + +_REGISTRY = "us-central1-docker.pkg.dev/testproj/binder" + + +@pytest.fixture() +def project(tmp_path: Path) -> Path: + proj = tmp_path / "proj" + proj.mkdir() + (proj / "Containerfile").write_text( + "FROM python:3.12-slim\nCOPY requirements.txt .\n" + "RUN pip install -r requirements.txt\n" + ) + (proj / "requirements.txt").write_text("numpy\n") + (proj / "unrelated.py").write_text("print('not part of the context')\n") + return proj + + +@pytest.fixture() +def hub_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LIGHTCONE_REGISTRY", _REGISTRY) + monkeypatch.setenv(BUCKET_ENV, "test-build-bucket") + monkeypatch.setattr( + "lightcone.engine.cloudbuild._metadata_access_token", lambda: "tok" + ) + + +class _FakeResponse: + def __init__(self, payload: object, status: int = 200) -> None: + self._body = ( + payload if isinstance(payload, bytes) else json.dumps(payload).encode() + ) + self.status = status + + def read(self) -> bytes: + return self._body + + def __enter__(self) -> _FakeResponse: + return self + + def __exit__(self, *args: object) -> bool: + return False + + +def _install_fake_gcp( + monkeypatch: pytest.MonkeyPatch, + *, + statuses: list[str], + log_tail: str = "", +) -> list[tuple[str, str, bytes | None]]: + """Fake GCS upload + Cloud Build create/poll; record requests.""" + calls: list[tuple[str, str, bytes | None]] = [] + polls = iter(statuses) + + def fake_urlopen(req: urllib.request.Request, timeout: float = 0): # noqa: ANN202 + url = req.full_url + calls.append((req.get_method(), url, req.data)) + if "upload/storage" in url: + return _FakeResponse({"name": "sources/x.tar.gz"}) + if url.endswith("/builds"): + return _FakeResponse({"metadata": {"build": {"id": "build-123"}}}) + if "/builds/build-123" in url: + return _FakeResponse({"status": next(polls)}) + if "alt=media" in url: + return _FakeResponse(log_tail.encode()) + raise AssertionError(f"unexpected URL {url}") + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + monkeypatch.setattr("lightcone.engine.cloudbuild._POLL_INTERVAL_S", 0.0) + return calls + + +def test_available_needs_bucket_and_registry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert not cloudbuild_available() + monkeypatch.setenv(BUCKET_ENV, "b") + assert not cloudbuild_available() + monkeypatch.setenv("LIGHTCONE_REGISTRY", _REGISTRY) + assert cloudbuild_available() + + +def test_ensure_builds_and_returns_content_addressed_ref( + monkeypatch: pytest.MonkeyPatch, project: Path, hub_env: None +) -> None: + monkeypatch.setattr( + "lightcone.engine.cloudbuild.registry_image_exists", lambda ref: False + ) + calls = _install_fake_gcp(monkeypatch, statuses=["QUEUED", "WORKING", "SUCCESS"]) + + phases: list[str] = [] + ref = ensure_worker_image( + project, + "Containerfile", + project_name="myproj", + on_progress=lambda p, m: phases.append(p), + ) + + assert ref.startswith(f"{_REGISTRY}/lc-myproj:") + # Build submission targeted the right project and Containerfile. + method, url, body = next(c for c in calls if c[1].endswith("/builds")) + assert "projects/testproj/builds" in url + build = json.loads(body.decode()) + step_args = build["steps"][0]["args"] + assert step_args[:2] == ["build", "-t"] and ref in step_args + assert "-f" in step_args and "Containerfile" in step_args + assert build["images"] == [ref] + assert build["logsBucket"] == "gs://test-build-bucket/logs" + assert phases == ["staging", "queued", "working", "success"] + + # The uploaded tarball is the *staged* context: hashed files only. + _, _, tar_bytes = next(c for c in calls if "upload/storage" in c[1]) + assert tar_bytes is not None + names = tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r:gz").getnames() + assert "Containerfile" in names and "requirements.txt" in names + assert "unrelated.py" not in names + + +def test_ensure_cached_image_skips_everything( + monkeypatch: pytest.MonkeyPatch, project: Path, hub_env: None +) -> None: + monkeypatch.setattr( + "lightcone.engine.cloudbuild.registry_image_exists", lambda ref: True + ) + + def boom(*a: object, **k: object) -> None: + raise AssertionError("no HTTP expected on the cached path") + + monkeypatch.setattr("urllib.request.urlopen", boom) + ref = ensure_worker_image(project, "Containerfile", project_name="myproj") + assert ref.startswith(f"{_REGISTRY}/lc-myproj:") + + +def test_ensure_failure_surfaces_log_tail( + monkeypatch: pytest.MonkeyPatch, project: Path, hub_env: None +) -> None: + monkeypatch.setattr( + "lightcone.engine.cloudbuild.registry_image_exists", lambda ref: False + ) + _install_fake_gcp( + monkeypatch, + statuses=["WORKING", "FAILURE"], + log_tail="Step 2/3: pip install\nERROR: no matching distribution", + ) + with pytest.raises(CloudBuildError, match="FAILURE") as exc: + ensure_worker_image(project, "Containerfile", project_name="myproj") + assert "no matching distribution" in str(exc.value) + + +def test_ensure_custom_service_account( + monkeypatch: pytest.MonkeyPatch, project: Path, hub_env: None +) -> None: + monkeypatch.setenv( + "LIGHTCONE_BUILD_SERVICE_ACCOUNT", "builder@testproj.iam.gserviceaccount.com" + ) + monkeypatch.setattr( + "lightcone.engine.cloudbuild.registry_image_exists", lambda ref: False + ) + calls = _install_fake_gcp(monkeypatch, statuses=["SUCCESS"]) + ensure_worker_image(project, "Containerfile", project_name="myproj") + _, _, body = next(c for c in calls if c[1].endswith("/builds")) + assert ( + json.loads(body.decode())["serviceAccount"] + == "projects/testproj/serviceAccounts/builder@testproj.iam.gserviceaccount.com" + ) + + +def test_ensure_requires_registry( + monkeypatch: pytest.MonkeyPatch, project: Path +) -> None: + monkeypatch.setenv(BUCKET_ENV, "b") + with pytest.raises(CloudBuildError, match="LIGHTCONE_REGISTRY"): + ensure_worker_image(project, "Containerfile", project_name="p") + + +def test_non_artifact_registry_rejected( + monkeypatch: pytest.MonkeyPatch, project: Path, hub_env: None +) -> None: + monkeypatch.setenv("LIGHTCONE_REGISTRY", "ghcr.io/org") + monkeypatch.setattr( + "lightcone.engine.cloudbuild.registry_image_exists", lambda ref: False + ) + with pytest.raises(CloudBuildError, match="Artifact Registry"): + ensure_worker_image(project, "Containerfile", project_name="p")