refactor: Make project column nullable and improve container regist…

…ry migration script (#2978) (#3038)

Co-authored-by: Gyubong Lee <jopemachine@naver.com>
Co-authored-by: Sanghun Lee <sanghun@lablup.com>

changes/2978.fix.md

@@ -0,0 +1 @@
+Make image, container_registry table's `project` column nullable and improve container registry storage config migration script.

src/ai/backend/common/docker.py

@@ -567,8 +567,8 @@ def merge_aliases(genned_aliases_1, genned_aliases_2) -> Mapping[str, "ImageRef"
 
     @property
     def canonical(self) -> str:
-        # e.g., cr.backend.ai/stable/python:3.9-ubuntu
         join = functools.partial(join_non_empty, sep="/")
+        # e.g., cr.backend.ai/stable/python:3.9-ubuntu
         return f"{join(self.registry, self.project, self.name)}:{self.tag}"
 
     @property

src/ai/backend/manager/models/alembic/versions/1d42c726d8a3_migrate_container_registries_from_etcd_to_psql.py

@@ -9,6 +9,7 @@
 import asyncio
 import enum
 import json
+import logging
 import os
 import sys
 import warnings
@@ -35,6 +36,8 @@
 branch_labels = None
 depends_on = None
 
+logger = logging.getLogger("alembic.runtime.migration")
+
 warnings.filterwarnings(
     "ignore",
     message="This declarative base already contains a class with the same class name and module name as .*",
@@ -60,6 +63,10 @@ class ContainerRegistryType(enum.StrEnum):
     DOCKER = "docker"
     HARBOR = "harbor"
     HARBOR2 = "harbor2"
+    GITHUB = "github"
+    GITLAB = "gitlab"
+    ECR = "ecr"
+    ECR_PUB = "ecr-public"
     LOCAL = "local"
 
 
@@ -109,7 +116,7 @@ class ContainerRegistryRow(Base):
             nullable=False,
             index=True,
         )
-        project = sa.Column("project", sa.String(length=255), nullable=False)
+        project = sa.Column("project", sa.String(length=255), nullable=True)
         username = sa.Column("username", sa.String(length=255), nullable=True)
         password = sa.Column("password", sa.String, nullable=True)
         ssl_verify = sa.Column("ssl_verify", sa.Boolean, server_default=sa.text("true"), index=True)
@@ -225,9 +232,13 @@ async def _get_container_registries():
             if registry_type == ContainerRegistryType.DOCKER and hostname == "index.docker.io":
                 input_config_template["project"] = "library"
             else:
-                raise Exception(
-                    f'ContainerRegistryRow.project is required! Please put the value to "{ETCD_CONTAINER_REGISTRY_KEY}/{hostname}/project" before migration.'
-                )
+                # ContainerRegistryRow with empty project is allowed only if local registry.
+                if registry_type == ContainerRegistryType.LOCAL:
+                    input_config_template["project"] = None
+                else:
+                    raise Exception(
+                        f'ContainerRegistryRow.project is required! Please put the value to "{ETCD_CONTAINER_REGISTRY_KEY}/{hostname}/project" before migration.'
+                    )
 
             input_configs.append(input_config_template)
 
@@ -246,7 +257,6 @@ def revert_data_psql_to_etcd() -> None:
 
     db_connection = op.get_bind()
 
-    # Prevent error from presence or absence of the extra column
     rows = db_connection.execute(
         sa.select([
             ContainerRegistryRow.url,
@@ -340,6 +350,9 @@ def insert_registry_id_to_images() -> None:
             if registry_name == "index.docker.io" and project is None:
                 project = "library"
 
+            if not project:
+                continue
+
             query_params.append({
                 "registry_id": registry_id,
                 "registry_name_and_project": f"{registry_name}/{project}",
@@ -351,6 +364,126 @@ def insert_registry_id_to_images() -> None:
         )
 
 
+def insert_registry_id_to_images_with_missing_registry_id() -> None:
+    """
+    If there are image rows with empty registry_id and the image's container registry name exists in the container_registries table,
+    Generate new container registry row corresponding to the image's project part.
+    """
+
+    db_connection = op.get_bind()
+    ContainerRegistryRow = get_container_registry_row_schema()
+
+    get_images_with_blank_registry_id = sa.select(images_table).where(
+        images_table.c.registry_id.is_(None)
+    )
+    images = db_connection.execute(get_images_with_blank_registry_id)
+
+    added_projects = []
+    for image in images:
+        namespace = (image.name.split("/"))[:2]
+        if len(namespace) < 2:
+            continue
+
+        registry_name, project = namespace
+
+        registry_info = db_connection.execute(
+            sa.select([
+                ContainerRegistryRow.url,
+                ContainerRegistryRow.registry_name,
+                ContainerRegistryRow.type,
+                ContainerRegistryRow.ssl_verify,
+            ]).where(ContainerRegistryRow.registry_name == registry_name)
+        ).fetchone()
+
+        if not registry_info:
+            continue
+
+        if project in added_projects:
+            continue
+        else:
+            added_projects.append(project)
+
+        registry_info = dict(registry_info)
+        registry_info["project"] = project
+
+        registry_id = db_connection.execute(
+            sa.insert(ContainerRegistryRow)
+            .values(**registry_info)
+            .returning(ContainerRegistryRow.id)
+        ).scalar_one()
+
+        db_connection.execute(
+            sa.update(images_table)
+            .values(registry_id=registry_id)
+            .where(
+                (
+                    images_table.c.name.startswith(f"{registry_name}/{project}")
+                    & (images_table.c.registry_id.is_(None))
+                )
+            )
+        )
+
+        logger.info(
+            f'Following container registry row auto-generated: "{registry_name}" with the project "{project}".'
+        )
+
+    if added_projects:
+        logger.info(
+            "If credential required for the auto-generated container registry rows, you should fill their credential columns manually."
+        )
+    else:
+        logger.info("No container registry row auto-generated.")
+
+
+def mark_local_images_with_missing_registry_id() -> None:
+    """
+    If there are image rows with empty registry_id and the image's container registry name does not exist in the container_registries table,
+    Generate a local type container registry record and then, consider and mark the images as local images.
+    """
+
+    db_connection = op.get_bind()
+    ContainerRegistryRow = get_container_registry_row_schema()
+
+    get_images_with_blank_registry_id = sa.select(images_table).where(
+        images_table.c.registry_id.is_(None)
+    )
+    images_with_missing_registry_id = db_connection.execute(get_images_with_blank_registry_id)
+
+    if not images_with_missing_registry_id:
+        return
+
+    local_registry_info = {
+        "type": ContainerRegistryType.LOCAL,
+        "registry_name": "local",
+        # url is not used for local registry.
+        # but it is required in the old schema (etcd),
+        # so, let's put a dummy value for compatibility purposes.
+        "url": "http://localhost",
+    }
+
+    local_registry_id = db_connection.execute(
+        sa.select([ContainerRegistryRow.id]).where(
+            ContainerRegistryRow.type == ContainerRegistryType.LOCAL
+        )
+    ).scalar_one_or_none()
+
+    if not local_registry_id:
+        local_registry_id = db_connection.execute(
+            sa.insert(ContainerRegistryRow)
+            .values(**local_registry_info)
+            .returning(ContainerRegistryRow.id)
+        ).scalar_one()
+
+    db_connection.execute(
+        sa.update(images_table)
+        .values(
+            registry_id=local_registry_id,
+            is_local=True,
+        )
+        .where((images_table.c.registry_id.is_(None)))
+    )
+
+
 def upgrade():
     metadata = sa.MetaData(naming_convention=convention)
     op.create_table(
@@ -367,7 +500,7 @@ def upgrade():
             nullable=False,
             index=True,
         ),
-        sa.Column("project", sa.String(length=255), nullable=False, index=True),
+        sa.Column("project", sa.String(length=255), nullable=True, index=True),
         sa.Column("username", sa.String(length=255), nullable=True),
         sa.Column("password", sa.String(), nullable=True),
         sa.Column(
@@ -387,10 +520,12 @@ def upgrade():
     )
 
     insert_registry_id_to_images()
-    delete_old_etcd_container_registries()
-
+    insert_registry_id_to_images_with_missing_registry_id()
+    mark_local_images_with_missing_registry_id()
     op.alter_column("images", "registry_id", nullable=False)
 
+    delete_old_etcd_container_registries()
+
 
 def downgrade():
     revert_data_psql_to_etcd()

src/ai/backend/manager/models/alembic/versions/7c8501cec07b_add_project_column_to_images_table.py

@@ -42,7 +42,7 @@ class ContainerRegistryRow(Base):
             nullable=False,
             index=True,
         )
-        project = sa.Column("project", sa.String(length=255), nullable=False)
+        project = sa.Column("project", sa.String(length=255), nullable=True)
         username = sa.Column("username", sa.String(length=255), nullable=True)
         password = sa.Column("password", sa.String, nullable=True)
         ssl_verify = sa.Column("ssl_verify", sa.Boolean, server_default=sa.text("true"), index=True)
@@ -62,7 +62,7 @@ class ImageRow(Base):
         __table_args__ = {"extend_existing": True}
         id = IDColumn("id")
         name = sa.Column("name", sa.String, nullable=False, index=True)
-        project = sa.Column("project", sa.String, nullable=False)
+        project = sa.Column("project", sa.String, nullable=True)
         image = sa.Column("image", sa.String, nullable=False, index=True)
         created_at = sa.Column(
             "created_at",
@@ -118,8 +118,6 @@ def upgrade():
 
     op.get_bind().execute(update_stmt)
 
-    op.alter_column("images", "project", nullable=False)
-
 
 def downgrade():
     op.drop_column("images", "project")

src/ai/backend/manager/models/container_registry.py

@@ -71,7 +71,7 @@ class ContainerRegistryRow(Base):
         nullable=False,
         index=True,
     )
-    project = sa.Column("project", sa.String(length=255), index=True, nullable=False)
+    project = sa.Column("project", sa.String(length=255), index=True, nullable=True)
     username = sa.Column("username", sa.String(length=255), nullable=True)
     password = sa.Column("password", sa.String, nullable=True)
     ssl_verify = sa.Column(

src/ai/backend/manager/models/image.py

@@ -157,7 +157,7 @@ class ImageRow(Base):
     __tablename__ = "images"
     id = IDColumn("id")
     name = sa.Column("name", sa.String, nullable=False, index=True)
-    project = sa.Column("project", sa.String, nullable=False)
+    project = sa.Column("project", sa.String, nullable=True)
     image = sa.Column("image", sa.String, nullable=False, index=True)
     created_at = sa.Column(
         "created_at",