docs: MkDocs site + gh-pages

.github/workflows/docs.yml

@@ -0,0 +1,26 @@
+name: Docs
+
+# Build the MkDocs site and publish it to the gh-pages branch (GitHub Pages).
+# Deploys only from main (i.e. on a release merge) or when run manually — never
+# from develop, so published docs always match a released state.
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Sync docs deps
+        run: uv sync --group docs
+      - name: Build and deploy
+        run: uv run mkdocs gh-deploy --force

docs/api.md

@@ -0,0 +1,28 @@
+# API reference
+
+The public Python API — what you import in your Dask jobs.
+
+## Capturing the source map
+
+::: daskgenie.track
+
+::: daskgenie.watch
+
+::: daskgenie.get_layer_map
+
+## Local profiling
+
+::: daskgenie.LocalProfiler
+    options:
+      members:
+        - __init__
+
+## Uploading to the collector
+
+::: daskgenie.create_run
+
+::: daskgenie.upload_graph
+
+## Distributed
+
+::: daskgenie.client.register

docs/changelog.md

@@ -0,0 +1 @@
+--8<-- "CHANGELOG.md"

docs/configuration.md

@@ -0,0 +1,35 @@
+# Configuration
+
+## Environment variables
+
+| Variable | Component | Purpose |
+| --- | --- | --- |
+| `DASKGENIE_DSN` | collector | Postgres/TimescaleDB DSN; selects the Timescale backend (default in Docker). |
+| `DASKGENIE_DB` | collector | SQLite path (or `:memory:`) when no DSN is set. |
+| `DASKGENIE_HOST` / `DASKGENIE_PORT` | collector | Bind address (default `127.0.0.1:8765`). |
+| `COLLECTOR_URL` | dashboard | Where the Next.js server proxies `/api` (e.g. `http://collector:8765`). |
+| `NEXT_PUBLIC_COLLECTOR_WS` | dashboard | WebSocket base the browser connects to (default `ws://<host>:8765`). |
+
+## Profiler knobs
+
+`register()` and `LocalProfiler` accept:
+
+| Argument | Default | Effect |
+| --- | --- | --- |
+| `deep` | `False` | Enable the memray deep-memory engine. |
+| `sample_interval` | `0.2` | Seconds between RSS/status samples. |
+| `flush_interval` | `0.5` | Seconds between pushes to the collector. Keep it well under how fast a worker OOMs so the killer chunk's metadata is sent before the process dies. |
+| `deep_epoch_seconds` | `5.0` | memray rotation window; smaller = finer time resolution, more overhead. |
+
+## Storage backend
+
+The collector defaults to **TimescaleDB** in Docker (set by `DASKGENIE_DSN` in
+`docker-compose.yml`); without a DSN it uses the self-contained SQLite store —
+which is also what the test suite runs on. Point it at any Postgres/Timescale
+instance with `DASKGENIE_DSN=postgresql://user:pass@host:5432/db`.
+
+## Prometheus
+
+The collector exposes `/metrics` with per-worker RSS and managed-memory gauges
+plus sample/death counters, so existing Grafana setups get value without the
+custom UI.

docs/contributing.md

@@ -0,0 +1 @@
+--8<-- "CONTRIBUTING.md"

docs/dashboard.md

@@ -0,0 +1,44 @@
+# The dashboard
+
+The dashboard lists every run (with the machine that opened it) and updates live
+over a WebSocket while a job runs. Open a run for these tabs:
+
+## Overview
+Live stats (workers, tasks, deaths, peak RSS), memory-over-time, and the hottest
+allocation line.
+
+## Timeline
+A large, zoomable memory-over-time chart. **Click any point** to pin that instant
+and see what was running (with source lines) and which lines were allocating in
+that window. Below it, a stacked **per-layer allocation timeline**.
+
+Navigation on every chart: scroll to zoom the time axis, drag to pan, shift-drag
+to box-zoom.
+
+## Workers
+A live, native-Dask-style table: per-worker RSS against the memory limit, managed
+memory, CPU, threads, and executing/ready counts.
+
+## Task stream
+Task rectangles coloured by layer on a zoomable time axis — an **ALL TASKS**
+global lane (all workers packed together) on top and **PER WORKER** lanes below.
+
+## Graph
+The real connected task DAG. Small graphs use an interactive view; large ones
+render the whole graph on a pan/zoom canvas (no collapsing to a blob). Hover a
+node to light up its edges; click for its source line and chunk sizes.
+
+## Memory
+The deep view (needs `deep=True`): an allocation **flamegraph** (per-worker,
+click to zoom, framework frames dimmed so your code stands out), a **peak bytes
+by source line** table, and **peak memory by task**.
+
+## Post-mortem
+Each worker death: the allocation lines at the high-water mark, the suspect task,
+its source line (syntax highlighted), and the chunk it was holding
+(`(4000, 4000) float64 = 128 MB`).
+
+!!! info "OOM is a heuristic"
+    The scheduler doesn't tell a plugin *why* a worker left. DaskGenie flags a
+    *suspected* OOM only when tasks were in flight at an unexpected removal, and
+    never over-claims.

docs/deep-memory.md

@@ -0,0 +1,50 @@
+# Deep memory (memray)
+
+Pass `deep=True` to profile *inside* the memory. On each worker DaskGenie drives
+[memray](https://github.com/bloomberg/memray) **as a library** — a short capture
+rotated every few seconds, read back and folded to the first line of *your* code
+responsible for each high-water-mark allocation. You never see or handle a
+capture file; only the aggregates are pushed.
+
+```python
+# distributed
+run_id = genie.register(client, "http://localhost:8765", deep=True)
+
+# local schedulers
+with dg.LocalProfiler(url, source_map=smap, collection=result, deep=True) as prof:
+    result.compute(scheduler="threads")
+```
+
+Needs the `deep` extra (`pip install 'daskgenie[deep]'`). Deep mode costs roughly
+1.5–2× runtime, so it is opt-in per run.
+
+## What you get
+
+- **Peak bytes by source line** — `job.py:42 build = 12.8 GB`.
+- **A per-worker flamegraph** — the full call stack (root → leaf) for each
+  allocation, aggregated into an icicle you can zoom into; framework frames are
+  dimmed so your code stands out. This is the memray "tree" read, on your code.
+- **Per-layer allocations over time** and **per-task peak memory**.
+- On a worker death, the allocation lines at the **high-water mark** in the
+  post-mortem.
+
+## How it works
+
+The engine runs a background thread of **epochs**: a memray `Tracker` captures to
+a throwaway temp file for a few seconds, then stops, is read back with
+`FileReader`, folded to the high-water-mark bytes per user source line (and full
+call stacks for the flamegraph), and the file is deleted before the next epoch
+starts. Rotation keeps each capture tiny and makes attribution time-resolved.
+Everything is guarded — a memray failure degrades to Tier-1 sampling and never
+crashes or OOMs the job.
+
+## Limitations
+
+- **memray** is Linux/macOS + CPython only; a single tracker runs per process.
+- Source-line attribution is strongest for `dask.array` and `dask.delayed`. For
+  `dask.dataframe` (dask-expr) and xarray the heavy work runs inside library/C
+  code, so per-line allocations fold to framework frames — the flamegraph still
+  shows the full tree.
+- On a **hard OOM kill** the in-flight epoch can die before it flushes, so a
+  specific post-mortem's line attribution is best-effort. The **Memory** tab
+  remains the reliable place to see the culprit line.

docs/examples.md

@@ -0,0 +1,37 @@
+# Examples
+
+Runnable scripts covering the breadth of Dask live in the
+[`examples/`](https://github.com/polymood/DaskGenie/tree/develop/examples)
+directory.
+
+| Script | Shows |
+| --- | --- |
+| `graph_source_map.py` | Source attribution only — layer → the line that built it. |
+| `local_scheduler.py` | Memory + chunks + graph for a local scheduler. |
+| `distributed_oom.py` | A real worker OOM and its post-mortem on `dask.distributed`. |
+| `deep_oom.py` | Deep memory (memray): the source line that allocated the OOM array. |
+| `big_pipeline_oom.py` | A minutes-long multi-stage pipeline, a large task graph, a real OOM. |
+| `threaded_big.py` | Threaded scheduler hoarding memory instead of spilling to Zarr. |
+| `threaded_crash.py` | Threaded scheduler that actually crashes (self-limiting `MemoryError`). |
+| `dask_delayed.py` | `dask.delayed` custom ETL DAG on a distributed cluster. |
+| `dask_dataframe.py` | `dask.dataframe` groupby + shuffle. |
+| `dask_bag.py` | `dask.bag` aggregation on the processes scheduler. |
+| `xarray_zarr.py` | xarray on Zarr — a chunked cube → anomaly. |
+| `xarray_netcdf.py` | xarray on NetCDF (HDF5) — rolling detrend. |
+
+## Running them
+
+Start the stack once, then run an example and open the dashboard:
+
+```bash
+docker compose up -d --build
+
+uv run --extra demo --extra deep python examples/deep_oom.py
+uv run --extra examples --extra deep python examples/xarray_zarr.py
+```
+
+!!! note "Attribution across collections"
+    Source-line attribution is strongest for `dask.array` and `dask.delayed`. For
+    `dask.dataframe` (dask-expr) and xarray the heavy work runs inside library C
+    code, so allocations fold to framework frames — the flamegraph still shows
+    the full call tree, and the graph / memory views work throughout.

docs/index.md

@@ -0,0 +1,40 @@
+# DaskGenie
+
+A memory profiler and live dashboard for Dask that ties a worker's memory back to
+the **line of your code** that caused it.
+
+DaskGenie fuses [memray](https://github.com/bloomberg/memray)-deep allocation
+tracing with Dask's task graph and streams it to a real-time dashboard. When a
+worker dies from an out-of-memory error, you can see the suspect task, the chunk
+it was holding, and the exact source line that allocated the array that killed
+it — not just "a worker disappeared". It works on `dask.distributed` and on the
+local (threaded / processes / synchronous) schedulers, across `dask.array`,
+`dask.dataframe`, `dask.bag`, `dask.delayed`, and xarray on Zarr/NetCDF.
+
+## Features
+
+- **Source attribution** — every task-graph layer mapped back to the user source
+  line that built it, never into `dask`/`numpy`/`xarray` internals.
+- **Deep memory (memray, as a library)** — per-line high-water-mark attribution
+  and a per-worker flamegraph / memray-style tree. Opt-in `deep=True`.
+- **Worker-death post-mortem** — which chunk, and which line, killed the worker.
+- **Real-time dashboard** — a Next.js app streaming over WebSocket: live workers,
+  a zoomable task stream, the whole task-graph DAG, memory-over-time with a
+  click-to-inspect spike explorer, per-layer allocations over time, and the
+  flamegraph.
+- **Any scheduler** — distributed plugins or the local-scheduler `LocalProfiler`.
+- **Team-friendly** — runs record the machine (hostname + IP) that opened them.
+- **Prometheus + TimescaleDB** — `/metrics` for Grafana, TimescaleDB or SQLite.
+
+## Where to next
+
+<div class="grid cards" markdown>
+
+- :material-download: **[Installation](installation.md)** — `pip install daskgenie` and the extras.
+- :material-rocket-launch: **[Quick start](quickstart.md)** — bring up the stack and profile a job.
+- :material-memory: **[Deep memory](deep-memory.md)** — the memray engine and flamegraph.
+- :material-view-dashboard: **[The dashboard](dashboard.md)** — every tab explained.
+- :material-cog: **[Configuration](configuration.md)** — environment variables and knobs.
+- :material-code-braces: **[API reference](api.md)** — the public Python API.
+
+</div>

docs/installation.md

@@ -0,0 +1,37 @@
+# Installation
+
+```bash
+pip install daskgenie
+```
+
+DaskGenie requires **Python 3.11 or newer**.
+
+## Optional extras
+
+```bash
+pip install 'daskgenie[deep]'       # memray-backed deep memory profiling
+pip install 'daskgenie[collector]'  # the FastAPI collector service (server side)
+pip install 'daskgenie[examples]'   # xarray, zarr, pandas, ... for the examples
+```
+
+| Extra | Pulls in | For |
+| --- | --- | --- |
+| `deep` | `memray` | Line-level allocation tracing and the flamegraph. |
+| `collector` | `fastapi`, `uvicorn`, `psycopg`, `websockets`, ... | Running the collector service yourself (usually you run it via Docker instead). |
+| `examples` | `xarray`, `zarr`, `pandas`, `pyarrow`, `h5netcdf`, ... | The runnable example scripts. |
+| `demo` | `distributed`, `numpy` | The minimal distributed demos. |
+
+!!! note "memray platform support"
+    Deep memory profiling needs memray, which is **Linux/macOS + CPython** only.
+    Where memray can't be imported, the profiler silently degrades to lightweight
+    RSS / managed-memory sampling — you still get memory-over-time and per-chunk
+    metadata, just not per-line attribution.
+
+## The dashboard stack
+
+The collector, TimescaleDB, and the web dashboard run as containers — you don't
+pip-install them. See [Quick start](quickstart.md):
+
+```bash
+docker compose up -d --build
+```

docs/local-schedulers.md

@@ -0,0 +1,33 @@
+# Local schedulers
+
+The distributed worker/scheduler plugins only exist on `dask.distributed`. For
+the local schedulers (`.compute(scheduler="threads"|"processes"|"synchronous")`,
+the default for bare dask collections) use `LocalProfiler`, which hooks Dask's
+callback API instead of installing cluster plugins.
+
+```python
+import daskgenie as dg
+
+with dg.track() as source_map:
+    result = build_pipeline()
+
+with dg.LocalProfiler(
+    "http://localhost:8765",
+    run_name="threaded job",
+    source_map=source_map,
+    collection=result,
+    deep=True,
+) as prof:
+    result.compute(scheduler="threads")
+
+# prof.run_id shows up in the dashboard like any other run
+```
+
+`LocalProfiler` samples process RSS over time, tags each sample with the tasks
+running at that instant, records the shape/dtype/nbytes of each task's output
+chunk, and — with `deep=True` — runs the same memray engine as the distributed
+path. There is no worker-death attribution here: a local scheduler has no workers
+to lose, so an OOM simply kills the process (the memory-over-time curve up to that
+point is still captured, flushed once per second).
+
+See [`examples/xarray_zarr.py`](examples.md) and the other threaded examples.