diff --git a/src/ai/backend/manager/container_registry/__init__.py b/src/ai/backend/manager/container_registry/__init__.py index b92e74f0393..7eb4b8f4335 100644 --- a/src/ai/backend/manager/container_registry/__init__.py +++ b/src/ai/backend/manager/container_registry/__init__.py @@ -32,6 +32,10 @@ def get_container_registry_cls(registry_info: Mapping[str, Any]) -> Type[BaseCon from .github import GitHubRegistry_v2 cr_cls = GitHubRegistry_v2 + elif registry_type == "gitlab": + from .gitlab import GitLabRegistry_v2 + + cr_cls = GitLabRegistry_v2 elif registry_type == "local": from .local import LocalRegistry diff --git a/src/ai/backend/manager/container_registry/gitlab.py b/src/ai/backend/manager/container_registry/gitlab.py new file mode 100644 index 00000000000..4ff9c9110c0 --- /dev/null +++ b/src/ai/backend/manager/container_registry/gitlab.py @@ -0,0 +1,50 @@ +import logging +import urllib.parse +from typing import AsyncIterator + +import aiohttp + +from ai.backend.common.logging import BraceStyleAdapter + +from .base import ( + BaseContainerRegistry, +) + +log = BraceStyleAdapter(logging.getLogger(__spec__.name)) # type: ignore[name-defined] + + +class GitLabRegistry_v2(BaseContainerRegistry): + async def fetch_repositories(self, sess: aiohttp.ClientSession) -> AsyncIterator[str]: + access_token, gitlab_project = ( + self.registry_info["password"], + self.registry_info["gitlab_project"], + ) + encoded_project_id = urllib.parse.quote(gitlab_project, safe="") + + base_url = f"https://gitlab.com/api/v4/projects/{encoded_project_id}/registry/repositories" + + headers = { + "Accept": "application/json", + "PRIVATE-TOKEN": access_token, + } + page = 1 + + while True: + async with sess.get( + base_url, + headers=headers, + params={"per_page": 30, "page": page, "membership": "true"}, + ) as response: + if response.status == 200: + data = await response.json() + + for repo in data: + yield repo["path"] + if "next" in response.headers.get("Link", ""): + page += 1 + else: + break + else: + raise RuntimeError( + f"Failed to fetch repositories! {response.status} error occurred." + )