-
Notifications
You must be signed in to change notification settings - Fork 6
feat: Python version-aware Docker image selection #261
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
6ff6fcd
57accb7
cea8fa4
1c85da4
e3f90cf
764413d
f3053c6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,15 +19,121 @@ def _endpoint_domain_from_base_url(base_url: str) -> str: | |
| ENDPOINT_DOMAIN = _endpoint_domain_from_base_url(runpod.endpoint_url_base) | ||
|
|
||
|
|
||
| # Python version support | ||
| SUPPORTED_PYTHON_VERSIONS: tuple[str, ...] = ("3.10", "3.11", "3.12") | ||
| GPU_PYTHON_VERSIONS: tuple[str, ...] = ("3.10", "3.11", "3.12") | ||
| CPU_PYTHON_VERSIONS: tuple[str, ...] = ("3.10", "3.11", "3.12") | ||
|
Comment on lines
+23
to
+25
|
||
| DEFAULT_PYTHON_VERSION: str = "3.11" | ||
|
|
||
|
|
||
| def local_python_version() -> str: | ||
| """Return the running interpreter's major.minor version string.""" | ||
| import sys | ||
|
|
||
| return f"{sys.version_info.major}.{sys.version_info.minor}" | ||
|
|
||
|
|
||
| # Image type to repository mapping | ||
| _IMAGE_REPOS: dict[str, str] = { | ||
| "gpu": "runpod/flash", | ||
| "cpu": "runpod/flash-cpu", | ||
| "lb": "runpod/flash-lb", | ||
| "lb-cpu": "runpod/flash-lb-cpu", | ||
| } | ||
|
|
||
| # Image types that require GPU-compatible Python versions | ||
| _GPU_IMAGE_TYPES: frozenset[str] = frozenset({"gpu", "lb"}) | ||
|
|
||
| # Image type to environment variable override mapping | ||
| _IMAGE_ENV_VARS: dict[str, str] = { | ||
| "gpu": "FLASH_GPU_IMAGE", | ||
| "cpu": "FLASH_CPU_IMAGE", | ||
| "lb": "FLASH_LB_IMAGE", | ||
| "lb-cpu": "FLASH_CPU_LB_IMAGE", | ||
| } | ||
|
|
||
|
|
||
| def validate_python_version(version: str) -> str: | ||
| """Validate that a Python version string is supported. | ||
|
|
||
| Args: | ||
| version: Python version string (e.g. "3.11"). | ||
|
|
||
| Returns: | ||
| The validated version string. | ||
|
|
||
| Raises: | ||
| ValueError: If version is not in SUPPORTED_PYTHON_VERSIONS. | ||
| """ | ||
| if version not in SUPPORTED_PYTHON_VERSIONS: | ||
| supported = ", ".join(SUPPORTED_PYTHON_VERSIONS) | ||
| raise ValueError( | ||
| f"Python {version} is not supported. Supported versions: {supported}" | ||
| ) | ||
| return version | ||
|
|
||
|
|
||
| def get_image_name( | ||
| image_type: str, | ||
| python_version: str, | ||
| *, | ||
| tag: str | None = None, | ||
| ) -> str: | ||
| """Resolve a versioned Docker image name for the given type and Python version. | ||
|
|
||
| Args: | ||
| image_type: One of 'gpu', 'cpu', 'lb', 'lb-cpu'. | ||
| python_version: Python version string (e.g. "3.11", "3.12"). | ||
| tag: Image tag suffix. Defaults to FLASH_IMAGE_TAG env var or "latest". | ||
|
|
||
| Returns: | ||
| Fully qualified image name, e.g. "runpod/flash:py3.12-latest". | ||
|
|
||
| Raises: | ||
| ValueError: If image_type is unknown, python_version is unsupported, | ||
| or a GPU image type is requested with a CPU-only Python version. | ||
| """ | ||
| if image_type not in _IMAGE_REPOS: | ||
| raise ValueError( | ||
| f"Unknown image type '{image_type}'. " | ||
| f"Valid types: {', '.join(sorted(_IMAGE_REPOS))}" | ||
| ) | ||
|
|
||
| # Environment variable override takes precedence, bypassing version validation | ||
| env_var = _IMAGE_ENV_VARS[image_type] | ||
| override = os.environ.get(env_var) | ||
| if override: | ||
| return override | ||
|
|
||
| validate_python_version(python_version) | ||
|
|
||
| if image_type in _GPU_IMAGE_TYPES and python_version not in GPU_PYTHON_VERSIONS: | ||
| gpu_versions = ", ".join(GPU_PYTHON_VERSIONS) | ||
| raise ValueError( | ||
| f"GPU endpoints require Python {gpu_versions}. Got Python {python_version}." | ||
| ) | ||
|
|
||
| resolved_tag = tag or os.environ.get("FLASH_IMAGE_TAG", "latest") | ||
| repo = _IMAGE_REPOS[image_type] | ||
| return f"{repo}:py{python_version}-{resolved_tag}" | ||
|
|
||
|
|
||
| # Docker image configuration | ||
| FLASH_IMAGE_TAG = os.environ.get("FLASH_IMAGE_TAG", "latest") | ||
| _RESOLVED_TAG = FLASH_IMAGE_TAG | ||
|
|
||
| FLASH_GPU_IMAGE = os.environ.get("FLASH_GPU_IMAGE", f"runpod/flash:{_RESOLVED_TAG}") | ||
| FLASH_CPU_IMAGE = os.environ.get("FLASH_CPU_IMAGE", f"runpod/flash-cpu:{_RESOLVED_TAG}") | ||
| FLASH_LB_IMAGE = os.environ.get("FLASH_LB_IMAGE", f"runpod/flash-lb:{_RESOLVED_TAG}") | ||
| FLASH_GPU_IMAGE = os.environ.get( | ||
| "FLASH_GPU_IMAGE", f"runpod/flash:py{DEFAULT_PYTHON_VERSION}-{_RESOLVED_TAG}" | ||
| ) | ||
| FLASH_CPU_IMAGE = os.environ.get( | ||
| "FLASH_CPU_IMAGE", f"runpod/flash-cpu:py{DEFAULT_PYTHON_VERSION}-{_RESOLVED_TAG}" | ||
| ) | ||
| FLASH_LB_IMAGE = os.environ.get( | ||
| "FLASH_LB_IMAGE", f"runpod/flash-lb:py{DEFAULT_PYTHON_VERSION}-{_RESOLVED_TAG}" | ||
| ) | ||
| FLASH_CPU_LB_IMAGE = os.environ.get( | ||
| "FLASH_CPU_LB_IMAGE", f"runpod/flash-lb-cpu:{_RESOLVED_TAG}" | ||
| "FLASH_CPU_LB_IMAGE", | ||
| f"runpod/flash-lb-cpu:py{DEFAULT_PYTHON_VERSION}-{_RESOLVED_TAG}", | ||
| ) | ||
|
|
||
| # Worker configuration defaults | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Narrowing requires-python to <3.13 will likely break the existing CI matrix entries that run on Python 3.13 and 3.14 (see .github/workflows/ci.yml). Either update the workflow matrix in this PR (preferred) or keep requires-python compatible with the versions CI is still running.