feat: Expose user info to session environment variables (#2043)

lablup · Apr 22, 2024 · 4f39048 · 4f39048
1 parent f08911e
commit 4f39048
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/changes/2043.feature.md b/changes/2043.feature.md
@@ -0,0 +1 @@
+Expose user info to environment variables
diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py
@@ -1449,7 +1449,20 @@ async def start_session(
                 Optional[ClusterSSHPortMapping], cluster_ssh_port_mapping
             ),
         )
+
+        async with self.db.begin_readonly_session() as db_sess:
+            uuid, email, username = (
+                await db_sess.execute(
+                    sa.select([UserRow.uuid, UserRow.email, UserRow.username]).where(
+                        UserRow.uuid == scheduled_session.user_uuid
+                    )
+                )
+            ).fetchone()
+
         scheduled_session.environ.update({
+            "BACKENDAI_USER_UUID": str(uuid),
+            "BACKENDAI_USER_EMAIL": email,
+            "BACKENDAI_USER_NAME": username,
             "BACKENDAI_SESSION_ID": str(scheduled_session.id),
             "BACKENDAI_SESSION_NAME": str(scheduled_session.name),
             "BACKENDAI_CLUSTER_SIZE": str(scheduled_session.cluster_size),