Build miles as a pure wheel that imports FSDP and torchtitan without Megatron

.github/actions/resolve-dependency-refs/action.yml

@@ -0,0 +1,63 @@
+name: Resolve dependency refs
+description: Pick the Megatron-LM and SGLang refs a CI run tests Miles against.
+inputs:
+  megatron_ref:
+    description: 'Explicit Megatron-LM ref (#N, branch or commit); wins over everything else.'
+    default: ''
+  sglang_ref:
+    description: 'Explicit SGLang ref (#N, branch or commit); wins over everything else.'
+    default: ''
+  pr_body:
+    description: 'PR description with optional ci-megatron-pr / ci-sglang-pr directives.'
+    default: ''
+outputs:
+  ci_megatron_pr:
+    description: 'Megatron-LM ref to fetch; #N becomes refs/pull/N/head.'
+    value: ${{ steps.resolve.outputs.ci_megatron_pr }}
+  ci_sglang_pr:
+    description: 'SGLang ref to fetch; #N becomes refs/pull/N/head.'
+    value: ${{ steps.resolve.outputs.ci_sglang_pr }}
+runs:
+  using: composite
+  steps:
+    - name: Resolve dependency refs
+      id: resolve
+      shell: bash
+      env:
+        PR_BODY: ${{ inputs.pr_body }}
+        INPUT_MEGATRON_PR: ${{ inputs.megatron_ref }}
+        INPUT_SGLANG_PR: ${{ inputs.sglang_ref }}
+      run: |
+        MEGATRON_PR="${INPUT_MEGATRON_PR}"
+        SGLANG_PR="${INPUT_SGLANG_PR}"
+        if [ -n "$PR_BODY" ]; then
+          PR_MEGATRON_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-megatron-pr:\s+\K\S+' || true)
+          PR_SGLANG_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-sglang-pr:\s+\K\S+' || true)
+          [ -z "$MEGATRON_PR" ] && [ -n "$PR_MEGATRON_PR" ] && MEGATRON_PR="$PR_MEGATRON_PR"
+          [ -z "$SGLANG_PR" ] && [ -n "$PR_SGLANG_PR" ] && SGLANG_PR="$PR_SGLANG_PR"
+        fi
+        # Release branches carry release-lock.json (written at branch cut) pinning
+        # the dependency SHAs the release was frozen with. Explicit overrides above
+        # still win; the moving branch heads are only the last resort.
+        if [ -f release-lock.json ]; then
+          LOCK_SGLANG=$(python3 -c "import json; print(json.load(open('release-lock.json'))['sglang_commit'])")
+          LOCK_MEGATRON=$(python3 -c "import json; print(json.load(open('release-lock.json'))['megatron_commit'])")
+          [ -z "$SGLANG_PR" ] && SGLANG_PR="$LOCK_SGLANG"
+          [ -z "$MEGATRON_PR" ] && MEGATRON_PR="$LOCK_MEGATRON"
+          echo "release-lock.json present: sglang=$LOCK_SGLANG megatron=$LOCK_MEGATRON"
+        fi
+        [ -z "$MEGATRON_PR" ] && MEGATRON_PR="miles-main"
+        [ -z "$SGLANG_PR" ] && SGLANG_PR="sglang-miles"
+        resolve_fetch_ref() {
+          local ref="$1"
+          if [[ "$ref" =~ ^#([0-9]+)$ ]]; then
+            echo "refs/pull/${BASH_REMATCH[1]}/head"
+          else
+            echo "$ref"
+          fi
+        }
+        MEGATRON_FETCH=$(resolve_fetch_ref "$MEGATRON_PR")
+        SGLANG_FETCH=$(resolve_fetch_ref "$SGLANG_PR")
+        echo "ci_megatron_pr=$MEGATRON_FETCH" >> $GITHUB_OUTPUT
+        echo "ci_sglang_pr=$SGLANG_FETCH" >> $GITHUB_OUTPUT
+        echo "Resolved: megatron=$MEGATRON_PR -> fetch=$MEGATRON_FETCH, sglang=$SGLANG_PR -> fetch=$SGLANG_FETCH"

.github/workflows/_run-ci.yml

@@ -126,6 +126,14 @@ jobs:
         with:
           ref: ${{ inputs.plan_already_resolved && inputs.ref || needs.plan.outputs.ref }}
 
+      # local actions load from the checked-out ref, which may predate them (/rerun-test, release CI)
+      - name: Checkout CI actions
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.sha }}
+          path: .ci-actions
+          sparse-checkout: .github/actions
+
       - name: Configure source paths
         shell: bash
         run: echo "PYTHONPATH=$GITHUB_WORKSPACE:$SGLANG_SOURCE_ROOT:$MEGATRON_SOURCE_ROOT" >> "$GITHUB_ENV"
@@ -192,45 +200,11 @@ jobs:
       - name: Resolve dependency refs
         if: ${{ !inputs.skip_dependency_install }}
         id: resolve-refs
-        shell: bash
-        env:
-          PR_BODY: ${{ github.event.pull_request.body || '' }}
-          INPUT_MEGATRON_PR: ${{ github.event.inputs.ci_megatron_pr || '' }}
-          INPUT_SGLANG_PR: ${{ github.event.inputs.ci_sglang_pr || '' }}
-        run: |
-          MEGATRON_PR="${INPUT_MEGATRON_PR}"
-          SGLANG_PR="${INPUT_SGLANG_PR}"
-          if [ -n "$PR_BODY" ]; then
-            PR_MEGATRON_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-megatron-pr:\s+\K\S+' || true)
-            PR_SGLANG_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-sglang-pr:\s+\K\S+' || true)
-            [ -z "$MEGATRON_PR" ] && [ -n "$PR_MEGATRON_PR" ] && MEGATRON_PR="$PR_MEGATRON_PR"
-            [ -z "$SGLANG_PR" ] && [ -n "$PR_SGLANG_PR" ] && SGLANG_PR="$PR_SGLANG_PR"
-          fi
-          # Release branches carry release-lock.json (written at branch cut) pinning
-          # the dependency SHAs the release was frozen with. Explicit overrides above
-          # still win; the moving branch heads are only the last resort.
-          if [ -f release-lock.json ]; then
-            LOCK_SGLANG=$(python3 -c "import json; print(json.load(open('release-lock.json'))['sglang_commit'])")
-            LOCK_MEGATRON=$(python3 -c "import json; print(json.load(open('release-lock.json'))['megatron_commit'])")
-            [ -z "$SGLANG_PR" ] && SGLANG_PR="$LOCK_SGLANG"
-            [ -z "$MEGATRON_PR" ] && MEGATRON_PR="$LOCK_MEGATRON"
-            echo "release-lock.json present: sglang=$LOCK_SGLANG megatron=$LOCK_MEGATRON"
-          fi
-          [ -z "$MEGATRON_PR" ] && MEGATRON_PR="miles-main"
-          [ -z "$SGLANG_PR" ] && SGLANG_PR="sglang-miles"
-          resolve_fetch_ref() {
-            local ref="$1"
-            if [[ "$ref" =~ ^#([0-9]+)$ ]]; then
-              echo "refs/pull/${BASH_REMATCH[1]}/head"
-            else
-              echo "$ref"
-            fi
-          }
-          MEGATRON_FETCH=$(resolve_fetch_ref "$MEGATRON_PR")
-          SGLANG_FETCH=$(resolve_fetch_ref "$SGLANG_PR")
-          echo "ci_megatron_pr=$MEGATRON_FETCH" >> $GITHUB_OUTPUT
-          echo "ci_sglang_pr=$SGLANG_FETCH" >> $GITHUB_OUTPUT
-          echo "Resolved: megatron=$MEGATRON_PR -> fetch=$MEGATRON_FETCH, sglang=$SGLANG_PR -> fetch=$SGLANG_FETCH"
+        uses: ./.ci-actions/.github/actions/resolve-dependency-refs
+        with:
+          megatron_ref: ${{ github.event.inputs.ci_megatron_pr || '' }}
+          sglang_ref: ${{ github.event.inputs.ci_sglang_pr || '' }}
+          pr_body: ${{ github.event.pull_request.body || '' }}
 
       - name: Reconcile Miles dependencies
         shell: bash
@@ -250,7 +224,7 @@ jobs:
 
       - name: Install Megatron Bridge override
         if: ${{ !inputs.skip_dependency_install && (github.event.inputs.ci_megatron_bridge_pr != '' || contains(github.event.pull_request.body, 'ci-megatron-bridge-pr:')) }}
-        uses: ./.github/actions/install-megatron-bridge
+        uses: ./.ci-actions/.github/actions/install-megatron-bridge
         with:
           ref: ${{ github.event.inputs.ci_megatron_bridge_pr || '' }}
           pr_body: ${{ github.event.pull_request.body || '' }}

.github/workflows/_run-cpu-ci.yml

@@ -41,6 +41,14 @@ jobs:
         with:
           ref: ${{ inputs.ref }}
 
+      # local actions load from the checked-out ref, which may predate them (/rerun-test, release CI)
+      - name: Checkout CI actions
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.sha }}
+          path: .ci-actions
+          sparse-checkout: .github/actions
+
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
@@ -52,45 +60,11 @@ jobs:
 
       - name: Resolve dependency refs
         id: resolve-refs
-        shell: bash
-        env:
-          PR_BODY: ${{ github.event.pull_request.body || '' }}
-          INPUT_MEGATRON_PR: ${{ github.event.inputs.ci_megatron_pr || '' }}
-          INPUT_SGLANG_PR: ${{ github.event.inputs.ci_sglang_pr || '' }}
-        run: |
-          MEGATRON_PR="${INPUT_MEGATRON_PR}"
-          SGLANG_PR="${INPUT_SGLANG_PR}"
-          if [ -n "$PR_BODY" ]; then
-            PR_MEGATRON_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-megatron-pr:\s+\K\S+' || true)
-            PR_SGLANG_PR=$(echo "$PR_BODY" | grep -m1 -oP '^ci-sglang-pr:\s+\K\S+' || true)
-            [ -z "$MEGATRON_PR" ] && [ -n "$PR_MEGATRON_PR" ] && MEGATRON_PR="$PR_MEGATRON_PR"
-            [ -z "$SGLANG_PR" ] && [ -n "$PR_SGLANG_PR" ] && SGLANG_PR="$PR_SGLANG_PR"
-          fi
-          # Release branches carry release-lock.json (written at branch cut) pinning
-          # the dependency SHAs the release was frozen with. Explicit overrides above
-          # still win; the moving branch heads are only the last resort.
-          if [ -f release-lock.json ]; then
-            LOCK_SGLANG=$(python3 -c "import json; print(json.load(open('release-lock.json'))['sglang_commit'])")
-            LOCK_MEGATRON=$(python3 -c "import json; print(json.load(open('release-lock.json'))['megatron_commit'])")
-            [ -z "$SGLANG_PR" ] && SGLANG_PR="$LOCK_SGLANG"
-            [ -z "$MEGATRON_PR" ] && MEGATRON_PR="$LOCK_MEGATRON"
-            echo "release-lock.json present: sglang=$LOCK_SGLANG megatron=$LOCK_MEGATRON"
-          fi
-          [ -z "$MEGATRON_PR" ] && MEGATRON_PR="miles-main"
-          [ -z "$SGLANG_PR" ] && SGLANG_PR="sglang-miles"
-          resolve_fetch_ref() {
-            local ref="$1"
-            if [[ "$ref" =~ ^#([0-9]+)$ ]]; then
-              echo "refs/pull/${BASH_REMATCH[1]}/head"
-            else
-              echo "$ref"
-            fi
-          }
-          MEGATRON_FETCH=$(resolve_fetch_ref "$MEGATRON_PR")
-          SGLANG_FETCH=$(resolve_fetch_ref "$SGLANG_PR")
-          echo "ci_megatron_pr=$MEGATRON_FETCH" >> $GITHUB_OUTPUT
-          echo "ci_sglang_pr=$SGLANG_FETCH" >> $GITHUB_OUTPUT
-          echo "Resolved: megatron=$MEGATRON_PR -> fetch=$MEGATRON_FETCH, sglang=$SGLANG_PR -> fetch=$SGLANG_FETCH"
+        uses: ./.ci-actions/.github/actions/resolve-dependency-refs
+        with:
+          megatron_ref: ${{ github.event.inputs.ci_megatron_pr || '' }}
+          sglang_ref: ${{ github.event.inputs.ci_sglang_pr || '' }}
+          pr_body: ${{ github.event.pull_request.body || '' }}
 
       - name: Checkout sglang
         uses: actions/checkout@v4
@@ -119,7 +93,7 @@ jobs:
 
       - name: Install Megatron Bridge override
         if: ${{ github.event.inputs.ci_megatron_bridge_pr != '' || contains(github.event.pull_request.body, 'ci-megatron-bridge-pr:') }}
-        uses: ./.github/actions/install-megatron-bridge
+        uses: ./.ci-actions/.github/actions/install-megatron-bridge
         with:
           ref: ${{ github.event.inputs.ci_megatron_bridge_pr || '' }}
           pr_body: ${{ github.event.pull_request.body || '' }}

.github/workflows/wheel-smoke.yml

@@ -0,0 +1,67 @@
+name: Wheel smoke
+
+# The built wheel, installed without Megatron, must import and parse args for FSDP and torchtitan.
+# SGLang stays a source tree on PYTHONPATH until a hosted sglang-miles wheel exists.
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+
+permissions:
+  contents: read
+
+jobs:
+  wheel-smoke:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+      - name: Checkout miles
+        uses: actions/checkout@v4
+
+      - name: Resolve dependency refs
+        id: resolve-refs
+        uses: ./.github/actions/resolve-dependency-refs
+        with:
+          pr_body: ${{ github.event.pull_request.body || '' }}
+
+      - name: Checkout sglang
+        uses: actions/checkout@v4
+        with:
+          repository: sgl-project/sglang
+          ref: ${{ steps.resolve-refs.outputs.ci_sglang_pr }}
+          path: sglang
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          # Match the Miles Docker runtime.
+          python-version: '3.12'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Build the sdist and the wheel from it
+        shell: bash
+        run: uv build --out-dir "$RUNNER_TEMP/dist" .
+
+      - name: Install the wheel into a clean environment
+        shell: bash
+        run: |
+          # torch follows sglang's pin (torchtitan needs >=2.13), torchvision follows torch
+          python3 -c "import tomllib; deps = tomllib.load(open('sglang/python/pyproject.toml', 'rb'))['project']['dependencies']; print(next(d for d in deps if d.startswith('torch==')))" \
+            > "$RUNNER_TEMP/overrides.txt"
+          echo torchvision >> "$RUNNER_TEMP/overrides.txt"
+          uv venv --python 3.12 "$RUNNER_TEMP/venv"
+          VIRTUAL_ENV="$RUNNER_TEMP/venv" uv pip install \
+            "$(ls "$RUNNER_TEMP"/dist/miles-*-py3-none-any.whl)[torchtitan]" \
+            -r tests/ci/requirements-ci-cpu.txt \
+            --override "$RUNNER_TEMP/overrides.txt"
+
+      - name: Run the smoke check from outside the checkout
+        shell: bash
+        working-directory: ${{ runner.temp }}
+        env:
+          PYTHONPATH: ${{ github.workspace }}/sglang/python
+        run: '"$RUNNER_TEMP/venv/bin/python" "$GITHUB_WORKSPACE/tests/ci/wheel_smoke.py" "$GITHUB_WORKSPACE"'

MANIFEST.in

@@ -0,0 +1 @@
+include requirements.txt

setup.py

@@ -1,54 +1,32 @@
-import sys
-import platform
-
 from setuptools import find_packages, setup
-from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
-
-
-def _get_platform_tag():
-    if platform.system() != "Linux":
-        return platform.system().lower()
-
-    machine = platform.machine().lower()
-    if machine in ("x86_64", "amd64"):
-        return "manylinux1_x86_64"
-    if machine in ("aarch64", "arm64"):
-        return "manylinux2014_aarch64"
-    return f"linux_{machine}"
 
 
 def _fetch_requirements(path):
     with open(path) as fd:
         return [r.strip() for r in fd.readlines() if r.strip() and not r.startswith("#")]
 
 
-# Custom wheel class to modify the wheel name
-class bdist_wheel(_bdist_wheel):
-    def finalize_options(self):
-        _bdist_wheel.finalize_options(self)
-        self.root_is_pure = False
-
-    def get_tag(self):
-        python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
-        abi_tag = f"{python_version}"
-        platform_tag = _get_platform_tag()
-
-        return python_version, abi_tag, platform_tag
-
-
 # Setup configuration
 setup(
     author="miles Team",
     name="miles",
     version="0.1.3",
     packages=find_packages(include=["miles*", "miles_plugins*"]),
     include_package_data=True,
-    package_data={"miles.dashboard": ["static/*"]},
+    package_data={
+        "miles.dashboard": ["static/*"],
+        "miles.utils.chat_template_utils.templates": ["*.jinja"],
+    },
     install_requires=_fetch_requirements("requirements.txt"),
     extras_require={
         "fsdp": [
             "torch>=2.0",
         ],
+        # torchtitan 0.3.0 needs torch>=2.13 but does not declare torch
+        "torchtitan": [
+            "torch>=2.13",
+            "torchtitan==0.3.0",
+        ],
         "mlflow": [
             "mlflow>=2.0",
         ],
@@ -70,14 +48,13 @@ def get_tag(self):
             "modal>=1.5.5",
         ],
     },
-    python_requires=">=3.10",
+    # sglang needs 3.11; nvidia-resiliency-ext 0.6 has no 3.13 wheel
+    python_requires=">=3.11",
     classifiers=[
-        "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3.12",
         "Environment :: GPU :: NVIDIA CUDA",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
         "Topic :: System :: Distributed Computing",
     ],
-    cmdclass={"bdist_wheel": bdist_wheel},
 )