electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first

Browse Source
This commit is contained in:
Evan
2026-03-19 11:10:58 +00:00
parent be731d3a85
commit a3ce437fd4
13 changed files with 517 additions and 3318 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flake.nix
+43 -109
View File
@@ -130,122 +130,56 @@
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
torch-cuda = pkgsCuda.python313Packages.torch;
vllm-cuda = pkgsCuda.python313Packages.vllm;
# Smoke test script for verifying vLLM + CUDA GPU setup
vllm-check = pkgs.writeShellApplication {
name = "vllm-check";
runtimeInputs = [
(pkgsCuda.python313.withPackages (ps: [ ps.torch ps.vllm ]))
];
# On non-NixOS hosts, NVIDIA driver libraries live in /usr/lib and must be
# LD_PRELOAD'd individually (adding the whole dir causes SIGILL from conflicts).
# These are: CUDA driver, NVML, and the PTX JIT compiler (for flash attention).
# libnvJitLink comes from the nix CUDA toolkit via LD_LIBRARY_PATH.
text = ''
for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
if [ -e "$dir/libcuda.so.1" ]; then
NVIDIA_LIBS="$dir/libcuda.so.1"
for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
[ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
done
export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
break
fi
done
export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
exec python ${inputs.self + /tests/test_vllm_smoke.py}
'';
};
# exo with CUDA torch + vLLM — wraps the uv2nix-built package with host driver libs
exo-cuda = pkgs.writeShellApplication {
name = "exo-cuda";
runtimeInputs = [ self'.packages.exo-cuda-unwrapped ];
text = ''
for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
if [ -e "$dir/libcuda.so.1" ]; then
NVIDIA_LIBS="$dir/libcuda.so.1"
for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
[ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
done
export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
break
fi
done
export LD_LIBRARY_PATH="${pkgsCuda.stdenv.cc.cc.lib}/lib:${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
exec exo-cuda "$@"
'';
};
};
);
# CUDA development shell with torch + vLLM (aarch64-linux only)
devShells = lib.optionalAttrs (pkgsCuda != null)
devShells =
{
cuda = pkgs.mkShell {
packages = [
(pkgsCuda.python313.withPackages (ps: [
ps.torch
ps.vllm
]))
pkgs.uv
pkgs.just
];
default = with pkgs; pkgs.mkShell {
inputsFrom = [ self'.checks.cargo-build ];
packages =
[
# FORMATTING
config.treefmt.build.wrapper
# PYTHON
python313
uv
ruff
basedpyright
# RUST
config.rust.toolchain
maturin
# NIX
nixpkgs-fmt
# SVELTE
nodejs
# MISC
just
jq
]
++ lib.optionals stdenv.isLinux [
unixtools.ifconfig
]
++ lib.optionals stdenv.isDarwin [
macmon
];
OPENSSL_NO_VENDOR = "1";
shellHook = ''
echo "CUDA dev shell with torch + vLLM"
python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')" 2>/dev/null || true
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
${lib.optionalString stdenv.isLinux ''
export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
''}
'';
};
} // {
default = with pkgs; pkgs.mkShell {
inputsFrom = [ self'.checks.cargo-build ];
packages =
[
# FORMATTING
config.treefmt.build.wrapper
# PYTHON
python313
uv
ruff
basedpyright
# RUST
config.rust.toolchain
maturin
# NIX
nixpkgs-fmt
# SVELTE
nodejs
# MISC
just
jq
]
++ lib.optionals stdenv.isLinux [
unixtools.ifconfig
]
++ lib.optionals stdenv.isDarwin [
macmon
];
OPENSSL_NO_VENDOR = "1";
shellHook = ''
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
${lib.optionalString stdenv.isLinux ''
export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
''}
'';
};
};
};
};
}
pyproject.toml
+5 -29
View File
@@ -46,12 +46,8 @@ dev = [
]
[project.optional-dependencies]
cuda = [
"torch>=2.10.0; sys_platform == 'linux'",
"vllm>=0.13.0; sys_platform == 'linux'",
"mlx-cuda-13==0.30.6; sys_platform == 'linux'",
"fastsafetensors>=0.1.10; sys_platform == 'linux'",
]
mlx = ["mlx_runner"]
cuda = ["vllm_runner"]
###
# workspace configuration
@@ -62,19 +58,14 @@ members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo_pyo3_bindings = { workspace = true }
mlx_runner = { path = "python/mlx_runner" }
vllm_runner = { path = "python/vllm_runner" }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
# Uncomment to use local mlx/mlx-lm development versions:
# mlx = { path = "/Users/Shared/mlx", editable=true }
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
[build-system]
requires = ["uv_build>=0.8.9,<0.9.0"]
build-backend = "uv_build"
@@ -130,22 +121,7 @@ root = "src"
[tool.uv]
required-version = ">=0.8.6"
prerelease = "allow"
environments = [
"sys_platform == 'darwin'",
"sys_platform == 'linux' and platform_machine == 'aarch64'",
]
no-binary-package = ["vllm"]
no-build-isolation-package = ["vllm"]
extra-build-dependencies = { vllm = [
"cmake>=3.26.1",
"ninja",
"packaging>=24.2",
"setuptools>=77.0.3,<81.0.0",
"setuptools-scm>=8.0",
"wheel",
"jinja2",
] }
conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
###
# ruff configuration
python/mlx_runner/README.md
View File
python/mlx_runner/main.py
+6
View File
@@ -0,0 +1,6 @@
def main():
print("Hello from mlx-runner!")
if __name__ == "__main__":
main()
python/mlx_runner/pyproject.toml
+11
View File
@@ -0,0 +1,11 @@
[project]
name = "mlx-runner"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = []
[build-system]
requires = ["uv_build>=0.8.9,<0.9.0"]
build-backend = "uv_build"
python/parts.nix
+39 -49
View File
@@ -3,7 +3,6 @@
perSystem =
{ config, self', pkgs, lib, system, ... }:
let
pkgsCuda = import ../nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
# Load workspace from uv.lock
workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
workspaceRoot = inputs.self;
@@ -34,6 +33,8 @@
};
};
inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux;
python = pkgs.python313;
# Overlay to provide build systems and custom packages
@@ -65,7 +66,7 @@
final.setuptools
];
});
} // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
} // lib.optionalAttrs isDarwin {
# Use our pure Nix-built MLX with Metal support (macOS only)
mlx = self'.packages.mlx;
};
@@ -77,7 +78,7 @@
ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
in
lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
lib.optionalAttrs isLinux (
(lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
mlx = ignoreMissing prev.mlx;
mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
@@ -100,39 +101,48 @@
}
);
baseOverlays = [
inputs.pyproject-build-systems.overlays.default
overlay
exoOverlay
buildSystemsOverlay
linuxOverlay
];
pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
inherit python;
}).overrideScope (
lib.composeManyExtensions baseOverlays
lib.composeManyExtensions [
inputs.pyproject-build-systems.overlays.default
overlay
exoOverlay
buildSystemsOverlay
linuxOverlay
]
);
# mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
# mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
venvCollisionPaths = lib.optionals isLinux [
"lib/python3.13/site-packages/mlx*"
"lib/python3.13/site-packages/nvidia*"
];
# Exclude bench deps from main env (bench has its own benchVenv)
exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
exoVenv = (pythonSet.mkVirtualEnv "exo-env" {
exo = lib.optionals isDarwin [ "mlx" ];
exo-pyo3-bindings = [ ];
}).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
exoCudaVenv = (pythonSet.mkVirtualEnv "exo-env" {
exo = lib.optionals isLinux [ "cuda" ];
exo-pyo3-bindings = [ ];
}).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
# Virtual environment with dev dependencies for testing
testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
exoDeps // {
testVenv = (pythonSet.mkVirtualEnv "exo-test-env"
{
exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
exo-pyo3-bindings = [ ];
}
)).overrideAttrs {
).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
@@ -173,52 +183,32 @@
makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
'';
vllmEnv = pkgsCuda.python313.withPackages (ps: [ ps.vllm ps.fastsafetensors ]);
vllmSite = pkgs.runCommand "vllm-site-filtered" { } ''
mkdir -p $out
for pkg in ${vllmEnv}/${python.sitePackages}/*; do
name=$(basename "$pkg")
case "$name" in
anyio*|pydantic*) ;;
*) ln -s "$pkg" "$out/$name" ;;
esac
done
'';
exoCudaDeps = exoDeps // {
mlx-cuda-13 = [ ];
};
exoCudaVenv = (pythonSet.mkVirtualEnv "exo-cuda-env" exoCudaDeps).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
exoCudaPackage = pkgs.runCommand "exo-cuda"
exoCudaPackage = pkgs.runCommand "exo"
{
nativeBuildInputs = [ pkgs.makeWrapper ];
}
''
mkdir -p $out/bin
makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo-cuda \
# Create wrapper script
makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
--prefix PYTHONPATH : "${vllmSite}"
${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
'';
in
{
# Python package only available on macOS (requires MLX/Metal)
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
packages = (lib.optionalAttrs isDarwin
{
exo = exoPackage;
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
exo-test-env = testVenv;
} // lib.optionalAttrs (pkgsCuda != null) {
exo-cuda-unwrapped = exoCudaPackage;
} // {
}) // {
exo = exoPackage;
exo-cuda = exoCudaPackage;
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
python/vllm_runner/.python-version
+1
View File
@@ -0,0 +1 @@
3.13
python/vllm_runner/README.md
View File
python/vllm_runner/main.py
+6
View File
@@ -0,0 +1,6 @@
def main():
print("Hello from vllm-runner!")
if __name__ == "__main__":
main()
python/vllm_runner/pyproject.toml
+25
View File
@@ -0,0 +1,25 @@
[project]
name = "vllm-runner"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"torch>=2.10.0; sys_platform == 'linux'",
"vllm>=0.13.0; sys_platform == 'linux'",
"mlx-cuda-13==0.30.6; sys_platform == 'linux'",
"fastsafetensors>=0.1.10; sys_platform == 'linux'",
]
[tool.uv.sources]
vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
[build-system]
requires = ["uv_build>=0.8.9,<0.9.0"]
build-backend = "uv_build"
src/exo/worker/runner/bootstrap.py
+2 -2
View File
@@ -97,8 +97,8 @@ def entrypoint(
else:
from exo.worker.runner.llm_inference.runner import MlxBuilder, Runner
builder = MlxBuilder(
model_id=bound_instance.bound_shard.model_card.model_id,
builder = MlxBuilder.create(
bound_instance,
event_sender=event_sender,
cancel_receiver=cancel_receiver,
)
src/exo/worker/runner/llm_inference/runner.py
+57 -16
View File
@@ -6,11 +6,9 @@ from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Self
import mlx.core as mx
from anyio import WouldBlock
from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.models.model_cards import ModelTask
from exo.shared.types.chunks import (
@@ -83,12 +81,20 @@ class ExitCode(str, Enum):
class Builder(ABC):
@abstractmethod
def connect(self, bound_instance: BoundInstance) -> None: ...
@classmethod
def create(
cls,
bound_instance: BoundInstance,
event_sender: MpSender[Event],
cancel_receiver: MpReceiver[TaskId],
) -> Self: ...
@abstractmethod
def connect(self) -> None: ...
@abstractmethod
def load(
self,
bound_instance: BoundInstance,
on_timeout: Callable[[], None],
on_layer_loaded: Callable[[int, int], None],
) -> None: ...
@@ -184,7 +190,7 @@ class Runner:
self.update_status(RunnerConnecting())
self.acknowledge_task(task)
self.generator.connect(self.bound_instance)
self.generator.connect()
self.send_task_status(task.task_id, TaskStatus.Complete)
self.update_status(RunnerConnected())
@@ -233,7 +239,6 @@ class Runner:
), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
self.generator.load(
self.bound_instance,
on_timeout=on_model_load_timeout,
on_layer_loaded=on_layer_loaded,
)
@@ -402,31 +407,47 @@ class Runner:
@dataclass
class MlxBuilder(Builder):
import mlx.core as mx
from mlx_lm.tokenizer_utils import TokenizerWrapper
model_id: ModelId
bound_instance: BoundInstance
event_sender: MpSender[Event]
cancel_receiver: MpReceiver[TaskId]
inference_model: Model | None = None
tokenizer: TokenizerWrapper | None = None
group: mx.distributed.Group | None = None
def connect(self, bound_instance: BoundInstance) -> None:
self.group = initialize_mlx(bound_instance)
@classmethod
def create(
cls,
bound_instance: BoundInstance,
event_sender: MpSender[Event],
cancel_receiver: MpReceiver[TaskId],
) -> Self:
return cls(
bound_instance.instance.shard_assignments.model_id,
bound_instance,
event_sender,
cancel_receiver,
)
def connect(self) -> None:
self.group = initialize_mlx(self.bound_instance)
def load(
self,
bound_instance: BoundInstance,
on_timeout: Callable[[], None],
on_layer_loaded: Callable[[int, int], None],
) -> None:
self.inference_model, self.tokenizer = load_mlx_items(
bound_instance,
self.bound_instance,
self.group,
on_timeout=on_timeout,
on_layer_loaded=on_layer_loaded,
)
def build(self) -> InferenceGenerator:
assert self.model_id
assert self.inference_model
assert self.tokenizer
@@ -514,16 +535,34 @@ class VllmBuilder(Builder):
trust_remote_code: bool
cancel_receiver: MpReceiver[TaskId]
event_sender: MpSender[Event]
group: mx.distributed.Group | None = None
bound_instance: BoundInstance
def connect(self, bound_instance: BoundInstance) -> None:
@classmethod
def create(
cls,
bound_instance: BoundInstance,
event_sender: MpSender[Event],
cancel_receiver: MpReceiver[TaskId],
) -> Self:
from exo.shared.constants import EXO_MODELS_DIR
mid = bound_instance.instance.shard_assignments.model_id
return cls(
mid,
str(EXO_MODELS_DIR / mid.normalize()),
bound_instance.bound_shard.model_card.trust_remote_code,
cancel_receiver,
event_sender,
bound_instance,
)
def connect(self) -> None:
raise NotImplementedError(
"Multiple node VLLM instances are not supported at the moment!"
)
def load(
self,
bound_instance: BoundInstance,
on_timeout: Callable[[], None],
on_layer_loaded: Callable[[int, int], None],
) -> None:
@@ -533,7 +572,7 @@ class VllmBuilder(Builder):
model_path=self.model_path,
model_id=self.model_id,
trust_remote_code=self.trust_remote_code,
n_layers=bound_instance.bound_shard.model_card.n_layers,
n_layers=self.bound_instance.bound_shard.model_card.n_layers,
on_layer_loaded=on_layer_loaded,
)
@@ -545,6 +584,8 @@ class VllmBuilder(Builder):
model_id=self.model_id,
prefix_cache=self._prefix_cache,
)
from mlx_lm.tokenizer_utils import TokenizerWrapper
tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
uv.lock
Generated
+322 -3113
View File
File diff suppressed because it is too large Load Diff