feat: Add uid and gid columns to users and kernels tables

Author: fregataa

src/ai/backend/manager/models/alembic/versions/f6ca2f2d04c1_add_uid_and_gid_columns.py

@@ -0,0 +1,43 @@
+"""add uid and gid columns
+
+Revision ID: f6ca2f2d04c1
+Revises: 0bb88d5a46bf
+Create Date: 2024-12-20 12:16:07.077845
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f6ca2f2d04c1"
+down_revision = "0bb88d5a46bf"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "kernels", sa.Column("uid", sa.Integer(), server_default=sa.text("NULL"), nullable=True)
+    )
+    op.add_column(
+        "kernels",
+        sa.Column("gids", sa.ARRAY(sa.Integer()), server_default=sa.text("NULL"), nullable=True),
+    )
+    op.add_column(
+        "users",
+        sa.Column("container_uid", sa.Integer(), server_default=sa.text("NULL"), nullable=True),
+    )
+    op.add_column(
+        "users",
+        sa.Column(
+            "container_gids", sa.ARRAY(sa.Integer()), server_default=sa.text("NULL"), nullable=True
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("users", "container_gids")
+    op.drop_column("users", "container_uid")
+    op.drop_column("kernels", "gids")
+    op.drop_column("kernels", "uid")

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

@@ -416,6 +416,9 @@ class KernelRow(Base):
     cluster_hostname = sa.Column(
         "cluster_hostname", sa.String(length=64), nullable=False, default=default_hostname
     )
+    uid = sa.Column("uid", sa.Integer, nullable=True, server_default=sa.null())
+    gids = sa.Column("gids", sa.ARRAY(sa.Integer), nullable=True, server_default=sa.null())
+
     # Resource ownership
     scaling_group = sa.Column(
         "scaling_group", sa.ForeignKey("scaling_groups.name"), index=True, nullable=True

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

@@ -159,6 +159,8 @@ class UserStatus(enum.StrEnum):
         sa.ForeignKey("keypairs.access_key", ondelete="SET NULL"),
         nullable=True,  # keypairs.user is non-nullable
     ),
+    sa.Column("container_uid", sa.Integer, nullable=True, server_default=sa.null()),
+    sa.Column("container_gids", sa.ARRAY(sa.Integer), nullable=True, server_default=sa.null()),
 )