jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.49k stars 2.8k forks source link

Please provide PEP 503 compliant indices for CUDA versions of packages #5410

Open Froskekongen opened 3 years ago

Froskekongen commented 3 years ago

At https://storage.googleapis.com/jax-releases/jax_releases.html, we can find cuda versions of jaxlib. This index is not PEP 503-compliant, and can therefore not be used in a package manager such as poetry.

Can you please provide an index that respects the requirements of PEP 503?

skye commented 3 years ago

Are you able to say more about what's not compliant about it, or what kind of error you're getting? I can believe we messed something up here, but it's a little hard to pinpoint exactly what, especially since the index works with pip.

hawkinsp commented 3 years ago

We didn't necessarily intend that our repository be PEP 503 compliant, just that pip understands it.

To be PEP 503 compliant this I think we'd need to serve our wheels a different way. e.g. "Within a repository, the root URL (/ for this PEP which represents the base URL) MUST be a valid HTML5 page with a single anchor element per project in the repository." is something that cannot be satisfied by a raw GCS bucket. This is certainly possible, just a bit of work to move our web hosting around.

Froskekongen commented 3 years ago

To be able to use gpu versions with package managers in a reasonable way, I think using a format that is compatible with pypi would be greatly appreciated. Here is an example where a third party provided something that is compatible for pytorch: https://vxlabs.com/pypi/

Having a simple page with this structure that refers to the already existing wheels should be enough, no?

I guess being pypi compatible is not necessarily the same as being PEP 503 compliant.

skye commented 3 years ago

@hawkinsp do you have any ideas as to where we could (easily) host such a page?

I'm gonna mark this contributions welcome for now, since it sounds like anyone can set up such a page and link to the wheels in our GCS bucket.

JanLuca commented 3 years ago

You could just host a GitHub pages with the corresponding PEP 503 structure. I just wrote a script to generate a index.html for this purpose:

#!/usr/bin/env python3

import re
import requests
import textwrap

TEMPLATE = """<!DOCTYPE html>
<html>
  <head>
    <meta name="pypi:repository-version" content="1.0">
    <title>Links for jaxlib</title>
  </head>
  <body>
    <h1>Links for jaxlib</h1>
{links}
    </body>
</html>
"""

TEMPLATE_HREF = '<a href="{url}">{name}</a>'

re_a_href = re.compile('<a[^>]*?href="(?P<url>[^"]*?)"[^>]*?>(?P<name>[^<]*?)</a>')

def main():
    r = requests.get("https://storage.googleapis.com/jax-releases/jax_releases.html")

    result = []

    for m in re_a_href.finditer(r.text):
        name = m.group("name").split("/")[-1]

        result.append(TEMPLATE_HREF.format(name=name, url=m.group("url")))

    links = textwrap.indent("\n".join(result), "    ")

    index = TEMPLATE.format(links=links)

    with open("index.html", "w") as f:
        f.write(index)

if __name__ == "__main__":
    main()

This can be used to update a simple pip index interface, I just tested it here: https://janluca.github.io/pypi/ The source code can be found in https://github.com/JanLuca/janluca.github.io

Would it be possible to just create some jax GitHub page in the organization for this purpose?

JanLuca commented 3 years ago

@skye Would it be possible to have such a GitHub page hosted along the jax repository?

skye commented 3 years ago

I thought GitHub pages were only per-user/organization, but it looks like you can have per-project pages as well. So yes, this does seem like it should be possible. I can look into this, thanks for providing the example script!

jorenham commented 2 years ago

I have created a daily-rebuilding repository for this, see https://github.com/jorenham/jax_pep503. With Poetry, you can use it like:

[tool.poetry.dependencies]
...
jax = {extras = ["cuda"], version = "^0.3.13"}

[[tool.poetry.source]]
name = "jorenham/jax_pep503"
url = "https://jorenham.github.io/jax_pep503/"
secondary = true
CactusPuppy commented 2 years ago

With pypa/pip#10825, the lack of a <!DOCTYPE html> tag will potentially start to prevent pip from understanding the index page. Please consider adding a DOCTYPE tag before the html tag to ensure operability with future versions of pip.

abn commented 2 years ago

Just a drive by comment on this coming from https://github.com/python-poetry/poetry/issues/5516. We have implmented single page support in https://github.com/python-poetry/poetry/pull/5517. However, for tools like Poetry, since metadata is important when resolving a project's universal dependencies, having package checksums (described in PEP 503) and PEP 658 support enabled would considerably help the dependency resolution and lockfile generation.