fix: Add missing indexes in the vfolders table (#1937)

src/ai/backend/manager/models/alembic/versions/41f332243bf9_add_missing_vfolder_indexes.py

@@ -0,0 +1,31 @@
+"""add-missing-vfolder-indexes
+
+Revision ID: 41f332243bf9
+Revises: 7ff52ff68bfc
+Create Date: 2024-02-28 17:27:40.387122
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "41f332243bf9"
+down_revision = "7ff52ff68bfc"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_index(op.f("ix_vfolders_host"), "vfolders", ["host"], unique=False)
+    op.create_index(
+        op.f("ix_vfolders_ownership_type"), "vfolders", ["ownership_type"], unique=False
+    )
+    op.create_index(op.f("ix_vfolders_status"), "vfolders", ["status"], unique=False)
+    op.create_index(op.f("ix_vfolders_usage_mode"), "vfolders", ["usage_mode"], unique=False)
+
+
+def downgrade():
+    op.drop_index(op.f("ix_vfolders_usage_mode"), table_name="vfolders")
+    op.drop_index(op.f("ix_vfolders_status"), table_name="vfolders")
+    op.drop_index(op.f("ix_vfolders_ownership_type"), table_name="vfolders")
+    op.drop_index(op.f("ix_vfolders_host"), table_name="vfolders")

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

@@ -469,6 +469,7 @@ async def handle_kernel_exception(
         default=KernelRole.COMPUTE,
         server_default=KernelRole.COMPUTE.name,
         nullable=False,
+        index=True,
     ),
     sa.Column("status_changed", sa.DateTime(timezone=True), nullable=True, index=True),
     sa.Column("status_info", sa.Unicode(), nullable=True, default=sa.null()),

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

@@ -223,14 +223,15 @@ class VFolderCloneInfo(NamedTuple):
     metadata,
     IDColumn("id"),
     # host will be '' if vFolder is unmanaged
-    sa.Column("host", sa.String(length=128), nullable=False),
+    sa.Column("host", sa.String(length=128), nullable=False, index=True),
     sa.Column("quota_scope_id", QuotaScopeIDType, nullable=False),
     sa.Column("name", sa.String(length=64), nullable=False, index=True),
     sa.Column(
         "usage_mode",
         EnumValueType(VFolderUsageMode),
         default=VFolderUsageMode.GENERAL,
         nullable=False,
+        index=True,
     ),
     sa.Column("permission", EnumValueType(VFolderPermission), default=VFolderPermission.READ_WRITE),
     sa.Column("max_files", sa.Integer(), default=1000),
@@ -248,6 +249,7 @@ class VFolderCloneInfo(NamedTuple):
         EnumValueType(VFolderOwnershipType),
         default=VFolderOwnershipType.USER,
         nullable=False,
+        index=True,
     ),
     sa.Column("user", GUID, sa.ForeignKey("users.uuid"), nullable=True),  # owner if user vfolder
     sa.Column("group", GUID, sa.ForeignKey("groups.id"), nullable=True),  # owner if project vfolder
@@ -258,6 +260,7 @@ class VFolderCloneInfo(NamedTuple):
         default=VFolderOperationStatus.READY,
         server_default=VFolderOperationStatus.READY,
         nullable=False,
+        index=True,
     ),
     # status_history records the most recent status changes for each status
     # e.g)