Jax fixes #323154

zimbatm · 2024-06-28T13:31:00Z

Description of changes

Doing a pass to get jax are related libraries to a good place.

Things done

Add a 👍 reaction to pull requests you find important.

mweinelt · 2024-06-28T16:57:52Z

Deferring to the CUDA team, they are much closer to Jax.

SomeoneSerge · 2024-06-29T02:23:35Z

pkgs/top-level/python-packages.nix

@@ -6114,13 +6114,13 @@ self: super: with self; {
    IOKit = pkgs.darwin.apple_sdk_11_0.IOKit;
  };

-  jaxlib = self.jaxlib-build;
+  jaxlib = if self.jaxlib-build.meta.broken then self.jaxlib-bin else self.jaxlib-build;


Ohhh I'm not sure how good of an idea this is: I can see myself tearing my hair out a few steps down the line (steps involving changing meta.broken) trying to figure out how I ended up with an unfree pypi wheel in the closure instead of a cpu-only source-build

I personally don't mind, but I get your point. Feel free to decide.

I agree that this feels like a footgun.

Another option: we alias jaxlib = jaxlib-bin and leave a comment linking to this PR and the hash mismatch issue.

Ok, I like samuela's simple approach. Fixed.

I'd suggest a more radical approach, which would probably cause frustration to other maintainers: we add a lib.warn to the if broken then jaxlib-bin branch saying that we're using a prebuilt and patchelfed jaxlib from pypi, potentially including unfree components

The "if broken then jaxlib-bin else jaxlib-build" logic existed in the jax package.

Ohhh, I had completely missed that

Ok, there was already an alias in place, and in this PR it place two roles:

Supporting darwin

Working around the hash issue

For the former, I'd replace the logic with if jaxlib.meta.unsupported then jaxlib-bin else jaxlib to look at platforms instead of broken. For the latter, I hope we can work around the instability soon.

And if the build version is broken, the packages should keep working.

Yes but whether shipping a substitute "product" counts as "it's still working" is a matter of setting expectations.

Anyway I see now that my comment was out of scope

Yes but whether shipping a substitute "product" counts as "it's still working" is a matter of setting expectations.

Agreed. One difference I glossed over is that the jaxlib-source also supports different mtl variants. If the user overrides that package, he could be surprised by the switch.

The "if broken then jaxlib-bin else jaxlib-build" logic existed in the jax package.

Yes, I wrote this in order to fix the build on darwin.

Ok, there was already an alias in place

Actually, there is/was currently no such alias in place. The "if broken then jaxlib-bin else jaxlib-build" was only logic local to the jax package so that it would use jaxlib in checkPhase, but its impact was limited to that. See

nixpkgs/pkgs/development/python-modules/jax/default.nix

Line 29 in 6a8b974

jaxlib' = if jaxlib.meta.broken then jaxlib-bin else jaxlib;

I think just writting if stdenv.isDarwin then A else B is way easier to understand and we should probably just go with that.

samuela · 2024-06-28T23:43:25Z

pkgs/development/python-modules/jax/default.nix

@@ -130,10 +127,14 @@ buildPythonPackage rec {
      "testQdwhWithOnRankDeficientInput5"
    ];

-  disabledTestPaths = lib.optionals (stdenv.isDarwin && stdenv.isAarch64) [
+  disabledTestPaths = [
+    # Segmentation fault


do we have an upstream issue for this?

Not yet. I don't know if this is an upstream or a packaging issue.

See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba

Looking at the log I'm 95% confident this is an upstream issue.

pkgs/development/python-modules/jaxlib/default.nix

samuela · 2024-06-29T23:07:19Z

pkgs/top-level/python-packages.nix

@@ -6114,13 +6114,13 @@ self: super: with self; {
    IOKit = pkgs.darwin.apple_sdk_11_0.IOKit;
  };

-  jaxlib = self.jaxlib-build;
+  jaxlib = if self.jaxlib-build.meta.broken then self.jaxlib-bin else self.jaxlib-build;


I agree that this feels like a footgun.

Another option: we alias jaxlib = jaxlib-bin and leave a comment linking to this PR and the hash mismatch issue.

pkgs/development/python-modules/jaxlib/default.nix

SomeoneSerge · 2024-07-02T17:07:39Z

pkgs/top-level/python-packages.nix

+  # Using the -bin version since the -source one is currently broken (see #323154)
+  jaxlib = self.jaxlib-bin;


Personally, I've no preference as to whether this is implemented in jax/default.nix or in python-packages.nix. Let's keep this as a source build on linux-*. Let's keep it at -bin for darwin-* but use meta.unsupported instead of broken. With a comment that this is a macos exception

Switched to this logic

So that packages depending on jaxlib can also build on macOS.

pkgs/development/python-modules/jax/default.nix

to the currently supported ones, so that the top-level jaxlib can fallback to jaxlib-bin

SomeoneSerge · 2024-07-03T15:17:11Z

pkgs/development/python-modules/jaxlib/default.nix

    # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136
    # however even with that fix applied, it doesn't work for everyone:
    # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129
-    # NOTE: We always build with NCCL; if it is unsupported, then our build is broken.
-    broken = effectiveStdenv.isDarwin || nccl.meta.unsupported;
+    platforms = platforms.linux;


This is to make the new meta.unsupported logic work. Honestly, I fail to predict what's this going to look like on search.nixos.org now

SomeoneSerge · 2024-07-03T15:18:42Z

pkgs/development/python-modules/jaxlib/default.nix

    # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136
    # however even with that fix applied, it doesn't work for everyone:
    # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129
-    # NOTE: We always build with NCCL; if it is unsupported, then our build is broken.


NOTE: We always build with NCCL; if it is unsupported, then our build is broken.

This is a peculiar comment, if the derivation depended on nccl and nccl was unsupported for the platform, the evaluation would fail saying exactly that: nccl is not supported on this platform. It also seems like the cpu-only jaxlib didn't actually have nccl in the inputs. Am I reading this right? CC @samuela

i'm not entirely sure... my read is that cudaSupport -> using nccl, and the logic that was added just happened to forget about the CPU-only case.

cc @ConnorBaker who wrote this in 9bebd9e

Co-authored-by: Sandro <[email protected]>

github-actions bot added the 6.topic: python label Jun 28, 2024

zimbatm force-pushed the jax-fixes branch 2 times, most recently from da16155 to b5e7b73 Compare June 28, 2024 15:18

ofborg bot added 8.has: clean-up 8.has: package (new) This PR adds a new package labels Jun 28, 2024

ofborg bot requested a review from samuela June 28, 2024 15:49

ofborg bot added 10.rebuild-darwin: 11-100 10.rebuild-linux: 11-100 labels Jun 28, 2024

zimbatm marked this pull request as ready for review June 28, 2024 16:21

zimbatm requested review from mweinelt and GaetanLepage June 28, 2024 16:51

mweinelt requested review from SomeoneSerge and removed request for mweinelt June 28, 2024 16:58

zimbatm requested a review from a team June 28, 2024 22:41

SomeoneSerge reviewed Jun 29, 2024

View reviewed changes

GaetanLepage approved these changes Jun 29, 2024

View reviewed changes

samuela reviewed Jun 29, 2024

View reviewed changes

wegank added the 12.approvals: 1 This PR was reviewed and approved by one reputable person label Jun 30, 2024

zimbatm force-pushed the jax-fixes branch from b5e7b73 to a785a1c Compare June 30, 2024 18:29

ofborg bot requested a review from samuela June 30, 2024 19:01

wegank removed the 12.approvals: 1 This PR was reviewed and approved by one reputable person label Jun 30, 2024

zimbatm force-pushed the jax-fixes branch from a785a1c to a128565 Compare July 1, 2024 07:44

SomeoneSerge reviewed Jul 2, 2024

View reviewed changes

pkgs/development/python-modules/jaxlib/default.nix Outdated Show resolved Hide resolved

SomeoneSerge reviewed Jul 2, 2024

View reviewed changes

zimbatm added 2 commits July 3, 2024 13:21

python3Packages.jaxlib: switch version based on support

62e36ed

So that packages depending on jaxlib can also build on macOS.

python3Packages.jax: disable segfaulting test

0ec9091

zimbatm force-pushed the jax-fixes branch from a128565 to 0ec9091 Compare July 3, 2024 11:22

SuperSandro2000 reviewed Jul 3, 2024

View reviewed changes

pkgs/development/python-modules/jax/default.nix Outdated Show resolved Hide resolved

SuperSandro2000 reviewed Jul 3, 2024

View reviewed changes

pkgs/development/python-modules/jax/default.nix Outdated Show resolved Hide resolved

ofborg bot added 10.rebuild-darwin: 0 This PR does not cause any packages to rebuild on Darwin and removed 10.rebuild-darwin: 11-100 labels Jul 3, 2024

python311Packages.jaxlib: limit platforms

cf4bfbb

to the currently supported ones, so that the top-level jaxlib can fallback to jaxlib-bin

SomeoneSerge requested a review from natsukium as a code owner July 3, 2024 14:49

python3Packages.jax{lib,lib-bin}: disambiguate meta.descriptions

291f236

SomeoneSerge reviewed Jul 3, 2024

View reviewed changes

SomeoneSerge approved these changes Jul 3, 2024

View reviewed changes

ofborg bot added 10.rebuild-darwin: 11-100 and removed 10.rebuild-darwin: 0 This PR does not cause any packages to rebuild on Darwin labels Jul 3, 2024

SomeoneSerge mentioned this pull request Jul 3, 2024

release-cuda: build with config.cudaSupport #324379

Merged

13 tasks

python3Packages.jax: drop redundant parens

039cf11

Co-authored-by: Sandro <[email protected]>

zimbatm merged commit 9f38e1c into NixOS:master Jul 4, 2024
26 checks passed

zimbatm deleted the jax-fixes branch July 4, 2024 09:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Jax fixes #323154

Jax fixes #323154

zimbatm commented Jun 28, 2024 •

edited

Loading

mweinelt commented Jun 28, 2024

SomeoneSerge Jun 29, 2024 •

edited

Loading

GaetanLepage Jun 29, 2024

samuela Jun 29, 2024

zimbatm Jun 30, 2024

SomeoneSerge Jun 30, 2024

SomeoneSerge Jul 1, 2024

SomeoneSerge Jul 1, 2024

zimbatm Jul 1, 2024

samuela Jul 1, 2024

SuperSandro2000 Jul 3, 2024

samuela Jun 28, 2024

zimbatm Jun 30, 2024

This comment was marked as outdated.

zimbatm Jul 1, 2024

samuela Jul 1, 2024

samuela Jun 29, 2024

SomeoneSerge Jul 2, 2024 •

edited

Loading

zimbatm Jul 3, 2024

SomeoneSerge Jul 3, 2024

SomeoneSerge Jul 3, 2024 •

edited

Loading

samuela Jul 4, 2024

		# Using the -bin version since the -source one is currently broken (see #323154)
		jaxlib = self.jaxlib-bin;

Jax fixes #323154

Jax fixes #323154

Conversation

zimbatm commented Jun 28, 2024 • edited Loading

Description of changes

Things done

mweinelt commented Jun 28, 2024

SomeoneSerge Jun 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment was marked as outdated.

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SomeoneSerge Jul 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SomeoneSerge Jul 3, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zimbatm commented Jun 28, 2024 •

edited

Loading

SomeoneSerge Jun 29, 2024 •

edited

Loading

SomeoneSerge Jul 2, 2024 •

edited

Loading

SomeoneSerge Jul 3, 2024 •

edited

Loading