From 04a2456ba451b0842097ba3544ec33e41ca333dc Mon Sep 17 00:00:00 2001 From: mindgitrwx Date: Mon, 20 Dec 2021 02:38:03 +0900 Subject: [PATCH 1/7] fix: add max_matches limitation on reenter_txn --- src/ai/backend/manager/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py index efa93139c..de152f321 100644 --- a/src/ai/backend/manager/registry.py +++ b/src/ai/backend/manager/registry.py @@ -666,6 +666,7 @@ async def get_session( for_update=for_update, extra_cond=extra_cond, db_connection=conn, + max_matches=1, ) if not session_infos: raise SessionNotFound() From feafdcfbf909c397feb5de6399fdab50ec5fa595 Mon Sep 17 00:00:00 2001 From: jhyeo Date: Mon, 20 Dec 2021 22:24:29 +0900 Subject: [PATCH 2/7] Revert "fix: add max_matches limitation on reenter_txn" This reverts commit 04a2456ba451b0842097ba3544ec33e41ca333dc. --- src/ai/backend/manager/registry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py index de152f321..efa93139c 100644 --- a/src/ai/backend/manager/registry.py +++ b/src/ai/backend/manager/registry.py @@ -666,7 +666,6 @@ async def get_session( for_update=for_update, extra_cond=extra_cond, db_connection=conn, - max_matches=1, ) if not session_infos: raise SessionNotFound() From e85ae95e5fd4cb44770790d4f4850359aed2147b Mon Sep 17 00:00:00 2001 From: jhyeo Date: Mon, 20 Dec 2021 22:29:35 +0900 Subject: [PATCH 3/7] Add cond_equal_name and cond_prefix_name --- src/ai/backend/manager/models/kernel.py | 39 ++++++++++++++++++++----- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/src/ai/backend/manager/models/kernel.py b/src/ai/backend/manager/models/kernel.py index f8724d527..c573a0bc2 100644 --- a/src/ai/backend/manager/models/kernel.py +++ b/src/ai/backend/manager/models/kernel.py @@ -301,7 +301,6 @@ class SessionInfo(TypedDict): status: KernelStatus created_at: datetime - async def match_session_ids( session_name_or_id: Union[str, UUID], access_key: AccessKey, @@ -321,12 +320,17 @@ async def match_session_ids( ) if extra_cond is not None: cond_id = cond_id & extra_cond - cond_name = ( + cond_equal_name = ( + (kernels.c.session_name == (f'{session_name_or_id}')) & + (kernels.c.access_key == access_key) + ) + cond_prefix_name = ( (kernels.c.session_name.like(f'{session_name_or_id}%')) & (kernels.c.access_key == access_key) ) if extra_cond is not None: - cond_name = cond_name & extra_cond + cond_equal_name = cond_equal_name & extra_cond + cond_prefix_name = cond_prefix_name & extra_cond cond_session_id = ( (sa.sql.expression.cast(kernels.c.session_id, sa.String).like(f'{session_name_or_id}%')) & (kernels.c.access_key == access_key) @@ -358,7 +362,24 @@ async def match_session_ids( ) if for_update: match_sid_by_id = match_sid_by_id.with_for_update() - match_sid_by_name = ( + match_sid_by_equal_name = ( + sa.select(info_cols) + .select_from(kernels) + .where( + (kernels.c.session_id.in_( + sa.select( + [kernels.c.session_id], + ) + .select_from(kernels) + .where(cond_equal_name) + .group_by(kernels.c.session_id) + .limit(max_matches).offset(0), + )) & + (kernels.c.cluster_role == DEFAULT_ROLE), + ) + .order_by(sa.desc(kernels.c.created_at)) + ) + match_sid_by_prefix_name = ( sa.select(info_cols) .select_from(kernels) .where( @@ -367,7 +388,7 @@ async def match_session_ids( [kernels.c.session_id], ) .select_from(kernels) - .where(cond_name) + .where(cond_prefix_name) .group_by(kernels.c.session_id) .limit(max_matches).offset(0), )) & @@ -376,7 +397,8 @@ async def match_session_ids( .order_by(sa.desc(kernels.c.created_at)) ) if for_update: - match_sid_by_name = match_sid_by_name.with_for_update() + match_sid_by_equal_name = match_sid_by_equal_name.with_for_update() + match_sid_by_prefix_name = match_sid_by_prefix_name.with_for_update() match_sid_by_session_id = ( sa.select(info_cols) .select_from(kernels) @@ -398,11 +420,14 @@ async def match_session_ids( match_sid_by_session_id = match_sid_by_session_id.with_for_update() for match_query in [ match_sid_by_session_id, - match_sid_by_name, + match_sid_by_equal_name, + match_sid_by_prefix_name, match_sid_by_id, ]: result = await db_connection.execute(match_query) rows = result.fetchall() + print("rows") + print("rows, len(lows)", rows, len(rows)) if not rows: continue return [ From b1fff9cd3e0b828019df6dfe310b657b06fa26cf Mon Sep 17 00:00:00 2001 From: jhyeo Date: Mon, 20 Dec 2021 22:38:07 +0900 Subject: [PATCH 4/7] Add blank line --- src/ai/backend/manager/models/kernel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ai/backend/manager/models/kernel.py b/src/ai/backend/manager/models/kernel.py index c573a0bc2..2073f93e0 100644 --- a/src/ai/backend/manager/models/kernel.py +++ b/src/ai/backend/manager/models/kernel.py @@ -301,6 +301,7 @@ class SessionInfo(TypedDict): status: KernelStatus created_at: datetime + async def match_session_ids( session_name_or_id: Union[str, UUID], access_key: AccessKey, From 0433f835191808d116a48f0f326d4ce6789a4e18 Mon Sep 17 00:00:00 2001 From: jhyeo Date: Mon, 20 Dec 2021 22:54:30 +0900 Subject: [PATCH 5/7] Add news fragment --- changes/506.fix | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/506.fix diff --git a/changes/506.fix b/changes/506.fix new file mode 100644 index 000000000..4ad93ef03 --- /dev/null +++ b/changes/506.fix @@ -0,0 +1 @@ +Fix too many session matching problem when exectue app by session_name From 13c491a613be27404f58ade48c88c5235237a05a Mon Sep 17 00:00:00 2001 From: jhyeo Date: Mon, 20 Dec 2021 23:00:19 +0900 Subject: [PATCH 6/7] Del test print --- src/ai/backend/manager/models/kernel.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ai/backend/manager/models/kernel.py b/src/ai/backend/manager/models/kernel.py index 2073f93e0..5b2d43b69 100644 --- a/src/ai/backend/manager/models/kernel.py +++ b/src/ai/backend/manager/models/kernel.py @@ -427,8 +427,6 @@ async def match_session_ids( ]: result = await db_connection.execute(match_query) rows = result.fetchall() - print("rows") - print("rows, len(lows)", rows, len(rows)) if not rows: continue return [ From 3213a75cd4b7451806e6da2fa10432b96443394c Mon Sep 17 00:00:00 2001 From: Joongi Kim Date: Tue, 21 Dec 2021 11:01:45 +0900 Subject: [PATCH 7/7] docs: Clarify the news fragment --- changes/506.fix | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changes/506.fix b/changes/506.fix index 4ad93ef03..4216a4358 100644 --- a/changes/506.fix +++ b/changes/506.fix @@ -1 +1 @@ -Fix too many session matching problem when exectue app by session_name +Fix "too many sessions matched" error when the given session name has an exact match with additional prefix matches