From 58fb80e41244b0c5895c6b2ede3b2cf1dfcb3059 Mon Sep 17 00:00:00 2001 From: Evan Date: Fri, 17 Apr 2026 18:12:48 +0100 Subject: [PATCH] format, comment --- python/parts.nix | 151 +++++++++++++++++++++++------------------------ 1 file changed, 75 insertions(+), 76 deletions(-) diff --git a/python/parts.nix b/python/parts.nix index b72e9d061..b5d6db622 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -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 =