[Python] Basic dependency management with pyproject and uv #6

rolfmorel · 2025-10-14T16:55:27Z

Sets us up to do a uv sync to set up the current Python environment with the correct dependencies. Note that this is just a facilitating mechanism for preparing the python environment; users are still free to set up the environment via any means they like.

By default the only dependency installed is the Python bindings (from llvm/eudsl) -- that is, enabling import mlir in Python. This is the common dependency for generating and running pipelines and for ingress. Optionally one can pass (a variant of) --extra ingress-torch to uv sync to also get the torch and torch-mlir dependencies -- that is, enabling import torch and import torch_mlir in Python. Per the included docs, it is recommended to first set up a fresh virtualenv, e.g. via uv venv.

The pyproject.toml comes with uv-specific sections so as to enable specifying from which "indices", i.e. Python package repositories, each specific package is expected to come from. This file pins particular nightlies/versions from these package repositories.

makslevental

I'm a simple man: I see red, I approve

dchigarev · 2025-10-17T07:34:48Z

@rolfmorel thanks for the PR! I'm good with the changes

banach-space

Makes sense, thanks!

Some minor comments/questions inline.

banach-space · 2025-10-17T10:21:56Z

pyproject.toml

+[project.optional-dependencies]
+ingress_torch_cpu = [
+    "torch==v2.8.0+cpu",
+    "torch-mlir==v20251004.590"


Where does this version come from? I've tried here, but no luck:

https://github.com/llvm/torch-mlir/releases?q=20251004.590&expanded=true

Presumably that's from https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels? Who manages that?

Indeed, fair point! The https://github.com/llvm/torch-mlir-release repo does state this in its README:

This repository houses release automation for torch-mlir. It is isolated from the main repository so that we can more strictly control people with write access.

As for "who manages that?" -- the contribution history (e.g. https://github.com/llvm/torch-mlir-release/graphs/contributors) shows @marbre, @sjain-stanford, @stellaraccident, and @sahas3. Maybe any of them can chime in on regarding who's managing torch-mlir's releases (in particular from this dedicated repo).

As for "who manages that?" -- the contribution history (e.g. https://github.com/llvm/torch-mlir-release/graphs/contributors) shows @marbre, @sjain-stanford, @stellaraccident, and @sahas3. Maybe any of them can chime in on regarding who's managing torch-mlir's releases (in particular from this dedicated repo).

I contributed to the project as I had a need for wheels earlier this year (and those where basically broken) but I am not actively managing releases for torch-mlir.

Thanks for the quick response, @marbre.

Could you let us know who to reach out to about these packages, e.g. regarding them inevitably having issues? For example the person who gave you access to repo or anybody else who is taking responsibility for this releases repo/torch-mlir's releases.

@sjarus, could you maybe comment on the infrastructure you have going that makes use of torch-mlir. In this particular case regarding getting the torch-mlir dependency to be available. Do you just build it yourself?

Not an official maintainer of the torch-mlir-release repo as well but I've vested interest in ensuring that the wheels are released properly from that repo as we do use a forked repo to build the python wheels with stable torch release to use in our work.

The main repo currently builds against nightly torch version which is pinned in https://github.com/llvm/torch-mlir/blob/main/pytorch-requirements.txt

banach-space · 2025-10-17T10:34:27Z

pyproject.toml

+    "torch-mlir==v20251004.590"
+]
+ingress_torch = [
+    "torch==v2.8.0",  # Nvidia-enabled version of torch


[nit] IIUC, v is not needed.

Suggested change

"torch==v2.8.0", # Nvidia-enabled version of torch

"torch==2.8.0", # Nvidia-enabled version of torch

Drive-by comment without looking deep into what is done here, so excuse if this seems noisy: You may want to consider the CPU-only version? I some of our projects we explicitly skipped setting a strict dependency on torch and leave it to the user to pick a suitable version (CPU-only, Nvidia-enabled, ROCm-enabled). This has among others the benefit of being able to use the way smaller CPU-ony packages in the CI.

Yep, that's correct. The leading vs were extraneous. They're gone now 👍

Regarding cpu-only: there's no default here, only different optional alternatives made available through an optional package management facilitator. Your point regarding the cpu-only version being lighter weight is well-noted though. I have moved the note that there are other versions to a collapsed section in the README.

pyproject.toml

banach-space · 2025-10-17T10:44:40Z

pyproject.toml

+pytorch_triton_xpu = { index = "pytorch" }
+pytorch_triton_rocm = { index = "pytorch" }


Why Triton? As in, what's the relevance here?

When using uv's explict = true flag, uv's dependency constraint solver needs to be told where to find the transitive dependencies (of torch). I think another option is to remove the explicit = true flags below, though that has the effect that we don't know from which non-Pypi index/package repository a dependency is going to be coming from (in effect, explicit = true makes it so that "dependencies that come from this repo need to explicitly state they will come from this repo").

banach-space · 2025-10-17T10:50:38Z

README.md

+$ uv sync --extra ingress-torch-cpu  # Optionally install the dependencies for torch ingress
+```
+
+For vendor-specific versions of `torch` use the targets `ingress-torch`, `ingress-torch-rocm` or `ingress-torch-xpu` for Nvidia, AMD, and Intel-enabled versions, respectively.


I have two points on "vendor specific" stuff:

Nothing against Nvidia, but why would we use the default name,ingress-torch, for Nvidia? Shouldn't that be e.g. ingress-torch-nvidia?

Will you be testing all the vendor specific stuff? Making this logic "generic" is a worthy goal, but if its not tested ...

I took the naming to be reflective of 1) the directory structure, i.e. for the ingress subfolder, in which you might like 2) the dependencies associated to the python packages with identifiers torch, or torch+cpu or torch+rocm (or actually, torch+rocm6.4) or torch+xpu, respectively (with the first supporting Nvidia by default -- I wonder why 🙃 ). I don't have a strong opinion on the naming of these targets (though they aren't allowed to contain "+"), so if you have an idea for a better scheme, do let me know!

I think for now we will indeed not bother with testing with all the different vendor-specific dependencies (at least not until these vendors themselves contribute to the repo). This is an argument for not bothering with adding these different targets now. I mostly added them so that we would have feature parity with the two install-virtualenv.sh scripts that I am removing in this PR. If you and @rengolin are fine with just an ingress-torch-cpu target for now, I can push a new commit removing the other targets and leave the current commit, showing how to get uv to play nice with the others, in this PR's history so we could revisit later. Just let me know!

I have now updated to the following pragmatic solution: the torch version for Nvidia now has a suffix -nvidia. As pointed out up above, the -cpu version is a lot lighter weight and hence will be the one the docs direct towards in the first instance (there's now a collapsed section explaining you can use a different target to obtain a different torch).

With regards to these targets getting tested: in the spirit of inclusivity, lets keep the different targets and take the position that them getting tested/integrated in CI is a "patches welcome"-type of situation.

Let me rephrase a bit but please excuse that I am lacking context of on the overall PR. What we, and others, do is to not introduce a dependency on torch via the package management system at all. Instead we have

https://github.com/nod-ai/shark-ai/blob/59259e974e02f25433df480cd52328b56c033752/sharktank/sharktank/__init__.py#L9-L12

to let the users decide and make them responsible for picking an appropriate version.

rolfmorel · 2025-10-17T14:06:00Z

Thanks for the review, @banach-space !

I hope your comments are sufficiently addressed for us to go ahead and merge. Do let me know if you feel otherwise!

banach-space · 2025-10-20T09:54:33Z

I hope your comments are sufficiently addressed for us to go ahead and merge. Do let me know if you feel otherwise!

Yes, thanks for the replies, that clarified the context 🙏🏻 Feel free to land this, my other comments are non-blockers.

rolfmorel added 2 commits October 14, 2025 09:46

Basic dependency management with pyproject and uv

bc3c6d2

Clean-up mlir-gen

e2accdc

rolfmorel requested review from Groverkss, banach-space, makslevental and rengolin October 15, 2025 17:58

rengolin approved these changes Oct 15, 2025

View reviewed changes

makslevental approved these changes Oct 15, 2025

View reviewed changes

banach-space approved these changes Oct 17, 2025

View reviewed changes

A number of small fixes prompted by Andrzej's review

3b18487

rolfmorel force-pushed the users/rolfmorel/pyproject branch from ac703fc to 3b18487 Compare October 17, 2025 13:48

rolfmorel merged commit 8d6bed1 into main Oct 20, 2025

dchigarev mentioned this pull request Oct 20, 2025

[ingress][torch-mlir][RFC] Initial version of fx-importer script using torch-mlir #4

Open

	"torch==v2.8.0", # Nvidia-enabled version of torch
	"torch==2.8.0", # Nvidia-enabled version of torch

		pytorch_triton_xpu = { index = "pytorch" }
		pytorch_triton_rocm = { index = "pytorch" }

[Python] Basic dependency management with pyproject and uv #6

[Python] Basic dependency management with pyproject and uv #6

Uh oh!

Conversation

rolfmorel commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

makslevental left a comment

Choose a reason for hiding this comment

Uh oh!

dchigarev commented Oct 17, 2025

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel commented Oct 17, 2025

Uh oh!

banach-space commented Oct 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

rolfmorel commented Oct 14, 2025 •

edited

Loading

rolfmorel Oct 17, 2025 •

edited

Loading

rolfmorel Oct 17, 2025 •

edited

Loading

rolfmorel Oct 17, 2025 •

edited

Loading

rolfmorel Oct 17, 2025 •

edited

Loading