diff --git a/changes/2614.feature.md b/changes/2614.feature.md new file mode 100644 index 0000000000..c78bff0ac4 --- /dev/null +++ b/changes/2614.feature.md @@ -0,0 +1 @@ +Allow regular users to assign agent manually if `hide-agent` configuration is disabled diff --git a/src/ai/backend/client/cli/session/execute.py b/src/ai/backend/client/cli/session/execute.py index a708e2f28b..26c1415619 100644 --- a/src/ai/backend/client/cli/session/execute.py +++ b/src/ai/backend/client/cli/session/execute.py @@ -384,7 +384,6 @@ def prepare_mount_arg( type=list_expr, help=( "Show mapping list of tuple which mapped containers with agent. " - "When user role is Super Admin. " "(e.g., --assign-agent agent_id_1,agent_id_2,...)" ), ) diff --git a/src/ai/backend/client/cli/session/lifecycle.py b/src/ai/backend/client/cli/session/lifecycle.py index 7cd192bdd9..4b4e819ab5 100644 --- a/src/ai/backend/client/cli/session/lifecycle.py +++ b/src/ai/backend/client/cli/session/lifecycle.py @@ -134,7 +134,6 @@ def _create_cmd(docs: Optional[str] = None): type=list_expr, help=( "Assign the session to specific agents. " - "This option is only applicable when the user role is Super Admin. " "(e.g., --assign-agent agent_id_1,agent_id_2,...)" ), ) diff --git a/src/ai/backend/manager/api/session.py b/src/ai/backend/manager/api/session.py index 915837356f..ef3b2b365b 100644 --- a/src/ai/backend/manager/api/session.py +++ b/src/ai/backend/manager/api/session.py @@ -25,6 +25,7 @@ List, Mapping, MutableMapping, + Optional, Set, Tuple, Union, @@ -654,30 +655,32 @@ async def create_from_params(request: web.Request, params: dict[str, Any]) -> we else: raise InvalidAPIParameters("API version not supported") params["config"] = creation_config - if params["config"]["agent_list"] is not None and request["user"]["role"] != ( - UserRole.SUPERADMIN - ): - raise InsufficientPrivilege( - "You are not allowed to manually assign agents for your session." - ) - if request["user"]["role"] == (UserRole.SUPERADMIN): - if not params["config"]["agent_list"]: - pass + + root_ctx: RootContext = request.app["_root.context"] + + agent_list = cast(Optional[list[str]], params["config"]["agent_list"]) + if agent_list is not None: + if ( + request["user"]["role"] != UserRole.SUPERADMIN + and root_ctx.local_config["manager"]["hide-agents"] + ): + raise InsufficientPrivilege( + "You are not allowed to manually assign agents for your session." + ) + agent_count = len(agent_list) + if params["cluster_mode"] == "multi-node": + if agent_count != params["cluster_size"]: + raise InvalidAPIParameters( + "For multi-node cluster sessions, the number of manually assigned" + " agents must be same to the cluster size. Note that you may specify" + " duplicate agents in the list.", + ) else: - agent_count = len(params["config"]["agent_list"]) - if params["cluster_mode"] == "multi-node": - if agent_count != params["cluster_size"]: - raise InvalidAPIParameters( - "For multi-node cluster sessions, the number of manually assigned" - " agents must be same to the cluster size. Note that you may specify" - " duplicate agents in the list.", - ) - else: - if agent_count != 1: - raise InvalidAPIParameters( - "For non-cluster sessions and single-node cluster sessions, " - "you may specify only one manually assigned agent.", - ) + if agent_count != 1: + raise InvalidAPIParameters( + "For non-cluster sessions and single-node cluster sessions, " + "you may specify only one manually assigned agent.", + ) return await _create(request, params)