From da571bc3e6119537d258b3a808f89684e83a3643 Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 12 Jul 2026 14:19:47 +0200 Subject: [PATCH 1/7] Add Dask Gateway support to lc run (JupyterHub deployments) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cluster_for_run gains a fourth branch, selected when a Gateway environment is ambient (DASK_GATEWAY__ADDRESS, as on a lightcone JupyterHub pod) or a cluster is named via LIGHTCONE_GATEWAY_CLUSTER. Gateway scheduler addresses use a gateway:// comm scheme that a bare distributed.Client cannot dial, so the parent/child rendezvous changes shape: cluster_for_run now yields the env overlay the child snakemake needs ({DASK_SCHEDULER_ADDRESS: ...} on existing branches, {LIGHTCONE_GATEWAY_CLUSTER: name} on the new one), and the executor rejoins the cluster by name through the authenticated Gateway API — the same rendezvous-by-ambient-context pattern as the SLURM branch. Ownership mirrors the address branch: clusters we create are adaptive (minimum=1, bounded by --jobs) and shut down on exit; clusters named by the user (e.g. created from the JupyterLab Dask sidebar) are left running with their scaling untouched. Startup fails fast when gateway workers don't advertise the cpus/memory/gpus resource contract, which otherwise hangs every task silently. Also: a lightcone-hub site entry detected via env markers (pods have generated hostnames, so detect_current_site now checks declared env markers before hostname patterns), and a [gateway] optional extra for the dask-gateway client. Validated against the lightcone-hub GKE staging deployment; design notes and the deployment procedure live in docs/design/jupyterhub-dask-gateway-gcp.md. Co-Authored-By: Claude Fable 5 --- docs/api/dask_cluster.md | 38 +- docs/contributing/backends.md | 5 +- docs/design/jupyterhub-dask-gateway-gcp.md | 526 ++++++++++++++++++ docs/user/cluster.md | 25 + pyproject.toml | 7 +- src/lightcone/cli/commands.py | 17 +- src/lightcone/engine/dask_cluster.py | 159 +++++- src/lightcone/engine/site_registry.py | 41 +- .../executor.py | 61 +- tests/test_dask_cluster.py | 208 ++++++- tests/test_site_registry.py | 29 + 11 files changed, 1061 insertions(+), 55 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..a3a59b67 100644 --- a/docs/api/dask_cluster.md +++ b/docs/api/dask_cluster.md @@ -1,24 +1,38 @@ # 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, 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 +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) → create an + adaptive Gateway cluster bounded by *max_workers*, or attach to the + named one; yield `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Gateway + scheduler addresses use a `gateway://` comm scheme a bare `Client` + cannot dial, so the child rejoins **by name** through the + authenticated Gateway API. Created clusters are shut down on exit; + attached ones are left running. Startup fails fast if workers don't + advertise the resource contract below. 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. ## Resource keys diff --git a/docs/contributing/backends.md b/docs/contributing/backends.md index 5c651e2d..b59f4ac2 100644 --- a/docs/contributing/backends.md +++ b/docs/contributing/backends.md @@ -23,8 +23,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..8bbb7d98 --- /dev/null +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -0,0 +1,526 @@ +# JupyterHub + Dask Gateway + kbatch on GCP — deployment procedure and lightcone-cli integration + +*Status: design proposal — researched 2026-07-12 against `main` (post-#157).* + +## 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. + +## 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..e3653adf 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -170,6 +170,31 @@ 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), `lc run` needs no configuration at +all: it requests an adaptive Dask Gateway cluster for the run, bounded +by `--jobs`, and shuts it down afterwards. Workers scale up from zero +on the compute pool, so the first rule may wait a couple of minutes +for node provisioning. + +To reuse a cluster you already created — for example from the +JupyterLab Dask sidebar, with its dashboard panels docked — name it: + +```bash +export LIGHTCONE_GATEWAY_CLUSTER= # e.g. hub.a1b2c3... +lc run +``` + +`lc run` attaches without changing the cluster's scaling and leaves it +running on exit, mirroring the `DASK_SCHEDULER_ADDRESS` convention. +Note that a Gateway scheduler's `gateway://` address cannot be used +with `DASK_SCHEDULER_ADDRESS` — always attach by name. + +Requires the optional gateway extra: +`pip install lightcone-cli[gateway]` (preinstalled on the hub image). + ## NERSC Perlmutter: site-specific notes !!! note "Setting up on Perlmutter for the first time?" diff --git a/pyproject.toml b/pyproject.toml index f7e5d297..51491bc1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,11 @@ dependencies = [ "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>=2024.1"] + [dependency-groups] eval = [ "anthropic>=0.90", @@ -100,7 +105,7 @@ 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 diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 259f8428..42d7849f 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -537,8 +537,9 @@ 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 - existing scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. + workstation, srun-launched workers inside a SLURM allocation, a + Dask Gateway cluster on a JupyterHub deployment, or an existing + scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. """ _abort_on_perlmutter_login() @@ -638,11 +639,17 @@ def run( raise click.ClickException(str(e)) with cluster_for_run( - verbose=verbose, local_directory=str(rundirs.dask_local) - ) as scheduler_addr: + verbose=verbose, + local_directory=str(rundirs.dask_local), + max_workers=int(n), + ) as cluster_env: env = { **os.environ, - "DASK_SCHEDULER_ADDRESS": scheduler_addr, + # 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 diff --git a/src/lightcone/engine/dask_cluster.py b/src/lightcone/engine/dask_cluster.py index 03049e5b..803465ff 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -1,19 +1,30 @@ # 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) → + create a Gateway cluster (or attach to a named one) via the + ``dask-gateway`` client. 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. Owned Gateway clusters are shut down on +exit; attached ones (named via ``LIGHTCONE_GATEWAY_CLUSTER``) are left +running, mirroring the ``DASK_SCHEDULER_ADDRESS`` convention. """ from __future__ import annotations @@ -34,6 +45,21 @@ RESOURCE_MEMORY = "memory" RESOURCE_GPUS = "gpus" +#: Env var carrying a Dask Gateway cluster name. Set by the user to make +#: ``lc run`` attach to an existing Gateway cluster (e.g. one created from +#: the JupyterLab Dask sidebar); 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" + +#: How long to wait for the first Gateway worker. Generous because a +#: scale-from-zero pool must provision a node and pull the worker image. +GATEWAY_WORKER_TIMEOUT = 600 + +#: Ceiling for adaptive scaling of an owned Gateway cluster when ``lc run`` +#: has no parallelism hint. +GATEWAY_DEFAULT_MAX_WORKERS = 8 + @dataclass class _NodeShape: @@ -89,32 +115,147 @@ 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`. + max_workers: int | 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). + + *max_workers* bounds adaptive scaling of an owned Gateway cluster; + the in-process branches size themselves to the node/allocation. """ 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, max_workers=max_workers) 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, max_workers: int | None +) -> Iterator[str]: + """Create (or attach to) a Dask Gateway cluster; yield its name. + + Ownership follows the same convention as the address branch: a + cluster we create is ours to shut down; a cluster named by the user + via ``LIGHTCONE_GATEWAY_CLUSTER`` (e.g. created from the JupyterLab + Dask sidebar, dashboard already docked) is left running on exit. + + 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 + + gateway = Gateway() + name = os.environ.get(GATEWAY_CLUSTER_ENV) + owned = name is None + if owned: + cluster = gateway.new_cluster(shutdown_on_close=False) + cluster.adapt( + minimum=1, maximum=max_workers or GATEWAY_DEFAULT_MAX_WORKERS + ) + else: + cluster = gateway.connect(name) + + if verbose: + mode = "started" if owned else "attached to" + print( + f"→ {mode} Dask Gateway cluster {cluster.name} " + f"(dashboard: {cluster.dashboard_link})" + ) + + try: + if owned: + # A worker is guaranteed to be coming (adaptive minimum=1), + # so wait for it and fail fast if the deployment forgot the + # resource contract — otherwise every task hangs silently. + client = cluster.get_client() + try: + client.wait_for_workers(1, timeout=GATEWAY_WORKER_TIMEOUT) + _assert_worker_resources(client) + finally: + client.close() + else: + # Don't touch the user's scaling; verify the contract only + # if workers are already up (an adaptive cluster at zero + # will scale once the executor submits tasks). + client = cluster.get_client() + try: + workers = client.scheduler_info().get("workers", {}) + if workers: + _assert_worker_resources(client) + finally: + client.close() + yield cluster.name + finally: + if owned: + cluster.shutdown() + else: + cluster.close() + + +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 + if any( + RESOURCE_CPUS in (w.get("resources") or {}) for w in workers.values() + ): + return + raise RuntimeError( + "Dask Gateway workers do not advertise the lightcone resource " + f"contract ({RESOURCE_CPUS}/{RESOURCE_MEMORY}/{RESOURCE_GPUS}); " + "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/site_registry.py b/src/lightcone/engine/site_registry.py index f4c00351..966b62a4 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,23 @@ "//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. + "container_runtime": "none", + # 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 +170,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/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index 41bf9c1e..fdd7b9b9 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -19,6 +19,7 @@ ) from lightcone.engine.dask_cluster import ( + GATEWAY_CLUSTER_ENV, RESOURCE_CPUS, RESOURCE_GPUS, RESOURCE_MEMORY, @@ -90,25 +91,53 @@ 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: a Gateway cluster is rejoined *by name* + through the authenticated Gateway API (its ``gateway://`` scheduler + address cannot be dialled by a bare ``Client``); anything else is a + plain address in ``DASK_SCHEDULER_ADDRESS``. + + Returns ``(client, gateway_cluster_or_None)`` — the cluster handle is + kept so ``shutdown()`` can release its local connections. It is never + shut down here: cluster lifetime belongs to the parent ``lc run``. + """ + if name := os.environ.get(GATEWAY_CLUSTER_ENV): try: - from dask.distributed import Client + from dask_gateway import Gateway except ImportError as exc: raise WorkflowError( - "dask.distributed is required for the dask executor " - "(`pip install distributed`)." + 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 + + try: + from dask.distributed import Client + except ImportError as exc: + raise WorkflowError( + "dask.distributed is required for the dask executor " + "(`pip install distributed`)." + ) from exc + + addr = os.environ.get("DASK_SCHEDULER_ADDRESS") + if not addr: + 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." + ) + return Client(addr), None - addr = os.environ.get("DASK_SCHEDULER_ADDRESS") - if not addr: - 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) + +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 run_job(self, job: JobExecutorInterface) -> None: cmd = self.format_job_exec(job) @@ -169,5 +198,9 @@ 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 is owned (and shut down) by the parent lc run. + self._gateway_cluster.close() finally: super().shutdown() diff --git a/tests/test_dask_cluster.py b/tests/test_dask_cluster.py index 4bb6b473..5e3866d2 100644 --- a/tests/test_dask_cluster.py +++ b/tests/test_dask_cluster.py @@ -28,6 +28,8 @@ def _clean_env(monkeypatch: pytest.MonkeyPatch) -> None: for var in ( "DASK_SCHEDULER_ADDRESS", + "DASK_GATEWAY__ADDRESS", + "LIGHTCONE_GATEWAY_CLUSTER", "SLURM_JOB_ID", "SLURM_NNODES", "SLURM_CPUS_ON_NODE", @@ -69,8 +71,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 +85,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 +100,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 +118,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 +278,196 @@ 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) -> None: + self._resources = resources + + def wait_for_workers(self, n: int, timeout: int) -> None: + pass + + def scheduler_info(self) -> dict[str, object]: + if self._resources is None: + return {"workers": {}} + return {"workers": {"w0": {"resources": self._resources}}} + + def close(self) -> None: + pass + + +class _FakeGatewayCluster: + def __init__( + self, name: str, log: list[str], resources: dict[str, float] | None + ) -> None: + self.name = name + self.dashboard_link = f"http://hub/services/dask-gateway/{name}/status" + self._log = log + self._resources = resources + + def adapt(self, minimum: int, maximum: int) -> None: + self._log.append(f"adapt({minimum},{maximum})") + + def get_client(self) -> _FakeGatewayClient: + return _FakeGatewayClient(self._resources) + + def shutdown(self) -> None: + self._log.append("shutdown") + + def close(self) -> None: + self._log.append("close") + + +def _install_fake_gateway( + monkeypatch: pytest.MonkeyPatch, + log: list[str], + resources: dict[str, float] | None = None, +): + """Inject a fake ``dask_gateway`` module; return it for inspection.""" + import sys + import types + + class _FakeGateway: + def new_cluster(self, **kwargs: object) -> _FakeGatewayCluster: + log.append(f"new_cluster({kwargs})") + return _FakeGatewayCluster("hub.new123", log, resources) + + def connect(self, name: str) -> _FakeGatewayCluster: + log.append(f"connect({name})") + return _FakeGatewayCluster(name, log, resources) + + 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_takes_gateway_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Ambient DASK_GATEWAY__ADDRESS (a hub pod) routes to the gateway branch.""" + monkeypatch.setenv("DASK_GATEWAY__ADDRESS", "http://proxy/services/dask-gateway/") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources={RESOURCE_CPUS: 2.0}) + + with cluster_for_run(max_workers=5) as env: + assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.new123"} + assert any(call.startswith("new_cluster") for call in log) + assert "adapt(1,5)" in log + + assert "shutdown" in log, "owned gateway cluster must be shut down on exit" + + +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={RESOURCE_CPUS: 2.0}) + + @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_attach_by_name_is_not_shut_down( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIGHTCONE_GATEWAY_CLUSTER attaches to an existing (e.g. sidebar-created) + cluster and must leave it running on exit — same convention as the + DASK_SCHEDULER_ADDRESS branch.""" + monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") + log: list[str] = [] + _install_fake_gateway(monkeypatch, log, resources={RESOURCE_CPUS: 2.0}) + + 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), ( + "attached clusters keep the user's scaling" + ) + + 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={}) # workers, no contract + + with pytest.raises(RuntimeError, match="resource contract"): + with cluster_for_run(): + pass # pragma: no cover + + assert "shutdown" in log, "failed owned cluster must still be cleaned up" + + +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={RESOURCE_CPUS: 2.0}) + + 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_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_site_registry.py b/tests/test_site_registry.py index 61396e1a..ebdca393 100644 --- a/tests/test_site_registry.py +++ b/tests/test_site_registry.py @@ -95,3 +95,32 @@ 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" + assert site.get("container_runtime") == "none" + 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 878d9e13b9d7597971f5325a16eb604d1c5b3070 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Sun, 12 Jul 2026 23:51:39 +0200 Subject: [PATCH 2/7] Attach-only Gateway + kubernetes as a first-class container runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reworks the Dask Gateway integration around one idea: on a hub deployment, the worker pod IS the container. The Containerfile keeps defining the recipe environment on every path; only its realization differs (OCI wrap locally/SLURM, pod image on the hub). - container.py: new "kubernetes" runtime — wrap_recipe no-ops, a site-declared non-OCI runtime resolves as explicit in load_runtime (kills the false "no container runtime found" warning on the hub; Perlmutter's OCI hint semantics unchanged). LIGHTCONE_REGISTRY env names the deployment registry; registry_image_ref maps the local lc-- tag to /lc-: (same content hash — same artifact). resolve_image_for_run stays site-agnostic so code_version is stable across laptop<->hub. Best-effort registry probe via the GKE metadata token. - dask_cluster.py: the Gateway branch is attach-only. lc run never creates clusters; it discovers the user's running cluster through the user-scoped Gateway API (one -> attach; zero -> error with a copy-pasteable new_cluster snippet, image pre-filled and shutdown_on_close=False; several -> LIGHTCONE_GATEWAY_CLUSTER disambiguates). Scaling untouched, cluster left running — the Gateway flavor of the DASK_SCHEDULER_ADDRESS convention. Verifies the resource contract on live workers (now cpus+memory together; cpus-only would still hang mem_mb rules) and the worker image via the deployment-injected LIGHTCONE_WORKER_IMAGE, warning on mismatch. Gateway()/connect failures translate to guided errors. - executor.py: _connect_client now mirrors the parent's branch priority (address first), and lc run strips the losing variable from the child env — a stale shell-exported gateway name can no longer redirect the child to an unverified cluster. - commands.py: lc build on kubernetes checks the registry and prints publish instructions that reproduce the hash-attested build context (lc build + docker tag/push, not a raw docker build). lc run cross-checks runtime vs cluster branch and warns on both mismatch directions. Manifests record the pod's actual image (worker_image). - tests: attach-only gateway suite (discovery, lifecycle, contract, image verification, stale-name guidance), registry helpers, runtime precedence; conftest neutralizes ambient hub env markers so the suite passes inside a hub pod. Validated end-to-end under a simulated hub environment; reviewed by a 25-finding multi-agent workflow with adversarial verification, all confirmed findings fixed. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRVPnWVUtpumnFsCzEEA8U --- docs/api/dask_cluster.md | 39 +- docs/contributing/backends.md | 12 +- docs/design/jupyterhub-dask-gateway-gcp.md | 83 +++++ docs/user/cluster.md | 85 ++++- pyproject.toml | 4 + src/lightcone/cli/commands.py | 350 +++++++++++++++--- src/lightcone/engine/container.py | 234 +++++++++++- src/lightcone/engine/dask_cluster.py | 253 +++++++++---- src/lightcone/engine/manifest.py | 6 + src/lightcone/engine/site_registry.py | 6 +- src/lightcone/engine/snakefile.py | 11 +- .../executor.py | 53 +-- tests/conftest.py | 19 + tests/test_container.py | 232 ++++++++++++ tests/test_dask_cluster.py | 335 +++++++++++++++-- tests/test_site_registry.py | 4 +- 16 files changed, 1499 insertions(+), 227 deletions(-) diff --git a/docs/api/dask_cluster.md b/docs/api/dask_cluster.md index a3a59b67..0784b5ea 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, max_workers=None) → Iterator[dict[str, str]]` +## `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 @@ -16,15 +16,24 @@ 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) → create an - adaptive Gateway cluster bounded by *max_workers*, or attach to the - named one; yield `{"LIGHTCONE_GATEWAY_CLUSTER": name}`. Gateway - scheduler addresses use a `gateway://` comm scheme a bare `Client` - cannot dial, so the child rejoins **by name** through the - authenticated Gateway API. Created clusters are shut down on exit; - attached ones are left running. Startup fails fast if workers don't - advertise the resource contract below. Requires the optional - dependency: `pip install lightcone-cli[gateway]`. + `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. @@ -32,7 +41,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. +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 @@ -87,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 b59f4ac2..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). diff --git a/docs/design/jupyterhub-dask-gateway-gcp.md b/docs/design/jupyterhub-dask-gateway-gcp.md index 8bbb7d98..6242d45b 100644 --- a/docs/design/jupyterhub-dask-gateway-gcp.md +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -1,6 +1,8 @@ # 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 @@ -513,6 +515,87 @@ single environment. 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. +## 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/ · diff --git a/docs/user/cluster.md b/docs/user/cluster.md index e3653adf..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 @@ -173,28 +177,81 @@ 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), `lc run` needs no configuration at -all: it requests an adaptive Dask Gateway cluster for the run, bounded -by `--jobs`, and shuts it down afterwards. Workers scale up from zero -on the compute pool, so the first rule may wait a couple of minutes -for node provisioning. - -To reuse a cluster you already created — for example from the -JupyterLab Dask sidebar, with its dashboard panels docked — name it: +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 ``` -`lc run` attaches without changing the cluster's scaling and leaves it -running on exit, mirroring the `DASK_SCHEDULER_ADDRESS` convention. -Note that a Gateway scheduler's `gateway://` address cannot be used -with `DASK_SCHEDULER_ADDRESS` — always attach by name. +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 51491bc1..36135a1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,6 +108,10 @@ mypy_path = "src" 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 42d7849f..7f5df5b2 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -537,14 +537,15 @@ 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, a - Dask Gateway cluster on a JupyterHub deployment, or an existing - scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. + 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 from lightcone.engine.scratch import ( RunLockBusyError, acquire_run_lock, @@ -570,6 +571,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 ) @@ -638,29 +648,48 @@ def run( except RunLockBusyError as e: raise click.ClickException(str(e)) - with cluster_for_run( - verbose=verbose, - local_directory=str(rundirs.dask_local), - max_workers=int(n), - ) as 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), - } - 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( @@ -897,7 +926,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. @@ -908,8 +940,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: @@ -917,6 +959,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 " @@ -930,34 +976,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 = ( @@ -965,9 +999,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 + 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 - seen.add(spec_str) + 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..450930d0 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( @@ -703,6 +744,156 @@ 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.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 + + req = urllib.request.Request( + f"https://{host}/v2/{path}/manifests/{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): + return None + + def wrap_recipe( recipe: str, *, @@ -718,18 +909,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 +968,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 803465ff..1ff1c71e 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -7,8 +7,12 @@ 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 Gateway cluster (or attach to a named one) via the - ``dask-gateway`` client. Gateway scheduler addresses use a custom + **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 — @@ -22,9 +26,9 @@ 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. Owned Gateway clusters are shut down on -exit; attached ones (named via ``LIGHTCONE_GATEWAY_CLUSTER``) are left -running, mirroring the ``DASK_SCHEDULER_ADDRESS`` convention. +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 @@ -45,20 +49,21 @@ RESOURCE_MEMORY = "memory" RESOURCE_GPUS = "gpus" -#: Env var carrying a Dask Gateway cluster name. Set by the user to make -#: ``lc run`` attach to an existing Gateway cluster (e.g. one created from -#: the JupyterLab Dask sidebar); 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. 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" -#: How long to wait for the first Gateway worker. Generous because a -#: scale-from-zero pool must provision a node and pull the worker image. -GATEWAY_WORKER_TIMEOUT = 600 - -#: Ceiling for adaptive scaling of an owned Gateway cluster when ``lc run`` -#: has no parallelism hint. -GATEWAY_DEFAULT_MAX_WORKERS = 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``). +WORKER_IMAGE_ENV = "LIGHTCONE_WORKER_IMAGE" @dataclass @@ -115,7 +120,7 @@ def cluster_for_run( *, verbose: bool = False, local_directory: str | None = None, - max_workers: int | None = None, + expected_worker_image: str | None = None, ) -> Iterator[dict[str, str]]: """Yield the env overlay the child snakemake needs to reach the cluster. @@ -132,8 +137,11 @@ def cluster_for_run( spill lands on Lustre instead of DVS-mounted home/CFS (where small- file I/O is slow and can pressure the gateway nodes). - *max_workers* bounds adaptive scaling of an owned Gateway cluster; - the in-process branches size themselves to the node/allocation. + *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: @@ -142,7 +150,9 @@ def cluster_for_run( return if GATEWAY_CLUSTER_ENV in os.environ or "DASK_GATEWAY__ADDRESS" in os.environ: - with _gateway_cluster(verbose=verbose, max_workers=max_workers) as name: + with _gateway_cluster( + verbose=verbose, expected_worker_image=expected_worker_image + ) as name: yield {GATEWAY_CLUSTER_ENV: name} return @@ -161,14 +171,19 @@ def cluster_for_run( @contextmanager def _gateway_cluster( - *, verbose: bool, max_workers: int | None + *, verbose: bool, expected_worker_image: str | None ) -> Iterator[str]: - """Create (or attach to) a Dask Gateway cluster; yield its name. - - Ownership follows the same convention as the address branch: a - cluster we create is ours to shut down; a cluster named by the user - via ``LIGHTCONE_GATEWAY_CLUSTER`` (e.g. created from the JupyterLab - Dask sidebar, dashboard already docked) is left running on exit. + """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 @@ -185,52 +200,141 @@ def _gateway_cluster( "`pip install lightcone-cli[gateway]`." ) from exc - gateway = Gateway() - name = os.environ.get(GATEWAY_CLUSTER_ENV) - owned = name is None - if owned: - cluster = gateway.new_cluster(shutdown_on_close=False) - cluster.adapt( - minimum=1, maximum=max_workers or GATEWAY_DEFAULT_MAX_WORKERS - ) - else: + 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: - mode = "started" if owned else "attached to" print( - f"→ {mode} Dask Gateway cluster {cluster.name} " + f"→ Attached to Dask Gateway cluster {cluster.name} " f"(dashboard: {cluster.dashboard_link})" ) try: - if owned: - # A worker is guaranteed to be coming (adaptive minimum=1), - # so wait for it and fail fast if the deployment forgot the - # resource contract — otherwise every task hangs silently. - client = cluster.get_client() - try: - client.wait_for_workers(1, timeout=GATEWAY_WORKER_TIMEOUT) - _assert_worker_resources(client) - finally: - client.close() - else: - # Don't touch the user's scaling; verify the contract only - # if workers are already up (an adaptive cluster at zero - # will scale once the executor submits tasks). - client = cluster.get_client() - try: - workers = client.scheduler_info().get("workers", {}) - if workers: - _assert_worker_resources(client) - finally: - client.close() + 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: - if owned: - cluster.shutdown() - else: - cluster.close() + # 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: @@ -245,16 +349,27 @@ def _assert_worker_resources(client: object) -> None: 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 (w.get("resources") or {}) for w in workers.values() + 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}/{RESOURCE_GPUS}); " - "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)." + 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)." ) 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 966b62a4..b3b7c644 100644 --- a/src/lightcone/engine/site_registry.py +++ b/src/lightcone/engine/site_registry.py @@ -88,8 +88,10 @@ "backend": "dask-gateway", "connection": {}, # Kubernetes is the container runtime: the Gateway worker pod - # image IS the recipe environment, so recipes run unwrapped. - "container_runtime": "none", + # 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. 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 fdd7b9b9..d40a6a5e 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -95,15 +95,32 @@ 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: a Gateway cluster is rejoined *by name* - through the authenticated Gateway API (its ``gateway://`` scheduler - address cannot be dialled by a bare ``Client``); anything else is a - plain address in ``DASK_SCHEDULER_ADDRESS``. + 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: cluster lifetime belongs to the parent ``lc run``. + 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: + raise WorkflowError( + "dask.distributed is required for the dask executor " + "(`pip install distributed`)." + ) from exc + return Client(addr), None + if name := os.environ.get(GATEWAY_CLUSTER_ENV): try: from dask_gateway import Gateway @@ -115,23 +132,12 @@ def _connect_client(): # type: ignore[no-untyped-def] cluster = Gateway().connect(name) return cluster.get_client(), cluster - try: - from dask.distributed import Client - except ImportError as exc: - raise WorkflowError( - "dask.distributed is required for the dask executor " - "(`pip install distributed`)." - ) from exc - - addr = os.environ.get("DASK_SCHEDULER_ADDRESS") - if not addr: - 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." - ) - return Client(addr), None + 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] @@ -200,7 +206,8 @@ def shutdown(self) -> None: self._client.close() if self._gateway_cluster is not None: # Releases this process's connections only; the cluster - # itself is owned (and shut down) by the parent lc run. + # 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_container.py b/tests/test_container.py index 2fb10a74..fc1ca814 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,167 @@ 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 + + +# ---- 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_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 +833,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 5e3866d2..904fece2 100644 --- a/tests/test_dask_cluster.py +++ b/tests/test_dask_cluster.py @@ -30,6 +30,7 @@ def _clean_env(monkeypatch: pytest.MonkeyPatch) -> None: "DASK_SCHEDULER_ADDRESS", "DASK_GATEWAY__ADDRESS", "LIGHTCONE_GATEWAY_CLUSTER", + "LIGHTCONE_WORKER_IMAGE", "SLURM_JOB_ID", "SLURM_NNODES", "SLURM_CPUS_ON_NODE", @@ -284,35 +285,56 @@ def close(self) -> None: class _FakeGatewayClient: - def __init__(self, resources: dict[str, float] | None) -> None: + def __init__( + self, + resources: dict[str, float] | None, + worker_image: str | None = None, + ) -> None: self._resources = resources - - def wait_for_workers(self, n: int, timeout: int) -> None: - pass + 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 + 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) + return _FakeGatewayClient(self._resources, self._worker_image) def shutdown(self) -> None: self._log.append("shutdown") @@ -321,23 +343,52 @@ 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.""" + """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 new_cluster(self, **kwargs: object) -> _FakeGatewayCluster: - log.append(f"new_cluster({kwargs})") - return _FakeGatewayCluster("hub.new123", log, resources) + 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})") - return _FakeGatewayCluster(name, log, resources) + 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] @@ -345,25 +396,79 @@ def connect(self, name: str) -> _FakeGatewayCluster: return module -def test_gateway_env_takes_gateway_path(monkeypatch: pytest.MonkeyPatch) -> None: - """Ambient DASK_GATEWAY__ADDRESS (a hub pod) routes to the gateway branch.""" +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, resources={RESOURCE_CPUS: 2.0}) + _install_fake_gateway(monkeypatch, log, running=[]) - with cluster_for_run(max_workers=5) as env: - assert env == {"LIGHTCONE_GATEWAY_CLUSTER": "hub.new123"} - assert any(call.startswith("new_cluster") for call in log) - assert "adapt(1,5)" in log + 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) - assert "shutdown" in log, "owned gateway cluster must be shut down on exit" + +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={RESOURCE_CPUS: 2.0}) + _install_fake_gateway( + monkeypatch, log, resources=dict(_GOOD_RESOURCES), running=["hub.run1"] + ) @contextmanager def _should_not_run(*, verbose: bool, local_directory: str | None = None): @@ -385,23 +490,21 @@ def test_explicit_address_wins_over_gateway( assert env == {"DASK_SCHEDULER_ADDRESS": "tcp://existing:8786"} -def test_gateway_attach_by_name_is_not_shut_down( +def test_gateway_named_cluster_skips_discovery( monkeypatch: pytest.MonkeyPatch, ) -> None: - """LIGHTCONE_GATEWAY_CLUSTER attaches to an existing (e.g. sidebar-created) - cluster and must leave it running on exit — same convention as the - DASK_SCHEDULER_ADDRESS branch.""" + """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={RESOURCE_CPUS: 2.0}) + _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), ( - "attached clusters keep the user's scaling" - ) + 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 @@ -427,13 +530,152 @@ def test_gateway_rejects_workers_without_resource_contract( 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={}) # workers, no contract + _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 "shutdown" in log, "failed owned cluster must still be cleaned up" + 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( @@ -444,7 +686,7 @@ def test_executor_connects_via_gateway_name( monkeypatch.setenv("LIGHTCONE_GATEWAY_CLUSTER", "hub.abc999") monkeypatch.delenv("DASK_SCHEDULER_ADDRESS", raising=False) log: list[str] = [] - _install_fake_gateway(monkeypatch, log, resources={RESOURCE_CPUS: 2.0}) + _install_fake_gateway(monkeypatch, log, resources=dict(_GOOD_RESOURCES)) from snakemake_executor_plugin_dask.executor import _connect_client @@ -454,6 +696,35 @@ def test_executor_connects_via_gateway_name( 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: diff --git a/tests/test_site_registry.py b/tests/test_site_registry.py index ebdca393..2d339059 100644 --- a/tests/test_site_registry.py +++ b/tests/test_site_registry.py @@ -110,7 +110,9 @@ def test_env_markers_detect_hub_despite_generated_hostname( ) site = detect_current_site() assert site.key == "lightcone-hub" - assert site.get("container_runtime") == "none" + # 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( From bce0908529161a5b18ca4e0f24a12652ee18a35b Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Mon, 13 Jul 2026 00:24:20 +0200 Subject: [PATCH 3/7] Fix gateway execution: worker pods don't share the driver's environment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validated live on the lightcone-hub GKE deployment; a two-rule pipeline now runs end to end on a Gateway cluster (attach, contract+image verification, DAG on spot workers, manifests with worker_image ground truth, lc verify clean, idempotent re-run). Three real-world fixes, all the same root cause — Gateway worker pods share the project volume with the driver but not its HOME, install prefix, or cwd, unlike local/SLURM workers which inherit the driver's environment by process ancestry: - Executor overrides get_python_executable(): the driver's sys.executable (conda notebook path) doesn't exist in a slim worker image — exit 127 before snakemake starts. The worker's own PATH resolves the interpreter; the worker image is the software deployment. - Executor pins --directory on the child snakemake: worker pods start in their own HOME, so relative outputs "succeeded" into the pod's ephemeral filesystem while the driver saw nothing. - lc run passes --shared-fs-usage input-output/persistence/sources/ storage-local-copies (not source-cache: the child mkdir'd the driver's ~/.cache path and died on PermissionError) and --latency-wait 60 (NFS attribute caching hid freshly written worker outputs from the driver past snakemake's 5s default) — gateway branch only; local/SLURM keep snakemake defaults. Also: _run_shell forwards a bounded stderr tail when the child fails (exit 127 with zero diagnostics is a debugging dead end), and gateway_branch_active() exposes the branch decision so lc run can shape the snakemake invocation before entering the cluster context. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRVPnWVUtpumnFsCzEEA8U --- src/lightcone/cli/commands.py | 33 ++++++++- src/lightcone/engine/dask_cluster.py | 16 ++++ .../executor.py | 53 +++++++++++++ tests/test_cli.py | 74 +++++++++++++++++++ 4 files changed, 175 insertions(+), 1 deletion(-) diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 7f5df5b2..b71df8f0 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -545,7 +545,11 @@ def run( _abort_on_perlmutter_login() from lightcone.engine.container import KUBERNETES_RUNTIME, load_runtime - from lightcone.engine.dask_cluster import GATEWAY_CLUSTER_ENV, cluster_for_run + from lightcone.engine.dask_cluster import ( + GATEWAY_CLUSTER_ENV, + cluster_for_run, + gateway_branch_active, + ) from lightcone.engine.scratch import ( RunLockBusyError, acquire_run_lock, @@ -635,6 +639,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 @@ -748,12 +753,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", @@ -770,6 +786,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 diff --git a/src/lightcone/engine/dask_cluster.py b/src/lightcone/engine/dask_cluster.py index 1ff1c71e..67bb9c89 100644 --- a/src/lightcone/engine/dask_cluster.py +++ b/src/lightcone/engine/dask_cluster.py @@ -115,6 +115,22 @@ 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( *, diff --git a/src/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index d40a6a5e..6dbb8387 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -17,6 +17,10 @@ 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, @@ -58,6 +62,16 @@ 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) if forwarded: block = "\n".join(forwarded) + "\n" lock_path = os.environ.get("LIGHTCONE_OUT_LOCK") @@ -145,6 +159,45 @@ 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) self.logger.debug(cmd) diff --git a/tests/test_cli.py b/tests/test_cli.py index a5c03fc1..0f1a2c74 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -213,6 +213,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 From 4b39e63781ffad2a96b6333d6599d514c1ef9370 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Mon, 13 Jul 2026 09:58:56 +0200 Subject: [PATCH 4/7] Sanitize project names into valid OCI image names; probe never raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found live on the hub: an ASTRA analysis named "Union2.1 Flat ΛCDM MAP Fit" produced the image name lc-union2.1-flat-λcdm-map-fit — not a valid OCI reference (registries and docker/podman require lowercase ASCII), and the λ reached http.client's ASCII-only request line, so `lc build` crashed with UnicodeEncodeError instead of reporting. - compute_image_tag routes through image_name_slug(): NFKD transliteration for accented latin, other non-ASCII dropped, invalid chars collapsed to "-". Already-valid names (gwtest, union2.1) pass through unchanged, so existing tags keep their names. - registry_image_exists percent-encodes the URL path and treats ValueError/UnicodeError as "cannot tell" (None) — the probe's contract is tri-state, never a traceback. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRVPnWVUtpumnFsCzEEA8U --- src/lightcone/engine/container.py | 44 +++++++++++++++++-- tests/test_container.py | 73 +++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 4 deletions(-) diff --git a/src/lightcone/engine/container.py b/src/lightcone/engine/container.py index 450930d0..3534fd7d 100644 --- a/src/lightcone/engine/container.py +++ b/src/lightcone/engine/container.py @@ -290,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] @@ -391,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: @@ -858,6 +882,7 @@ def registry_image_exists(ref: str) -> bool | None: report status, never to gate execution. """ import urllib.error + import urllib.parse import urllib.request if "/" not in ref: @@ -875,8 +900,16 @@ def registry_image_exists(ref: str) -> bool | None: 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/{path}/manifests/{tag}", + f"https://{host}/v2/{quoted_path}/manifests/{quoted_tag}", method="HEAD", headers={ "Authorization": f"Bearer {token}", @@ -890,7 +923,10 @@ def registry_image_exists(ref: str) -> bool | None: if exc.code == 404: return False return None - except (urllib.error.URLError, OSError): + 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 diff --git a/tests/test_container.py b/tests/test_container.py index fc1ca814..b1f2655b 100644 --- a/tests/test_container.py +++ b/tests/test_container.py @@ -675,6 +675,43 @@ def test_resolution_ignores_deployment_registry( 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) ------------------------------ @@ -772,6 +809,42 @@ def test_no_credentials_returns_unknown( 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 From 2d4081885741f13ec959fe9bb78ef332beedcf75 Mon Sep 17 00:00:00 2001 From: Alexandre Boucaud Date: Wed, 15 Jul 2026 22:54:06 +0200 Subject: [PATCH 5/7] Surface Gateway worker stderr to the driver (LCR-175) --- .../executor.py | 30 +++++++++++++--- tests/test_dask_plugin.py | 34 +++++++++++++++++-- 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/src/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index 6dbb8387..86001115 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -31,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 @@ -52,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 @@ -72,6 +77,7 @@ def _run_shell(cmd: str) -> int: 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") @@ -87,7 +93,21 @@ 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 _build_resources(job: JobExecutorInterface) -> dict[str, float]: @@ -231,7 +251,9 @@ async def check_active_jobs( ) continue - exit_code = future.result() + exit_code, block = 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}." diff --git a/tests/test_dask_plugin.py b/tests/test_dask_plugin.py index ba1ad686..b16d1095 100644 --- a/tests/test_dask_plugin.py +++ b/tests/test_dask_plugin.py @@ -12,6 +12,7 @@ from snakemake_executor_plugin_dask.executor import ( _build_resources, + _emit_block, _run_shell, ) @@ -21,13 +22,40 @@ 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_build_resources_default_uses_threads() -> None: From 8113c556c442456317df6711678fe684ff1a0609 Mon Sep 17 00:00:00 2001 From: Alexandre Boucaud Date: Wed, 15 Jul 2026 23:38:00 +0200 Subject: [PATCH 6/7] Tolerate legacy int worker result; add _unpack_result (LCR-175) --- src/snakemake_executor_plugin_dask/executor.py | 18 +++++++++++++++++- tests/test_dask_plugin.py | 9 +++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/snakemake_executor_plugin_dask/executor.py b/src/snakemake_executor_plugin_dask/executor.py index 86001115..97bfeb3f 100644 --- a/src/snakemake_executor_plugin_dask/executor.py +++ b/src/snakemake_executor_plugin_dask/executor.py @@ -110,6 +110,22 @@ def _emit_block(block: str) -> None: 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) + return int(result), "" + + def _build_resources(job: JobExecutorInterface) -> dict[str, float]: """Translate Snakemake resources to Dask abstract resource units.""" res: dict[str, float] = {} @@ -251,7 +267,7 @@ async def check_active_jobs( ) continue - exit_code, block = future.result() + exit_code, block = _unpack_result(future.result()) if self._gateway_cluster is not None: _emit_block(block) if exit_code != 0: diff --git a/tests/test_dask_plugin.py b/tests/test_dask_plugin.py index b16d1095..172f8c64 100644 --- a/tests/test_dask_plugin.py +++ b/tests/test_dask_plugin.py @@ -14,6 +14,7 @@ _build_resources, _emit_block, _run_shell, + _unpack_result, ) @@ -58,6 +59,14 @@ def test_emit_block_writes_to_stdout(capsys) -> None: # type: ignore[no-untyped 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: res = _build_resources(_job(threads=4)) assert res == {"cpus": 4.0} From 50389ee513b26960731c4265681e729175accdfc Mon Sep 17 00:00:00 2001 From: Alexandre Boucaud Date: Thu, 16 Jul 2026 10:55:15 +0200 Subject: [PATCH 7/7] Add open question regarding build path to the PRD (LCR-175) --- docs/design/jupyterhub-dask-gateway-gcp.md | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/design/jupyterhub-dask-gateway-gcp.md b/docs/design/jupyterhub-dask-gateway-gcp.md index 6242d45b..3d299b06 100644 --- a/docs/design/jupyterhub-dask-gateway-gcp.md +++ b/docs/design/jupyterhub-dask-gateway-gcp.md @@ -514,6 +514,32 @@ 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).** + §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