Handle cyclic dependency in pip wheels, e.g. PyTorch 2.0

[siddharthab]

🚀 feature request

Relevant Rules

pip_repository, whl_library and py_library

Description

rules_python does dependency resolution at the level of a wheel and sets up the entire wheel as one py_library. This difference in the level of granularity here is why we can not handle cyclic dependencies in wheels. Some pip wheels have cyclic dependencies which is valid Python, and pip handles them well. For example, the upcoming PyTorch 2.0 release has a cyclic dependency between the torch (.whl file) and pytorch-triton (.whl file) wheels.

Describe the solution you'd like

The proper solution for this would be to have more granular build units which would also be better from a sandboxing and build caching efficiency perspective. But as an intermediate solution, maybe we could install the wheels with cyclical dependencies between them into one repo (through one whl_library call) and set up the composite of such wheels as one py_library. For gazelle, a resolve directive can ensure that all packages point to the same py_library target.

Describe alternatives you've considered

An alternative could be that I manually build a composite wheel and set it up as a whl_library in my WORKSPACE.

[arrdem]

Does Gazelle actually have access to this information? Consider a package with platform-dependent dependencies and packaging. Gazelle can't reasonably be made to determine that 1) there is a conditional dependency 2) on some platform+architecture combination it results in a cycle. The dependency graph is only really available after the pip install (or ideally pip download) has occurred.

Maybe one could make pip_parse install requirements, create a py_library/whl_library implementation target pair with no dependencies, collect dependency information, compute connected components within the dependencies, create library/whl targets which wrap connected components into a single dependency and generate alias targets for everything else.

[siddharthab]

In my intermediate solution, I am not suggesting that Gazelle do anything. The only comment I made regarding Gazelle is that an extra resolve directive would be enough to make things work there without any changes to what it has currently. In fact, looking at how it works more closely, I think Gazelle will work for my suggestion without any changes to it.

My suggestion is all this could be handled at the pip_parse level. The requirements.txt file from pip-compile has enough information to construct a directional graph and detect cycles. And even if we do not want to perform cycle detection in pip_parse, we could manually provide another attribute to it to create composite wheels. This new attribute could be a list of lists or a dict of strings to lists.

For PyTorch 2.0, one could imagine a new pip_parse invocation (without automatic cycle detection) to look like this:

pip_parse(
  name = "pip",
  requirements_lock = "//:requirements_lock.txt",
  composite_wheels = {
    "torch_group": ["torch", "pytorch-triton"],
  },
)

With automatic cycle detection, the invocation would stay as-is.

Going back to the more long-term proper solution, I was suggesting something akin to what Gazelle does for Golang, where dependencies are specified per Go package, and not per Go module. And Gazelle accounts for platform specific dependencies seamlessly there. It is able to do this by running Gazelle inside each go_repository after downloading and extracting the contents, as part of the implementation of go_repository. Note that as per the specification, Go modules are allowed to have cyclical dependencies, but Go packages are not. So Gazelle's behavior for Golang is appropriate.

[aignas]

At least pipdeptree is printing warnings when such things occur: https://github.com/tox-dev/pipdeptree/blob/1fed5d4928489d7f2d05c89ce5857a80ae6ca50e/README.md#warnings-about-circular-dependencies

[arrdem]

That makes good sense. Even being able to manually specify composite libraries would be a big help. At present my deployment is maintaining some forked packages with "fixed" dependencies to resolve cycles that way.

[arrdem]

I've pulled together a draft patch which implements the above sketched composite_deps flow since it blocks some stuff for us, working with my employer to get approval and a CLA to release the code for review.

[Ubehebe]

I ran into the same problem with the same dependency (PyTorch 2.0). This is an important package, so I think we need to resolve this sooner rather than later. A few observations:

• The cycle only exists on linux. The macos wheel does not have a dependency on triton at all.
• An immediate solution is to patch wheel_installer.py here to omit the back-dependency from triton to torch:

if whl.name == "triton":
  sanitised_dependencies = [s for s in sanitised_dependencies if s != '"@pip_torch//:pkg"']
  sanitised_wheel_file_dependencies = [s for s in sanitised_wheel_file_dependencies if s != '"@pip_torch//:whl"' 

The patch can be maintained locally by changing your rules_python http_archive setup to a local_repository.

At a higher level, I'd like to learn more about how the Python ecosystem uses cyclic dependencies, and how we should translate that into a bazel target graph. A quick grep through the triton source code shows that the back-dependency to torch is not easily removed; it's imported from many places.

A couple questions for the OP @siddharthab:

Some pip wheels have cyclic dependencies which is valid Python, and pip handles them well.

Do you know if any PEP specifies what should happen in the presence of cyclic imports?

The proper solution for this would be to have more granular build units

It sounds like you have an idea for how this might be done. Could you be more specific? When a .py file does from torch import cuda, would that correspond to a BUILD dependency like @pip_torch//cuda:pkg?

My concern with the composite_wheels approach discussed above is that it leaks implementation details (the presence of a cyclic dep) into the BUILD-level API. For example, if you wanted to use the same bazel repo on macos and linux, you'd have to list the triton dep in requirements.in and use the torch-triton composite wheel everywhere, even though triton has no purpose on macos.

[arrdem]

I'd argue that all of pip_parse and even Golang's equivalent repos incantation is "leaking" specifics of the package fetching implementation. You already have to specify your interpreter, setup.py env, the requirements and soforth in the pip_parse incantation. This seems entirely consistent with that pattern, although perhaps the config could/should come from a file rather than being directly written into the WORKSPACE.

Re: my patchset, still chasing our lawyers.

[siddharthab]

Do you know if any PEP specifies what should happen in the presence of cyclic imports?

A point to be noted here is that we are talking about cyclic dedendencies in wheels and not python packages. I did not find something specifically about cyclic imports in wheels, but it is to be expected from the specification of file contents in a wheel. It simply states that the root contains all files to be installed in purelib or platlib, which are usually both the site-packages directory. So any random subset of python packages can be made a wheel, and there can obviously be cyclic dependencies between these random subsets of packages. So it all comes down to pip and whether it wants to handle cyclic wheels. From what I could see, it forms a transitive closure of the requirements and "installs" them all, so cycles are OK. This is different from other packaging systems (e.g. for R) where dependencies need to be installed in order.

A note on circular Python package imports. While I could not find a PEP or reference doc talking about this specifically either, it is possible from the reference on the import mechanism, because an import statement does not really execute any code in the imported module (but import-from does). This can also be tested with a small example program, and there are several community discussions about this. The mechanism to allow circular imports is what allows submodules in a package to implicitly depend on __init__.py and for __init__.py to explicitly depend on the submodules.

It sounds like you have an idea for how this might be done. Could you be more specific? When a .py file does from torch import cuda, would that correspond to a BUILD dependency like @pip_torch//cuda:pkg?

Yes, I gave a little more thought to this since I posted the issue and I am leaning towards this solution as the best way forward. I also want to implement proper granularity in Gazelle's Python extension, for the benefit of test caching in our internal codebase. Please bear with me if I don't use the right terms or get something wrong because I am new to Python and its packaging system.

As per the spec, there are no restrictions on what file/directory names may exist in a wheel (only setup.py and setup.cfg are explicitly prohibited). All that the spec says is that all the files in a wheel are installed to the site-packages directory, with some notion of a separate data directory that is spread out.

1. At the top-level, subdirectories can be regular or namespace packages, and .py files can be Python modules. It will be unusual to have files that are typically associated with a package, e.g. __init__.py. These files/directories also need to have some global uniqueness to their names to avoid conflict between wheels meant to be installed to the same location.
2. There are two special directories - dist-info and data. Specifically, the data directory may have runtime dependencies used by the package. From the language in the spec, I think it will be unusual for a wheel to have non-Python files outside one of the package directories or the data directories.
3. Unless a Python module depends on a file (Python or not Python) using paths derived from __file__, we can rely on Python's import mechanism to build a dependency graph between modules, and assume that all runtime dependencies are in the data directory.

In my ideal scenario, a Bazel repository initialized from a Python wheel is "spread" out as defined in the spec, and then we run Gazelle to build a finer dependency graph. Gazelle's Python extension currently does not do some of these things correctly.

1. Python modules can have circular dependencies, and because all packages are modules, packages may also have circular dependencies. Although I would imagine that this is not very common outside the context of a single package.
2. Within a single package, all modules implicitly depend on __init__.py and __init__.py may also depend on some of these modules. The transitive closure of the modules that __init__.py needs, forms the default py_library target of the package. All modules outside of this transitive closure get their own py_library target that depends on the default target. Note that these other targets can also have transitive closures of their own because of circular imports. We can have a Gazelle option to not decompose explicitly specified packages if it does not get them right by itself.
3. Any non-Python files in a package may be part of a data filegroup and all modules in the package can depend on it. This filegroup can also include subdirectories if there are no Python files inside them.
4. Initialization of a package depends on its parent package's initialization (see import mechanism for regular packages). So the default target of a package also depends on the parent's default target.
5. Every target in a package may also depend on the data files in the wheel from which the package was obtained. But this is optional.

Short of running Gazelle inside an unpacked wheel (with or without my suggested changes), I think a composite wheel is the only way forward to keep everything functional (or you could just omit dependencies because you know you won't use that code). For example, if we just descended into the immediate top-level subdirectories inside a wheel and defined py_library targets there, we would still have cycles in the Bazel action graph. To verify, see that torch._inductor imports triton and triton imports torch, so having torch and triton trees as separate py_library targets won't be enough.

[siddharthab]

Edited the last paragraph on my previous comment.

[groodt]

Cyclic dependencies are invalid / unsupported in the rest of the PyPA ecosystem as well. It’s just that pip has historically not enforced it. Poetry, conda and others are a bit stricter.

The upstream issue with OpenAI/triton has been fixed triton-lang/triton#1374

[siddharthab]

In the issue you linked above, triton has moved torch to be a test dependency in its metadata file, but it still maintains imports to torch in its modules, and likewise, torch continues to maintain imports to triton in its modules.

Because this is valid Python, this will continue to work as long as the user installs both torch and triton, manually installing the other if needed. This means that Bazel users will also have to manually specify the other dependency in their BUILD file. This solution just shifts the burden to the users, which may be OK as a special case here. I suppose there are some other packages which do the same.

Thanks for looking into this everyone. I think I am OK with the solution (triton deleting the link in its metadata) for now.

[ccutch]

Is this being maintained or is this being ignored? It is clear that including both dependencies does not work and if this rules package is no longer being maintained what is the correct package?

[siddharthab]

The idea here is that it will work once the torch and triton wheels have been updated to remove the cyclic dependency. I am not sure what the plan there is, we may have to wait a few more days.

[ccutch]

I see. I have found a solution in the mean time of running pip at runtime with

import pip

pip.main(['install', 'torch'])