Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

python311Packages.orbax-checkpoint: init at 0.5.3 #266467

Closed
wants to merge 2 commits into from

Conversation

GaetanLepage
Copy link
Contributor

@GaetanLepage GaetanLepage commented Nov 9, 2023

Description of changes

Add orbax-checkpoint.

Things done

  • Built on platform(s)
    • x86_64-linux
    • aarch64-linux
    • x86_64-darwin
    • aarch64-darwin
  • For non-Linux: Is sandboxing enabled in nix.conf? (See Nix manual)
    • sandbox = relaxed
    • sandbox = true
  • Tested, as applicable:
  • Tested compilation of all packages that depend on this change using nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD". Note: all changes have to be committed, also see nixpkgs-review usage
  • Tested basic functionality of all binary files (usually in ./result/bin/)
  • 23.11 Release Notes (or backporting 23.05 Release notes)
    • (Package updates) Added a release notes entry if the change is major or breaking
    • (Module updates) Added a release notes entry if the change is significant
    • (Module addition) Added a release notes entry if adding a new NixOS module
  • Fits CONTRIBUTING.md.

@GaetanLepage GaetanLepage marked this pull request as draft November 9, 2023 14:21
@ofborg ofborg bot added 8.has: package (new) This PR adds a new package 11.by: package-maintainer This PR was created by the maintainer of the package it changes 10.rebuild-darwin: 0 This PR does not cause any packages to rebuild on Darwin 10.rebuild-linux: 1-10 labels Nov 9, 2023
@GaetanLepage GaetanLepage marked this pull request as ready for review November 16, 2023 10:11
@GaetanLepage
Copy link
Contributor Author

Result of nixpkgs-review pr 266467 run on x86_64-linux 1

4 packages built:
  • python310Packages.orbax-checkpoint
  • python310Packages.orbax-checkpoint.dist
  • python311Packages.orbax-checkpoint
  • python311Packages.orbax-checkpoint.dist

@GaetanLepage GaetanLepage force-pushed the orbax-checkpoint branch 4 times, most recently from 6804158 to 4af94bc Compare January 15, 2024 09:36
@ofborg ofborg bot requested review from samuela and ndl January 15, 2024 13:35
Comment on lines +62 to +64
preCheck = ''
export HOME=$(mktemp -d)
'';
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this necessary? also why not put it in the overridePythonAttrs below?

absl-py
etils
jax
jaxlib
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

jaxlib should never be in propagatedBuildInputs. if setup.py asks for it, remove with pythonRemoveDeps

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is upgrading tensorstore required by orbax-checkpoint?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After orbax-checkpoint version 0.4.3 I run into failures at runtime with the current version of tensorstore. I see despite the title that the orbax version is at 0.4.8 here.

@GaetanLepage GaetanLepage changed the title python311Packages.orbax-checkpoint: init at 0.4.2 python311Packages.orbax-checkpoint: init at 0.5.3 Feb 8, 2024
@wegank wegank added the 2.status: merge conflict This PR has merge conflicts with the target branch label Mar 20, 2024
@@ -21,7 +22,7 @@
, tensorflow
}:

buildPythonPackage rec {
let flax = buildPythonPackage rec {
pname = "flax";
version = "0.7.5";
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Version 0.7.5 of flax fails to build because jax.config.define_*_state was removed in jax 0.4.25.

nix build .#python311Packages.flax
error: builder for '/nix/store/w04njz8gwd8q2qa2giwlfg7yicb8jjhi-python3.11-flax-0.7.5.drv' failed with exit code 1;
       last 10 log lines:
       >   File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
       >   File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/__init__.py", line 19, in <module>
       >     from .configurations import (
       >   File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/configurations.py", line 93, in <module>
       >     flax_filter_frames = define_bool_state(
       >                          ^^^^^^^^^^^^^^^^^^
       >   File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/configurations.py", line 42, in define_bool_state
       >     return jax_config.define_bool_state('flax_' + name, default, help)
       >            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
       > AttributeError: 'Config' object has no attribute 'define_bool_state'
       For full logs, run 'nix-store -l /nix/store/w04njz8gwd8q2qa2giwlfg7yicb8jjhi-python3.11-flax-0.7.5.drv'.

Bumping to the latest version 0.8.2 fixes the issue (it was fixed in 0.8.0).


let orbax-checkpoint = buildPythonPackage rec {
pname = "orbax-checkpoint";
version = "0.5.3";
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The latest version is 0.5.7.

@@ -8892,6 +8892,8 @@ self: super: with self; {

oras = callPackage ../development/python-modules/oras { };

orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint { };
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@GaetanLepage perhaps out of scope for this PR, but orbax-checkpoint's dependency on protobuf causes a conflict with tensorflow since tensorflow has its own protobuf version. This means python packages can't depend on orbax-checkpoint and tensorflow at the same time. See tensorflow-build in pkgs/top-level/python-packages.nix.

  tensorflow-build = let
    compat = rec {
      protobufTF = pkgs.protobuf_21.override {
        abseil-cpp = pkgs.abseil-cpp_202301;
      };
  ...

I first noticed this when trying to package gpjax with the following code.

nix why-depends --all .#gpjax /nix/store/cb1gs888vfqxawvc65q1dk6jzbayh3wz-source#protobuf_21
/nix/store/icb03rmfzgph4bhxm4wic934qwx7y8j7-python3.11-gpjax-0.8.2
├───/nix/store/n8blzy5lqs9gl7yi4m8q5c44gnfkdgcc-python3.11-orbax-checkpoint-0.5.7
│   └───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
│       └───/nix/store/swx2q314if8j37ryn35nr2yb60axmrxz-protobuf-21.12
└───/nix/store/3smllqadx197wmspwd3zkyrxz2rrj01b-python3.11-tensorflow-probability-0.21.0
    └───/nix/store/b6fvycgjpin2g72mfv5fw3f9agqhn697-python3.11-tensorflow-2.13.0
        ├───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
        ├───/nix/store/api4113pl9qg9296cf15igrfnxyifjhl-python3.11-tensorboard-2.16.2
        │   ├───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
        │   └───/nix/store/d4dyq3y1scsaprv1y72fvz7sbn31dqjb-python3.11-tensorboard_plugin_profile-2.11.1
        │       └───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
        └───/nix/store/d4dyq3y1scsaprv1y72fvz7sbn31dqjb-python3.11-tensorboard_plugin_profile-2.11.1

It doesn't seem tensorflow-build exposes protobuf-pythonTF ergonomically but

orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint { };

could still be replaced with

orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint {
  protobuf = pkgs.protobuf_21.override {
    abseil-cpp = pkgs.abseil-cpp_202301;
  };
}

with the caveat that the two have to stay in sync. I'm a bit confused what the current situation is, for example, pkgs.python3Packages.tensorboard also conflicts with tensorflow which seems strange to me. Why does tensorflow carry its own versions of packages that are also exposed at the top level?

@stephen-huan
Copy link
Member

@GaetanLepage sorry for the noise, but this PR seems to be superseded by

@samuela
Copy link
Member

samuela commented Apr 2, 2024

i'll close for now since it looks like these changes have been superseded by other PRs, and we can always reopen or open follow up PRs to resurrect anything from this one

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
2.status: merge conflict This PR has merge conflicts with the target branch 6.topic: python 8.has: package (new) This PR adds a new package 10.rebuild-darwin: 0 This PR does not cause any packages to rebuild on Darwin 10.rebuild-linux: 1-10 11.by: package-maintainer This PR was created by the maintainer of the package it changes
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants