Skip to content

Commit d95ce09

Browse files
jopemachineclaude
andcommitted
fix(BA-5650-I): resolve cascaded rebase conflicts and align tests
- Resolve conflict markers in scheduler db_source by keeping the user-table-joined ``main_access_key`` lookup (single source of truth) rather than the obsolete ``session.access_key`` snapshot. - Drop duplicate ``resolve_main_access_keys`` definition introduced by the rebase. - Use ``spec.access_key`` (SessionCreationSpec) where applicable — the spec dataclass keeps ``access_key`` for now. - Test fixture ``_make_session_data`` now passes ``owner_id`` instead of relying on attribute assignment after construction. - Use ``main_access_key`` on PendingSessionData in provisioner test. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 95b1bba commit d95ce09

4 files changed

Lines changed: 6 additions & 40 deletions

File tree

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

Lines changed: 4 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1215,8 +1215,7 @@ async def get_pending_timeout_sessions_by_ids(
12151215
SweptSessionInfo(
12161216
session_id=row.id,
12171217
creation_id=row.creation_id,
1218-
main_access_key=row.access_key,
1219-
>>>>>>> e8b82576d (refactor(BA-5650-I): test and remaining ORM updates)
1218+
main_access_key=row.main_access_key,
12201219
)
12211220
)
12221221

@@ -1489,7 +1488,7 @@ async def fetch_session_creation_data(
14891488
db_sess,
14901489
spec.user_scope.domain_name,
14911490
str(spec.user_scope.group_id),
1492-
spec.main_access_key,
1491+
spec.access_key,
14931492
)
14941493
image_infos = await self._resolve_image_info(db_sess, image_refs)
14951494

@@ -1510,7 +1509,7 @@ async def fetch_session_creation_data(
15101509
dotfile_data = await self._fetch_dotfiles(
15111510
db_sess,
15121511
spec.user_scope,
1513-
spec.main_access_key,
1512+
spec.access_key,
15141513
vfolder_mounts,
15151514
)
15161515

@@ -1554,7 +1553,7 @@ async def fetch_session_creation_context(
15541553
db_sess,
15551554
spec.user_scope.domain_name,
15561555
str(spec.user_scope.group_id),
1557-
spec.main_access_key,
1556+
spec.access_key,
15581557
)
15591558
image_infos = await self._resolve_image_info(db_sess, image_refs)
15601559

@@ -2927,12 +2926,7 @@ async def _get_sessions_by_statuses(
29272926
scheduled_session = ScheduledSessionData(
29282927
session_id=session.id,
29292928
creation_id=session.creation_id or "",
2930-
main_access_key=AccessKey(session.access_key)
2931-
if session.access_key
2932-
else AccessKey(""),
2933-
=======
29342929
main_access_key=AccessKey(owner_main_ak) if owner_main_ak else AccessKey(""),
2935-
>>>>>>> 4ac120c73 (refactor(BA-5650-I): drop stray non-BA-5650 changes from slice)
29362930
reason="triggered-by-scheduler",
29372931
)
29382932
scheduled_sessions.append(scheduled_session)
@@ -2979,12 +2973,7 @@ async def _get_scheduled_sessions(self, db_sess: SASession) -> list[ScheduledSes
29792973
ScheduledSessionData(
29802974
session_id=session.id,
29812975
creation_id=session.creation_id or "",
2982-
main_access_key=AccessKey(session.access_key)
2983-
if session.access_key
2984-
else AccessKey(""),
2985-
=======
29862976
main_access_key=AccessKey(owner_main_ak) if owner_main_ak else AccessKey(""),
2987-
>>>>>>> 4ac120c73 (refactor(BA-5650-I): drop stray non-BA-5650 changes from slice)
29882977
reason="triggered-by-scheduler",
29892978
)
29902979
)
@@ -4828,21 +4817,3 @@ async def _get_db_now_in_session(self, db_sess: SASession) -> datetime:
48284817
"""
48294818
result = await db_sess.execute(sa.select(sa.func.now()))
48304819
return result.scalar_one()
4831-
4832-
async def resolve_main_access_keys(
4833-
self, session_ids: Sequence[SessionId]
4834-
) -> dict[SessionId, AccessKey]:
4835-
"""Resolve the main access key for each session's owner.
4836-
4837-
Joins sessions → users to look up the owner's main_access_key.
4838-
"""
4839-
if not session_ids:
4840-
return {}
4841-
async with self._db.begin_readonly_session() as db_sess:
4842-
stmt = (
4843-
sa.select(SessionRow.id, UserRow.main_access_key)
4844-
.join(UserRow, SessionRow.user_uuid == UserRow.uuid)
4845-
.where(SessionRow.id.in_([sid for sid in session_ids]))
4846-
)
4847-
rows = (await db_sess.execute(stmt)).all()
4848-
return {SessionId(row[0]): AccessKey(row[1]) for row in rows if row[1] is not None}

tests/unit/manager/services/session/test_session_lifecycle_service.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ def _make_session_data(
210210
) -> SessionData:
211211
return SessionData(
212212
id=session_id,
213+
owner_id=user_id,
213214
creation_id="test-creation-id",
214215
name=name,
215216
session_type=session_type,

tests/unit/manager/services/session/test_session_service.py

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -245,10 +245,7 @@ def sample_session_data(
245245
agent_ids=["i-ubuntu"],
246246
domain_name="default",
247247
group_id=sample_group_id,
248-
<<<<<<< HEAD
249248
owner_id=sample_user_id,
250-
=======
251-
>>>>>>> d59487b17 (refactor(BA-5650-I): test and remaining ORM updates)
252249
images=["cr.backend.ai/stable/python:latest"],
253250
tag=None,
254251
occupying_slots=ResourceSlot({"cpu": 1, "mem": 1024}),
@@ -349,10 +346,7 @@ async def test_multiple_matches(
349346
agent_ids=[],
350347
domain_name="default",
351348
group_id=sample_group_id,
352-
<<<<<<< HEAD
353349
owner_id=sample_user_id,
354-
=======
355-
>>>>>>> d59487b17 (refactor(BA-5650-I): test and remaining ORM updates)
356350
images=["python:latest"],
357351
tag=None,
358352
occupying_slots=ResourceSlot({}),

tests/unit/manager/sokovan/scheduler/provisioner/test_provisioner.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def _create_scheduling_data_with_strategy(
6666
# Create one pending session
6767
session = PendingSessionData(
6868
id=SessionId(uuid.uuid4()),
69-
access_key=AccessKey("test-key"),
69+
main_access_key=AccessKey("test-key"),
7070
requested_slots=ResourceSlot({"cpu": Decimal("1"), "mem": Decimal("1024")}),
7171
owner_id=uuid.uuid4(),
7272
group_id=uuid.uuid4(),

0 commit comments

Comments
 (0)