Skip to content

Commit 350294e

Browse files
committed
refactor(BA-5714): resolve rebase conflicts and add resolve_main_access_keys
- Resolve merge conflicts from cascading onto the updated slice E (drf.py keeps ``existing_sess.user_uuid``; predicates.py keeps the ``main_ak is None`` early returns; scheduler types stay on ``main_access_key``; stream db_source keeps ``UserNotFound``). - Add ``SchedulerRepository.resolve_main_access_keys`` and the matching ``ScheduleDBSource.resolve_main_access_keys`` so the coordinator's cache-invalidation step can look up each session's owner main_access_key in a single query. Required by sokovan/scheduler/coordinator.py and lifecycle/deprioritize_sessions.py call sites that were previously reading ``session_info.metadata.access_key``.
1 parent 23db4d0 commit 350294e

3 files changed

Lines changed: 26 additions & 1 deletion

File tree

changes/BA-5650-E.misc.md

Lines changed: 0 additions & 1 deletion
This file was deleted.

src/ai/backend/manager/repositories/scheduler/db_source/db_source.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4774,6 +4774,26 @@ async def get_db_now(self) -> datetime:
47744774
result = await conn.execute(sa.select(sa.func.now()))
47754775
return result.scalar_one()
47764776

4777+
async def resolve_main_access_keys(
4778+
self, session_ids: Sequence[SessionId]
4779+
) -> dict[SessionId, AccessKey]:
4780+
"""Resolve the main access key for each session's owner.
4781+
4782+
Joins ``sessions`` → ``users`` to look up the owner's
4783+
``main_access_key``. Sessions whose owner has no configured
4784+
main access key are omitted from the returned mapping.
4785+
"""
4786+
if not session_ids:
4787+
return {}
4788+
async with self._db.begin_readonly_session() as db_sess:
4789+
stmt = (
4790+
sa.select(SessionRow.id, UserRow.main_access_key)
4791+
.join(UserRow, SessionRow.user_uuid == UserRow.uuid)
4792+
.where(SessionRow.id.in_([sid for sid in session_ids]))
4793+
)
4794+
rows = (await db_sess.execute(stmt)).all()
4795+
return {SessionId(row[0]): AccessKey(row[1]) for row in rows if row[1] is not None}
4796+
47774797
async def _get_db_now_in_session(self, db_sess: SASession) -> datetime:
47784798
"""Get the current timestamp from the database within an existing session.
47794799

src/ai/backend/manager/repositories/scheduler/repository.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -959,3 +959,9 @@ async def get_db_now(self) -> datetime:
959959
Current database timestamp with timezone
960960
"""
961961
return await self._db_source.get_db_now()
962+
963+
async def resolve_main_access_keys(
964+
self, session_ids: Sequence[SessionId]
965+
) -> dict[SessionId, AccessKey]:
966+
"""Resolve the main access key for each session's owner."""
967+
return await self._db_source.resolve_main_access_keys(session_ids)

0 commit comments

Comments
 (0)