diff --git a/source/examples/index.md b/source/examples/index.md index 4a5843ad..36ebb259 100644 --- a/source/examples/index.md +++ b/source/examples/index.md @@ -27,4 +27,5 @@ rapids-morpheus-pipeline/notebook fraud-detection-mlops-pipeline/notebook lulc-classification-gpu/notebook cuml-ray-hpo/notebook +rapids-topic-modeling-slurm/notebook ``` diff --git a/source/examples/rapids-topic-modeling-slurm/notebook.ipynb b/source/examples/rapids-topic-modeling-slurm/notebook.ipynb new file mode 100644 index 00000000..6de8aa03 --- /dev/null +++ b/source/examples/rapids-topic-modeling-slurm/notebook.ipynb @@ -0,0 +1,564 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": { + "tags": [ + "platforms/hpc", + "library/cuml", + "library/dagster-slurm", + "library/metaxy", + "library/scikit-learn", + "data-format/parquet", + "workflow/topic-modeling" + ] + }, + "source": [ + "# GPU Topic Modeling on HPC with dagster-slurm, cuML UMAP, and HDBSCAN\n", + "\n", + "_July, 2026_\n", + "\n", + "In this example workflow, we run a topic-modeling pipeline on a Slurm HPC cluster with [dagster-slurm](https://github.com/ascii-supply-networks/dagster-slurm), an open source integration that runs [Dagster](https://dagster.io/) assets as Slurm jobs directly from your laptop.\n", + "CPU stages train [scikit-learn](https://scikit-learn.org/) LDA models as a partitioned fan-out of `sbatch` jobs; GPU stages reduce and cluster the resulting topic vectors with [RAPIDS cuML](https://docs.rapids.ai/api/cuml/stable/) UMAP and HDBSCAN on a GPU node; the final asset streams a labeled meta-topic map back into the Dagster UI.\n", + "\n", + "This workflow targets classic HPC infrastructure: a shared Slurm cluster with no container runtime, where Python environments arrive as [pixi-pack](https://github.com/Quantco/pixi-pack) archives over SSH.\n", + "The same asset code runs in three settings without modification:\n", + "\n", + "- **local mode** on a laptop (no Slurm, no GPU),\n", + "- a **docker Slurm cluster** for CI and development (CPU fallback for the RAPIDS stages),\n", + "- a **real HPC cluster**, where the UMAP and HDBSCAN stages run on cuML with `gpus_per_node: 1`.\n", + "\n", + "The example is a small-scale, public-data version of a production Common Crawl topic-modeling pipeline.\n", + "It deliberately does not show GPU speedup benchmarks: the corpus is small, and the point is the orchestration mechanics, which are identical at production scale.\n", + "\n", + "## Quickstart\n", + "\n", + "```bash\n", + "git clone https://github.com/ascii-supply-networks/dagster-slurm.git\n", + "cd dagster-slurm && docker compose up -d # local Slurm cluster\n", + "cd examples && pixi run start-staging\n", + "# open http://localhost:3000 and materialize the rapids_topics group\n", + "```\n", + "\n", + "That is the whole interface.\n", + "The rest of this page explains what happens when you press the button, and how to point the same button at a real HPC cluster.\n", + "\n", + "### What you would normally do, and what happens here\n", + "\n", + "The usual workflow for a pipeline like this is: SSH to the login node, `module load` a Python that is almost the right version, hand-maintain a venv or conda environment on the cluster, write an `sbatch` script per stage, submit, poll `squeue`, and grep `slurm-*.out` files when something fails.\n", + "Here, none of that is manual: dagster-slurm packs the exact environment from your repo's lockfile into a content-hashed archive, ships it and the payload script over SSH, generates and submits the `sbatch` jobs, streams the logs back live, and records structured results per stage.\n", + "The cluster needs nothing beyond SSH and `sbatch`; your laptop needs nothing beyond pixi and docker." + ] + }, + { + "cell_type": "markdown", + "id": "pipeline", + "metadata": {}, + "source": [ + "## Overview\n", + "\n", + "This workflow shows how to:\n", + "\n", + "- Fan out partitioned model training as independent `sbatch` jobs (one Slurm job per partition) from a Dagster backfill\n", + "- Target **different packed environments per asset**: CPU stages run in a scikit-learn environment, GPU stages in a self-contained RAPIDS environment\n", + "- Request GPUs per asset (`gpus_per_node: 1`) only on deployments that have them, with an automatic CPU fallback (umap-learn / hdbscan) elsewhere\n", + "- Ship Python environments to clusters **without container runtimes** using `pixi-pack`\n", + "- Stream Slurm job logs, structured metadata, and inline plots back into the Dagster UI over SSH\n", + "- Override Slurm sizing (CPUs, memory, wall time, GPU count) per run from the Dagster launchpad\n", + "\n", + "## The pipeline\n", + "\n", + "Six assets in the `rapids_topics` group:\n", + "\n", + "```{note}\n", + "Every stage in this example runs as an `sbatch` job by choice, to keep the demo uniform.\n", + "That is not a constraint: Dagster orchestrates before, during, and after the cluster, so ingest or publishing assets can run in the Dagster process (or anywhere else) while only the heavy middle targets Slurm, all in one lineage graph.\n", + "On sites where compute nodes have no internet access you would want exactly that for `reuters_corpus`: run the download as a local or login-node asset instead of an `sbatch` job.\n", + "```\n", + "\n", + "```{mermaid}\n", + ":config: {\"flowchart\": {\"htmlLabels\": false, \"subGraphTitleMargin\": {\"top\": 4, \"bottom\": 10}}}\n", + "\n", + "flowchart LR\n", + " subgraph sk[\"workload-topic-modeling env\"]\n", + " direction TB\n", + " corpus[\"`**reuters_corpus**\n", + "CPU sbatch\n", + "download + SGML parse + shared vocabulary`\"]\n", + " lda[\"`**lda_models**\n", + "CPU sbatch, one job per (month, seed) partition\n", + "5 months x 3 seeds = 15 independent Slurm jobs`\"]\n", + " ttm[\"`**topic_term_matrix**\n", + "CPU sbatch\n", + "stack topic-term vectors of all models`\"]\n", + " corpus --> lda --> ttm\n", + " end\n", + " subgraph rapids[\"packaged-cluster-rapids env\"]\n", + " direction TB\n", + " umap[\"`**umap_embedding**\n", + "GPU sbatch\n", + "cuML UMAP (umap-learn fallback on CPU)`\"]\n", + " hdb[\"`**hdbscan_meta_topics**\n", + "GPU sbatch\n", + "cuML HDBSCAN (hdbscan fallback on CPU)`\"]\n", + " tmap[\"`**topic_map**\n", + "CPU sbatch\n", + "labeled meta-topic scatter + JSON summary`\"]\n", + " umap --> hdb --> tmap\n", + " end\n", + " sk --> rapids\n", + "```\n", + "\n", + "![The rapids_topics asset group in the Dagster lineage view, fully materialized against a real Slurm cluster](../../images/dagster-slurm-topics-lineage-overview.png)\n", + "\n", + "### Why cluster topic-term vectors?\n", + "\n", + "Doc-topic vectors from independently trained LDA models are not comparable:\n", + "topic 7 in the February model has nothing to do with topic 7 in the March model.\n", + "The chain therefore clusters **topic-term vectors** (normalized rows of the topic-word matrix over a vocabulary shared by all models) into meta-topics.\n", + "Every topic from every `(month, seed)` model becomes one point in vocabulary space; UMAP reduces those points to 2D and HDBSCAN groups recurring themes across months and seeds into meta-topics.\n", + "\n", + "### Dataset\n", + "\n", + "[Reuters-21578](https://kdd.ics.uci.edu/databases/reuters21578/reuters21578.html), the classic 1987 newswire research corpus (~21k documents, ~18k with usable body text).\n", + "It is small, quick to download, and dated, which preserves the temporal partitioning story of the production pipeline: the corpus asset buckets documents by month, and LDA training fans out over `(month, seed)`.\n", + "\n", + "### Code layout\n", + "\n", + "Dagster-slurm separates orchestration from computation.\n", + "The *assets* only pick a payload script, an environment, and Slurm resources; the *payloads* are plain Python scripts that talk to Dagster through [Dagster Pipes](https://docs.dagster.io/guides/build/external-pipelines) and also run standalone.\n", + "\n", + "| Piece | Path in the [dagster-slurm repo](https://github.com/ascii-supply-networks/dagster-slurm) |\n", + "|---|---|\n", + "| Asset definitions | [`examples/.../defs/rapids_topics/topic_assets.py`](https://github.com/ascii-supply-networks/dagster-slurm/blob/main/examples/projects/dagster-slurm-example/dagster_slurm_example/defs/rapids_topics/topic_assets.py) |\n", + "| Payload scripts | [`examples/.../dagster_slurm_example_hpc_workload/rapids_topics/`](https://github.com/ascii-supply-networks/dagster-slurm/tree/main/examples/projects/dagster-slurm-example-hpc-workload/dagster_slurm_example_hpc_workload/rapids_topics) |\n", + "| Environments | [`examples/pyproject.toml`](https://github.com/ascii-supply-networks/dagster-slurm/blob/main/examples/pyproject.toml) (`workload-topic-modeling`, `packaged-cluster-rapids`) |\n" + ] + }, + { + "cell_type": "markdown", + "id": "environments", + "metadata": {}, + "source": [ + "## One packed environment per stage\n", + "\n", + "The pipeline uses two pixi environments, declared per asset via `slurm_pack_cmd` metadata:\n", + "\n", + "- **`workload-topic-modeling`**: the standard cluster stack plus scikit-learn.\n", + "Used by the CPU stages (corpus, LDA, aggregation).\n", + "- **`packaged-cluster-rapids`**: a self-contained environment (Python 3.12, `numpy<2.3`) with `cuml` from the `rapidsai` conda channel on linux-64, and the CPU fallback libraries (umap-learn, hdbscan, matplotlib) co-installed.\n", + "Used by the UMAP, HDBSCAN, and report stages on every deployment, CPU fallback included.\n", + "\n", + "```toml\n", + "# examples/pyproject.toml (excerpt)\n", + "[tool.pixi.feature.cluster-rapids]\n", + "channels = [{ channel = \"rapidsai\", priority = 1 }]\n", + "\n", + "[tool.pixi.feature.cluster-rapids.dependencies]\n", + "python = \"3.12.*\"\n", + "numpy = \">=2.0,<2.3\"\n", + "umap-learn = \">=0.5,<1\"\n", + "hdbscan = \">=0.8,<1\"\n", + "matplotlib = \">=3.9,<4\"\n", + "\n", + "[tool.pixi.feature.cluster-rapids.target.linux-64.dependencies]\n", + "cuml = \">=25.10,<26\"\n", + "```\n", + "\n", + "Why a separate solve-group instead of adding cuML to the main cluster environment? RAPIDS pins numba, and numba pins numpy below what the rest of the cluster stack wants.\n", + "Co-locating umap-learn with the main stack makes the resolver backtrack into unbuildable sdists.\n", + "Giving RAPIDS its own solve-group keeps both environments installable, and dagster-slurm makes running each asset in its own environment a one-line metadata declaration:\n", + "\n", + "```python\n", + "_RAPIDS_PACK_METADATA = {\n", + " \"slurm_pack_cmd\": [\n", + " \"pixi\", \"run\", \"-e\", \"opstooling\", \"--frozen\",\n", + " \"python\", \"scripts/pack_environment.py\",\n", + " \"--env\", \"packaged-cluster-rapids\", \"--build-missing\",\n", + " ],\n", + "}\n", + "\n", + "@dg.asset(group_name=\"rapids_topics\", metadata=_RAPIDS_PACK_METADATA, ...)\n", + "def umap_embedding(...): ...\n", + "```\n", + "\n", + "`pixi-pack` builds the environment into a single self-extracting archive, cached on the cluster by content hash.\n", + "Where it builds is configurable: by default dagster-slurm packs on the cluster's login node, which is faster because only the lockfile and a few small inputs travel over SSH; if remote packing fails, it falls back to packing locally and uploading the multi-gigabyte archive.\n", + "When you know remote packing cannot work, for example because the login node has no internet access, force the local path with `SLURM_PACK_ON_REMOTE=0`.\n", + "Either way the archive carries every conda and PyPI package, so extracting it needs no internet access on the compute nodes.\n", + "\n", + "## GPU on the cluster, CPU everywhere else\n", + "\n", + "The payloads select their backend at import time: if cuML imports, they use the GPU implementation; otherwise the CPU library.\n", + "One payload serves both the docker Slurm cluster and a GPU node:\n", + "\n", + "```python\n", + "# umap_reduce.py (excerpt)\n", + "try:\n", + " import cuml\n", + " cuml.set_global_output_type(\"numpy\")\n", + " _HAS_CUML = True\n", + "except ImportError:\n", + " _HAS_CUML = False\n", + "\n", + "\n", + "def make_umap(*, n_components, n_neighbors, min_dist, metric,\n", + " random_state, build_algo):\n", + " if _HAS_CUML:\n", + " from cuml.manifold import UMAP as _UMAP\n", + " return _UMAP(\n", + " n_components=n_components, n_neighbors=n_neighbors,\n", + " min_dist=min_dist, metric=metric,\n", + " random_state=random_state, build_algo=build_algo,\n", + " verbose=True,\n", + " )\n", + " from umap import UMAP as _UMAP\n", + " return _UMAP(\n", + " n_components=n_components, n_neighbors=n_neighbors,\n", + " min_dist=min_dist, metric=metric,\n", + " random_state=random_state, low_memory=True, verbose=True,\n", + " )\n", + "```\n", + "\n", + "The matching asset requests a GPU only when the deployment is a real supercomputer:\n", + "\n", + "```python\n", + "def _gpu_slurm_opts() -> dict:\n", + " if _is_supercomputer():\n", + " return {\"nodes\": 1, \"cpus_per_task\": 8, \"mem\": \"32G\",\n", + " \"gpus_per_node\": 1}\n", + " return {\"nodes\": 1, \"cpus_per_task\": 2, \"mem\": \"4G\",\n", + " \"gpus_per_node\": 0}\n", + "```\n", + "\n", + "Each materialization reports which backend actually ran (`backend: cuml (GPU)` or `backend: umap-learn (CPU)`) in its asset metadata, so a misconfigured deployment is visible in the UI rather than silent.\n", + "\n", + "### Practical notes on cuML on HPC\n", + "\n", + "Hard-won details baked into the example, worth knowing before you adapt it:\n", + "\n", + "- **`cuml.set_global_output_type(\"numpy\")`** keeps the rest of the payload backend-agnostic: downstream code sees numpy arrays whether cuML or the CPU library produced them.\n", + "- **UMAP `build_algo` stays on `\"auto\"`.** cuML picks brute-force kNN for small inputs and GPU nn-descent at scale.\n", + "Forcing `nn_descent` on a small input (fewer than ~150 rows) crashes cuML with a CUDA invalid-argument error.\n", + "Set it explicitly only for large corpora.\n", + "- **cuML HDBSCAN caps `min_samples` at 1023.** The payload clamps the value on the GPU branch only.\n", + "- **cuML HDBSCAN labels more points as noise** than the CPU library at identical settings.\n", + "Tune `min_cluster_size` / `min_samples` against your real data, not against the CPU fallback.\n", + "- **cuML is linux-64 only** in this setup, declared under `[tool.pixi.feature.cluster-rapids.target.linux-64.dependencies]`, so the same pixi environment still solves on a macOS laptop (CPU libraries only)." + ] + }, + { + "cell_type": "markdown", + "id": "running", + "metadata": {}, + "source": [ + "## Running it\n", + "\n", + "### Prerequisites\n", + "\n", + "- [pixi](https://pixi.sh/) installed locally\n", + "- A clone of the [dagster-slurm repository](https://github.com/ascii-supply-networks/dagster-slurm)\n", + "- For the docker mode: docker compose\n", + "- For the HPC mode: SSH access to a Slurm cluster (a login node you can `sbatch` from); GPUs optional but required for the cuML path\n", + "\n", + "```bash\n", + "git clone https://github.com/ascii-supply-networks/dagster-slurm.git\n", + "cd dagster-slurm/examples\n", + "```\n", + "\n", + "### 1. Local mode (laptop, no Slurm)\n", + "\n", + "```bash\n", + "pixi run start\n", + "# open http://localhost:3000, materialize assets in the rapids_topics group\n", + "```\n", + "\n", + "```{tip}\n", + "Local mode materializes **3 of the 6 assets**: `reuters_corpus`, `lda_models`, and `topic_term_matrix` run directly on your machine.\n", + "`umap_embedding`, `hdbscan_meta_topics`, and `topic_map` need a Slurm deployment with the rapids environment and will fail locally: the dev environment deliberately excludes umap-learn/hdbscan/matplotlib because of the numba/numpy pin conflict described above.\n", + "Use the docker Slurm cluster below for the full chain.\n", + "```\n", + "\n", + "### 2. Docker Slurm cluster (full chain, CPU fallback)\n", + "\n", + "Start the dockerized Slurm cluster that ships with the repo and run in staging mode, which packs and deploys environments on demand:\n", + "\n", + "```bash\n", + "docker compose up -d # repo root: slurmctld + compute nodes\n", + "cd examples\n", + "pixi run start-staging\n", + "```\n", + "\n", + "Materialize the whole `rapids_topics` group.\n", + "Every asset becomes an `sbatch` job inside the docker cluster; the UMAP/HDBSCAN payloads log `backend: umap-learn (CPU)` and produce the same artifact shapes as the GPU path.\n", + "This is also what CI exercises.\n", + "\n", + "What to expect: the Reuters-21578 download is about 8 MB, and with cached environments the full chain completes in a few minutes (the individual jobs take seconds to ~1 minute each).\n", + "The first run also packs the two environments, which dominates wall-clock: the RAPIDS environment in\n", + "particular is large, so expect the initial pack to take on the order of tens\n", + "of minutes depending on your machine and network.\n", + "\n", + "### 3. Real HPC cluster with GPUs\n", + "\n", + "Point dagster-slurm at your cluster and start in supercomputer mode:\n", + "\n", + "```bash\n", + "export SLURM_EDGE_NODE_HOST=login.your-cluster.example\n", + "export SLURM_EDGE_NODE_USER=your-user\n", + "export SLURM_EDGE_NODE_KEY_PATH=~/your/key/path\n", + "\n", + "cd examples\n", + "pixi run start-staging-supercomputer # pack + deploy envs on demand\n", + "# or, with pre-deployed environments:\n", + "pixi run start-production-supercomputer\n", + "```\n", + "\n", + "In supercomputer deployments the GPU assets submit with `gpus_per_node: 1` and the payloads log `backend: cuml (GPU)`.\n", + "\n", + "Two things an experienced Slurm user will ask:\n", + "\n", + "- **\"A 15-job fan-out is 15 queue waits, can I run those inside one allocation?\"** dagster-slurm has session and heterogeneous-job modes for exactly this, running multiple assets inside a single Slurm allocation to amortize queueing.\n", + "They are experimental at the time of writing, which is why this example sticks to one `sbatch` per asset; see [execution modes](https://dagster-slurm.geoheil.com/docs/how-to/execution-modes) for status.\n", + "- **\"Can I target more than one cluster?\"** Deployments are configuration, not code: each deployment names its own edge node, so the same asset graph can run against your institute cluster in one deployment and a national system in another (the dagster-slurm docs ship site notes for several European HPC systems).\n", + "Even in the same graph you could override the default cluster config at execution time from the UI and have a multi-cluster pipeline.\n", + "\n", + "Optional environment variables:\n", + "\n", + "| Variable | Effect |\n", + "|---|---|\n", + "| `RAPIDS_TOPICS_BASE` | Base output directory on the cluster (default `$HOME/rapids_topics`) |\n", + "| `RAPIDS_TOPICS_CPU_ENV` | Path to an already-extracted CPU env on the cluster; skips packing |\n", + "| `RAPIDS_TOPICS_GPU_ENV` | Same, for the rapids environment |\n", + "\n", + "```{note}\n", + "Packing the rapids environment takes a while the first time (cuML is large).\n", + "For iterating on a real cluster, extract it once and set `RAPIDS_TOPICS_GPU_ENV`, or use the launchpad override below.\n", + "```\n", + "\n", + "### Per-run overrides from the launchpad\n", + "\n", + "All six assets share a config schema whose fields default to \"use the deployment-aware defaults\".\n", + "From the Dagster launchpad you can override `cpus_per_task`, `mem`, `time_limit`, `gpus_per_node`, and `pre_deployed_env_path` for a single run without touching code, plus the modeling knobs (topic count, UMAP neighbors, HDBSCAN cluster sizes) each stage exposes." + ] + }, + { + "cell_type": "markdown", + "id": "ui-walkthrough", + "metadata": {}, + "source": [ + "## What you see in the UI\n", + "\n", + "The screenshots below are from a run against a real Slurm cluster (staging supercomputer mode, one A100 for the GPU stages).\n", + "\n", + "**Partitioned fan-out.** `lda_models` is partitioned by `(month, seed)`; a backfill dispatches each partition as its own `sbatch` job.\n", + "The lineage view shows the fan-out filling up while downstream assets wait:\n", + "\n", + "![lda_models mid-backfill: partitions filling while topic_term_matrix and umap_embedding wait](../../images/dagster-slurm-topics-backfill-fanout.png)\n", + "\n", + "**One Slurm job per run.** The backfill's run list tags every run with its Slurm job id (`dagster_slurm/job_id`), so you can correlate Dagster runs with `sacct`/`squeue` output directly:\n", + "\n", + "![Backfill run list: each lda_models partition is a separate run with its own Slurm job id](../../images/dagster-slurm-topics-backfill-runs.png)\n", + "\n", + "**Live event log, including environment packing.** The event log shows the full lifecycle: cache miss on the environment hash, the reproducible `pixi-pack` command, job submission, and live log streaming over SSH:\n", + "\n", + "![Event log of a downstream run: environment packing, submission, and state transitions](../../images/dagster-slurm-topics-run-event-log.png)\n", + "\n", + "**Raw Slurm stdout, streamed.** The `stdout` tab shows exactly what ran on the compute node: working directory, payload path, which Python the packed environment resolved to:\n", + "\n", + "![Slurm job stdout streamed into the Dagster UI: environment activation and payload launch](../../images/dagster-slurm-topics-slurm-stdout.png)\n", + "\n", + "**Structured results per stage.** Every payload reports metadata through Pipes: row counts, output paths, the Slurm job id, plus scheduler-derived efficiency numbers (`node_hours`, `cpu_efficiency_pct`, `max_memory_mb`):\n", + "\n", + "![STEP_OUTPUT of topic_term_matrix: model count, output path, Slurm job id, efficiency metrics](../../images/dagster-slurm-topics-step-output-metadata.png)\n", + "\n", + "The corpus asset does the same at the head of the pipeline (document counts per month, vocabulary size, output directory):\n", + "\n", + "![reuters_corpus materialization metadata: 17893 documents, 5 months, shared vocabulary](../../images/dagster-slurm-topics-corpus-asset-metadata.png)\n", + "\n", + "**Metrics over time.** Because efficiency numbers are numeric metadata, Dagster plots them across materializations for free.\n", + "Cost regressions in a pipeline stage show up as a line going the wrong way:\n", + "\n", + "![Metadata plots: cpu_efficiency_pct and elapsed_seconds across materializations](../../images/dagster-slurm-topics-metadata-plots.png)\n", + "\n", + "**The result.** The terminal `topic_map` asset reports everything a reader of the pipeline needs as materialization metadata: cluster counts (`n_meta_topics: 7`, `n_noise_topics: 2`), the plot and summary paths on\n", + "the cluster filesystem, the labeled cluster summary as JSON, and a markdown preview of the plot itself:\n", + "\n", + "![topic_map run view: all downstream stages green, materialization metadata with cluster counts, artifact paths, and Slurm job id](../../images/dagster-slurm-topics-topic-map-run-success.png)\n", + "\n", + "That preview means the topic map renders directly inside the Dagster UI, no `scp` of PNGs off the cluster required:\n", + "\n", + "![topic_map_preview rendered inline in the Dagster UI: the meta-topic scatter inside the run view](../../images/dagster-slurm-topics-topic-map-inline-preview.png)\n", + "\n", + "The map itself, as written to the cluster filesystem: 45 topic-term vectors from the LDA models, UMAP-reduced and HDBSCAN-clustered into 7 meta-topics, each labeled with its top shared terms.\n", + "At this toy scale some clusters collapse onto newswire boilerplate and function words (\"vs, mln, loss\" is the earnings-report cluster; \"that, will, be\" is not a topic anyone would publish), which is exactly the honest output of 1987 newswire at 15 topics per month; the production-scale version of this chain uses far larger corpora and vocabulary filtering:\n", + "\n", + "![Labeled meta-topic map: UMAP scatter of topic-term vectors, colored by HDBSCAN cluster](../../images/dagster-slurm-topics-topic-map.png)\n", + "\n", + "**The whole chain.** A complete backfill of the group on the real cluster, three LDA partitions plus the four surrounding stages, finished in 19m28s end to end: the corpus stage (download + parse) took just under 12 minutes, each LDA partition under a minute, and the downstream aggregation-UMAP-HDBSCAN-report run just under five.\n", + "The backfill overview is the at-a-glance version:\n", + "\n", + "![Backfill overview: all six assets at 100 percent, every stage succeeded](../../images/dagster-slurm-topics-backfill-complete.png)\n" + ] + }, + { + "cell_type": "markdown", + "id": "metaxy", + "metadata": {}, + "source": [ + "## Refined example: incremental reprocessing with metaxy\n", + "\n", + "The basic pipeline recomputes everything downstream of a change: re-train one month's LDA models and the aggregation, UMAP, and HDBSCAN stages all run again over the full topic set.\n", + "At toy scale that is fine.\n", + "In production, where the fan-out is hundreds of partitions and model retrains arrive continuously, you want to know *which topic vectors actually changed* and skip the rest.\n", + "\n", + "[metaxy](https://github.com/anam-org/metaxy) adds sample-level incremental tracking on top of the same pipeline.\n", + "The dagster-slurm examples ship two working metaxy integrations (the `metaxy_simple` and `metaxy_ray` asset groups); the refinement below applies the identical pattern to the topic chain.\n", + "\n", + "### Feature specs\n", + "\n", + "Each topic-term vector is a tracked sample, keyed by a stable id.\n", + "Meta-topic assignments depend on them:\n", + "\n", + "```python\n", + "import metaxy as mx\n", + "\n", + "class TopicTermVectors(\n", + " mx.BaseFeature,\n", + " spec=mx.FeatureSpec(\n", + " key=\"rapids_topics/topic_term_vectors\",\n", + " id_columns=[\"topic_uid\"], # e.g. \"1987-02/seed=1/topic=7\"\n", + " fields=[\"month\", \"seed\", \"topic_id\", \"vector\"],\n", + " ),\n", + "):\n", + " topic_uid: str\n", + " month: str\n", + " seed: int\n", + " topic_id: int\n", + " vector: list[float]\n", + "\n", + "\n", + "class MetaTopics(\n", + " mx.BaseFeature,\n", + " spec=mx.FeatureSpec(\n", + " key=\"rapids_topics/meta_topics\",\n", + " id_columns=[\"topic_uid\"],\n", + " fields=[\"embedding\", \"cluster\"],\n", + " deps=[TopicTermVectors],\n", + " ),\n", + "):\n", + " topic_uid: str\n", + " embedding: list[float]\n", + " cluster: int\n", + "```\n", + "\n", + "### Asset side\n", + "\n", + "The aggregation asset registers vectors in a metaxy store instead of only writing a parquet file.\n", + "`metaxy.toml` selects the store per deployment (DuckDB locally, a Delta table on the shared cluster filesystem in\n", + "production), and the config file ships to the cluster with the payload via `extra_files`:\n", + "\n", + "```python\n", + "import metaxy.ext.dagster as mxd\n", + "\n", + "@mxd.metaxify\n", + "@dg.asset(\n", + " metadata={\"metaxy/feature\": \"rapids_topics/topic_term_vectors\",\n", + " **_CPU_PACK_METADATA},\n", + " group_name=\"rapids_topics_metaxy\",\n", + " deps=[lda_models],\n", + ")\n", + "def topic_term_matrix(context, compute: ComputeResource, config: TopicSlurmConfig):\n", + " metaxy_config = dg.file_relative_path(__file__, \"../../../../../metaxy.toml\")\n", + " return compute.run(\n", + " context=context,\n", + " payload_path=_payload(\"aggregate_topics_metaxy.py\"),\n", + " config=config,\n", + " extra_files=[metaxy_config],\n", + " extra_env={\"METAXY_STORE\": _store_for_deployment(), **_base_env()},\n", + " extra_slurm_opts=_merged_slurm_opts(_cpu_slurm_opts(), config),\n", + " ).get_results()\n", + "```\n", + "\n", + "### Payload side\n", + "\n", + "Inside the Slurm job, the payload asks the store what changed and processes only that increment:\n", + "\n", + "```python\n", + "# aggregate_topics_metaxy.py (core of the payload)\n", + "import metaxy as mx\n", + "\n", + "cfg = mx.init() # reads the shipped metaxy.toml\n", + "store = cfg.get_store()\n", + "\n", + "with store:\n", + " increment = store.resolve_update(\"rapids_topics/topic_term_vectors\",\n", + " samples=stacked_vectors)\n", + "\n", + "to_write = increment.new.to_polars()\n", + "stale = increment.stale.to_polars() # vectors whose upstream model changed\n", + "context.log.info(f\"{len(to_write)} new + {len(stale)} stale topic vectors\")\n", + "\n", + "with store.open(mode=\"w\"):\n", + " if len(to_write) > 0:\n", + " store.write(\"rapids_topics/topic_term_vectors\", to_write)\n", + " if len(stale) > 0:\n", + " store.write(\"rapids_topics/topic_term_vectors\", stale)\n", + "```\n", + "\n", + "The UMAP/HDBSCAN payload then resolves the increment for\n", + "`rapids_topics/meta_topics`.\n", + "If the increment is empty, it reports `status: up_to_date` and exits without touching the GPU, which on a busy cluster means the job releases its allocation in seconds.\n", + "\n", + "### What this buys you\n", + "\n", + "UMAP and HDBSCAN are global models: when the topic set does change, the reduction and clustering rerun over the full set, because a partial re-embed is not meaningful.\n", + "The incremental win for the GPU stages is therefore **change detection** (skip the whole GPU job when nothing upstream changed, for example after a partial backfill retry) and **provenance** (every meta-topic assignment is traceable to the exact model version that produced its topic vector).\n", + "For the fan-in stage the win is the classic one: the payload re-registers only new or stale vectors.\n", + "\n", + "Re-materializing a single month's `lda_models` partitions now results in:\n", + "\n", + "1. `topic_term_matrix` registers only that month's ~45 vectors as stale; everything else is untouched.\n", + "2. `umap_embedding` sees a non-empty increment and reruns (global model).\n", + "3. A second materialization with no upstream changes reports `status: up_to_date` at every stage and submits no compute-heavy work.\n", + "\n", + "For a complete, runnable reference of the store setup, `@metaxify` wiring, and `MetaxyDatasource`/`MetaxyDatasink` inside distributed payloads, see the `metaxy_simple` and `metaxy_ray` groups in the [dagster-slurm examples](https://github.com/ascii-supply-networks/dagster-slurm/tree/main/examples/projects/dagster-slurm-example/dagster_slurm_example/defs).\n", + "\n", + "## Conclusion\n", + "\n", + "The pattern shown here, per-asset Slurm sizing and per-asset packed environments around a CPU fan-out plus GPU reduction, is the shape of a large class of scientific workloads:\n", + "embarrassingly parallel training or extraction, followed by accelerated aggregation.\n", + "dagster-slurm contributes the orchestration ergonomics (lineage, backfills, live logs, structured metadata) without asking the HPC site for anything beyond SSH and `sbatch`, and RAPIDS contributes drop-in GPU acceleration for the reduction stages with a clean CPU fallback for development and CI.\n", + "\n", + "The workflow argument underneath all of it is the iteration loop.\n", + "Because the same asset code runs in local mode, on the docker cluster, and on the real cluster, you (or a coding agent, or CI) iterate at seconds-scale on your laptop, promote to the dockerized Slurm cluster to check scheduling behavior, and only then spend queue time and GPU hours, without rewriting anything between settings.\n", + "Cluster time becomes something you spend on purpose, not something you burn debugging environment drift.\n", + "\n", + "- **dagster-slurm**: [repository](https://github.com/ascii-supply-networks/dagster-slurm), [documentation](https://dagster-slurm.geoheil.com)\n", + "- **This example in the dagster-slurm repo**: [assets](https://github.com/ascii-supply-networks/dagster-slurm/tree/main/examples/projects/dagster-slurm-example/dagster_slurm_example/defs/rapids_topics) and [payloads](https://github.com/ascii-supply-networks/dagster-slurm/tree/main/examples/projects/dagster-slurm-example-hpc-workload/dagster_slurm_example_hpc_workload/rapids_topics)\n", + "- **Environment packaging details**: [Packaging dependencies](https://dagster-slurm.geoheil.com/docs/how-to/environment-packaging)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/source/hpc.md b/source/hpc.md index b918194d..3642fc79 100644 --- a/source/hpc.md +++ b/source/hpc.md @@ -334,6 +334,12 @@ python -m cudf.pandas my_script.py sbatch rapids_job.sh ``` + ## Orchestrating Pipelines + + The workflows above submit one job at a time by hand. For multi-stage pipelines, for example CPU preprocessing followed by GPU training, [dagster-slurm](https://github.com/ascii-supply-networks/dagster-slurm) runs [Dagster](https://dagster.io/) assets as `sbatch` jobs over SSH: it packs your Python environment with [pixi-pack](https://github.com/Quantco/pixi-pack), ships it to the cluster, submits each stage with its own resources (`gpus_per_node`, memory, wall time), and streams logs and results back into the Dagster UI. + It needs nothing on the cluster beyond SSH and `sbatch`, so it works without containers or admin access. + See the [GPU topic modeling on HPC](/examples/rapids-topic-modeling-slurm/notebook) example for a complete RAPIDS pipeline with cuML UMAP and HDBSCAN stages. + ```{relatedexamples} ``` diff --git a/source/images/dagster-slurm-topics-backfill-complete.png b/source/images/dagster-slurm-topics-backfill-complete.png new file mode 100644 index 00000000..9c3436fb Binary files /dev/null and b/source/images/dagster-slurm-topics-backfill-complete.png differ diff --git a/source/images/dagster-slurm-topics-backfill-fanout.png b/source/images/dagster-slurm-topics-backfill-fanout.png new file mode 100644 index 00000000..cc19ede9 Binary files /dev/null and b/source/images/dagster-slurm-topics-backfill-fanout.png differ diff --git a/source/images/dagster-slurm-topics-backfill-runs.png b/source/images/dagster-slurm-topics-backfill-runs.png new file mode 100644 index 00000000..f1053d01 Binary files /dev/null and b/source/images/dagster-slurm-topics-backfill-runs.png differ diff --git a/source/images/dagster-slurm-topics-corpus-asset-metadata.png b/source/images/dagster-slurm-topics-corpus-asset-metadata.png new file mode 100644 index 00000000..d5d682c1 Binary files /dev/null and b/source/images/dagster-slurm-topics-corpus-asset-metadata.png differ diff --git a/source/images/dagster-slurm-topics-lineage-overview.png b/source/images/dagster-slurm-topics-lineage-overview.png new file mode 100644 index 00000000..0829ea11 Binary files /dev/null and b/source/images/dagster-slurm-topics-lineage-overview.png differ diff --git a/source/images/dagster-slurm-topics-metadata-plots.png b/source/images/dagster-slurm-topics-metadata-plots.png new file mode 100644 index 00000000..3b9bb233 Binary files /dev/null and b/source/images/dagster-slurm-topics-metadata-plots.png differ diff --git a/source/images/dagster-slurm-topics-run-event-log.png b/source/images/dagster-slurm-topics-run-event-log.png new file mode 100644 index 00000000..5dbbeaf6 Binary files /dev/null and b/source/images/dagster-slurm-topics-run-event-log.png differ diff --git a/source/images/dagster-slurm-topics-slurm-stdout.png b/source/images/dagster-slurm-topics-slurm-stdout.png new file mode 100644 index 00000000..f476d799 Binary files /dev/null and b/source/images/dagster-slurm-topics-slurm-stdout.png differ diff --git a/source/images/dagster-slurm-topics-step-output-metadata.png b/source/images/dagster-slurm-topics-step-output-metadata.png new file mode 100644 index 00000000..1f70125b Binary files /dev/null and b/source/images/dagster-slurm-topics-step-output-metadata.png differ diff --git a/source/images/dagster-slurm-topics-topic-map-inline-preview.png b/source/images/dagster-slurm-topics-topic-map-inline-preview.png new file mode 100644 index 00000000..cbdc7bf2 Binary files /dev/null and b/source/images/dagster-slurm-topics-topic-map-inline-preview.png differ diff --git a/source/images/dagster-slurm-topics-topic-map-run-success.png b/source/images/dagster-slurm-topics-topic-map-run-success.png new file mode 100644 index 00000000..4abea8a7 Binary files /dev/null and b/source/images/dagster-slurm-topics-topic-map-run-success.png differ diff --git a/source/images/dagster-slurm-topics-topic-map.png b/source/images/dagster-slurm-topics-topic-map.png new file mode 100644 index 00000000..07713f3a Binary files /dev/null and b/source/images/dagster-slurm-topics-topic-map.png differ