feat: custom docker image namespace

Signed-off-by: thxCode <[email protected]>

Author: thxCode

gpustack_runtime/deployer/__types__.py

@@ -1023,13 +1023,12 @@ def validate_and_default(self):
                 and envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY
                 not in ["docker.io", "index.docker.io"]
             ):
+                image_registry = envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_NAMESPACE
                 image_split = c.image.split("/")
                 if len(image_split) == 1:
-                    c.image = f"{envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY}/library/{c.image}"
+                    c.image = f"{image_registry}/library/{c.image}"
                 elif len(image_split) == 2:
-                    c.image = (
-                        f"{envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY}/{c.image}"
-                    )
+                    c.image = f"{image_registry}/{c.image}"
             # Correct runner image if needed.
             if envs.GPUSTACK_RUNTIME_DEPLOY_CORRECT_RUNNER_IMAGE:
                 c.image, ok = correct_runner_image(c.image)

gpustack_runtime/deployer/__utils__.py

@@ -651,3 +651,127 @@ def bytes_to_human_readable(size_in_bytes: int) -> str:
     if size_in_bytes >= _KiB:
         return f"{size_in_bytes / _KiB:.2f} KiB"
     return f"{size_in_bytes} B"
+
+
+def replace_image_with(
+    image: str,
+    registry: str | None = None,
+    namespace: str | None = None,
+    repository: str | None = None,
+) -> str:
+    """
+    Replace the registry, namespace, and repository of a Docker image string.
+
+    The given image string is parsed into its components (registry, namespace, repository, tag),
+    and the specified components are replaced with the provided values.
+
+    The format of a Docker image string is:
+    [registry/][namespace/]repository[:tag|@digest]
+
+    Args:
+        image:
+            The original Docker image string.
+        registry:
+            The new registry to use. If None, keep the original registry.
+        namespace:
+            The new namespace to use. If None, keep the original namespace.
+        repository:
+            The new repository to use. If None, keep the original repository.
+
+    Returns:
+        The modified Docker image string.
+
+    """
+    if not image or (not registry and not namespace and not repository):
+        return image
+
+    registry = registry.strip() if registry else None
+    namespace = namespace.strip() if namespace else None
+    repository = repository.strip() if repository else None
+
+    image_reg, image_ns, image_repo, image_tag = (
+        None,
+        None,
+        None,
+        None,
+    )
+    image_rest = image.strip()
+
+    # Get tag.
+    parts = image_rest.rsplit("@", maxsplit=1)
+    if len(parts) == 2:
+        image_rest, image_tag = parts
+    else:
+        parts = image_rest.rsplit(":", maxsplit=1)
+        if len(parts) == 2 and "/" not in parts[1]:
+            image_rest, image_tag = parts
+    if not image_rest:
+        return image
+
+    # Get repository.
+    parts = image_rest.rsplit("/", maxsplit=1)
+    if len(parts) == 2:
+        image_rest, image_repo = parts
+    else:
+        image_rest, image_repo = None, image_rest
+
+    # Get namespace.
+    if image_rest:
+        parts = image_rest.rsplit("/", maxsplit=1)
+        if len(parts) == 2:
+            image_reg, image_ns = parts
+        else:
+            image_reg, image_ns = None, image_rest
+
+    return make_image_with(
+        repository=repository or image_repo,
+        registry=registry or image_reg,
+        namespace=namespace or image_ns,
+        tag=image_tag,
+    )
+
+
+def make_image_with(
+    repository: str,
+    registry: str | None = None,
+    namespace: str | None = None,
+    tag: str | None = None,
+) -> str:
+    """
+    Make a Docker image string from the given registry, namespace, repository, and tag.
+
+    The format of a Docker image string is:
+    [registry/][namespace/]repository[:tag|@digest]
+
+    Args:
+        repository:
+            The repository name.
+        registry:
+            The registry to use. If None, no registry will be used.
+        namespace:
+            The namespace to use. If None, no namespace will be used.
+        tag:
+            The tag to use. If None, no tag will be used.
+
+    Returns:
+        The Docker image string.
+
+    """
+    if not repository or (not registry and not namespace and not tag):
+        return repository
+
+    image = ""
+    if registry:
+        image += f"{registry}/"
+    if namespace:
+        image += f"{namespace}/"
+    elif registry:
+        image += "library/"
+    image += repository
+    if not tag:
+        return image
+    if tag.startswith("sha256:"):
+        image += f"@{tag}"
+    else:
+        image += f":{tag}"
+    return image

gpustack_runtime/deployer/docker.py

@@ -44,7 +44,7 @@
     WorkloadStatusOperation,
     WorkloadStatusStateEnum,
 )
-from .__utils__ import _MiB, bytes_to_human_readable, safe_json
+from .__utils__ import _MiB, bytes_to_human_readable, replace_image_with, safe_json
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -138,6 +138,17 @@ def validate_and_default(self):
         # Default and validate in the base class.
         super().validate_and_default()
 
+        # Adjust default image namespace if needed.
+        if namespace := envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_NAMESPACE:
+            self.pause_image = replace_image_with(
+                image=self.pause_image,
+                namespace=namespace,
+            )
+            self.unhealthy_restart_image = replace_image_with(
+                image=self.unhealthy_restart_image,
+                namespace=namespace,
+            )
+
 
 @dataclass_json
 @dataclass

gpustack_runtime/envs.py

@@ -68,6 +68,11 @@
     If not set, it should be "docker.io".
     If the image name already contains a registry, this setting will be ignored.
     """
+    GPUSTACK_RUNTIME_DEPLOY_DEFAULT_NAMESPACE: str | None = None
+    """
+    Namespace for default runner images.
+    If not set, it should be "gpustack".
+    """
     GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY_USERNAME: str | None = None
     """
     Username for the default container registry.
@@ -282,7 +287,15 @@
         "Auto",
     ),
     "GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY": lambda: trim_str(
-        getenv("GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY"),
+        getenvs(
+            keys=[
+                "GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY",
+                "GPUSTACK_SYSTEM_DEFAULT_CONTAINER_REGISTRY",  # Compatible with gpustack/gpustack.
+            ],
+        ),
+    ),
+    "GPUSTACK_RUNTIME_DEPLOY_DEFAULT_NAMESPACE": lambda: trim_str(
+        getenv("GPUSTACK_RUNTIME_DEPLOY_DEFAULT_NAMESPACE"),
     ),
     "GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY_USERNAME": lambda: trim_str(
         getenv("GPUSTACK_RUNTIME_DEPLOY_DEFAULT_REGISTRY_USERNAME"),

tests/gpustack_runtime/deployer/fixtures/test_make_image_with.json

@@ -0,0 +1,69 @@
+[
+  [
+    "1: empty",
+    {
+      "repository": ""
+    },
+    ""
+  ],
+  [
+    "2: empty",
+    {
+      "repository": "nginx"
+    },
+    "nginx"
+  ],
+  [
+    "nginx:latest",
+    {
+      "repository": "nginx",
+      "tag": "latest"
+    },
+    "nginx:latest"
+  ],
+  [
+    "website.com/abc/library/nginx",
+    {
+      "repository": "nginx",
+      "registry": "website.com/abc"
+    },
+    "website.com/abc/library/nginx"
+  ],
+  [
+    "xyz/nginx",
+    {
+      "repository": "nginx",
+      "namespace": "xyz"
+    },
+    "xyz/nginx"
+  ],
+  [
+    "1: website.com/abc/xyz/nginx:v0.1.0",
+    {
+      "repository": "nginx",
+      "registry": "website.com/abc",
+      "namespace": "xyz",
+      "tag": "v0.1.0"
+    },
+    "website.com/abc/xyz/nginx:v0.1.0"
+  ],
+  [
+    "2: website.com/abc/xyz/nginx:v0.1.0",
+    {
+      "repository": "nginx",
+      "registry": "website.com",
+      "namespace": "abc/xyz",
+      "tag": "v0.1.0"
+    },
+    "website.com/abc/xyz/nginx:v0.1.0"
+  ],
+  [
+    "xyz/nginx@sha256:abcdef1234567890",
+    {
+      "repository": "nginx",
+      "namespace": "xyz",
+      "tag": "sha256:abcdef1234567890"
+    },
+    "xyz/nginx@sha256:abcdef1234567890"
+  ]
+]

tests/gpustack_runtime/deployer/fixtures/test_replace_image_with.json

@@ -0,0 +1,58 @@
+[
+  [
+    "1: empty",
+    {
+      "image": ""
+    },
+    ""
+  ],
+  [
+    "2: empty",
+    {
+      "image": "nginx"
+    },
+    "nginx"
+  ],
+  [
+    "nginx",
+    {
+      "image": "nginx",
+      "registry": "quay.io"
+    },
+    "quay.io/library/nginx"
+  ],
+  [
+    "nginxx:latest",
+    {
+      "image": "nginxx:latest",
+      "registry": "quay.io",
+      "repository": "nginx"
+    },
+    "quay.io/library/nginx:latest"
+  ],
+  [
+    "xyz/nginx@sha256:abcdef1234567890",
+    {
+      "image": "xyz/nginx@sha256:abcdef1234567890",
+      "namespace": "lmn"
+    },
+    "lmn/nginx@sha256:abcdef1234567890"
+  ],
+  [
+    "docker.io/library/redis:6.0",
+    {
+      "image": "docker.io/library/redis:6.0",
+      "registry": "myregistry.com/mysuperrepo",
+      "repository": "myredis"
+    },
+    "myregistry.com/mysuperrepo/library/myredis:6.0"
+  ],
+  [
+    "docker.io/abc/xyz/redis:6.0",
+    {
+      "image": "docker.io/abc/xyz/redis:6.0",
+      "namespace": "lmn"
+    },
+    "docker.io/abc/lmn/redis:6.0"
+  ]
+]

tests/gpustack_runtime/deployer/test_utils.py

@@ -1,7 +1,12 @@
 import pytest
 from fixtures import load
 
-from gpustack_runtime.deployer.__utils__ import compare_versions, correct_runner_image
+from gpustack_runtime.deployer.__utils__ import (
+    compare_versions,
+    correct_runner_image,
+    make_image_with,
+    replace_image_with,
+)
 
 
 @pytest.mark.parametrize(
@@ -28,3 +33,29 @@ def test_correct_runner_image(name, kwargs, expected):
     assert list(actual) == expected, (
         f"case {name} expected {expected}, but got {actual} for kwargs: {kwargs}"
     )
+
+
+@pytest.mark.parametrize(
+    "name, kwargs, expected",
+    load(
+        "test_replace_image_with.json",
+    ),
+)
+def test_replace_image_with(name, kwargs, expected):
+    actual = replace_image_with(**kwargs)
+    assert actual == expected, (
+        f"case {name} expected {expected}, but got {actual} for kwargs: {kwargs}"
+    )
+
+
+@pytest.mark.parametrize(
+    "name, kwargs, expected",
+    load(
+        "test_make_image_with.json",
+    ),
+)
+def test_make_image_with(name, kwargs, expected):
+    actual = make_image_with(**kwargs)
+    assert actual == expected, (
+        f"case {name} expected {expected}, but got {actual} for kwargs: {kwargs}"
+    )