From f80a229b76ac2f7ca2975761f2759401e68b5ec7 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Wed, 25 Oct 2023 12:43:24 -0700 Subject: [PATCH] python3Packages.jaxlib-bin: support Python 3.9-3.12 --- .../development/python-modules/jaxlib/bin.nix | 93 ++++++++++++++++--- .../python-modules/jaxlib/prefetch.sh | 15 +-- 2 files changed, 87 insertions(+), 21 deletions(-) diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 8b673d6040d5..ba3a6adffb74 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -47,38 +47,101 @@ let # official instructions recommend installing CPU-only versions via PyPI. cpuSrcs = let - getSrcFromPypi = { platform, hash }: fetchPypi { - inherit version platform hash; + getSrcFromPypi = { platform, dist, hash }: fetchPypi { + inherit version platform dist hash; pname = "jaxlib"; format = "wheel"; # See the `disabled` attr comment below. - dist = "cp310"; - python = "cp310"; - abi = "cp310"; + python = dist; + abi = dist; }; in { - "x86_64-linux" = getSrcFromPypi { + "3.9-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; + dist = "cp39"; + hash = "sha256-8bTrWutuK0qVnbkcwMfgBf414YdaLc3GK5IsCm/JNPE="; + }; + "3.9-aarch64-darwin" = getSrcFromPypi { + platform = "macosx_11_0_arm64"; + dist = "cp39"; + hash = "sha256-Tmv2iOqlNbZqw/rYjef6GmM0N18EA5JTt6T3lQe+4Rs="; + }; + "3.9-x86_64-darwin" = getSrcFromPypi { + platform = "macosx_10_14_x86_64"; + dist = "cp39"; + hash = "sha256-mDT1INLqPdCkxtMMFR0qHLOIZdWEy8Iuzw1/vOoECsA="; + }; + + "3.10-x86_64-linux" = getSrcFromPypi { + platform = "manylinux2014_x86_64"; + dist = "cp310"; hash = "sha256-ksnY+CPEstact5lKjbSg+ZSPJtSt0Y0NFWEFufBCByk="; }; - "aarch64-darwin" = getSrcFromPypi { + "3.10-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; + dist = "cp310"; hash = "sha256-O7dHvdKLKfNELGfF4TKy7N5EX6Ca7Zu8OtLXWvFykR8="; }; - "x86_64-darwin" = getSrcFromPypi { + "3.10-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; + dist = "cp310"; hash = "sha256-gqKMUZSXrt8sQtTAoQbzAfCzO8gM9Y1/tZpuJVWyN0Y="; }; - }; + "3.11-x86_64-linux" = getSrcFromPypi { + platform = "manylinux2014_x86_64"; + dist = "cp311"; + hash = "sha256-m+NDzwXMNboNjDl2nLY+vqAoN2dQJZVWb1UQDpqqDPw="; + }; + "3.11-aarch64-darwin" = getSrcFromPypi { + platform = "macosx_11_0_arm64"; + dist = "cp311"; + hash = "sha256-zCOAjaWWCQT9Jnm1jjc1Rh5gemqy7ACtTKLM0MqSJzM="; + }; + "3.11-x86_64-darwin" = getSrcFromPypi { + platform = "macosx_10_14_x86_64"; + dist = "cp311"; + hash = "sha256-gOLIxkk+2hew2GqWu1WgMVEx1YEutx7Zod7QbwsuUVQ="; + }; + + "3.12-x86_64-linux" = getSrcFromPypi { + platform = "manylinux2014_x86_64"; + dist = "cp312"; + hash = "sha256-BZTmkgNuV4nWtfbY4t/19aP43szZQEdgpFXh5qwGRXk="; + }; + "3.12-aarch64-darwin" = getSrcFromPypi { + platform = "macosx_11_0_arm64"; + dist = "cp312"; + hash = "sha256-aAMTrLXU9EYwPv+kdeyI88/D7b4NANB39Fn8vuXUqFA="; + }; + "3.12-x86_64-darwin" = getSrcFromPypi { + platform = "macosx_10_14_x86_64"; + dist = "cp312"; + hash = "sha256-KHzlIfa9KtYcHX+i/F/SKaYTpD4/XjHVu5j3BdRTUmc="; + }; + }; # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html. # When upgrading, you can get these hashes from prefetch.sh. See # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. - gpuSrc = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; - hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o="; + gpuSrcs = { + "3.9" = fetchurl { + url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl"; + hash = "sha256-WB5Vbr/XeYKXCP/3DIXF20jR6/1xE3huX1h5ow8ETl0="; + }; + "3.10" = fetchurl { + url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; + hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o="; + }; + "3.11" = fetchurl { + url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl"; + hash = "sha256-Q8ZtF2GCrG30GFbCeCZTWPmW2TBybeXzh2u+NRiYpx4="; + }; + "3.12" = fetchurl { + url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl"; + hash = "sha256-lphkSDOJ9SwbO0hp/xC1bYn5fWgth9A9Iwsc9zV0buI="; + }; }; in @@ -87,15 +150,15 @@ buildPythonPackage { inherit version; format = "wheel"; - disabled = !(pythonVersion == "3.10"); + disabled = !(pythonVersion == "3.9" || pythonVersion == "3.10" || pythonVersion == "3.11" || pythonVersion == "3.12"); # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6. src = if !cudaSupport then ( - cpuSrcs."${stdenv.hostPlatform.system}" + cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}" or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}") - ) else gpuSrc; + ) else gpuSrcs."${pythonVersion}"; # Prebuilt wheels are dynamically linked against things that nix can't find. # Run `autoPatchelfHook` to automagically fix them. diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh index 3362e2d0b781..46b2573e3c30 100755 --- a/pkgs/development/python-modules/jaxlib/prefetch.sh +++ b/pkgs/development/python-modules/jaxlib/prefetch.sh @@ -1,7 +1,8 @@ -#!/usr/bin/env bash +#! /usr/bin/env nix-shell +#! nix-shell -i sh -p jq prefetch () { - expr="(import { system = \"$1\"; config.cudaSupport = $2; }).python3.pkgs.jaxlib-bin.src.url" + expr="(import { system = \"$2\"; config.cudaSupport = $3; }).python$1.pkgs.jaxlib-bin.src.url" url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r) echo "$url" sha256=$(nix-prefetch-url "$url") @@ -9,7 +10,9 @@ prefetch () { echo } -prefetch "x86_64-linux" "false" -prefetch "aarch64-darwin" "false" -prefetch "x86_64-darwin" "false" -prefetch "x86_64-linux" "true" +for py in "39" "310" "311" "312"; do + prefetch "$py" "x86_64-linux" "false" + prefetch "$py" "aarch64-darwin" "false" + prefetch "$py" "x86_64-darwin" "false" + prefetch "$py" "x86_64-linux" "true" +done