diff --git a/.github/workflows/build-app.yml b/.github/workflows/build-app.yml index 44e9c77e..542dc800 100644 --- a/.github/workflows/build-app.yml +++ b/.github/workflows/build-app.yml @@ -113,11 +113,22 @@ jobs: uv python install uv sync --locked + - name: Install Nix + uses: cachix/install-nix-action@v31 + with: + nix_path: nixpkgs=channel:nixos-unstable + + - name: Configure Cachix + uses: cachix/cachix-action@v14 + with: + name: exo + authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" + - name: Build dashboard run: | - cd dashboard - npm ci - npm run build + DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths) + mkdir -p dashboard/build + cp -r "$DASHBOARD_OUT"/* dashboard/build/ - name: Install Sparkle CLI run: | diff --git a/dashboard/dashboard.nix b/dashboard/dashboard.nix new file mode 100644 index 00000000..6ecd1490 --- /dev/null +++ b/dashboard/dashboard.nix @@ -0,0 +1,60 @@ +{ lib +, config +, dream2nix +, ... +}: +let + # Read and parse the lock file + rawLockFile = builtins.fromJSON (builtins.readFile "${config.deps.dashboardSrc}/package-lock.json"); + + # For packages with bundleDependencies, filter out deps that are bundled + # (bundled deps are inside the tarball, not separate lockfile entries) + fixedPackages = lib.mapAttrs + (path: entry: + if entry ? bundleDependencies && entry.bundleDependencies != [ ] + then entry // { + dependencies = lib.filterAttrs + (name: _: !(lib.elem name entry.bundleDependencies)) + (entry.dependencies or { }); + } + else entry + ) + (rawLockFile.packages or { }); + + fixedLockFile = rawLockFile // { packages = fixedPackages; }; +in +{ + imports = [ + dream2nix.modules.dream2nix.nodejs-package-lock-v3 + dream2nix.modules.dream2nix.nodejs-granular-v3 + ]; + + name = "exo-dashboard"; + version = "1.0.0"; + + mkDerivation = { + src = config.deps.dashboardSrc; + + buildPhase = '' + runHook preBuild + npm run build + runHook postBuild + ''; + + installPhase = '' + runHook preInstall + cp -r build $out/build + runHook postInstall + ''; + }; + + deps = { nixpkgs, ... }: { + inherit (nixpkgs) stdenv; + dashboardSrc = null; # Injected by parts.nix + }; + + nodejs-package-lock-v3 = { + # Don't use packageLockFile - provide the fixed lock content directly + packageLock = fixedLockFile; + }; +} diff --git a/dashboard/parts.nix b/dashboard/parts.nix new file mode 100644 index 00000000..80fbb156 --- /dev/null +++ b/dashboard/parts.nix @@ -0,0 +1,44 @@ +{ inputs, ... }: +{ + perSystem = + { pkgs, lib, ... }: + let + # Filter source to only include dashboard directory + src = lib.cleanSourceWith { + src = inputs.self; + filter = + path: type: + let + baseName = builtins.baseNameOf path; + inDashboardDir = + (lib.hasInfix "/dashboard/" path) + || (lib.hasSuffix "/dashboard" (builtins.dirOf path)) + || (baseName == "dashboard" && type == "directory"); + in + inDashboardDir; + }; + + # Build the dashboard with dream2nix (includes node_modules in output) + dashboardFull = inputs.dream2nix.lib.evalModules { + packageSets.nixpkgs = pkgs; + modules = [ + ./dashboard.nix + { + paths.projectRoot = inputs.self; + paths.projectRootFile = "flake.nix"; + paths.package = inputs.self + "/dashboard"; + } + # Inject the filtered source + { + deps.dashboardSrc = lib.mkForce "${src}/dashboard"; + } + ]; + }; + in + { + # Extract just the static site from the full build + packages.dashboard = pkgs.runCommand "exo-dashboard" { } '' + cp -r ${dashboardFull}/build $out + ''; + }; +} diff --git a/flake.lock b/flake.lock index d45f6f50..3c7b28e1 100644 --- a/flake.lock +++ b/flake.lock @@ -1,5 +1,42 @@ { "nodes": { + "crane": { + "locked": { + "lastModified": 1767744144, + "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=", + "owner": "ipetkov", + "repo": "crane", + "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e", + "type": "github" + }, + "original": { + "owner": "ipetkov", + "repo": "crane", + "type": "github" + } + }, + "dream2nix": { + "inputs": { + "nixpkgs": [ + "nixpkgs" + ], + "purescript-overlay": "purescript-overlay", + "pyproject-nix": "pyproject-nix" + }, + "locked": { + "lastModified": 1765953015, + "narHash": "sha256-5FBZbbWR1Csp3Y2icfRkxMJw/a/5FGg8hCXej2//bbI=", + "owner": "nix-community", + "repo": "dream2nix", + "rev": "69eb01fa0995e1e90add49d8ca5bcba213b0416f", + "type": "github" + }, + "original": { + "owner": "nix-community", + "repo": "dream2nix", + "type": "github" + } + }, "fenix": { "inputs": { "nixpkgs": [ @@ -21,6 +58,22 @@ "type": "github" } }, + "flake-compat": { + "flake": false, + "locked": { + "lastModified": 1696426674, + "narHash": "sha256-kvjfFW7WAETZlt09AgDn1MrtKzP7t90Vf7vypd3OL1U=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "0f9255e01c2351cc7d116c072cb317785dd33b33", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, "flake-parts": { "inputs": { "nixpkgs-lib": [ @@ -73,8 +126,54 @@ "type": "github" } }, + "purescript-overlay": { + "inputs": { + "flake-compat": "flake-compat", + "nixpkgs": [ + "dream2nix", + "nixpkgs" + ], + "slimlock": "slimlock" + }, + "locked": { + "lastModified": 1728546539, + "narHash": "sha256-Sws7w0tlnjD+Bjck1nv29NjC5DbL6nH5auL9Ex9Iz2A=", + "owner": "thomashoneyman", + "repo": "purescript-overlay", + "rev": "4ad4c15d07bd899d7346b331f377606631eb0ee4", + "type": "github" + }, + "original": { + "owner": "thomashoneyman", + "repo": "purescript-overlay", + "type": "github" + } + }, + "pyproject-nix": { + "inputs": { + "nixpkgs": [ + "dream2nix", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1763017646, + "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=", + "owner": "pyproject-nix", + "repo": "pyproject.nix", + "rev": "47bd6f296502842643078d66128f7b5e5370790c", + "type": "github" + }, + "original": { + "owner": "pyproject-nix", + "repo": "pyproject.nix", + "type": "github" + } + }, "root": { "inputs": { + "crane": "crane", + "dream2nix": "dream2nix", "fenix": "fenix", "flake-parts": "flake-parts", "nixpkgs": "nixpkgs", @@ -99,6 +198,28 @@ "type": "github" } }, + "slimlock": { + "inputs": { + "nixpkgs": [ + "dream2nix", + "purescript-overlay", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1688756706, + "narHash": "sha256-xzkkMv3neJJJ89zo3o2ojp7nFeaZc2G0fYwNXNJRFlo=", + "owner": "thomashoneyman", + "repo": "slimlock", + "rev": "cf72723f59e2340d24881fd7bf61cb113b4c407c", + "type": "github" + }, + "original": { + "owner": "thomashoneyman", + "repo": "slimlock", + "type": "github" + } + }, "treefmt-nix": { "inputs": { "nixpkgs": [ diff --git a/flake.nix b/flake.nix index 5c7a56cf..1c049537 100644 --- a/flake.nix +++ b/flake.nix @@ -9,6 +9,8 @@ inputs.nixpkgs-lib.follows = "nixpkgs"; }; + crane.url = "github:ipetkov/crane"; + fenix = { url = "github:nix-community/fenix"; inputs.nixpkgs.follows = "nixpkgs"; @@ -19,6 +21,11 @@ inputs.nixpkgs.follows = "nixpkgs"; }; + dream2nix = { + url = "github:nix-community/dream2nix"; + inputs.nixpkgs.follows = "nixpkgs"; + }; + # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) nixpkgs-swift.url = "github:NixOS/nixpkgs/08dacfca559e1d7da38f3cf05f1f45ee9bfd213c"; }; @@ -39,10 +46,12 @@ imports = [ inputs.treefmt-nix.flakeModule + ./dashboard/parts.nix + ./rust/parts.nix ]; perSystem = - { config, inputs', pkgs, lib, system, ... }: + { config, self', inputs', pkgs, lib, system, ... }: let fenixToolchain = inputs'.fenix.packages.complete; # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) @@ -59,7 +68,7 @@ }; rustfmt = { enable = true; - package = fenixToolchain.rustfmt; + package = config.rust.toolchain; }; prettier = { enable = true; @@ -79,6 +88,8 @@ ''; devShells.default = with pkgs; pkgs.mkShell { + inputsFrom = [ self'.checks.cargo-build ]; + packages = [ # FORMATTING @@ -91,14 +102,8 @@ basedpyright # RUST - (fenixToolchain.withComponents [ - "cargo" - "rustc" - "clippy" - "rustfmt" - "rust-src" - ]) - rustup # Just here to make RustRover happy + config.rust.toolchain + maturin # NIX nixpkgs-fmt @@ -110,30 +115,20 @@ just jq ] - ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [ - # IFCONFIG + ++ lib.optionals stdenv.isLinux [ unixtools.ifconfig - - # Build dependencies for Linux - pkg-config - openssl - ]) - ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [ - # MACMON + ] + ++ lib.optionals stdenv.isDarwin [ macmon - ]); + ]; + + OPENSSL_NO_VENDOR = "1"; shellHook = '' - # PYTHON - export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib" - ${lib.optionalString pkgs.stdenv.isLinux '' - # Build environment for Linux - export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH" - export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH" + export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib" + ${lib.optionalString stdenv.isLinux '' + export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH" ''} - echo - echo "🍎🍎 Run 'just ' to get started" - just --list ''; }; }; diff --git a/rust/parts.nix b/rust/parts.nix new file mode 100644 index 00000000..e00b27c6 --- /dev/null +++ b/rust/parts.nix @@ -0,0 +1,159 @@ +{ inputs, ... }: +{ + perSystem = + { config, self', inputs', pkgs, lib, ... }: + let + # Fenix nightly toolchain with all components + fenixPkgs = inputs'.fenix.packages; + rustToolchain = fenixPkgs.complete.withComponents [ + "cargo" + "rustc" + "clippy" + "rustfmt" + "rust-src" + "rust-analyzer" + ]; + + # Crane with fenix toolchain + craneLib = (inputs.crane.mkLib pkgs).overrideToolchain rustToolchain; + + # Source filtering - only include rust/ directory and root Cargo files + # This ensures changes to Python/docs/etc don't trigger Rust rebuilds + src = lib.cleanSourceWith { + src = inputs.self; + filter = + path: type: + let + baseName = builtins.baseNameOf path; + parentDir = builtins.dirOf path; + inRustDir = + (lib.hasInfix "/rust/" path) + || (lib.hasSuffix "/rust" parentDir) + || (baseName == "rust" && type == "directory"); + isRootCargoFile = + (baseName == "Cargo.toml" || baseName == "Cargo.lock") + && (builtins.dirOf path == toString inputs.self); + in + isRootCargoFile + || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path)); + }; + + # Common arguments for all Rust builds + commonArgs = { + inherit src; + pname = "exo-rust"; + version = "0.0.1"; + strictDeps = true; + + nativeBuildInputs = [ + pkgs.pkg-config + pkgs.python313 # Required for pyo3-build-config + ]; + + buildInputs = [ + pkgs.openssl + pkgs.python313 # Required for pyo3 tests + ]; + + OPENSSL_NO_VENDOR = "1"; + + # Required for pyo3 tests to find libpython + LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ]; + }; + + # Build dependencies once for caching + cargoArtifacts = craneLib.buildDepsOnly ( + commonArgs + // { + cargoExtraArgs = "--workspace"; + } + ); + in + { + # Export toolchain for use in treefmt and devShell + options.rust = { + toolchain = lib.mkOption { + type = lib.types.package; + default = rustToolchain; + description = "The Rust toolchain to use"; + }; + }; + + config = { + packages = { + # The system_custodian binary + system_custodian = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "-p system_custodian"; + + meta = { + description = "System custodian daemon for exo"; + mainProgram = "system_custodian"; + }; + } + ); + + # Python bindings wheel via maturin + exo_pyo3_bindings = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + pname = "exo_pyo3_bindings"; + + nativeBuildInputs = commonArgs.nativeBuildInputs ++ [ + pkgs.maturin + ]; + + buildPhaseCargoCommand = '' + maturin build \ + --release \ + --manylinux off \ + --manifest-path rust/exo_pyo3_bindings/Cargo.toml \ + --features "pyo3/extension-module,pyo3/experimental-async" \ + --interpreter ${pkgs.python313}/bin/python \ + --out dist + ''; + + # Don't use crane's default install behavior + doNotPostBuildInstallCargoBinaries = true; + + installPhaseCommand = '' + mkdir -p $out + cp dist/*.whl $out/ + ''; + } + ); + }; + + checks = { + # Full workspace build (all crates) + cargo-build = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + # Run tests with nextest + cargo-nextest = craneLib.cargoNextest ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + + # Build documentation + cargo-doc = craneLib.cargoDoc ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + }; + }; + }; +} diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index fd4f9003..d46a8115 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -82,6 +82,7 @@ MODEL_CARDS: dict[str, ModelCard] = { # storage_size=Memory.from_kb(754706307), # n_layers=61, # hidden_size=7168, + # supports_tensor=True, # ), # ), # "deepseek-v3.2-4bit": ModelCard( @@ -96,6 +97,7 @@ MODEL_CARDS: dict[str, ModelCard] = { # storage_size=Memory.from_kb(754706307 // 2), # TODO !!!!! # n_layers=61, # hidden_size=7168, + # supports_tensor=True, # ), # ), # deepseek r1 @@ -554,6 +556,81 @@ MODEL_CARDS: dict[str, ModelCard] = { supports_tensor=True, ), ), + "glm-4.7-4bit": ModelCard( + short_id="glm-4.7-4bit", + model_id=ModelId("mlx-community/GLM-4.7-4bit"), + name="GLM 4.7 4bit", + description="GLM 4.7 4bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-4bit"), + pretty_name="GLM 4.7 4bit", + storage_size=Memory.from_bytes(198556925568), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "glm-4.7-6bit": ModelCard( + short_id="glm-4.7-6bit", + model_id=ModelId("mlx-community/GLM-4.7-6bit"), + name="GLM 4.7 6bit", + description="GLM 4.7 6bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-6bit"), + pretty_name="GLM 4.7 6bit", + storage_size=Memory.from_bytes(286737579648), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "glm-4.7-8bit-gs32": ModelCard( + short_id="glm-4.7-8bit-gs32", + model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"), + name="GLM 4.7 8bit (gs32)", + description="GLM 4.7 8bit (gs32)", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"), + pretty_name="GLM 4.7 8bit (gs32)", + storage_size=Memory.from_bytes(396963397248), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "minimax-m2.1-8bit": ModelCard( + short_id="minimax-m2.1-8bit", + model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"), + name="MiniMax M2.1 8bit", + description="MiniMax M2.1 8bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"), + pretty_name="MiniMax M2.1 8bit", + storage_size=Memory.from_bytes(242986745856), + n_layers=61, + hidden_size=3072, + supports_tensor=True, + ), + ), + "minimax-m2.1-3bit": ModelCard( + short_id="minimax-m2.1-3bit", + model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"), + name="MiniMax M2.1 3bit", + description="MiniMax M2.1 3bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"), + pretty_name="MiniMax M2.1 3bit", + storage_size=Memory.from_bytes(100086644736), + n_layers=61, + hidden_size=3072, + supports_tensor=True, + ), + ), # "devstral-2-123b-instruct-2512-8bit": ModelCard( # short_id="devstral-2-123b-instruct-2512-8bit", # model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"), diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 55a549fd..46d812e3 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -19,6 +19,7 @@ from mlx_lm.models.glm4_moe import MoE from mlx_lm.models.gpt_oss import GptOssMoeModel from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.models.llama import Model as LlamaModel +from mlx_lm.models.minimax import Model as MiniMaxModel from mlx_lm.models.ministral3 import Model as Ministral3Model from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock @@ -252,6 +253,14 @@ def tensor_auto_parallel( all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) + elif isinstance(model, MiniMaxModel): + tensor_parallel_sharding_strategy = MiniMaxShardingStrategy( + group, + all_to_sharded_linear, + sharded_to_all_linear, + all_to_sharded_linear_in_place, + sharded_to_all_linear_in_place, + ) elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)): tensor_parallel_sharding_strategy = QwenShardingStrategy( group, @@ -394,6 +403,35 @@ class ShardedDeepseekV3MoE(CustomMlxLayer): return y +class MiniMaxShardingStrategy(TensorParallelShardingStrategy): + def shard_model(self, model: nn.Module) -> nn.Module: + model = cast(MiniMaxModel, model) + for layer in model.layers: + # Shard the self attention + layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) + layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) + layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) + layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) + layer.self_attn.num_attention_heads //= self.N + layer.self_attn.num_key_value_heads //= self.N + + # Shard the MoE. Shard in place since the MoE should be responsible + # for aggregating the results. + self.all_to_sharded_linear_in_place( + layer.block_sparse_moe.switch_mlp.gate_proj + ) + self.sharded_to_all_linear_in_place( + layer.block_sparse_moe.switch_mlp.down_proj + ) + self.all_to_sharded_linear_in_place( + layer.block_sparse_moe.switch_mlp.up_proj + ) + layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + layer.block_sparse_moe.sharding_group = self.group + + return model + + class QwenShardingStrategy(TensorParallelShardingStrategy): def shard_model(self, model: nn.Module) -> nn.Module: model = cast(Qwen3MoeModel, model) @@ -414,7 +452,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy): self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) - layer.mlp = ShardedQwenMoE(layer.mlp) # type: ignore + layer.mlp = ShardedQwenMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] layer.mlp.sharding_group = self.group # Shard the MLP