refactor: enhance mirrored deployment self container detection

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

Author: thxCode

.github/workflows/pack.yml

@@ -90,6 +90,20 @@ jobs:
         uses: docker/metadata-action@v5
         with:
           images: "${{ env.INPUT_NAMESPACE }}/${{ env.INPUT_REPOSITORY }}"
+      - name: Get Image Labels
+        run: |
+          #!/usr/bin/env bash
+
+          set -eo pipefail
+
+          LABELS=(
+            "org.opencontainers.image.source=${{ fromJSON(steps.metadata.outputs.json).labels['org.opencontainers.image.source'] }}"
+            "org.opencontainers.image.version=${{ fromJSON(steps.metadata.outputs.json).labels['org.opencontainers.image.version'] }}"
+            "org.opencontainers.image.revision=${{ fromJSON(steps.metadata.outputs.json).labels['org.opencontainers.image.revision'] }}"
+            "org.opencontainers.image.created=${{ fromJSON(steps.metadata.outputs.json).labels['org.opencontainers.image.created'] }}"
+          )
+          INPUT_DOCKER_MIRRORED_NAME_FILTER_LABELS=$(printf "%s;" "${LABELS[@]}")
+          echo "INPUT_DOCKER_MIRRORED_NAME_FILTER_LABELS=${INPUT_DOCKER_MIRRORED_NAME_FILTER_LABELS}" >> $GITHUB_ENV
       - name: Package
         uses: docker/build-push-action@v6
         with:
@@ -107,6 +121,7 @@ jobs:
           platforms: "linux/amd64,linux/arm64"
           build-args: |
             PYTHON_VERSION=${{ env.INPUT_PYTHON_VERSION }}
+            GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS=${{ env.INPUT_DOCKER_MIRRORED_NAME_FILTER_LABELS }}
           tags: |
             ${{ steps.metadata.outputs.tags }}
           labels: |

Makefile

@@ -94,6 +94,7 @@ package:
 			--buildkitd-flags "--allow-insecure-entitlement=security.insecure --allow-insecure-entitlement=network.host --oci-worker-net=host --oci-worker-gc-keepstorage=204800" \
 			--bootstrap; \
 	fi
+	LABELS=("org.opencontainers.image.source=https://github.com/gpustack/runtime" "org.opencontainers.image.version=main" "org.opencontainers.image.revision=$(GIT_COMMIT)" "org.opencontainers.image.created=$$(date +"%Y-%m-%dT%H:%M:%S.%s")"); \
 	TAG=$(PACKAGE_NAMESPACE)/$(PACKAGE_REPOSITORY):$(PACKAGE_TAG); \
 	EXTRA_ARGS=(); \
 	if [[ "$(PACKAGE_WITH_CACHE)" == "true" ]]; then \
@@ -116,6 +117,7 @@ package:
 		--ulimit nofile=65536:65536 \
 		--shm-size 16G \
 		--progress plain \
+		--build-arg "GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS=$$(printf "%s;" "$${LABELS[@]}")" \
 		$${EXTRA_ARGS[@]} \
 		$(SRCDIR); \
 	set +x

gpustack_runtime/deployer/docker.py

@@ -1161,30 +1161,14 @@ def _prepare_create(self):
                 self_container_id,
             )
         try:
-            if envs.GPUSTACK_RUNTIME_DEPLOY_MIRRORED_NAME:
-                # Directly get container by name or ID.
-                self_container = self._client.containers.get(self_container_id)
-            else:
-                # Find containers that matches the hostname.
-                containers = self._client.containers.list()
-                containers = [
-                    c
-                    for c in containers
-                    if c.attrs["Config"].get("Hostname", "") == self_container_id
-                ]
-                if len(containers) != 1:
-                    msg = f"Container with name {self_container_id} not found"
-                    raise docker.errors.NotFound(
-                        msg,
-                    )
-                self_container = containers[0]
+            self_container = self._find_self_container(self_container_id)
             self_image = self_container.image
-        except docker.errors.APIError:
+        except docker.errors.APIError as e:
             output_log = logger.warning
             if logger.isEnabledFor(logging.DEBUG):
                 output_log = logger.exception
             output_log(
-                f"Mirrored deployment enabled, but failed to get self Container {self_container_id}, skipping",
+                f"Mirrored deployment enabled, but failed to get self Container {self_container_id}, skipping: {e}",
             )
             return
 
@@ -1394,6 +1378,55 @@ def mutate_create_options(create_options: dict[str, Any]) -> dict[str, Any]:
                     result = result.joinpath(b_subpath.lstrip("/"))
                 self._container_ephemeral_files_dir = result
 
+    def _find_self_container(
+        self,
+        self_container_id: str,
+    ) -> docker.models.containers.Container:
+        """
+        Find the current container if running inside a Docker container.
+
+        Args:
+            self_container_id:
+                The container name or ID to find.
+
+        Returns:
+            The Docker container if found, None otherwise.
+
+        Raises:
+            If failed to find itself.
+
+        """
+        if envs.GPUSTACK_RUNTIME_DEPLOY_MIRRORED_NAME:
+            # Directly get container by name or ID.
+            return self._client.containers.get(self_container_id)
+
+        # Find containers that matches the hostname.
+        containers: list[docker.models.containers.Container] = []
+        for c in self._client.containers.list():
+            # Ignore workload containers with host network enabled.
+            if _LABEL_WORKLOAD in c.labels:
+                continue
+            # Ignore containers that do not match the hostname.
+            if c.attrs["Config"].get("Hostname", "") != self_container_id:
+                continue
+            # Ignore containers that do not match the filter labels.
+            if envs.GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS and any(
+                c.labels.get(k) != v
+                for k, v in envs.GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS.items()
+            ):
+                continue
+            containers.append(c)
+
+        # Validate found containers.
+        if len(containers) != 1:
+            msg = (
+                f"Found multiple Containers with the same hostname {self_container_id}, "
+                "please use `--env GPUSTACK_RUNTIME_DEPLOY_MIRRORED_NAME=...` to specify the exact container name"
+            )
+            raise docker.errors.NotFound(msg)
+
+        return containers[0]
+
     @_supported
     def _create(self, workload: WorkloadPlan):
         """

gpustack_runtime/envs.py

@@ -132,6 +132,12 @@
     # Deployer
 
     ## Docker
+    GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS: dict[str, str] | None = None
+    """
+    Filter labels for selecting the mirrored deployer container in Docker.
+    Only works when `GPUSTACK_RUNTIME_DEPLOY_MIRRORED_NAME` is not set.
+    Normally, it should be injected automatically via CI without any manual configuration.
+    """
     GPUSTACK_RUNTIME_DOCKER_PAUSE_IMAGE: str | None = None
     """
     Docker image used for the pause container.
@@ -284,6 +290,12 @@
         getenv("GPUSTACK_RUNTIME_DETECT_PHYSICAL_INDEX_PRIORITY", "1"),
     ),
     # Deployer
+    "GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS": lambda: to_dict(
+        getenv(
+            "GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS",
+        ),
+        sep=";",
+    ),
     "GPUSTACK_RUNTIME_DOCKER_PAUSE_IMAGE": lambda: getenv(
         "GPUSTACK_RUNTIME_DOCKER_PAUSE_IMAGE",
         "gpustack/runtime:pause",

pack/Dockerfile

@@ -15,9 +15,11 @@
 # - PYTHON_VERSION: Version of Python to use.
 # - GPUSTACK_RUNTIME_BASE_IMAGE: Base image for the gpustack-runtime stage.
 # - GPUSTACK_RUNTIME_ROCM_VERSION: Version of ROCm to vendor libraries from, update this if project dependencies has changed.
+# - GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS: Line-separated list of labels to filter mirrored images when deploying mirrored deployment.
 ARG PYTHON_VERSION=3.11
 ARG GPUSTACK_RUNTIME_BASE_IMAGE=runtime
 ARG GPUSTACK_RUNTIME_ROCM_VERSION=6.2.4
+ARG GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS
 
 FROM ubuntu:22.04@sha256:3c61d3759c2639d4b836d32a2d3c83fa0214e36f195a3421018dbaaf79cbe37f AS runtime
 SHELL ["/bin/bash", "-eo", "pipefail", "-c"]
@@ -471,5 +473,9 @@ ENV NVIDIA_DISABLE_REQUIRE="true" \
     NVIDIA_VISIBLE_DEVICES="all" \
     NVIDIA_DRIVER_CAPABILITIES="compute,utility"
 
+ARG GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS
+
+ENV GPUSTACK_RUNTIME_DEPLOY_MIRRORED_DEPLOYMENT="true" \
+    GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS="${GPUSTACK_RUNTIME_DOCKER_MIRRORED_NAME_FILTER_LABELS}"
 WORKDIR /
 ENTRYPOINT [ "tini", "--" ]