orbax-0.1.8: Bad package release?

[andsteing]

Quick fix: Use pip install orbax==0.1.7 until this issue is fixed.

Something seems to have gone wrong:

$ !pip download orbax
Collecting orbax==0.1.8
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
  error: subprocess-exited-with-error
  
  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  Preparing metadata (setup.py) ... error
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

Downloading the package manually and inspecting its contents (from PyPI: https://pypi.org/project/orbax/0.1.8/#files):

$ wget https://files.pythonhosted.org/packages/1a/b2/bf13b6c0f73952d4ce11c13f439b1fbd688ad9a9a875ab5cf6106412644d/orbax-0.1.8.tar.gz
$ tar xzvf orbax-0.1.8.tar.gz
orbax-0.1.8/
orbax-0.1.8/PKG-INFO
orbax-0.1.8/README.md
orbax-0.1.8/orbax.egg-info/
orbax-0.1.8/orbax.egg-info/PKG-INFO
orbax-0.1.8/orbax.egg-info/SOURCES.txt
orbax-0.1.8/orbax.egg-info/dependency_links.txt
orbax-0.1.8/orbax.egg-info/top_level.txt
orbax-0.1.8/setup.cfg
orbax-0.1.8/setup.py

It should look something like the previous release:

$ wget https://files.pythonhosted.org/packages/15/44/d8a13c81c47302440e861d755e496742df25d293d49c935aeb3e05f04ce3/orbax-0.1.7.tar.gz
$ tar xzvf orbax-0.1.7.tar.gz
orbax-0.1.7/LICENSE
orbax-0.1.7/README.md
orbax-0.1.7/orbax/__init__.py
orbax-0.1.7/orbax/checkpoint/README.md
orbax-0.1.7/orbax/checkpoint/__init__.py
orbax-0.1.7/orbax/checkpoint/abstract_checkpointer.py
orbax-0.1.7/orbax/checkpoint/aggregate_handlers.py
orbax-0.1.7/orbax/checkpoint/array_checkpoint_handler.py
orbax-0.1.7/orbax/checkpoint/async_checkpoint_handler.py
orbax-0.1.7/orbax/checkpoint/async_checkpointer.py
orbax-0.1.7/orbax/checkpoint/checkpoint_handler.py
orbax-0.1.7/orbax/checkpoint/checkpoint_manager.py
orbax-0.1.7/orbax/checkpoint/checkpoint_utils.py
orbax-0.1.7/orbax/checkpoint/checkpoint_utils_test.py
orbax-0.1.7/orbax/checkpoint/checkpointer.py
orbax-0.1.7/orbax/checkpoint/future.py
orbax-0.1.7/orbax/checkpoint/json_checkpoint_handler.py
orbax-0.1.7/orbax/checkpoint/json_checkpoint_handler_test.py
orbax-0.1.7/orbax/checkpoint/lazy_utils.py
orbax-0.1.7/orbax/checkpoint/msgpack_utils.py
orbax-0.1.7/orbax/checkpoint/orbax_checkpoint.ipynb
orbax-0.1.7/orbax/checkpoint/pytree_checkpoint_handler.py
orbax-0.1.7/orbax/checkpoint/test_utils.py
orbax-0.1.7/orbax/checkpoint/transform_utils.py
orbax-0.1.7/orbax/checkpoint/transform_utils_test.py
orbax-0.1.7/orbax/checkpoint/type_handlers.py
orbax-0.1.7/orbax/checkpoint/utils.py
orbax-0.1.7/orbax/checkpoint/utils_test.py
orbax-0.1.7/orbax/conftest.py
orbax-0.1.7/pyproject.toml
orbax-0.1.7/PKG-INFO

=> The orbax-0.1.8 package should probably be yanked, and a new package orbax-0.1.9 should be released.

Ideally, the release process would be automated via Github action, like for example here:
https://github.com/google/flax/blob/main/.github/workflows/pythonpublish.yml

[ajain-23]

Hi Andreas,

Thanks for the note, and apologies for any confusion! orbax is no longer a standalone package, as it includes related functionalities for model checkpointing and exporting that have different dependencies (e.g. only the latter requires tensorflow). In order to reduce dependency bloat, installing Orbax's features is now done directly, with the packages orbax-checkpoint and orbax-export. Releases for these are automated :)

Import statements in actual code can remain unchanged, as orbax functions as a standard Python namespace.

This information should be in the error that appears on pip install orbax, and on the orbax-0.1.8 PyPi page (akin to pip install sklearn and pip install pytorch): https://pypi.org/project/orbax/0.1.8/#description

[andsteing]

Hi Ayush.

Oh, I missed that renaming indeed!

This is what I see on a fresh Colab VM when trying to install orbax:

>>> !pip install orbax
Collecting orbax
  Downloading orbax-0.1.8.tar.gz (1.6 kB)
  error: subprocess-exited-with-error
  
  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  Preparing metadata (setup.py) ... error
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

And indeed, I see similar error messages when trying to pip install pytorch and pip install sklearn.

So it seems this is not orbax specific. And I will think of navigating to the PyPI index page and reading the README next time before opening an issue :)

Thanks, Andreas

[andsteing]

Oh, actually if one things of adding the -v flag, then the error message appears:

>>> !pip install -v orbax
Collecting orbax
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
  Running command python setup.py egg_info

  *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***

  error: subprocess-exited-with-error

...

[andsteing]

Though I'm wondering: why push a new orbax-0.1.8 to force everybody to update to the new package?

This will break lots of other libs that specify orbax as a simple dependency (example: flax-0.6.4).

For example, the following command breaks: pip install flax==0.6.4

As does the following command: pip install jax==0.3.25 jaxlib==0.3.25 flax – which is required to pin the JAX version in Colab TPU runtimes.

Maybe it would be better to have an empty orbax package that has the both orbax-checkpoint and orbax-export as a dependency?

[ajain-23]

We did seriously consider an empty package dependent on both of the "child" packages, but ended up deciding that a clean break would be better in the long-term, better aligned with precedent from other packages. If a user must pin another package to an older version and for some reason miss the switch to e.g. orbax-checkpoint, they can opt to pin to earlier versions of orbax too.

cc @cpgaffney1

[andsteing]

When discussing these two options, what argument convinced you that a "clean break" would be the better option?

I think the "empty package with child dependencies" options would only have the minor disadvantage that some user that erroneously depends on orbax rather than orbax-checkpoint (assuming they don't use orbax-export) has some unused dependencies. This could be fixed at any point later.

On the other hand the "clean break" option might result in many breakages (e.g. all cases that need to pin JAX for some reason, like above young-geng/EasyLM#84, or the regular TPU Colab runtimes). These errors are easy to fix, but it still adds an overhead everytime it occurs.

[cpgaffney1]

One factor that we considered was that orbax was always actually just the checkpointing component, and did not really incorporate any elements of export (or perhaps towards the end a few primitive, un-developed pieces of the future export library).

Since orbax is conceptually equivalent to orbax-checkpoint, we considered that it would be better to simply leave it as-is instead of making a unified package. orbax-checkpoint and orbax-export occupy conceptually related, but ultimately distinct, roles in a vertical chain of JAX ecosystem tools, and in practice, there's not a whole lot of overlap for individual users. That was part of what prompted us to maintain a clean separation.

Ultimately, it can be argued that we have taken a less-than-optimal approach by doing this separation rather than an approach like that of the etils library, for example, and I'm definitely receptive to that line of thinking.

[andsteing]

Renaming is always painful :-) It's great though to do the effort and establish a well structured PyPI naming scheme!

If we go with the analogy of etils, then maybe^1 at some point it was a single package etils but then later it was split into multiple subcomponents that could be installed separately, like etils[ecolab] etc. But the default package is still etils[all] that includes all the subcomponents. This way users of old etils without a version number would still be compatible with the new naming scheme, at the price of maybe including more dependencies than they strictly need to.

^1: Note that everything that follows is an assumption – I didn't actually check if etils was a single package in the beginning.

[MaxBoyman]

tflite-model-maker won't install with pip because of the change mantioned above. what should we do? there is a dependency with flax and the tensorflowjs which breaks lots of packages installation with pypi. what is the solution?

[ajain-23]

The latest version of tensorflowjs allowed as a dependency for tflite-model-maker (https://github.com/tensorflow/examples/blob/master/tensorflow_examples/lite/model_maker/requirements.txt#L14) doesn't have any flax dependency: https://github.com/tensorflow/tfjs/blob/tfjs-v3.18.0/tfjs-converter/python/requirements.txt.

Even the most recent tensorflowjs still depends on a version of flax that has no orbax dependency (https://github.com/tensorflow/tfjs/blob/tfjs-v4.10.0/tfjs-converter/python/requirements.txt).

Can you tell us more about your issue?

[koktavy]

Also coming here from trying to install tflite-model-maker. Here's the error in full. I'm not typically working in python so I apologize if this would be a simple version change in the requirements.txt or such.

I've tried installing both orbax-checkpoint and orbax-export with no change.

Collecting orbax (from flax>=0.5.3->tensorflowjs>=2.4.0->tflite-model-maker)
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [3 lines of output]

      *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***

      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

[andsteing]

I would still vote for publishing a new orbax-0.1.9 to PyPI that includes orbax-checkpoint and emits a warnings.warn(..., DeprecationWarning) upon import.

This would effectively solve all these issues, while still pointing new users to the new updated package naming.

[koktavy]

I was able to successfully install tflite-model-maker using a conda environment set to python 3.9 instead of 3.8. This approach gathered orbax-checkpoint automatically.

[ajain-23]

That may be a happy coincidence, just pushed an orbax release :)

(though orbax-checkpoint does require a minimum Python version of 3.9)

After some brief discussion, we've decided to have orbax install orbax-checkpoint automatically to preserve existing usage (as Andreas also suggested). https://pypi.org/project/orbax/0.1.9/

Given that, will close this issue. Please reopen if any concerns arise!