From a2dec4f45b5620974df2d271d8bbd48072c01e5b Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 27 Mar 2024 03:33:17 +0000 Subject: [PATCH] Format with ruff --- src/ai/backend/common/types.py | 8 +++---- .../backend/manager/models/scaling_group.py | 24 +++++++++---------- src/ai/backend/manager/scheduler/types.py | 8 +++---- 3 files changed, 18 insertions(+), 22 deletions(-) diff --git a/src/ai/backend/common/types.py b/src/ai/backend/common/types.py index fa9685e0c1..4f5b602861 100644 --- a/src/ai/backend/common/types.py +++ b/src/ai/backend/common/types.py @@ -1190,11 +1190,9 @@ def from_json(cls, obj: Mapping[str, Any]) -> RoundRobinState: @classmethod def as_trafaret(cls) -> t.Trafaret: - return t.Dict( - { - t.Key("next_index"): t.Int, - } - ) + return t.Dict({ + t.Key("next_index"): t.Int, + }) # This is only used when AgentSelectionStrategy is ROUNDROBIN. diff --git a/src/ai/backend/manager/models/scaling_group.py b/src/ai/backend/manager/models/scaling_group.py index d10a060e5b..c28b4e9111 100644 --- a/src/ai/backend/manager/models/scaling_group.py +++ b/src/ai/backend/manager/models/scaling_group.py @@ -95,19 +95,17 @@ def from_json(cls, obj: Mapping[str, Any]) -> ScalingGroupOpts: @classmethod def as_trafaret(cls) -> t.Trafaret: - return t.Dict( - { - t.Key("allowed_session_types", default=["interactive", "batch"]): t.List( - tx.Enum(SessionTypes), min_length=1 - ), - t.Key("pending_timeout", default=0): tx.TimeDuration(allow_negative=False), - # Each scheduler impl refers an additional "config" key. - t.Key("config", default={}): t.Mapping(t.String, t.Any), - t.Key( - "agent_selection_strategy", default=AgentSelectionStrategy.DISPERSED - ): tx.Enum(AgentSelectionStrategy), - } - ).allow_extra("*") + return t.Dict({ + t.Key("allowed_session_types", default=["interactive", "batch"]): t.List( + tx.Enum(SessionTypes), min_length=1 + ), + t.Key("pending_timeout", default=0): tx.TimeDuration(allow_negative=False), + # Each scheduler impl refers an additional "config" key. + t.Key("config", default={}): t.Mapping(t.String, t.Any), + t.Key("agent_selection_strategy", default=AgentSelectionStrategy.DISPERSED): tx.Enum( + AgentSelectionStrategy + ), + }).allow_extra("*") scaling_groups = sa.Table( diff --git a/src/ai/backend/manager/scheduler/types.py b/src/ai/backend/manager/scheduler/types.py index 498ce73dbd..ecb5afcc33 100644 --- a/src/ai/backend/manager/scheduler/types.py +++ b/src/ai/backend/manager/scheduler/types.py @@ -512,10 +512,10 @@ async def select_agent( sgroup_name = roundrobin_context.sgroup_name requested_architecture = get_requested_architecture(pending_session_or_kernel) - rr_state: RoundRobinState | None = ( - await sched_ctx.registry.shared_config.get_roundrobin_state( - sgroup_name, requested_architecture - ) + rr_state: ( + RoundRobinState | None + ) = await sched_ctx.registry.shared_config.get_roundrobin_state( + sgroup_name, requested_architecture ) if rr_state is None: