mirror/exo
1
0
Fork 0
mirror of https://github.com/exo-explore/exo.git synced 2026-04-19 13:27:55 -04:00
Code Issues Packages Projects Releases Wiki Activity

format, comment

Browse Source
This commit is contained in:
Evan
2026-04-17 18:12:48 +01:00
parent 62590737eb
commit 58fb80e412
1 changed files with 75 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
python/parts.nix
View File

@@ -32,10 +32,6 @@ let
libnvshmem
nccl
];
pyprojectOverlay = workspace.mkPyprojectOverlay {
sourcePreference = "wheel";
dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
};
exoOverlay = final: prev: {
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
@@ -55,12 +51,68 @@ let
'';
};
};
buildSystemsOverlay = final: prev: { }
// lib.optionalAttrs isLinux {
"${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
});
buildSystemsOverlay = final: prev: { } //
lib.optionalAttrs isDarwin
{
mlx = prev.mlx.overrideAttrs (old:
let
# Static dependencies included directly during compilation
gguf-tools = pkgs.fetchFromGitHub {
owner = "antirez";
repo = "gguf-tools";
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
};
metal_cpp = pkgs.fetchzip {
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
};
nanobind = pkgs.fetchFromGitHub {
owner = "wjakob";
repo = "nanobind";
rev = "v2.10.2";
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
fetchSubmodules = true;
};
in
{
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
# TODO: non-sdk_26 support
buildInputs = (old.buildInputs or [ ])
++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
patches = [
(pkgs.replaceVars ../nix/darwin-build-fixes.patch {
sdkVersion = pkgs.apple-sdk_26.version;
inherit (self'.packages.metal-toolchain) metalVersion;
})
];
postPatch = ''
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
--replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
'';
DEV_RELEASE = 1;
CMAKE_ARGS = toString ([
(lib.cmakeBool "USE_SYSTEM_FMT" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
(lib.cmakeBool "MLX_BUILD_CPU" true)
(lib.cmakeBool "MLX_BUILD_METAL" true)
(lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
(lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${pkgs.apple-sdk_26.version}")
(lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${pkgs.apple-sdk_26.passthru.sdkroot}")
] ++ lib.optionals (isDarwin && isx86_64) [
(lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
]);
SDKROOT = pkgs.apple-sdk_26.passthru.sdkroot;
MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
});
} // lib.optionalAttrs isLinux {
mlx = prev.mlx.overrideAttrs (old: {
buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
@@ -68,67 +120,11 @@ let
cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
'';
});
}
// lib.optionalAttrs isDarwin {
mlx = prev.mlx.overrideAttrs (old:
let
# Static dependencies included directly during compilation
gguf-tools = pkgs.fetchFromGitHub {
owner = "antirez";
repo = "gguf-tools";
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
};
metal_cpp = pkgs.fetchzip {
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
};
nanobind = pkgs.fetchFromGitHub {
owner = "wjakob";
repo = "nanobind";
rev = "v2.10.2";
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
fetchSubmodules = true;
};
in
{
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
# TODO: non-sdk_26 support
buildInputs = (old.buildInputs or [ ])
++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
patches = [
(pkgs.replaceVars ../nix/darwin-build-fixes.patch {
sdkVersion = pkgs.apple-sdk_26.version;
inherit (self'.packages.metal-toolchain) metalVersion;
})
];
postPatch = ''
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
--replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
'';
DEV_RELEASE = 1;
CMAKE_ARGS = toString ([
(lib.cmakeBool "USE_SYSTEM_FMT" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
(lib.cmakeBool "MLX_BUILD_CPU" true)
(lib.cmakeBool "MLX_BUILD_METAL" true)
(lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
(lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${pkgs.apple-sdk_26.version}")
(lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${pkgs.apple-sdk_26.passthru.sdkroot}")
] ++ lib.optionals (isDarwin && isx86_64) [
(lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
]);
SDKROOT = pkgs.apple-sdk_26.passthru.sdkroot;
MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
});
} // lib.optionalAttrs cudaSupport {
"${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
@@ -150,15 +146,18 @@ let
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
};
basePythonSet = pkgs.callPackage inputs.pyproject-nix.build.packages {
inherit python;
pyprojectOverlay = workspace.mkPyprojectOverlay {
sourcePreference = "wheel";
dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
};
editableOverlay = workspace.mkEditablePyprojectOverlay {
# Use environment variable pointing to editable root directory
root = "$REPO_ROOT";
members = [ "exo" "exo-bench" ];
};
pythonSet = basePythonSet.overrideScope (
pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
inherit python;
}).overrideScope (
lib.composeManyExtensions [
inputs.pyproject-build-systems.overlays.default
pyprojectOverlay
@@ -166,7 +165,6 @@ let
buildSystemsOverlay
]
);
editablePythonSet = pythonSet.overrideScope editableOverlay;
mkApp = cmd: name: members: pkgs.writeShellApplication {
inherit name;
@@ -175,6 +173,7 @@ let
EXO_RESOURCES_DIR = inputs.self + /resources;
};
runtimeInputs = [
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
((pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; }))
]
++ lib.optionals isDarwin [ pkgs.macmon ];
@@ -182,11 +181,11 @@ let
};
in
{
inherit pythonSet editablePythonSet;
inherit pythonSet;
editablePythonSet = pythonSet.overrideScope editableOverlay;
mkPythonScript = members: name: path: mkApp ''python ${path} "$@"'' name members;
mkExo = name: members: mkApp ''exo "$@"'' name members;
};
in
{
perSystem =