Skip to content

SIZE_PROHIBITIVE_PACKAGES unconditionally strips torchvision, but it's neither present on all GPU base images nor actually large #353

Description

@orincolor

Summary

torchvision is currently hardcoded into SIZE_PROHIBITIVE_PACKAGES (cli/commands/build.py, lines 92–98), with the comment:

"~50 MB, requires torch"

and the assumption that it is already provided by the GPU base images (runpod/pytorch:*).

There are two issues with this approach.

1. The base image assumption is incorrect

On the GPU worker image provisioned for a GpuType.NVIDIA_GEFORCE_RTX_4090 endpoint, torch is installed, but torchvision is not.

However, any user-specified dependency such as:

dependencies=["torchvision"]

is silently discarded because of should_exclude_package() together with the unconditional merge of SIZE_PROHIBITIVE_PACKAGES (around line 260).

As a result, there is currently no way to bundle torchvision, even when it is explicitly requested.

Deployments that depend on it fail at runtime with an ImportError (for example, during trust_remote_code model loading), with no indication that the dependency was removed during the build.

2. The size justification is no longer accurate

The comment claims that torchvision is approximately 50 MB, but the current wheel:

torchvision-0.28.0-cp312-cp312-manylinux_2_28_x86_64.whl

is only 7.3 MB, according to the PyPI package metadata.

Including it would have a negligible impact on build or upload times.

Reproduction

Deploy an endpoint similar to:

@Endpoint(
    gpu=GpuType.NVIDIA_GEFORCE_RTX_4090,
    dependencies=["torchvision"],
)

and import torchvision in your application.

The deployment succeeds, but torchvision is missing at runtime despite being explicitly listed as a dependency.

Expected behavior

If a dependency is explicitly listed by the user, it should either:

  • be included in the deployment, or
  • fail the build with a clear explanation that it was intentionally excluded.

It should not be silently removed.

Suggested fixes

  • Remove torchvision from SIZE_PROHIBITIVE_PACKAGES, since its current wheel size no longer justifies automatic exclusion.
  • If the intention is to rely on GPU base images, verify that torchvision is actually installed on those images before excluding it.
  • Emit a build-time warning (or error) whenever a user-specified dependency is automatically excluded, rather than silently stripping it. This would make debugging significantly easier.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions