From af42d4ad357be3fdb2d9285fbf420b615faa473d Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Mon, 8 Apr 2024 14:54:24 +0900
Subject: [PATCH] fix: inconsistency of `kernels.session_name` column (#1985)

Co-authored-by: Sion Kang <siontama@gmail.com>
Backported-from: main (24.09)
Backported-to: 24.03
---
 changes/1985.fix.md                           |  1 +
 src/ai/backend/manager/api/session.py         |  3 ++
 ...ync_kernelssessionname_to_real_session_.py | 36 +++++++++++++++++++
 3 files changed, 40 insertions(+)
 create mode 100644 changes/1985.fix.md
 create mode 100644 src/ai/backend/manager/models/alembic/versions/dddf9be580f5_sync_kernelssessionname_to_real_session_.py

diff --git a/changes/1985.fix.md b/changes/1985.fix.md
new file mode 100644
index 0000000000..baaffe86fd
--- /dev/null
+++ b/changes/1985.fix.md
@@ -0,0 +1 @@
+Sync mismatch between `kernels.session_name` and `sessions.name` and fix session-rename API to update `session_name` of sibling kernels atomically.
\ No newline at end of file
diff --git a/src/ai/backend/manager/api/session.py b/src/ai/backend/manager/api/session.py
index 3ae5ab6967..7cf7587ecc 100644
--- a/src/ai/backend/manager/api/session.py
+++ b/src/ai/backend/manager/api/session.py
@@ -1315,10 +1315,13 @@ async def rename_session(request: web.Request, params: Any) -> web.Response:
             owner_access_key,
             allow_stale=True,
             for_update=True,
+            kernel_loading_strategy=KernelLoadingStrategy.ALL_KERNELS,
         )
         if compute_session.status != SessionStatus.RUNNING:
             raise InvalidAPIParameters("Can't change name of not running session")
         compute_session.name = new_name
+        for kernel in compute_session.kernels:
+            kernel.session_name = new_name
         await db_sess.commit()
 
     return web.Response(status=204)
diff --git a/src/ai/backend/manager/models/alembic/versions/dddf9be580f5_sync_kernelssessionname_to_real_session_.py b/src/ai/backend/manager/models/alembic/versions/dddf9be580f5_sync_kernelssessionname_to_real_session_.py
new file mode 100644
index 0000000000..0797b3ca06
--- /dev/null
+++ b/src/ai/backend/manager/models/alembic/versions/dddf9be580f5_sync_kernelssessionname_to_real_session_.py
@@ -0,0 +1,36 @@
+"""sync_kernelssessionname_to_real_session_name
+
+Revision ID: dddf9be580f5
+Revises: 857b763b8618
+Create Date: 2024-04-01 16:58:14.341114
+
+"""
+
+import textwrap
+
+from alembic import op
+from sqlalchemy.sql import text
+
+# revision identifiers, used by Alembic.
+revision = "dddf9be580f5"
+down_revision = "857b763b8618"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    conn = op.get_bind()
+    sync_stmt = textwrap.dedent(
+        """
+        UPDATE kernels
+        SET session_name = sessions.name
+        FROM sessions
+        WHERE kernels.session_id = sessions.id
+        AND kernels.session_name <> sessions.name;
+        """
+    )
+    conn.execute(text(sync_stmt))
+
+
+def downgrade():
+    pass