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.

Sign up for GitHub

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

Jump to bottom

No module named 'orbax.checkpoint._src.path' when used as bazel pip dependency #1429

Open
hartikainen opened this issue Dec 7, 2024 · 4 comments · May be fixed by #1458
Open

No module named 'orbax.checkpoint._src.path' when used as bazel pip dependency #1429

hartikainen opened this issue Dec 7, 2024 · 4 comments · May be fixed by #1458

Comments

@hartikainen
Copy link

hartikainen commented Dec 7, 2024 •
edited
Loading

I'm trying to use orbax-checkpoint as a transitive pip dependency of brax within my bazel setup. I've declared brax, and thus also orbax-checkpoint, as a dependency via pip/pypi. When running any of my bazel targets that depend on the pip-installed orbax-checkpoint, I get No module named 'orbax.checkpoint._src.path'.

I created a minimal reproduction here: main...hartikainen:orbax:bazel-pip-dependency-error. And here's the error I get running it:

$ cd ./pypi_dependency_test
$ bazel run //:pypi_dependency_test                                  2.336s
INFO: Analyzed target //:pypi_dependency_test (0 packages loaded, 0 targets configured).
INFO: Found 1 target...
Target //:pypi_dependency_test up-to-date:
  bazel-bin/pypi_dependency_test
INFO: Elapsed time: 0.147s, Critical Path: 0.01s
INFO: 1 process: 1 internal.
INFO: Build completed successfully, 1 total action
INFO: Running command line: external/bazel_tools/tools/test/test-setup.sh ./pypi_dependency_test
exec ${PAGER:-/usr/bin/less} "$0" || exit 1
Executing tests from //:pypi_dependency_test
-----------------------------------------------------------------------------
Traceback (most recent call last):
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/_main/pypi_dependency_test.py", line 3, in <module>
    import orbax.checkpoint
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python~~pip~pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/__init__.py", line 23, in <module>
    from orbax.checkpoint import aggregate_handlers
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python~~pip~pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/aggregate_handlers.py", line 25, in <module>
    from orbax.checkpoint import utils
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python~~pip~pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/utils.py", line 30, in <module>
    from orbax.checkpoint._src.path import async_utils
ModuleNotFoundError: No module named 'orbax.checkpoint._src.path'

Does anyone have an idea what might be causing this? I can't reproduce this with other pip systems like vanilla venv or even when creating the venv with rules_uv and running python manually through it.
@hartikainen
Copy link
Author

hartikainen commented Dec 15, 2024 •
edited
Loading

I managed to solve the most immediate issue by removing the BUILD files from orbax-checkpoint wheel using Bazel patches as shown here: main...hartikainen:orbax:bazel-pip-dependency-error. I guess the BUILD files somehow trip rules_python.

Is it intentional that only the following two folders have BUILD-files in them?

Should they be excluded from the wheels? Are they even supposed to exist in the repository?

Unfortunately, I'm still unable to use orbax-checkpoint as a dependency because (I believe) something similar is going on with tensorstore, which is a dependency of orbax-checkpoint:

$ bazel run //:pypi_dependency_test
WARNING: Couldn't auto load rules or symbols, because no dependency on module/repository 'rules_android' found. This will result in a failure if there's a reference to those rules or symbols.
INFO: Analyzed target //:pypi_dependency_test (117 packages loaded, 6407 targets configured).
INFO: Found 1 target...
Target //:pypi_dependency_test up-to-date:
  bazel-bin/pypi_dependency_test
INFO: Elapsed time: 0.451s, Critical Path: 0.02s
INFO: 1 process: 5 action cache hit, 1 internal.
INFO: Build completed successfully, 1 total action
INFO: Running command line: external/bazel_tools/tools/test/test-setup.sh ./pypi_dependency_test
exec ${PAGER:-/usr/bin/less} "$0" || exit 1
Executing tests from //:pypi_dependency_test
-----------------------------------------------------------------------------
Traceback (most recent call last):
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/_main/pypi_dependency_test.py", line 3, in <module>
    from orbax import checkpoint as ocp
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/__init__.py", line 23, in <module>
    from orbax.checkpoint import aggregate_handlers
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/aggregate_handlers.py", line 26, in <module>
    from orbax.checkpoint._src.metadata import tree as tree_metadata
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/metadata/tree.py", line 33, in <module>
    from orbax.checkpoint._src.metadata import tree_rich_types
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/metadata/tree_rich_types.py", line 25, in <module>
    from orbax.checkpoint._src.metadata import value_metadata_entry
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/metadata/value_metadata_entry.py", line 24, in <module>
    from orbax.checkpoint._src.serialization import types
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/serialization/types.py", line 32, in <module>
    from orbax.checkpoint._src.serialization import serialization
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 36, in <module>
    from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
  File "/private/var/tmp/_bazel/aae0c90911bf9b49b7fec6ebc38b486f/execroot/_main/bazel-out/darwin_arm64-fastbuild/bin/pypi_dependency_test.runfiles/rules_python++pip+pypi_312_orbax_checkpoint/site-packages/orbax/checkpoint/_src/serialization/tensorstore_utils.py", line 28, in <module>
    import tensorstore as ts
ModuleNotFoundError: No module named 'tensorstore'

@hartikainen
Copy link
Author

hartikainen commented Dec 15, 2024 •
edited
Loading

I managed fix the tensorstore issue by downgrading from tensorstore==0.1.71 to tensorstore==0.1.69. The diff between tensorstore==0.1.69 and tensorstore==0.1.71 is quite large and I don't know exactly what's gone wrong there.

@cpgaffney1
Copy link
Collaborator

There are only two BUILD files because we are in the process of adding more for testing purposes. Probably we could just exclude these from the package published to PyPI since they aren't really necessary for users.

@hartikainen
Copy link
Author

Excluding the BUILD files from the wheel would be awesome. That would fix the issue.

@hartikainen hartikainen linked a pull request Jan 2, 2025 that will close this issue
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Exclude BUILD files from flit builds hartikainen/orbax
2 participants
@hartikainen @cpgaffney1