pants + torch woes

3rdparty BUILD

TORCH_VERSION = "2.3.1"

CUDA_VERSION = "cu118"

TORCH_CUDA_WHEEL_URL = "https://download.pytorch.org/whl/cu118/torch-2.3.1%2Bcu118-cp311-cp311-linux_x86_64.whl#sha256=5b0d531814886573cbe8c8ca91d17676f96bbaa33b569dd37ea235f124314e97"

TORCH_LINUX_CPU_WHEEL_URL = "https://download.pytorch.org/whl/cpu/torch-2.3.1%2Bcpu-cp311-cp311-linux_x86_64.whl#sha256=a272defe305dbd944aa28a91cc3db0f0149495b3ebec2e39723a7224fa05dc57"



python_requirement(
    name = "torch_manual",
    modules = ["_dummy_torch"],
    **parametrize(
        "python_linux_cuda",
        requirements = ["torch @ {}".format(TORCH_CUDA_WHEEL_URL)],
        resolve = "python_linux_cuda",
    ),
    **parametrize(
        "python_linux",
        requirements = ["torch @ {}".format(TORCH_LINUX_CPU_WHEEL_URL)],
        resolve = "python_linux",
    ),
    **parametrize(
        "python_darwin",
        requirements = [
            "torch=={v},!={v}+cpu,!={v}+{cuda}".format(cuda = CUDA_VERSION, v = TORCH_VERSION),
        ],
        resolve = "python_darwin",
    )
)


python_requirements(
    name = "3rdparty",
    overrides = {
        "pex": {"dependencies": [":3rdparty#setuptools"]},
        "sqlalchemy": {"dependencies": [":3rdparty#psycopg2-binary"]},
        "torch": {"dependencies": [":torch_manual"]},
    },
    source = "pyproject.toml",
    **parametrize(
        "python_linux_cuda",
        resolve = "python_linux_cuda",
    ),
    **parametrize(
        "python_linux",
        resolve = "python_linux",
    ),
    **parametrize(
        "python_darwin",
        resolve = "python_darwin",
    )
)

BUILD file for the repo root

__defaults__(
    {
        docker_image: dict(build_platform=["linux/amd64"]),
        (python_sources, python_source, python_tests, python_test, pex_binary, pex_binaries): {
            **parametrize(
                "python_linux_cuda",
                resolve="python_linux_cuda",
            ),
            **parametrize(
                "python_linux",
                resolve="python_linux",
            ),
            **parametrize(
                "python_darwin",
                resolve="python_darwin",
            ),
        },
    },
)

local_environment(
    name="local_linux",
    compatible_platforms=["linux_x86_64"],
    fallback_environment="docker_linux",

)

local_environment(
    name="local_darwin",
    compatible_platforms=["macos_arm64"],
    fallback_environment="docker_linux",
    subprocess_environment_env_vars=["PANTS_PYTHON_RESOLVES=python_darwin"],  # pragma: allowlist secret
)

docker_environment(
    name="docker_linux",
    docker_env_vars=["PANTS_PYTHON_RESOLVES=python_linux"],  # pragma: allowlist secret
    image="mzdotai/golden:base_latest",
    platform="linux_x86_64",
    python_bootstrap_search_path=["<PATH>"]
)

docker_environment(
    name="docker_linux_cuda",
    docker_env_vars=["PANTS_PYTHON_RESOLVES=python_linux_cuda"],  # pragma: allowlist secret
    image="mzdotai/golden:base_gpu_latest",
    platform="linux_x86_64",
    python_bootstrap_search_path=["<PATH>"]
)

BUILD file with 'normal' dep specification

python_sources(
    name = "backend_sources",
    sources = [
        "**/*.py",
        "!tests",
    ],
)

pex_binary(
    name = "backend_app",
    dependencies = [":backend_sources"],  ## EXPLICIT DEPS INCLUDED
    entry_point = "backend",
    execution_mode = "venv",
    include_requirements = True,
    include_sources = True,
    **parametrize( "python_linux_cuda", complete_platforms = [ "//3rdparty/python:py311_linux_pex_platform_tags", ],
        environment = "docker_linux_cuda", resolve = "python_linux_cuda", ),
    **parametrize( "python_linux", complete_platforms = [ "//3rdparty/python:py311_linux_pex_platform_tags", ],
        environment = "docker_linux", resolve = "python_linux", ),
    **parametrize( "python_darwin", complete_platforms = [ "//3rdparty/python:py311_macos_14_pex_platform_tags", ],
        environment = "local_darwin", resolve = "python_darwin", )
)

BUILD file with explicitly parametrized deps

pex_binary(
    name = "backend_app",
    entry_point = "mzai.backend.main",
    execution_mode = "venv",
    include_requirements = True,
    include_sources = True,
    **parametrize(
        "python_linux_cuda",
        complete_platforms = ["//3rdparty/python:py311_linux_pex_platform_tags"],
        dependencies = [":backend_sources@parametrize=python_linux_cuda"],
        environment = "docker_linux_cuda", resolve = "python_linux_cuda",
    ),
    **parametrize(
        "python_linux",
        complete_platforms = [
            "//3rdparty/python:py311_linux_pex_platform_tags",
        ],
        dependencies = [":backend_sources@parametrize=python_linux"],
        environment = "docker_linux", resolve = "python_linux",
    ),
    **parametrize(
        "python_darwin",
        complete_platforms = ["//3rdparty/python:py311_macos_14_pex_platform_tags",],
        dependencies = [":backend_sources@parametrize=python_darwin"],
        environment = "local_darwin", resolve = "python_darwin",
    )
)

notes.md

the project BUILD results in pants not quite getting the parametrization correct. if you look at the transitive deps of backend_sources, it's accurate: - python_darwin, etc are correct.

pants dependencies --transitive --dependencies-format=json lumigator/python/mzai/backend:backend_sources | jq .
{
  "lumigator/python/mzai/backend:backend_sources@parametrize=python_darwin": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_darwin",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_darwin",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_darwin",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_darwin",
    "3rdparty/python:3rdparty#torch@parametrize=python_darwin",
    "3rdparty/python:torch_manual@parametrize=python_darwin",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_darwin",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_darwin",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_darwin"
  ],
  "lumigator/python/mzai/backend:backend_sources@parametrize=python_linux": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux.lock:_python_linux_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux",
    "3rdparty/python:torch_manual@parametrize=python_linux",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux"
  ],
  "lumigator/python/mzai/backend:backend_sources@parametrize=python_linux_cuda": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux_cuda",
    "3rdparty/python:torch_manual@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux_cuda"
  ]
}

Transitive deps of the pex, which is somehow pulling in the python_linux_cuda deps to both python_darwin and python_linux

pants dependencies --transitive --dependencies-format=json lumigator/python/mzai/backend:backend_app | jq .
{
  "lumigator/python/mzai/backend:backend_app@parametrize=python_darwin": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux_cuda",
    "3rdparty/python:py311_macos_14_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_darwin",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux_cuda"
  ],
  "lumigator/python/mzai/backend:backend_app@parametrize=python_linux": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux_cuda",
    "3rdparty/python:py311_linux_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux_cuda"
  ],
  "lumigator/python/mzai/backend:backend_app@parametrize=python_linux_cuda": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux_cuda",
    "3rdparty/python:py311_linux_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux_cuda"
  ]
}

When run with the explicitly laid out deps in the uupdated build file it works:

pants dependencies --transitive --dependencies-format=json lumigator/python/mzai/backend:backend_app | jq .
{
  "lumigator/python/mzai/backend:backend_app@parametrize=python_darwin": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_darwin",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_darwin",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_darwin",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_darwin",
    "3rdparty/python:3rdparty#torch@parametrize=python_darwin",
    "3rdparty/python:py311_macos_14_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_darwin",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_darwin",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_darwin",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_darwin"
  ],
  "lumigator/python/mzai/backend:backend_app@parametrize=python_linux": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_darwin.lock:_python_darwin_lockfile",
    "3rdparty/python/python_linux.lock:_python_linux_lockfile",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux",
    "3rdparty/python:py311_linux_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_linux",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux"
  ],
  "lumigator/python/mzai/backend:backend_app@parametrize=python_linux_cuda": [
    "3rdparty/python/pyproject.toml:3rdparty",
    "3rdparty/python/python_linux_cuda.lock:_python_linux_cuda_lockfile",
    "3rdparty/python:3rdparty#psycopg2-binary@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic-settings@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#pydantic@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#sqlalchemy@parametrize=python_linux_cuda",
    "3rdparty/python:3rdparty#torch@parametrize=python_linux_cuda",
    "3rdparty/python:py311_linux_pex_platform_tags",
    "3rdparty/python:torch_manual@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/__init__.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/main.py:backend_sources@parametrize=python_linux_cuda",
    "lumigator/python/mzai/backend/settings.py:backend_sources@parametrize=python_linux_cuda"
  ]
}

pants.toml

[GLOBAL]
pants_version = "2.21.0"
colors = true

backend_packages = [
  "pants.backend.docker",
  "pants.backend.shell",
  "pants.backend.python",
  "pants.backend.shell.lint.shfmt",
  "pants.backend.shell.lint.shellcheck",
  "pants.backend.experimental.python.lint.ruff.check",
  "pants.backend.experimental.python.lint.ruff.format",
  "pants.backend.build_files.fmt.ruff",
  "pants.backend.experimental.helm"
]

[environments-preview.names]
local_darwin = "//:local_darwin"
local_linux = "//:local_linux"
docker_linux = "//:docker_linux"
docker_linux_cuda = "//:docker_linux_cuda"


[cli.alias]
--debug-build = "--print-stacktrace -ldebug --keep-sandboxes=on_failure"

[generate-lockfiles]
diff = true


[source]
# Defines the "roots" for subprojects. anything with folders like
# //**/[python|docker|k8s|infra]/<name>/
# Will be importable as `<name>.whatever` in python/etc.
root_patterns = ["python", "infra"]

[tailor]
build_file_name = "BUILD.pants"

[python]
interpreter_constraints = ['CPython==3.11.9']
# Enable the "resolves" mechanism, which turns on lockfiles for user code. See
# https://www.pantsbuild.org/docs/python-third-party-dependencies. 
# This also adds the `generate-lockfiles` goal for Pants to generate the lockfile for you.
enable_resolves = true
resolver_manylinux = "manylinux2014"
default_resolve = "python_darwin"

[python-repos]
# these are required for nvidia's cuda libs etc.
indexes.add = [
"https://pypi.ngc.nvidia.com",
]


find_links = ["%(env.HOME)s/workspace/wheelhouse"]
path_mappings = ["WHEELS_DIR|%(env.HOME)s/workspace/wheelhouse"]

[python.resolves]
python_darwin = "3rdparty/python/python_darwin.lock"
python_linux = "3rdparty/python/python_linux.lock"
python_linux_cuda = "3rdparty/python/python_linux_cuda.lock"


[python-bootstrap]
search_path = [
  "%(buildroot)s/.python/python3.11.9/python/install/bin",
  "<PYENV>",
  "<PATH>",
]

[python-infer]
use_rust_parser = true

[ruff]
config = "ruff.toml"

[docker]
use_buildx = true
#### This is a terrible hidden edge -
#### the `--output=type=docker` is required to make buildx, which
# is sandboxed, push the built image to the local docker registry.
# without it, it requires you to add an explict
# local registry to pants.
build_args = ["--output=type=docker", "REPOVERSION"]
env_vars = [
  "REPOVERSION=%(env.REPOVERSION)s",
  "DOCKER_BUILDKIT=1",
  "DOCKER_CONFIG=%(homedir)s/.docker"
]
tools = [
  "dirname",
  "readlink",
  "python3",
  "cut",
  "sed",
  "bash",
]


[pytest]
args = ["-v"]