79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
"""Shared helpers for the Modal launcher scripts in this directory.
|
|
|
|
Standalone module: the launchers import it by file path (sys.path trick), not
|
|
through the ``fastvideo`` package, so they keep working on thin CI hosts that
|
|
have ``modal`` but not torch.
|
|
"""
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
|
|
_REGISTRY = "ghcr.io"
|
|
_ACCEPT_MANIFEST_TYPES = ", ".join([
|
|
"application/vnd.oci.image.index.v1+json",
|
|
"application/vnd.docker.distribution.manifest.list.v2+json",
|
|
"application/vnd.oci.image.manifest.v1+json",
|
|
"application/vnd.docker.distribution.manifest.v2+json",
|
|
])
|
|
|
|
|
|
def resolve_image_ref(image_ref: str) -> str:
|
|
"""Pin a mutable ghcr.io tag to its current digest (``tag@sha256:...``).
|
|
|
|
Modal caches registry images by reference string and never re-pulls a
|
|
moved tag, so a mutable tag like ``py3.12-latest`` silently serves stale
|
|
bytes after a new image is published. Resolving the tag to its digest at
|
|
launch time makes Modal's cache key track the actual image: a new publish
|
|
changes the digest (fresh pull), an unchanged image keeps the cache hit.
|
|
|
|
Falls back to the unresolved tag, with a warning, if the registry cannot
|
|
be reached (e.g. offline) — that preserves today's behavior.
|
|
"""
|
|
if "@" in image_ref or not image_ref.startswith(f"{_REGISTRY}/"):
|
|
return image_ref
|
|
repo_path, _, tag = image_ref[len(_REGISTRY) + 1:].partition(":")
|
|
if not tag:
|
|
return image_ref
|
|
try:
|
|
token_url = f"https://{_REGISTRY}/token?scope=repository:{repo_path}:pull"
|
|
with urllib.request.urlopen(token_url, timeout=10) as response:
|
|
token = json.load(response)["token"]
|
|
manifest_request = urllib.request.Request(
|
|
f"https://{_REGISTRY}/v2/{repo_path}/manifests/{tag}",
|
|
method="HEAD",
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Accept": _ACCEPT_MANIFEST_TYPES,
|
|
},
|
|
)
|
|
with urllib.request.urlopen(manifest_request, timeout=10) as response:
|
|
digest = response.headers.get("Docker-Content-Digest", "")
|
|
if digest.startswith("sha256:"):
|
|
return f"{image_ref}@{digest}"
|
|
print(f"WARNING: registry returned no digest for {image_ref}; "
|
|
"Modal may reuse a stale cached image for this tag.")
|
|
except Exception as error: # noqa: BLE001 - never fail the launch over this
|
|
print(f"WARNING: could not resolve {image_ref} to a digest ({error}); "
|
|
"Modal may reuse a stale cached image for this tag.")
|
|
return image_ref
|
|
|
|
|
|
def resolve_uv_torch_backend(image_tag: str) -> str | None:
|
|
"""UV_TORCH_BACKEND for a launcher image tag.
|
|
|
|
A caller-set UV_TORCH_BACKEND always wins. Otherwise sniff the CUDA
|
|
version from an explicit image tag (cuda13 -> cu130, cuda12.6 -> cu126)
|
|
so uv resolves torch against the image's toolkit. Mutable tags (e.g.
|
|
py3.12-latest) return None and inherit the registry image's baked
|
|
backend, which keeps a latest-tag CUDA transition safe.
|
|
"""
|
|
override = os.environ.get("UV_TORCH_BACKEND")
|
|
if override:
|
|
return override
|
|
tag = image_tag.lower()
|
|
if "cuda13" in tag:
|
|
return "cu130"
|
|
if "cuda12.6" in tag:
|
|
return "cu126"
|
|
return None
|