Add proper support for DGX spark with easy manager. Thank you NVIDIA!

README.md

@@ -102,7 +102,7 @@ cd ai-toolkit
 
 Then start the manager with the script for your platform:
 
-Linux:
+Linux (x86_64 and ARM64, including DGX Spark / DGX OS):
 ```bash
 chmod +x run_linux.sh
 ./run_linux.sh
@@ -146,7 +146,7 @@ pip3 install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0
 pip3 install -r requirements.txt
 ```
 
-For devices running **DGX OS** (including DGX Spark), follow [these](dgx_instructions.md) instructions.
+These steps also work on ARM64 Linux, including DGX Spark / DGX OS.
 
 
 Windows:

dgx_instructions.md

@@ -1,84 +0,0 @@
-# AI Toolkit by Ostris
-
-## DGX OS installation instructions
-
-You need to use Python 3.11 to run AI Toolkit on DGX OS. The easiest way to do this without affecting the system installation of Python is to create a virtual environment with **miniconda**, which allows you to specify the version of Python to use in the environment.
-
-This guide will assume you have a fresh installation of DGX OS, and will guide you through the installation of all requirements.
-
-### Installation instructions for DGX OS:
-
-**1) Get Python 3.11 (via miniconda)**
-
-Install the latest version of miniconda:
-```
-wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
-chmod u+x Miniconda3-latest-Linux-aarch64.sh
-./Miniconda3-latest-Linux-aarch64.sh
-```
-
-Restart your bash or ssh session. If miniconda was installed successfully, it will automatically load the 'base' environment by default. If you want to disable this behaviour, run:
-```
-conda config --set auto_activate_base false
-```
-
-Now you can create a Python 3.11 environment for ai-toolkit:
-```
-conda create --name ai-toolkit python=3.11
-```
-
-Then activate the environment with:
-
-```
-conda activate ai-toolkit
-```
-
-
-**2) Install PyTorch**
-
-```
-pip3 install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
-```
-
-
-**3) Install the remaining requirements (dgx_requirements.txt)**
-
-```
-pip3 install -r dgx_requirements.txt
-```
-
-### Running the UI on DGX OS:
-
-Running the UI is not that different from doing it on other systems, however, you need to install the ARM64 version of NodeJS for Linux, which is compatible with the NVIDIA Grace CPU.
-
-
-**1) Install Node.js**
-
-Download a Linux ARM64 build of Node.js from: https://nodejs.org (for example: https://nodejs.org/dist/v24.11.1/node-v24.11.1-linux-arm64.tar.xz)
-
-Extract it and add the bin directory to your path. I extracted it to **/opt** and added the following to my ~/.bashrc file:
-```
-export PATH=“/opt/node-v24.11.1-linux-arm64/bin:$PATH”
-```
-
-
-**2) Compile and run the Node.js UI**
-
-Change to the ui directory, then build and run the UI:
-```
-cd ui
-npm run build_and_start
-```
-
-If all went well, you’ll be able to access the UI on port 8675 and start training.
-
-
-<details>
-  <summary>Troubleshooting issues</summary>
-If you’re not getting any output when starting a training job from the UI, it’s probably crashing before the process started, the best way to debug these issues is to run the python training script directly (which is normally started by the UI). To do this, set up a training job in the UI, go to the advanced config screen, copy and paste the configuration into a file like train.yaml, then run the training script like this with the conda virtual environment active:
-
-```
-python run.py path/to/train.yaml
-```
-</details>
-<br>

dgx_requirements.txt

@@ -1,13 +0,0 @@
-# You need to use Python 3.11, the easiest way to get this on DGX OS without impacting the system version of Python is to create an environment with miniconda.
-
-# specific dependency versions needed on DGX OS devices:
-scipy==1.16.0
-tifffile==2025.6.11
-imageio==2.37.0
-scikit_image==0.25.2
-clean_fid==0.1.35
-pywavelets==1.9.0
-contourpy==1.3.3
-opencv_python_headless==4.11.0.86
-
--r requirements_base.txt

manager/README.md

@@ -21,11 +21,11 @@ python3 -m manager doctor      # diagnose problems
 - **Hardware → spec mapping** is in [spec.py](spec.py). One universal torch
   pin (2.13.0 / torchvision 0.28.0 / torchaudio 2.11.0) on every platform:
   cu130 wheels when the driver supports CUDA 13 (cu126 fallback for older
-  drivers, refused outright on Blackwell GPUs which need cu130), same stack +
+  drivers, refused outright on Blackwell GPUs which need cu130), the same
-  Python 3.11 + `dgx_requirements.txt` on DGX/Grace, PyPI wheels on Mac,
+  stack on Linux aarch64 (DGX Spark / Grace), PyPI wheels on Mac,
   rocm7.1 (experimental) for AMD, `--cpu` to force a CPU install. **Torch
   pins there must be updated together with the README install instructions,
-  run_mac.zsh, and dgx_instructions.md.**
+  and run_mac.zsh.**
 - **Accelerators everywhere wheels exist**, via per-spec `extra_packages`
   (installed after requirements with `--upgrade` so they override pins) and
   `optional_packages` (installed one-by-one, warn-only on failure):

manager/detect.py

@@ -108,7 +108,4 @@ def detect():
 
     if IS_MAC:
         result["backend"] = "mps" if arch == "arm64" else "cpu"
-
-    # DGX OS / Grace (GB10, DGX Spark): NVIDIA GPU on aarch64 Linux
-    result["is_dgx"] = bool(result["os"] == "linux" and arch == "aarch64" and nvidia)
     return result

manager/env.py

@@ -81,6 +81,34 @@ def _venv_platform():
         return None
 
 
+def _venv_has_headers():
+    """True if the venv's base interpreter ships Python.h.
+
+    triton JIT-compiles a small C launcher (cuda_utils.c) on first GPU use and
+    needs the headers; a venv built on a distro Python without the -dev
+    package fails there with "Python.h: No such file or directory".
+    """
+    if not venv_exists():
+        return False
+    try:
+        out = subprocess.run(
+            [
+                venv_python(),
+                "-c",
+                "import os, sysconfig; print(os.path.isfile(os.path.join("
+                "sysconfig.get_paths()['include'], 'Python.h')))",
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            timeout=30,
+            env=clean_env(),
+        )
+        # unknown -> assume fine rather than recreate on a flaky probe
+        return out.returncode != 0 or out.stdout.decode().strip() == "True"
+    except (OSError, subprocess.TimeoutExpired):
+        return True
+
+
 def _uv_python_platform(uv_python):
     """'win-arm64' / 'win-amd64' expected for a pinned uv interpreter request."""
     if not uv_python:
@@ -113,6 +141,21 @@ def ensure_venv(spec, dry_run=False):
             )
             import shutil
 
+            shutil.rmtree(venv_dir(), ignore_errors=True)
+        elif find_uv() and not _venv_has_headers():
+            if dry_run:
+                info(
+                    "[dry-run] venv Python has no C headers — would recreate "
+                    "it on a uv-managed Python."
+                )
+                return venv_python()
+            warn(
+                "Existing venv is built on a Python without C headers "
+                "(Python.h), which triton needs — recreating it on a "
+                "uv-managed Python (all packages will be reinstalled)."
+            )
+            import shutil
+
             shutil.rmtree(venv_dir(), ignore_errors=True)
         else:
             return venv_python()
@@ -134,7 +177,17 @@ def ensure_venv(spec, dry_run=False):
     if uv:
         info("Creating venv with uv (python %s) at %s" % (python_request, target))
         run(
-            [uv, "venv", target, "--python", python_request, "--seed"],
+            # --managed-python: never build on a distro Python, which may
+            # lack the C headers triton needs (see _venv_has_headers)
+            [
+                uv,
+                "venv",
+                target,
+                "--python",
+                python_request,
+                "--managed-python",
+                "--seed",
+            ],
             env=clean_env(),
         )
     else:
@@ -382,7 +435,6 @@ def requirements_hash(spec):
         for f in os.listdir(REPO_ROOT)
         if f.startswith("requirements") and f.endswith(".txt")
     ]
-    req_files.append(os.path.join(REPO_ROOT, "dgx_requirements.txt"))
     base = file_hash(req_files)
     import hashlib
 

manager/nodejs.py

@@ -26,7 +26,7 @@ from .util import (
 )
 
 NODE_DIR = os.path.join(REPO_ROOT, ".node")
-# Node 24 is the current LTS line and matches the dgx_instructions.md guidance.
+# Node 24 is the current LTS line.
 NODE_VERSION = "24.11.1"
 MIN_NODE_MAJOR = 20
 

manager/spec.py

@@ -306,16 +306,6 @@ def _cuda_spec(detection):
     arch = detection["arch"]
     flavor, notes = _cuda_flavor(detection)
     python_version = "3.12"
-    requirements = (
-        "dgx_requirements.txt" if detection.get("is_dgx") else "requirements.txt"
-    )
-    if detection.get("is_dgx"):
-        # the old "Python 3.11 on DGX OS" constraint was for conda/system
-        # installs; uv provisions 3.12 and all aarch64 cp312 wheels exist now
-        notes = notes + [
-            "DGX OS / Grace detected: using %s wheels and dgx_requirements.txt."
-            % flavor
-        ]
 
     extras = [TORCHCODEC]
     optional = [FLA]
@@ -340,7 +330,6 @@ def _cuda_spec(detection):
         TORCH,
         torch_index=PYTORCH_INDEX + flavor,
         python_version=python_version,
-        requirements_file=requirements,
         extra_packages=extras,
         optional_packages=optional,
         find_links=find_links,

requirements_base.txt

@@ -40,6 +40,6 @@ pytorch-wavelets==1.3.0
 matplotlib==3.10.1
 setuptools>=77.0.3
 av==16.0.1
-torchcodec==0.9.1
+torchcodec==0.15.0
 librosa==0.11.0
 mutagen==1.47.0