From d50ca7d46ba424344088adc3534fb52986bee119 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Thu, 2 Jul 2026 10:46:17 +0900 Subject: [PATCH 1/2] chore: enforce @override via mypy explicit-override rule Enable mypy's explicit-override error code and add the 2672 missing @override decorators across 609 files, driven by mypy's own diagnostics. Every method overriding a base-class method now carries @typing.override. ai.backend.client.cli.pretty is exempted per-module: its `total` property overrides tqdm's `total` instance attribute (not a method), which mypy cannot recognise @override on regardless of decorator placement. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 6 +- src/ai/backend/account_manager/config.py | 4 +- src/ai/backend/account_manager/exceptions.py | 5 +- src/ai/backend/account_manager/models/base.py | 9 +++ src/ai/backend/agent/agent.py | 2 + src/ai/backend/agent/alloc_map.py | 7 ++ src/ai/backend/agent/cli/context.py | 4 +- .../agent/dependencies/bootstrap/composer.py | 3 + .../agent/dependencies/bootstrap/config.py | 3 + .../agent/dependencies/bootstrap/etcd.py | 4 + .../dependencies/bootstrap/redis_config.py | 3 + src/ai/backend/agent/dependencies/composer.py | 3 + .../dependencies/infrastructure/composer.py | 3 + .../dependencies/infrastructure/docker.py | 4 + .../dependencies/infrastructure/redis.py | 4 + src/ai/backend/agent/docker/agent.py | 2 + src/ai/backend/agent/docker/intrinsic.py | 52 +++++++++++- src/ai/backend/agent/docker/kernel.py | 2 + src/ai/backend/agent/docker/metadata/root.py | 7 +- .../backend/agent/docker/metadata/server.py | 3 +- src/ai/backend/agent/dummy/intrinsic.py | 38 ++++++++- src/ai/backend/agent/dummy/kernel.py | 2 + src/ai/backend/agent/errors/agent.py | 32 +++++++- src/ai/backend/agent/errors/kernel.py | 8 ++ src/ai/backend/agent/errors/resources.py | 17 +++- src/ai/backend/agent/health/docker.py | 4 + src/ai/backend/agent/health/heartbeat.py | 6 +- src/ai/backend/agent/kernel.py | 17 ++++ src/ai/backend/agent/kubernetes/agent.py | 2 + src/ai/backend/agent/kubernetes/intrinsic.py | 38 ++++++++- src/ai/backend/agent/resources.py | 6 ++ src/ai/backend/agent/server.py | 2 + .../stage/kernel_lifecycle/docker/image.py | 3 + src/ai/backend/agent/stats.py | 2 + .../agent/tasks/cleanup_reported_kernels.py | 6 +- .../agent/tasks/collect_container_stat.py | 6 +- .../backend/agent/tasks/collect_node_stat.py | 6 +- .../agent/tasks/collect_process_stat.py | 6 +- .../tasks/report_kernel_commit_status.py | 6 +- src/ai/backend/agent/tasks/scan_images.py | 6 +- .../agent/tasks/sync_container_lifecycles.py | 6 +- src/ai/backend/agent/tasks/update_slots.py | 6 +- src/ai/backend/agent/types.py | 3 +- src/ai/backend/agent/utils.py | 3 + src/ai/backend/appproxy/common/config.py | 4 +- .../backend/appproxy/common/errors/config.py | 8 ++ src/ai/backend/appproxy/common/errors/http.py | 20 ++++- .../backend/appproxy/common/errors/service.py | 9 +++ src/ai/backend/appproxy/common/etcd.py | 4 +- src/ai/backend/appproxy/common/events.py | 18 ++++- src/ai/backend/appproxy/common/types.py | 2 + src/ai/backend/appproxy/common/utils.py | 2 + .../dependencies/bootstrap/composer.py | 3 + .../dependencies/bootstrap/config.py | 3 + .../coordinator/dependencies/composer.py | 3 + .../dependencies/infrastructure/composer.py | 3 + .../dependencies/infrastructure/database.py | 4 + .../dependencies/infrastructure/etcd.py | 4 + .../dependencies/infrastructure/redis.py | 5 +- .../appproxy/coordinator/errors/circuit.py | 5 ++ .../appproxy/coordinator/errors/config.py | 15 ++++ .../appproxy/coordinator/health/database.py | 4 + .../appproxy/coordinator/models/base.py | 30 +++++++ .../appproxy/coordinator/models/utils.py | 2 + src/ai/backend/appproxy/coordinator/pglock.py | 4 +- src/ai/backend/appproxy/worker/cli/context.py | 4 +- .../worker/dependencies/bootstrap/composer.py | 3 + .../worker/dependencies/bootstrap/config.py | 3 + .../appproxy/worker/dependencies/composer.py | 3 + .../dependencies/infrastructure/composer.py | 3 + .../dependencies/infrastructure/redis.py | 4 + .../backend/appproxy/worker/errors/circuit.py | 4 + .../backend/appproxy/worker/errors/config.py | 8 ++ .../backend/appproxy/worker/errors/process.py | 3 + .../proxy/backend/last_access_marker.py | 6 +- .../appproxy/worker/proxy/frontend/h2/base.py | 3 +- .../appproxy/worker/proxy/frontend/h2/port.py | 8 +- .../worker/proxy/frontend/h2/subdomain.py | 8 +- .../worker/proxy/frontend/http/base.py | 5 ++ .../worker/proxy/frontend/http/port.py | 5 +- .../worker/proxy/frontend/http/subdomain.py | 5 +- .../appproxy/worker/proxy/frontend/tcp.py | 9 ++- .../appproxy/worker/proxy/frontend/traefik.py | 19 ++++- .../appproxy/worker/tasks/heartbeat.py | 6 +- src/ai/backend/cli/extensions.py | 7 +- src/ai/backend/cli/params.py | 2 + .../v2/prometheus_query_preset/commands.py | 3 +- src/ai/backend/client/exceptions.py | 3 +- src/ai/backend/client/load_balancing.py | 4 + src/ai/backend/client/output/console.py | 11 ++- src/ai/backend/client/output/formatters.py | 29 ++++++- src/ai/backend/client/output/json.py | 11 ++- src/ai/backend/client/request.py | 5 ++ src/ai/backend/client/session.py | 10 +++ src/ai/backend/client/utils.py | 8 +- src/ai/backend/common/asyncio.py | 3 + src/ai/backend/common/auth/__init__.py | 8 +- src/ai/backend/common/bgtask/bgtask.py | 10 +++ .../common/bgtask/hooks/composite_hook.py | 2 + .../backend/common/bgtask/hooks/event_hook.py | 2 + .../common/bgtask/hooks/metric_hook.py | 3 +- .../common/bgtask/hooks/valkey_hook.py | 2 + .../backend/common/bgtask/tasks/heartbeat.py | 6 +- src/ai/backend/common/bgtask/tasks/retry.py | 6 +- src/ai/backend/common/cli.py | 6 +- .../common/clients/http_client/client_pool.py | 3 +- .../common/clients/valkey_client/client.py | 20 ++++- .../valkey_client/valkey_tus/exceptions.py | 5 ++ .../common/configs/loader/config_overrider.py | 1 + .../common/configs/loader/env_loader.py | 1 + .../common/configs/loader/etcd_loader.py | 1 + .../common/configs/loader/etcd_watcher.py | 3 +- .../common/configs/loader/loader_chain.py | 1 + .../common/configs/loader/toml_loader.py | 1 + src/ai/backend/common/cron/local_cron.py | 4 +- src/ai/backend/common/data/endpoint/types.py | 3 +- .../common/data/model_deployment/types.py | 3 +- .../common/data/notification/messages.py | 7 +- src/ai/backend/common/dependencies/base.py | 4 +- .../common/dependencies/stacks/builder.py | 4 +- .../common/dependencies/stacks/visualizing.py | 4 +- src/ai/backend/common/docker.py | 11 +++ src/ai/backend/common/dto/manager/context.py | 4 +- src/ai/backend/common/dto/manager/query.py | 3 +- src/ai/backend/common/enum_extension.py | 5 +- src/ai/backend/common/etcd.py | 12 +++ .../events/event_types/artifact/anycast.py | 6 ++ .../events/event_types/bgtask/broadcast.py | 2 + .../events/event_types/kernel/anycast.py | 3 + .../common/events/event_types/log/anycast.py | 7 +- .../event_types/notification/anycast.py | 2 + .../events/event_types/session/anycast.py | 7 ++ .../events/event_types/session/broadcast.py | 3 + .../events/event_types/vfolder/anycast.py | 12 +++ .../events/event_types/volume/broadcast.py | 10 +++ .../common/events/hub/propagators/bypass.py | 4 + .../common/events/hub/propagators/cache.py | 4 + src/ai/backend/common/events/types.py | 1 + src/ai/backend/common/exception.py | 79 ++++++++++++++++++- .../common/health_checker/checkers/etcd.py | 4 + .../common/health_checker/checkers/http.py | 4 + .../health_checker/checkers/prometheus.py | 4 + .../common/health_checker/checkers/valkey.py | 5 +- .../common/health_checker/exceptions.py | 9 +++ src/ai/backend/common/json.py | 3 +- src/ai/backend/common/jwt/exceptions.py | 4 + src/ai/backend/common/leader/exceptions.py | 4 + .../backend/common/leader/tasks/event_task.py | 6 +- .../common/leader/tasks/leader_cron.py | 4 +- .../common/leader/valkey_leader_election.py | 7 +- src/ai/backend/common/lock.py | 8 +- .../common/message_queue/hiredis_queue.py | 9 ++- .../message_queue/redis_queue/exceptions.py | 3 + src/ai/backend/common/metrics/safe.py | 10 ++- src/ai/backend/common/plugin/event.py | 4 +- src/ai/backend/common/plugin/monitor.py | 6 +- src/ai/backend/common/redis_client.py | 4 +- .../common/resilience/policies/metrics.py | 3 +- .../common/resilience/policies/retry.py | 4 +- .../common/resilience/policies/timeout.py | 4 +- .../etcd_discovery/service_discovery.py | 9 ++- .../service_discovery/event_publisher.py | 6 +- .../redis_discovery/service_discovery.py | 9 ++- .../service_discovery/service_discovery.py | 3 +- src/ai/backend/common/stage/types.py | 3 + src/ai/backend/common/typed_validators.py | 4 +- src/ai/backend/common/types.py | 26 ++++++ src/ai/backend/common/utils.py | 3 + src/ai/backend/common/validators.py | 2 + src/ai/backend/common/web/session/__init__.py | 8 ++ .../common/web/session/redis_storage.py | 1 + src/ai/backend/install/app.py | 9 ++- src/ai/backend/install/context.py | 10 ++- src/ai/backend/install/types.py | 3 + src/ai/backend/install/widgets.py | 7 +- src/ai/backend/kernel/app/__init__.py | 6 +- src/ai/backend/kernel/logging.py | 6 +- src/ai/backend/kernel/python/__init__.py | 6 +- .../backend/kernel/python/drawing/turtle.py | 6 +- src/ai/backend/kernel/utils.py | 3 +- src/ai/backend/logging/config.py | 4 +- src/ai/backend/logging/formatter.py | 8 +- src/ai/backend/logging/handler/intrinsic.py | 1 + src/ai/backend/logging/utils.py | 1 + src/ai/backend/manager/actions/action/base.py | 3 +- .../actions/validators/auth_validator.py | 3 + src/ai/backend/manager/api/gql/agent/types.py | 3 +- .../api/gql/app_config_allow_list/types.py | 3 +- .../api/gql/app_config_definition/types.py | 3 +- .../backend/manager/api/gql/artifact/types.py | 3 +- .../manager/api/gql/audit_log/types/node.py | 3 +- .../api/gql/container_registry/types.py | 3 +- .../api/gql/deployment/types/access_token.py | 3 +- .../api/gql/deployment/types/auto_scaling.py | 3 +- .../api/gql/deployment/types/deployment.py | 3 +- .../api/gql/deployment/types/replica.py | 3 +- .../api/gql/deployment/types/revision.py | 3 +- .../manager/api/gql/deployment/types/route.py | 3 +- src/ai/backend/manager/api/gql/domain.py | 2 + .../manager/api/gql/domain_v2/types/node.py | 3 +- .../api/gql/extensions/exception_handler.py | 3 +- .../manager/api/gql/extensions/logging_ext.py | 3 +- .../manager/api/gql/extensions/metric.py | 3 +- .../manager/api/gql/extensions/validation.py | 2 + .../manager/api/gql/huggingface_registry.py | 3 +- src/ai/backend/manager/api/gql/image/types.py | 4 +- .../manager/api/gql/image_federation.py | 2 + .../backend/manager/api/gql/kernel/types.py | 3 +- .../manager/api/gql/legacy_node_stubs.py | 5 ++ .../manager/api/gql/notification/types.py | 4 +- .../backend/manager/api/gql/object_storage.py | 3 +- src/ai/backend/manager/api/gql/project.py | 2 + .../manager/api/gql/project_v2/types/node.py | 3 +- .../manager/api/gql/pydantic_compat.py | 14 +++- .../manager/api/gql/rbac/types/entity.py | 3 +- .../manager/api/gql/rbac/types/permission.py | 3 +- .../manager/api/gql/rbac/types/role.py | 4 +- .../manager/api/gql/reservoir_registry.py | 3 +- .../api/gql/resource_group/federation.py | 2 + .../manager/api/gql/resource_group/types.py | 3 +- .../manager/api/gql/resource_slot/types.py | 5 +- .../api/gql/scheduling_history/types.py | 5 +- src/ai/backend/manager/api/gql/schema.py | 2 + .../backend/manager/api/gql/session/types.py | 3 +- .../manager/api/gql/session_federation.py | 2 + .../manager/api/gql/storage_namespace.py | 3 +- .../manager/api/gql/user/types/node.py | 3 +- .../manager/api/gql/user_federation.py | 2 + src/ai/backend/manager/api/gql/vfolder.py | 2 + .../manager/api/gql/vfolder_v2/types/node.py | 3 +- src/ai/backend/manager/api/gql/vfs_storage.py | 3 +- .../manager/api/rest/manager/handler.py | 16 +++- src/ai/backend/manager/api/utils.py | 2 + src/ai/backend/manager/api/wsproxy.py | 3 +- .../backend/manager/clients/agent/client.py | 4 +- .../manager/clients/prometheus/querier.py | 2 + src/ai/backend/manager/config/bootstrap.py | 3 +- .../config/loader/legacy_etcd_loader.py | 1 + src/ai/backend/manager/config/unified.py | 3 +- src/ai/backend/manager/config_legacy.py | 2 + .../manager/container_registry/local.py | 1 + .../backend/manager/data/common/sentinel.py | 3 +- src/ai/backend/manager/data/group/types.py | 1 + .../dependencies/agents/agent_client_pool.py | 3 + .../agents/appproxy_client_pool.py | 3 + .../manager/dependencies/agents/composer.py | 4 +- .../agents/deployment_controller.py | 4 +- .../manager/dependencies/agents/registry.py | 3 + .../dependencies/agents/route_controller.py | 3 + .../agents/scheduling_controller.py | 3 + .../dependencies/bootstrap/composer.py | 3 + .../manager/dependencies/bootstrap/config.py | 3 + .../manager/dependencies/bootstrap/etcd.py | 4 + .../dependencies/components/agent_cache.py | 3 + .../dependencies/components/composer.py | 3 + .../dependencies/components/storage.py | 3 + .../backend/manager/dependencies/composer.py | 3 + .../manager/dependencies/config/provider.py | 4 +- .../manager/dependencies/domain/composer.py | 4 +- .../dependencies/domain/distributed_lock.py | 4 +- .../dependencies/domain/notification.py | 3 + .../dependencies/domain/repositories.py | 4 +- .../manager/dependencies/domain/services.py | 4 +- src/ai/backend/manager/dependencies/errors.py | 4 + .../dependencies/infrastructure/composer.py | 3 + .../dependencies/infrastructure/database.py | 4 + .../dependencies/infrastructure/redis.py | 4 + .../dependencies/messaging/composer.py | 3 + .../dependencies/messaging/event_fetcher.py | 3 + .../dependencies/messaging/event_hub.py | 3 + .../dependencies/messaging/event_producer.py | 3 + .../dependencies/messaging/message_queue.py | 4 +- .../dependencies/orchestration/composer.py | 3 + .../orchestration/idle_checker.py | 3 + .../orchestration/leader_election.py | 3 + .../dependencies/orchestration/sokovan.py | 3 + .../manager/dependencies/plugins/composer.py | 3 + .../dependencies/plugins/event_dispatcher.py | 3 + .../manager/dependencies/plugins/hook.py | 3 + .../dependencies/plugins/monitoring.py | 6 +- .../manager/dependencies/plugins/network.py | 3 + .../processing/agent_lost_checker.py | 4 +- .../processing/bgtask_registry.py | 3 + .../dependencies/processing/composer.py | 3 + .../processing/event_dispatcher.py | 3 + .../processing/log_cleanup_timer.py | 3 + .../processing/manager_status_watcher.py | 4 +- .../dependencies/processing/processors.py | 3 + .../dependencies/processing/stats_reporter.py | 4 +- .../system/background_task_manager.py | 3 + .../manager/dependencies/system/composer.py | 3 + .../dependencies/system/cors_options.py | 3 + .../dependencies/system/gql_adapter.py | 3 + .../dependencies/system/health_probe.py | 3 + .../dependencies/system/jwt_validator.py | 3 + .../manager/dependencies/system/metrics.py | 3 + .../dependencies/system/prometheus_client.py | 3 + .../dependencies/system/service_discovery.py | 3 + src/ai/backend/manager/errors/agent.py | 3 + src/ai/backend/manager/errors/api.py | 9 +++ src/ai/backend/manager/errors/appproxy.py | 4 + src/ai/backend/manager/errors/artifact.py | 18 +++++ .../manager/errors/artifact_registry.py | 7 ++ src/ai/backend/manager/errors/auth.py | 19 +++++ src/ai/backend/manager/errors/bgtask.py | 3 + src/ai/backend/manager/errors/common.py | 13 ++- .../manager/errors/container_registry.py | 4 + src/ai/backend/manager/errors/deployment.py | 15 ++++ src/ai/backend/manager/errors/export.py | 5 ++ src/ai/backend/manager/errors/fair_share.py | 5 ++ src/ai/backend/manager/errors/image.py | 16 ++++ src/ai/backend/manager/errors/kernel.py | 26 +++++- src/ai/backend/manager/errors/keypair.py | 5 ++ src/ai/backend/manager/errors/notification.py | 10 +++ .../backend/manager/errors/object_storage.py | 4 + src/ai/backend/manager/errors/permission.py | 10 +++ src/ai/backend/manager/errors/repository.py | 12 ++- src/ai/backend/manager/errors/resource.py | 48 +++++++++++ .../backend/manager/errors/resource_slot.py | 6 ++ .../backend/manager/errors/role_invitation.py | 5 ++ src/ai/backend/manager/errors/role_preset.py | 4 + src/ai/backend/manager/errors/service.py | 17 ++++ src/ai/backend/manager/errors/storage.py | 29 ++++++- src/ai/backend/manager/errors/user.py | 11 +++ src/ai/backend/manager/errors/vfs_storage.py | 3 + .../manager/event_dispatcher/reporters.py | 4 + .../manager/events/hub/propagators/session.py | 5 +- src/ai/backend/manager/health/database.py | 4 + src/ai/backend/manager/health/rpc.py | 5 ++ src/ai/backend/manager/idle.py | 13 +++ src/ai/backend/manager/models/agent/row.py | 2 + src/ai/backend/manager/models/artifact/row.py | 4 +- .../manager/models/artifact_registries/row.py | 3 +- .../manager/models/artifact_revision/row.py | 3 +- .../backend/manager/models/audit_log/row.py | 3 + src/ai/backend/manager/models/base.py | 63 +++++++++++++++ src/ai/backend/manager/models/domain/row.py | 3 + src/ai/backend/manager/models/group/row.py | 3 + .../backend/manager/models/hasher/bcrypt.py | 4 + .../backend/manager/models/hasher/pbkdf2.py | 5 ++ src/ai/backend/manager/models/hasher/sha.py | 5 ++ src/ai/backend/manager/models/hasher/types.py | 3 +- .../models/huggingface_registry/row.py | 4 +- src/ai/backend/manager/models/image/row.py | 4 + .../manager/models/object_storage/row.py | 4 +- .../backend/manager/models/rbac/__init__.py | 17 +++- .../backend/manager/models/rbac/exceptions.py | 5 ++ .../manager/models/reservoir_registry/row.py | 4 +- .../manager/models/scaling_group/row.py | 3 + src/ai/backend/manager/models/session/row.py | 2 + src/ai/backend/manager/models/storage.py | 2 + src/ai/backend/manager/models/utils.py | 2 + src/ai/backend/manager/models/vfolder/row.py | 2 + .../backend/manager/models/vfs_storage/row.py | 4 +- src/ai/backend/manager/network/overlay.py | 7 +- .../notification/channels/email/channel.py | 2 + .../notification/channels/webhook/channel.py | 2 + src/ai/backend/manager/pglock.py | 4 +- .../backend/manager/plugin/error_monitor.py | 5 ++ src/ai/backend/manager/plugin/keypair/hook.py | 6 +- .../backend/manager/plugin/keypair/webapp.py | 6 +- .../manager/plugin/openid/exceptions.py | 3 + src/ai/backend/manager/plugin/openid/hook.py | 5 ++ .../backend/manager/plugin/openid/webapp.py | 9 +++ .../backend/manager/plugin/totp/exception.py | 5 ++ src/ai/backend/manager/plugin/totp/hook.py | 6 +- src/ai/backend/manager/plugin/totp/webapp.py | 6 +- .../manager/repositories/auth/types.py | 5 ++ .../manager/repositories/base/pagination.py | 14 +++- .../repositories/deployment/types/endpoint.py | 4 +- .../manager/repositories/domain/types.py | 3 + .../manager/repositories/fair_share/types.py | 7 ++ .../manager/repositories/group/types.py | 5 ++ .../manager/repositories/image/creators.py | 1 + .../manager/repositories/keypair/types.py | 5 ++ .../manager/repositories/model_card/types.py | 5 ++ .../manager/repositories/ops/rbac/provider.py | 2 + .../repositories/ops/sokovan/provider.py | 3 +- .../permission_controller/types.py | 8 +- .../resource_preset/db_source/types.py | 5 +- .../resource_usage_history/types.py | 8 +- .../repositories/role_invitation/types.py | 7 ++ .../repositories/scheduling_history/types.py | 8 +- .../manager/repositories/session/types.py | 3 + .../manager/repositories/user/types.py | 7 ++ .../manager/repositories/vfolder/types.py | 5 ++ .../services/image/actions/modify_image.py | 1 + .../model_serving/actions/force_sync.py | 2 + .../actions/scale_service_replicas.py | 1 + .../manager/sokovan/deployment/exceptions.py | 7 ++ .../deployment/group/lifecycle/source.py | 4 + .../deployment/group/lifecycle/types.py | 12 +++ .../deployment/group/scaling/source.py | 3 + .../sokovan/deployment/group/scaling/types.py | 10 +++ .../sokovan/deployment/handlers/destroying.py | 8 ++ .../sokovan/deployment/handlers/replica.py | 8 ++ .../route/handlers/appproxy_sync.py | 8 ++ .../deployment/route/handlers/draining.py | 8 ++ .../deployment/route/handlers/health_check.py | 8 ++ .../route/handlers/observer/health_check.py | 3 + .../route/handlers/probe_target_sync.py | 8 ++ .../deployment/route/handlers/provisioning.py | 8 ++ .../route/handlers/route_eviction.py | 8 ++ .../deployment/route/handlers/running.py | 8 ++ .../route/handlers/service_discovery_sync.py | 8 ++ .../deployment/route/handlers/starting.py | 8 ++ .../deployment/route/handlers/terminating.py | 8 ++ .../deployment/route/handlers/warming_up.py | 8 ++ .../manager/sokovan/reconciler/base.py | 3 + .../manager/sokovan/reconciler/flag.py | 3 + .../manager/sokovan/recorder/exceptions.py | 5 ++ .../manager/sokovan/scheduler/exceptions.py | 2 + .../handlers/cleanup/force_terminated.py | 5 +- .../handlers/kernel/sweep_stale_kernels.py | 7 +- .../handlers/lifecycle/check_precondition.py | 8 +- .../lifecycle/deprioritize_sessions.py | 7 ++ .../handlers/lifecycle/schedule_sessions.py | 8 +- .../handlers/lifecycle/start_sessions.py | 8 +- .../handlers/lifecycle/terminate_sessions.py | 8 +- .../handlers/maintenance/sweep_sessions.py | 7 ++ .../manager/sokovan/scheduler/hooks/status.py | 3 + .../post_processors/cache_invalidation.py | 3 +- .../kernel_schedule_marking.py | 3 +- .../post_processors/schedule_marking.py | 3 +- .../allocators/repository_allocator.py | 5 ++ .../provisioner/selectors/concentrated.py | 4 + .../provisioner/selectors/dispersed.py | 4 + .../scheduler/provisioner/selectors/legacy.py | 4 + .../provisioner/selectors/roundrobin.py | 4 + .../provisioner/validators/concurrency.py | 5 ++ .../provisioner/validators/dependencies.py | 5 ++ .../validators/domain_resource_limit.py | 5 ++ .../provisioner/validators/exceptions.py | 24 ++++++ .../validators/group_resource_limit.py | 5 ++ .../validators/keypair_resource_limit.py | 5 ++ .../validators/pending_session_count_limit.py | 5 ++ .../pending_session_resource_limit.py | 5 ++ .../provisioner/validators/reserved_batch.py | 4 + .../validators/user_resource_limit.py | 5 ++ .../assign_container_user_mapping_rule.py | 4 + .../preparers/assign_network_config_rule.py | 4 + .../preparers/assign_user_identity_rule.py | 4 + .../preparers/build_internal_data_rule.py | 4 +- .../compute_kernel_resources_rule.py | 4 +- .../preparers/expand_kernel_groups_rule.py | 4 + .../preparers/inject_session_environ_rule.py | 4 + .../merge_resource_group_defaults_rule.py | 4 + .../preparers/resolve_vfolder_mounts_rule.py | 4 + .../concurrent_session_limit_rule.py | 4 + .../validators/container_limit_rule.py | 4 + .../dotfile_vfolder_conflict_rule.py | 3 + .../validators/image_slot_type_rule.py | 3 + .../validators/inference_model_folder_rule.py | 4 + .../validators/mount_name_validation_rule.py | 4 + .../validators/requested_slot_type_rule.py | 3 + .../validators/required_resource_slot_rule.py | 3 + .../validators/resource_limit_rule.py | 4 +- .../validators/service_port_rule.py | 4 +- .../manager/tasks/agent_lost_checker.py | 6 +- .../backend/manager/tasks/stats_reporter.py | 6 +- src/ai/backend/storage/cli/context.py | 4 +- src/ai/backend/storage/data/storage/types.py | 4 +- .../dependencies/bootstrap/composer.py | 3 + .../storage/dependencies/bootstrap/config.py | 3 + .../backend/storage/dependencies/composer.py | 3 + .../dependencies/infrastructure/composer.py | 3 + .../dependencies/infrastructure/etcd.py | 4 + .../dependencies/infrastructure/redis.py | 4 + .../storage/dependencies/storage/composer.py | 3 + .../dependencies/storage/storage_pool.py | 3 + src/ai/backend/storage/errors/base.py | 6 ++ src/ai/backend/storage/errors/common.py | 14 ++++ src/ai/backend/storage/errors/object.py | 21 +++++ src/ai/backend/storage/errors/process.py | 10 +++ src/ai/backend/storage/errors/quota.py | 10 +++ src/ai/backend/storage/errors/vfolder.py | 3 + src/ai/backend/storage/errors/volume.py | 5 ++ src/ai/backend/storage/exception.py | 36 +++++++++ src/ai/backend/storage/plugin.py | 4 +- .../storage/services/artifacts/common.py | 4 + .../storage/services/artifacts/huggingface.py | 1 + .../storage/services/artifacts/reservoir.py | 1 + .../backend/storage/storages/vfs_storage.py | 1 + .../storage/volumes/cephfs/__init__.py | 13 ++- .../backend/storage/volumes/ddn/__init__.py | 9 ++- .../storage/volumes/dellemc/__init__.py | 13 ++- .../storage/volumes/dellemc/exceptions.py | 4 + .../backend/storage/volumes/gpfs/__init__.py | 15 +++- .../storage/volumes/gpfs/exceptions.py | 13 +++ .../storage/volumes/hammerspace/exception.py | 7 ++ .../storage/volumes/netapp/__init__.py | 20 ++++- .../backend/storage/volumes/noop/__init__.py | 36 +++++++++ .../storage/volumes/purestorage/__init__.py | 9 ++- .../storage/volumes/purestorage/exceptions.py | 3 + .../storage/volumes/purestorage/rapidfiles.py | 6 ++ .../volumes/purestorage/rapidfiles_v2.py | 4 + src/ai/backend/storage/volumes/types.py | 7 +- .../backend/storage/volumes/vast/__init__.py | 13 ++- .../storage/volumes/vast/exceptions.py | 4 + .../backend/storage/volumes/vfs/__init__.py | 41 ++++++++++ .../backend/storage/volumes/weka/__init__.py | 13 ++- .../storage/volumes/weka/exceptions.py | 9 +++ .../backend/storage/volumes/xfs/__init__.py | 10 ++- src/ai/backend/storage/watcher.py | 14 +++- src/ai/backend/test/cli/utils.py | 4 +- src/ai/backend/test/contexts/auth.py | 5 ++ .../backend/test/contexts/client_session.py | 4 + src/ai/backend/test/contexts/context.py | 3 +- src/ai/backend/test/contexts/domain.py | 3 + src/ai/backend/test/contexts/group.py | 4 + src/ai/backend/test/contexts/model_service.py | 5 ++ .../backend/test/contexts/resource_policy.py | 3 + src/ai/backend/test/contexts/scaling_group.py | 3 + src/ai/backend/test/contexts/tester.py | 3 + src/ai/backend/test/templates/auth/keypair.py | 2 + src/ai/backend/test/templates/auth/login.py | 1 + .../test/templates/group/add_group_to_user.py | 1 + src/ai/backend/test/templates/group/group.py | 1 + .../test/templates/model_service/endpoint.py | 3 + .../test/templates/model_service/jwt_token.py | 1 + .../test/templates/session/batch_session.py | 1 + .../templates/session/dependent_session.py | 1 + .../templates/session/interactive_session.py | 3 + .../templates/session/session_template.py | 3 + .../vfolder_mounted_interactive_session.py | 1 + src/ai/backend/test/templates/template.py | 7 +- src/ai/backend/test/templates/user/user.py | 1 + .../test/templates/vfolder/file_uploader.py | 1 + .../test/templates/vfolder/general_vfolder.py | 2 + .../backend/test/templates/vfolder/invite.py | 3 + .../test/templates/vfolder/model_vfolder.py | 1 + .../backend/test/templates/vfolder/share.py | 1 + .../backend/test/testcases/session/commit.py | 3 + .../creation_failure_command_timeout.py | 2 + .../session/creation_failure_low_resources.py | 3 + .../creation_failure_schedule_timeout.py | 3 + .../creation_failure_too_many_container.py | 3 + .../session/creation_failure_wrong_command.py | 2 + .../test/testcases/session/execution.py | 5 ++ .../test/testcases/session/filecheck.py | 2 + src/ai/backend/test/testcases/spec_manager.py | 4 +- .../backend/test/testcases/vfolder/access.py | 4 + .../backend/test/testcases/vfolder/clone.py | 3 + .../test/testcases/vfolder/delete_files.py | 4 + .../test/testcases/vfolder/download.py | 2 + .../test/testcases/vfolder/invite_failures.py | 4 + .../test/testcases/vfolder/list_files.py | 3 + src/ai/backend/test/testcases/vfolder/move.py | 3 + .../backend/test/testcases/vfolder/purge.py | 3 + .../test/testcases/vfolder/rename_file.py | 3 + .../backend/test/testcases/vfolder/restore.py | 3 + .../backend/test/testcases/vfolder/share.py | 3 + src/ai/backend/test/tester/exporter.py | 7 ++ src/ai/backend/testutils/db.py | 3 +- src/ai/backend/web/cli/context.py | 4 +- .../apollo_router_pool/health_checker.py | 7 ++ .../clients/manager_pool/health_checker.py | 7 ++ .../web/dependencies/components/composer.py | 3 + .../components/hive_router_client.py | 3 + .../dependencies/components/manager_client.py | 3 + src/ai/backend/web/dependencies/composer.py | 3 + .../dependencies/infrastructure/composer.py | 3 + .../web/dependencies/infrastructure/redis.py | 4 + src/ai/backend/web/errors.py | 5 ++ src/ai/backend/web/logging.py | 4 +- src/ai/backend/web/template.py | 4 +- stubs/trafaret/base.pyi | 3 +- stubs/trafaret/constructor.pyi | 2 + tests/unit/agent/test_affinity_map.py | 6 +- tests/unit/client/v2/test_anonymous_client.py | 2 + tests/unit/client/v2/test_infra_client.py | 3 +- tests/unit/client/v2/test_session_client.py | 3 +- tests/unit/client_v2/conftest.py | 2 + tests/unit/client_v2/test_base_client.py | 2 + .../configs/loader/test_loader_chain.py | 3 +- .../common/configs/loader/test_source_name.py | 3 +- tests/unit/common/cron/test_local_cron.py | 7 +- .../dependencies/stacks/test_visualizing.py | 12 ++- tests/unit/common/dependencies/test_base.py | 5 ++ tests/unit/common/events/test_dispatcher.py | 14 +++- tests/unit/common/events/test_hub.py | 13 ++- .../resilience/test_policy_composition.py | 9 ++- tests/unit/common/test_api_handlers.py | 4 +- tests/unit/common/test_distributed.py | 9 ++- tests/unit/common/test_events.py | 8 +- tests/unit/common/test_logging.py | 3 + tests/unit/common/test_plugin.py | 20 ++++- tests/unit/common/test_types.py | 4 +- .../common/valkey_leader/test_leader_cron.py | 6 +- .../test_leader_task_integration.py | 4 +- .../test_valkey_leader_client.py | 2 + tests/unit/manager/actions/test_processor.py | 8 +- tests/unit/manager/api/test_exceptions.py | 2 + tests/unit/manager/api/test_routing.py | 3 +- tests/unit/manager/models/test_base.py | 5 +- .../base/rbac/test_entity_creator.py | 7 +- .../base/rbac/test_entity_purger.py | 10 ++- .../base/rbac/test_scope_binder.py | 3 +- .../base/rbac/test_scope_unbinder.py | 7 +- .../manager/repositories/base/test_creator.py | 12 ++- .../manager/repositories/base/test_purger.py | 5 +- .../manager/repositories/base/test_querier.py | 6 +- .../manager/repositories/base/test_updater.py | 20 ++++- .../repositories/base/test_upserter.py | 5 +- .../manager/repositories/ops/test_provider.py | 8 +- .../deployment/test_deployment_service.py | 4 +- .../sokovan/reconciler/test_coordinator.py | 12 ++- .../preparers/test_session_spec_preparer.py | 3 + .../storage-proxy/storages/test_s3_client.py | 3 + .../storage-proxy/storages/test_service.py | 4 +- 610 files changed, 3337 insertions(+), 247 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3f0c239d570..4d82a3f8a03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -334,7 +334,7 @@ namespace_packages = true explicit_package_bases = true python_executable = "dist/export/python/virtualenvs/python-default/3.13.7/bin/python" disable_error_code = ["typeddict-unknown-key", "import-untyped"] -enable_error_code = ["possibly-undefined"] +enable_error_code = ["possibly-undefined", "explicit-override"] warn_unreachable = true # Disable errors in client CLI, web server, and tests @@ -377,6 +377,10 @@ module = [ ] ignore_errors = true +[[tool.mypy.overrides]] +module = "ai.backend.client.cli.pretty" +disable_error_code = ["explicit-override"] + [tool.pyright] pythonVersion = "3.13" extraPaths = ["src"] diff --git a/src/ai/backend/account_manager/config.py b/src/ai/backend/account_manager/config.py index 1e883467cea..dfd4ff33275 100644 --- a/src/ai/backend/account_manager/config.py +++ b/src/ai/backend/account_manager/config.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from pathlib import Path from pprint import pformat -from typing import Annotated, Any +from typing import Annotated, Any, override import click from pydantic import ( @@ -59,9 +59,11 @@ class HostPortPair(BaseSchema): host: Annotated[str, Field(examples=["127.0.0.1"])] port: Annotated[int, Field(gt=0, lt=65536, examples=[8201])] + @override def __repr__(self) -> str: return f"{self.host}:{self.port}" + @override def __str__(self) -> str: return self.__repr__() diff --git a/src/ai/backend/account_manager/exceptions.py b/src/ai/backend/account_manager/exceptions.py index 762c36373fc..4e72f5336b6 100644 --- a/src/ai/backend/account_manager/exceptions.py +++ b/src/ai/backend/account_manager/exceptions.py @@ -12,7 +12,7 @@ from __future__ import annotations import json -from typing import Any +from typing import Any, override from aiohttp import web @@ -51,6 +51,7 @@ def __init__( self.body_dict = body self.body = json.dumps(body).encode() + @override def __str__(self) -> str: lines = [] if self.extra_msg: @@ -61,6 +62,7 @@ def __str__(self) -> str: lines.append(" -> extra_data: " + repr(self.extra_data)) return "\n".join(lines) + @override def __repr__(self) -> str: lines = [] if self.extra_msg: @@ -73,6 +75,7 @@ def __repr__(self) -> str: lines.append(" -> extra_data: " + repr(self.extra_data)) return "\n".join(lines) + @override def __reduce__(self) -> tuple[type[BackendError], tuple[Any, ...], dict[str, Any]]: return ( type(self), diff --git a/src/ai/backend/account_manager/models/base.py b/src/ai/backend/account_manager/models/base.py index a0724bf8490..43b536cae4d 100644 --- a/src/ai/backend/account_manager/models/base.py +++ b/src/ai/backend/account_manager/models/base.py @@ -8,6 +8,7 @@ Self, TypeVar, cast, + override, ) import sqlalchemy as sa @@ -55,11 +56,13 @@ class GUID[UUID_SubType: uuid.UUID](TypeDecorator[uuid.UUID]): uuid_subtype_func: ClassVar[Callable[[Any], Any]] = lambda v: v cache_ok = True + @override def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator[Any]: if dialect.name == "postgresql": return cast(TypeDecorator[Any], dialect.type_descriptor(UUID())) return cast(TypeDecorator[Any], dialect.type_descriptor(CHAR(16))) + @override def process_bind_param( self, value: UUID_SubType | uuid.UUID | None, dialect: Dialect ) -> str | bytes | None: @@ -73,6 +76,7 @@ def process_bind_param( return str(value) return value.bytes + @override def process_result_value(self, value: Any, dialect: Dialect) -> UUID_SubType | None: if value is None: return value @@ -100,6 +104,7 @@ def __init__(self, enum_cls: type[T_StrEnum], **opts: Any) -> None: super().__init__(length=64, **opts) self._enum_cls = enum_cls + @override def process_bind_param( # type: ignore[override] self, value: T_StrEnum | None, @@ -107,6 +112,7 @@ def process_bind_param( # type: ignore[override] ) -> str | None: return value.value if value is not None else None + @override def process_result_value( # type: ignore[override] self, value: Any | None, @@ -114,10 +120,12 @@ def process_result_value( # type: ignore[override] ) -> T_StrEnum | None: return self._enum_cls(value) if value is not None else None + @override def copy(self, **kw: Any) -> Self: return StrEnumType(self._enum_cls, **self._opts) # type: ignore[return-value] @property + @override def python_type(self) -> type[T_StrEnum]: return self._enum_cls @@ -125,6 +133,7 @@ def python_type(self) -> type[T_StrEnum]: class PasswordColumn(TypeDecorator[str]): impl = VARCHAR + @override def process_bind_param(self, value: Any, dialect: Dialect) -> str: return hash_password(value) diff --git a/src/ai/backend/agent/agent.py b/src/ai/backend/agent/agent.py index 7f4c30e2620..ddd5087d280 100644 --- a/src/ai/backend/agent/agent.py +++ b/src/ai/backend/agent/agent.py @@ -42,6 +42,7 @@ Literal, ParamSpec, cast, + override, ) from uuid import UUID @@ -950,6 +951,7 @@ def __init__( self._sync_container_lifecycle_observer = SyncContainerLifecycleObserver.instance() self._clean_kernel_registry_task = asyncio.create_task(self._clean_kernel_registry_loop()) + @override async def __ainit__(self) -> None: """ An implementation of AbstractAgent would define its own ``__ainit__()`` method. diff --git a/src/ai/backend/agent/alloc_map.py b/src/ai/backend/agent/alloc_map.py index eaf446c7575..7ae26222575 100644 --- a/src/ai/backend/agent/alloc_map.py +++ b/src/ai/backend/agent/alloc_map.py @@ -15,6 +15,7 @@ Any, TypeVar, final, + override, ) import attr @@ -299,6 +300,7 @@ def __init__( } super().__init__(*args, **kwargs) + @override def allocate( self, slots: Mapping[SlotName, Decimal], @@ -479,6 +481,7 @@ def _allocate_evenly( return allocation + @override def apply_allocation( self, existing_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -487,6 +490,7 @@ def apply_allocation( for device_id, alloc in per_device_alloc.items(): self.allocations[slot_name][device_id] += alloc + @override def free( self, existing_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -516,6 +520,7 @@ def __init__( self.digits = Decimal(10) ** -2 # decimal points that is supported by agent self.powers = Decimal(100) # reciprocal of self.digits + @override def allocate( self, slots: Mapping[SlotName, Decimal], @@ -917,6 +922,7 @@ def allocate_across_devices( self.update_affinity_hint(slot_allocation, affinity_hint) return allocation + @override def apply_allocation( self, existing_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -925,6 +931,7 @@ def apply_allocation( for device_id, alloc in per_device_alloc.items(): self.allocations[slot_name][device_id] += alloc + @override def free( self, existing_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], diff --git a/src/ai/backend/agent/cli/context.py b/src/ai/backend/agent/cli/context.py index 8bfa409a561..343bf5a8e9b 100644 --- a/src/ai/backend/agent/cli/context.py +++ b/src/ai/backend/agent/cli/context.py @@ -2,7 +2,7 @@ from contextlib import AbstractContextManager from pathlib import Path -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, Self, override import click @@ -19,6 +19,7 @@ def __init__(self, log_level: LogLevel, config_path: Path | None = None) -> None self.config_path = config_path self.log_level = log_level + @override def __enter__(self) -> Self: from ai.backend.logging import LocalLogger @@ -31,6 +32,7 @@ def __enter__(self) -> Self: self._logger.__enter__() return self + @override def __exit__(self, *exc_info: Any) -> None: click_ctx = click.get_current_context() if click_ctx.invoked_subcommand != "start-server": diff --git a/src/ai/backend/agent/dependencies/bootstrap/composer.py b/src/ai/backend/agent/dependencies/bootstrap/composer.py index e9dda52c940..35b0d76e919 100644 --- a/src/ai/backend/agent/dependencies/bootstrap/composer.py +++ b/src/ai/backend/agent/dependencies/bootstrap/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.agent.config.unified import AgentUnifiedConfig from ai.backend.common.configs.redis import RedisConfig @@ -50,10 +51,12 @@ class AgentBootstrapComposer(DependencyComposer[AgentBootstrapInput, AgentBootst """ @property + @override def stage_name(self) -> str: return "bootstrap" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/agent/dependencies/bootstrap/config.py b/src/ai/backend/agent/dependencies/bootstrap/config.py index 694fde08cc1..bae799cd43e 100644 --- a/src/ai/backend/agent/dependencies/bootstrap/config.py +++ b/src/ai/backend/agent/dependencies/bootstrap/config.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.agent.config.unified import AgentConfigValidationContext, AgentUnifiedConfig from ai.backend.common import config as common_config @@ -35,10 +36,12 @@ class AgentConfigLoaderDependency( """ @property + @override def stage_name(self) -> str: return "config-loader" @asynccontextmanager + @override async def provide( self, setup_input: AgentConfigLoaderInput ) -> AsyncIterator[AgentUnifiedConfig]: diff --git a/src/ai/backend/agent/dependencies/bootstrap/etcd.py b/src/ai/backend/agent/dependencies/bootstrap/etcd.py index 2728b918ba8..c16dd875409 100644 --- a/src/ai/backend/agent/dependencies/bootstrap/etcd.py +++ b/src/ai/backend/agent/dependencies/bootstrap/etcd.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.agent.config.unified import AgentUnifiedConfig from ai.backend.common.dependencies import DependencyProvider @@ -18,10 +19,12 @@ class AgentEtcdDependency(DependencyProvider[AgentUnifiedConfig, AsyncEtcd]): """ @property + @override def stage_name(self) -> str: return "etcd" @asynccontextmanager + @override async def provide(self, setup_input: AgentUnifiedConfig) -> AsyncIterator[AsyncEtcd]: """Initialize and provide etcd client. @@ -60,6 +63,7 @@ async def provide(self, setup_input: AgentUnifiedConfig) -> AsyncIterator[AsyncE ) as etcd: yield etcd + @override def gen_liveness_checker(self, resource: AsyncEtcd) -> ServiceHealthChecker: """Liveness — stuck etcd connection observed; restart is the recovery path.""" return EtcdHealthChecker(etcd=resource) diff --git a/src/ai/backend/agent/dependencies/bootstrap/redis_config.py b/src/ai/backend/agent/dependencies/bootstrap/redis_config.py index 64d38edb83d..04ac7bd7794 100644 --- a/src/ai/backend/agent/dependencies/bootstrap/redis_config.py +++ b/src/ai/backend/agent/dependencies/bootstrap/redis_config.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common import config as common_config from ai.backend.common.configs.redis import RedisConfig @@ -16,10 +17,12 @@ class RedisConfigDependency(NonMonitorableDependencyProvider[AsyncEtcd, RedisCon """ @property + @override def stage_name(self) -> str: return "redis-config (temp)" @asynccontextmanager + @override async def provide(self, setup_input: AsyncEtcd) -> AsyncIterator[RedisConfig]: """Read redis config from etcd. diff --git a/src/ai/backend/agent/dependencies/composer.py b/src/ai/backend/agent/dependencies/composer.py index ff6fb99c693..cc335c212a8 100644 --- a/src/ai/backend/agent/dependencies/composer.py +++ b/src/ai/backend/agent/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.logging.types import LogLevel @@ -48,10 +49,12 @@ class AgentDependencyComposer(DependencyComposer[AgentDependencyInput, AgentDepe """ @property + @override def stage_name(self) -> str: return "agent-dependencies" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/agent/dependencies/infrastructure/composer.py b/src/ai/backend/agent/dependencies/infrastructure/composer.py index af03045a6c3..e8f1783547b 100644 --- a/src/ai/backend/agent/dependencies/infrastructure/composer.py +++ b/src/ai/backend/agent/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from aiodocker import Docker @@ -51,10 +52,12 @@ class AgentInfrastructureComposer( """ @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/agent/dependencies/infrastructure/docker.py b/src/ai/backend/agent/dependencies/infrastructure/docker.py index 9da3c8f0f90..6d9905cbb08 100644 --- a/src/ai/backend/agent/dependencies/infrastructure/docker.py +++ b/src/ai/backend/agent/dependencies/infrastructure/docker.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from aiodocker import Docker @@ -19,10 +20,12 @@ class DockerDependency(DependencyProvider[AgentUnifiedConfig, Docker]): """ @property + @override def stage_name(self) -> str: return "docker" @asynccontextmanager + @override async def provide(self, setup_input: AgentUnifiedConfig) -> AsyncIterator[Docker]: """Initialize and provide a Docker client. @@ -38,6 +41,7 @@ async def provide(self, setup_input: AgentUnifiedConfig) -> AsyncIterator[Docker finally: await docker.close() + @override def gen_liveness_checker( self, resource: Docker, diff --git a/src/ai/backend/agent/dependencies/infrastructure/redis.py b/src/ai/backend/agent/dependencies/infrastructure/redis.py index 2bf4a279780..830aa404302 100644 --- a/src/ai/backend/agent/dependencies/infrastructure/redis.py +++ b/src/ai/backend/agent/dependencies/infrastructure/redis.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_bgtask.client import ValkeyBgtaskClient from ai.backend.common.clients.valkey_client.valkey_container_log.client import ( @@ -53,10 +54,12 @@ class AgentValkeyDependency(DependencyProvider[RedisConfig, AgentValkeyClients]) """Provides lifecycle management for 4 agent-specific Valkey clients.""" @property + @override def stage_name(self) -> str: return "valkey" @asynccontextmanager + @override async def provide(self, setup_input: RedisConfig) -> AsyncIterator[AgentValkeyClients]: """Initialize and provide all agent Valkey clients. @@ -97,6 +100,7 @@ async def provide(self, setup_input: RedisConfig) -> AsyncIterator[AgentValkeyCl finally: await clients.close() + @override def gen_liveness_checker(self, resource: AgentValkeyClients) -> ServiceHealthChecker: """Liveness — Valkey connection-stuck observed; restart recovers.""" return ValkeyHealthChecker( diff --git a/src/ai/backend/agent/docker/agent.py b/src/ai/backend/agent/docker/agent.py index 89e541ee758..bbc2eae7db0 100644 --- a/src/ai/backend/agent/docker/agent.py +++ b/src/ai/backend/agent/docker/agent.py @@ -1516,6 +1516,7 @@ def __init__( [KernelRecoveryDataAdapterTarget(container_loader, container_writer)], ) + @override async def __ainit__(self) -> None: async with closing_async(Docker()) as docker: docker_host = "" @@ -1603,6 +1604,7 @@ async def __ainit__(self) -> None: blocklist=self.local_config.agent.block_network_plugins, ) + @override async def shutdown(self, stop_signal: signal.Signals) -> None: # Stop handling agent sock. if self.agent_sock_task is not None: diff --git a/src/ai/backend/agent/docker/intrinsic.py b/src/ai/backend/agent/docker/intrinsic.py index 1ff43b48754..653791eaf75 100644 --- a/src/ai/backend/agent/docker/intrinsic.py +++ b/src/ai/backend/agent/docker/intrinsic.py @@ -8,7 +8,7 @@ from collections.abc import Collection, Iterable, Mapping, Sequence from decimal import Decimal from pathlib import Path -from typing import Any, cast +from typing import Any, cast, override import aiohttp import psutil @@ -200,15 +200,19 @@ class CPUPlugin(AbstractComputePlugin): _docker: Docker + @override async def init(self, context: Any | None = None) -> None: self._docker = Docker() + @override async def cleanup(self) -> None: await self._docker.close() + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override async def list_devices(self) -> Collection[CPUDevice]: cores = await libnuma.get_available_cores() overcommit_factor = int(os.environ.get("BACKEND_CPU_OVERCOMMIT_FACTOR", "1")) @@ -227,15 +231,18 @@ async def list_devices(self) -> Collection[CPUDevice]: for core_idx in sorted(cores) ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() return { SlotName("cpu"): Decimal(sum(dev.processing_units for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return { "agent_version": __version__, @@ -243,6 +250,7 @@ async def extra_info(self) -> Mapping[str, str]: "os_type": platform.system(), } + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: _cstat = psutil.cpu_times(True) q = Decimal("0.000") @@ -269,6 +277,7 @@ async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasureme ), ] + @override async def gather_container_measures( self, ctx: StatContext, @@ -356,6 +365,7 @@ async def api_impl(container_id: str) -> float | None: ), ] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: @@ -421,6 +431,7 @@ async def api_impl(_cid: str, _pids: list[int]) -> list[Decimal | None]: ), ] + @override async def create_alloc_map(self) -> AbstractAllocMap: devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -432,10 +443,12 @@ async def create_alloc_map(self) -> AbstractAllocMap: }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: # TODO: move the sysconf hook in libbaihook.so here return [] + @override async def generate_docker_args( self, docker: Docker, @@ -451,6 +464,7 @@ async def generate_docker_args( }, } + @override async def restore_from_container( self, container: Container, @@ -469,6 +483,7 @@ async def restore_from_container( SlotName("cpu"): resource_spec.allocations[DeviceName("cpu")][SlotName("cpu")], }) + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -485,16 +500,19 @@ async def get_attached_devices( }) return attached_devices + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: return [] + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "cpu", @@ -527,15 +545,19 @@ class MemoryPlugin(AbstractComputePlugin): _docker: Docker + @override async def init(self, context: Any | None = None) -> None: self._docker = Docker() + @override async def cleanup(self) -> None: await self._docker.close() + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override async def list_devices(self) -> Collection[MemoryDevice]: memory_size = psutil.virtual_memory().total overcommit_factor = int(os.environ.get("BACKEND_MEM_OVERCOMMIT_FACTOR", "1")) @@ -550,18 +572,22 @@ async def list_devices(self) -> Collection[MemoryDevice]: ), ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() return { SlotName("mem"): Decimal(sum(dev.memory_size for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return {} + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: _mstat = psutil.virtual_memory() total_mem_used_bytes = Decimal(_mstat.total - _mstat.available) @@ -632,6 +658,7 @@ def get_disk_stat() -> tuple[Decimal, Decimal, dict[DeviceId, Measurement]]: ), ] + @override async def gather_container_measures( self, ctx: StatContext, container_ids: Sequence[str] ) -> Sequence[ContainerMeasurement]: @@ -912,6 +939,7 @@ async def api_impl( ), ] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: @@ -998,6 +1026,7 @@ async def api_impl( ), ] + @override async def create_alloc_map(self) -> AbstractAllocMap: devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -1010,9 +1039,11 @@ async def create_alloc_map(self) -> AbstractAllocMap: }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: return [] + @override async def generate_docker_args( self, docker: Docker, @@ -1026,6 +1057,7 @@ async def generate_docker_args( }, } + @override async def restore_from_container( self, container: Container, @@ -1040,6 +1072,7 @@ async def restore_from_container( SlotName("mem"): {DeviceId("root"): memory_limit}, }) + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -1056,16 +1089,19 @@ async def get_attached_devices( }) return attached_devices + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: return [] + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "ram", @@ -1078,18 +1114,23 @@ def get_metadata(self) -> AcceleratorMetadata: class OverlayNetworkPlugin(AbstractNetworkAgentPlugin[DockerKernel]): + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: return await super().update_plugin_config(plugin_config) + @override async def get_capabilities(self) -> set[ContainerNetworkCapability]: return set() + @override async def join_network( self, kernel_config: KernelCreationConfig, @@ -1111,23 +1152,29 @@ async def join_network( }, } + @override async def leave_network(self, kernel: DockerKernel) -> None: pass class HostNetworkPlugin(AbstractNetworkAgentPlugin[DockerKernel]): + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: return await super().update_plugin_config(plugin_config) + @override async def get_capabilities(self) -> set[ContainerNetworkCapability]: return {ContainerNetworkCapability.GLOBAL} + @override async def join_network( self, kernel_config: KernelCreationConfig, cluster_info: ClusterInfo, **kwargs: Any ) -> dict[str, Any]: @@ -1147,9 +1194,11 @@ async def join_network( }, } + @override async def leave_network(self, kernel: DockerKernel) -> None: pass + @override async def prepare_port_forward( self, kernel: DockerKernel, bind_host: str, ports: Iterable[tuple[int, int]], **kwargs: Any ) -> None: @@ -1173,6 +1222,7 @@ async def prepare_port_forward( lambda: (config_dir / "intrinsic-ports.json").write_bytes(dump_json(intrinsic_ports)), ) + @override async def expose_ports( self, kernel: DockerKernel, bind_host: str, ports: Iterable[tuple[int, int]], **kwargs: Any ) -> ContainerNetworkInfo: diff --git a/src/ai/backend/agent/docker/kernel.py b/src/ai/backend/agent/docker/kernel.py index 6b2d96e1419..cf79b2e41e0 100644 --- a/src/ai/backend/agent/docker/kernel.py +++ b/src/ai/backend/agent/docker/kernel.py @@ -79,9 +79,11 @@ def __init__( async def close(self) -> None: pass + @override def __getstate__(self) -> Mapping[str, Any]: return super().__getstate__() + @override def __setstate__(self, props: MutableMapping[str, Any]) -> None: if "network_driver" not in props: props["network_driver"] = "bridge" diff --git a/src/ai/backend/agent/docker/metadata/root.py b/src/ai/backend/agent/docker/metadata/root.py index bd90cab9fc3..fb616eed15e 100644 --- a/src/ai/backend/agent/docker/metadata/root.py +++ b/src/ai/backend/agent/docker/metadata/root.py @@ -1,6 +1,6 @@ from collections.abc import Mapping, Sequence from http import HTTPStatus -from typing import Any +from typing import Any, override from aiohttp import web @@ -10,19 +10,24 @@ class ContainerMetadataPlugin(MetadataPlugin): + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: pass + @override async def prepare_app(self) -> NewMetadataPluginResponse: app = web.Application() return NewMetadataPluginResponse(app, []) + @override async def routes(self) -> Sequence[MetadataPluginRoute]: return [ MetadataPluginRoute("GET", "/envs", self.get_envs, None), diff --git a/src/ai/backend/agent/docker/metadata/server.py b/src/ai/backend/agent/docker/metadata/server.py index 59eade29b07..166b56203f1 100644 --- a/src/ai/backend/agent/docker/metadata/server.py +++ b/src/ai/backend/agent/docker/metadata/server.py @@ -1,7 +1,7 @@ import logging from collections.abc import Mapping, MutableMapping, Sequence from http import HTTPStatus -from typing import Any, cast +from typing import Any, cast, override from uuid import UUID import attr @@ -113,6 +113,7 @@ def __init__( self.loaded_apps = [] self.route_structure = {"latest": {"extension": {}}} + @override async def __ainit__(self) -> None: local_config = cast(AgentUnifiedConfig, self.app["_root.context"].local_config) await prepare_kernel_metadata_uri_handling(local_config) diff --git a/src/ai/backend/agent/dummy/intrinsic.py b/src/ai/backend/agent/dummy/intrinsic.py index 9a5576ac7d6..43506eaab53 100644 --- a/src/ai/backend/agent/dummy/intrinsic.py +++ b/src/ai/backend/agent/dummy/intrinsic.py @@ -1,7 +1,7 @@ from collections.abc import Collection, Mapping, Sequence from decimal import Decimal from pathlib import Path -from typing import Any +from typing import Any, override import aiodocker @@ -59,15 +59,19 @@ def __init__( super().__init__(plugin_config, local_config) self.resource_config = dummy_config["agent"]["resource"] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "cpu", @@ -78,6 +82,7 @@ def get_metadata(self) -> AcceleratorMetadata: "display_icon": "cpu", } + @override async def list_devices(self) -> Collection[AbstractComputeDevice]: num_core: int = self.resource_config["cpu"]["num-core"] return [ @@ -91,21 +96,26 @@ async def list_devices(self) -> Collection[AbstractComputeDevice]: for core_idx in range(num_core) ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() return { SlotName("cpu"): Decimal(sum(dev.processing_units for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return {} + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: return [] + @override async def gather_container_measures( self, ctx: StatContext, @@ -113,11 +123,13 @@ async def gather_container_measures( ) -> Sequence[ContainerMeasurement]: return [] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: return [] + @override async def create_alloc_map(self) -> "AbstractAllocMap": devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -129,9 +141,11 @@ async def create_alloc_map(self) -> "AbstractAllocMap": }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: return [] + @override async def generate_docker_args( self, docker: aiodocker.docker.Docker, @@ -149,6 +163,7 @@ async def generate_docker_args( }, } + @override async def restore_from_container( self, container: Container, @@ -156,6 +171,7 @@ async def restore_from_container( ) -> None: return None + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -172,11 +188,13 @@ async def get_attached_devices( }) return attached_devices + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: @@ -210,15 +228,19 @@ def __init__( super().__init__(plugin_config, local_config) self.resource_config = dummy_config["agent"]["resource"] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "ram", @@ -229,6 +251,7 @@ def get_metadata(self) -> AcceleratorMetadata: "display_icon": "cpu", } + @override async def list_devices(self) -> Collection[AbstractComputeDevice]: memory_size = self.resource_config["memory"]["size"] return [ @@ -242,21 +265,26 @@ async def list_devices(self) -> Collection[AbstractComputeDevice]: ), ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() return { SlotName("mem"): Decimal(sum(dev.memory_size for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return {} + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: return [] + @override async def gather_container_measures( self, ctx: StatContext, @@ -264,11 +292,13 @@ async def gather_container_measures( ) -> Sequence[ContainerMeasurement]: return [] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: return [] + @override async def create_alloc_map(self) -> "AbstractAllocMap": devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -281,9 +311,11 @@ async def create_alloc_map(self) -> "AbstractAllocMap": }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: return [] + @override async def generate_docker_args( self, docker: aiodocker.docker.Docker, @@ -291,6 +323,7 @@ async def generate_docker_args( ) -> Mapping[str, Any]: return {} + @override async def restore_from_container( self, container: Container, @@ -298,6 +331,7 @@ async def restore_from_container( ) -> None: return None + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -314,11 +348,13 @@ async def get_attached_devices( }) return attached_devices + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: diff --git a/src/ai/backend/agent/dummy/kernel.py b/src/ai/backend/agent/dummy/kernel.py index a71ffe6ab67..a64f0a71722 100644 --- a/src/ai/backend/agent/dummy/kernel.py +++ b/src/ai/backend/agent/dummy/kernel.py @@ -263,9 +263,11 @@ def __init__( self.event_producer = event_producer + @override async def __ainit__(self) -> None: return + @override def __setstate__(self, props: MutableMapping[str, Any]) -> None: self.__dict__.update(props) self.zctx = None diff --git a/src/ai/backend/agent/errors/agent.py b/src/ai/backend/agent/errors/agent.py index db980e8d602..b9ae12ea491 100644 --- a/src/ai/backend/agent/errors/agent.py +++ b/src/ai/backend/agent/errors/agent.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, override from aiohttp import web @@ -23,6 +23,7 @@ class ImageArchitectureMismatchError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/image-architecture-mismatch" error_title = "Image architecture mismatch." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -37,6 +38,7 @@ class ImagePullTimeoutError(BackendAIError, web.HTTPGatewayTimeout): error_type = "https://api.backend.ai/probs/agent/image-pull-timeout" error_title = "Image pull timeout." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -51,6 +53,7 @@ class ContainerCreationFailedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/container-creation-failed" error_title = "Container creation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -65,6 +68,7 @@ class ContainerStartupTimeoutError(BackendAIError, web.HTTPGatewayTimeout): error_type = "https://api.backend.ai/probs/agent/container-startup-timeout" error_title = "Container startup timeout." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -79,6 +83,7 @@ class ContainerStartupCancelledError(BackendAIError, web.HTTPInternalServerError error_type = "https://api.backend.ai/probs/agent/container-startup-cancelled" error_title = "Container startup cancelled." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -93,6 +98,7 @@ class ContainerStartupFailedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/container-startup-failed" error_title = "Container startup failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -107,6 +113,7 @@ class ReservedPortError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/reserved-port" error_title = "Reserved port cannot be used." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -121,6 +128,7 @@ class PortConflictError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/port-conflict" error_title = "Port conflicts with existing service." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -135,6 +143,7 @@ class ModelDefinitionNotFoundError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/model-definition-not-found" error_title = "Model definition file not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -149,6 +158,7 @@ class ModelDefinitionEmptyError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/model-definition-empty" error_title = "Model definition is empty." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -163,6 +173,7 @@ class ModelDefinitionInvalidYAMLError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/model-definition-invalid-yaml" error_title = "Invalid YAML syntax in model definition." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -177,6 +188,7 @@ class ModelFolderNotSpecifiedError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/model-folder-not-specified" error_title = "Model folder not specified." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -191,6 +203,7 @@ class ImageCommandRequiredError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/image-command-required" error_title = "Image command is required." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -205,6 +218,7 @@ class AgentInitializationError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/initialization-failed" error_title = "Agent initialization failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -219,6 +233,7 @@ class InvalidAgentConfigError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-config" error_title = "Invalid agent configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -233,6 +248,7 @@ class InvalidMountPathError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-mount-path" error_title = "Invalid mount path." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -247,6 +263,7 @@ class InvalidLoggingConfigError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-logging-config" error_title = "Invalid logging configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -261,6 +278,7 @@ class InvalidChunkSizeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-chunk-size" error_title = "Invalid chunk size." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -275,6 +293,7 @@ class AsyncioTaskNotAvailableError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/asyncio-task-not-available" error_title = "Asyncio task context not available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -289,6 +308,7 @@ class KernelNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/agent/kernel-not-found" error_title = "Kernel not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -303,6 +323,7 @@ class InitializationError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/initialization-error" error_title = "Agent initialization error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -317,6 +338,7 @@ class InvalidArgumentError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-argument" error_title = "Invalid argument." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -331,6 +353,7 @@ class UnsupportedBaseDistroError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/unsupported-base-distro" error_title = "Unsupported base distribution." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -352,9 +375,11 @@ def __init__( self.message = message super().__init__(extra_msg=message) + @override def __reduce__(self) -> tuple[type[BackendAIError], tuple[Any, ...], dict[str, Any]]: return (self.__class__, (self.container_id, self.message), self.__dict__) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -373,9 +398,11 @@ def __init__(self, message: str) -> None: self.message = message super().__init__(extra_msg=message) + @override def __reduce__(self) -> tuple[type[BackendAIError], tuple[Any, ...], dict[str, Any]]: return (self.__class__, (self.message,), self.__dict__) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -390,6 +417,7 @@ class NetworkPluginNotFound(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/network-plugin-not-found" error_title = "Network plugin not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -404,6 +432,7 @@ class ServicePortAlreadyUsedError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/agent/service-port-already-used" error_title = "Service port already used." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -418,6 +447,7 @@ class InvalidSocket(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/invalid-socket" error_title = "Invalid socket." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/agent/errors/kernel.py b/src/ai/backend/agent/errors/kernel.py index cd8d326bc07..32687a7cef6 100644 --- a/src/ai/backend/agent/errors/kernel.py +++ b/src/ai/backend/agent/errors/kernel.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class KernelRunnerNotInitializedError(BackendAIError, web.HTTPInternalServerErro error_type = "https://api.backend.ai/probs/agent/kernel-runner-not-initialized" error_title = "Kernel runner is not initialized." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -35,6 +38,7 @@ class AsyncioContextError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/asyncio-context-error" error_title = "Asyncio context is invalid or unavailable." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -49,6 +53,7 @@ class SubprocessStreamError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/subprocess-stream-error" error_title = "Subprocess stream is not available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -63,6 +68,7 @@ class OutputQueueNotInitializedError(BackendAIError, web.HTTPInternalServerError error_type = "https://api.backend.ai/probs/agent/output-queue-not-initialized" error_title = "Output queue is not initialized." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -77,6 +83,7 @@ class OutputQueueMismatchError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/output-queue-mismatch" error_title = "Output queue mismatch." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -91,6 +98,7 @@ class RunIdNotSetError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/run-id-not-set" error_title = "Run ID is not set." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, diff --git a/src/ai/backend/agent/errors/resources.py b/src/ai/backend/agent/errors/resources.py index 98f13008179..2ed44ccdc2d 100644 --- a/src/ai/backend/agent/errors/resources.py +++ b/src/ai/backend/agent/errors/resources.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from decimal import Decimal -from typing import Any +from typing import Any, override from aiohttp import web @@ -26,6 +26,7 @@ class ResourceError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/resource-error" error_title = "Resource error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -85,6 +86,7 @@ def __init__( self.context_tag = context_tag super().__init__(extra_msg=str(self)) + @override def __str__(self) -> str: return ( f"InsufficientResource: {self.msg} ({self.slot_name}" @@ -92,6 +94,7 @@ def __str__(self) -> str: + f"allocating {self.requested_alloc} out of {self.total_allocatable})" ) + @override def __reduce__(self) -> tuple[type[BackendAIError], tuple[Any, ...], dict[str, Any]]: return ( self.__class__, @@ -128,6 +131,7 @@ def __init__( self.context_tag = context_tag super().__init__(extra_msg=str(self)) + @override def __str__(self) -> str: return ( f"FractionalResourceFragmented: {self.msg} ({self.slot_name}" @@ -135,6 +139,7 @@ def __str__(self) -> str: + f"allocating {self.requested_alloc} from {self.dev_allocs})" ) + @override def __reduce__(self) -> tuple[type[BackendAIError], tuple[Any, ...], dict[str, Any]]: return ( self.__class__, @@ -155,6 +160,7 @@ class InvalidResourceConfigError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-resource-config" error_title = "Invalid resource configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -169,6 +175,7 @@ class AgentIdNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/agent/agent-id-not-found" error_title = "Agent ID not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -183,6 +190,7 @@ class ResourceOverAllocatedError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/resource-over-allocated" error_title = "Resources are over-allocated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -197,6 +205,7 @@ class ResourceAllocationError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/resource-allocation-error" error_title = "Resource allocation error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -211,6 +220,7 @@ class InvalidOvercommitFactorError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-overcommit-factor" error_title = "Invalid overcommit factor." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -225,6 +235,7 @@ class InvalidAllocMapTypeError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/invalid-alloc-map-type" error_title = "Invalid allocation map type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -239,6 +250,7 @@ class InvalidMeasurementError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent/invalid-measurement" error_title = "Invalid measurement data." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -253,6 +265,7 @@ class InvalidConfigFormatError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/agent/invalid-config-format" error_title = "Invalid configuration format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -267,6 +280,7 @@ class InvalidContainerMeasurementError(BackendAIError, web.HTTPInternalServerErr error_type = "https://api.backend.ai/probs/agent/invalid-container-measurement" error_title = "Invalid container measurement type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -285,6 +299,7 @@ class PortPoolExhaustedError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/agent/port-pool-exhausted" error_title = "Host port pool exhausted." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/agent/health/docker.py b/src/ai/backend/agent/health/docker.py index b9f8fde4b4d..5ff11685fe3 100644 --- a/src/ai/backend/agent/health/docker.py +++ b/src/ai/backend/agent/health/docker.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from aiodocker import Docker @@ -36,10 +37,12 @@ def __init__(self, docker: Docker, timeout: float = 5.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return CONTAINER + @override async def check_service(self) -> ServiceHealth: """ Check Docker connection health by getting version info. @@ -66,6 +69,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={CID_DOCKER: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/agent/health/heartbeat.py b/src/ai/backend/agent/health/heartbeat.py index de264dc3f09..8caf91ff461 100644 --- a/src/ai/backend/agent/health/heartbeat.py +++ b/src/ai/backend/agent/health/heartbeat.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -19,16 +19,20 @@ def __init__(self, agent: AbstractAgent[Any, Any]) -> None: self._agent = agent @property + @override def name(self) -> str: return "heartbeat" @property + @override def interval(self) -> float: return self._agent.local_config.debug.heartbeat_interval @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.heartbeat() diff --git a/src/ai/backend/agent/kernel.py b/src/ai/backend/agent/kernel.py index 53cb2c06a19..717f02db657 100644 --- a/src/ai/backend/agent/kernel.py +++ b/src/ai/backend/agent/kernel.py @@ -26,6 +26,7 @@ TypedDict, cast, overload, + override, ) import zmq @@ -254,6 +255,7 @@ async def init(self, event_producer: EventProducer) -> None: self.runner = None raise + @override def __getstate__(self) -> Mapping[str, Any]: props = self.__dict__.copy() del props["agent_config"] @@ -482,20 +484,25 @@ def __init__(self, kernel_registry: KernelRegistry, agent_id: AgentId) -> None: self._registry = kernel_registry self._agent_id = agent_id + @override def __getitem__(self, key: KernelId) -> AbstractKernel: return self._registry[AgentKernelRegistryKey(self._agent_id, key)] + @override def __setitem__(self, key: KernelId, value: AbstractKernel) -> None: self._registry[AgentKernelRegistryKey(self._agent_id, key)] = value + @override def __delitem__(self, key: KernelId) -> None: del self._registry[AgentKernelRegistryKey(self._agent_id, key)] + @override def __iter__(self) -> Iterator[KernelId]: for registry_key in self._registry: if registry_key.agent_id == self._agent_id: yield registry_key.kernel_id + @override def __len__(self) -> int: return sum(1 for key in self._registry if key.agent_id == self._agent_id) @@ -508,13 +515,16 @@ def __init__(self, kernel_registry: KernelRegistry) -> None: self._registry = kernel_registry + @override def __getitem__(self, key: KernelId) -> AbstractKernel: return self._registry[key] + @override def __iter__(self) -> Iterator[KernelId]: for registry_key in self._registry: yield registry_key.kernel_id + @override def __len__(self) -> int: return len(self._registry) @@ -541,22 +551,27 @@ def __getitem__(self, key: KernelId) -> AbstractKernel: ... @overload def __getitem__(self, key: AgentKernelRegistryKey) -> AbstractKernel: ... + @override def __getitem__(self, key: KernelId | AgentKernelRegistryKey) -> AbstractKernel: if isinstance(key, AgentKernelRegistryKey): return self._registry[key] return self._global_registry[key] + @override def __setitem__(self, key: AgentKernelRegistryKey, value: AbstractKernel) -> None: self._registry[key] = value self._global_registry[key.kernel_id] = value + @override def __delitem__(self, key: AgentKernelRegistryKey) -> None: del self._registry[key] del self._global_registry[key.kernel_id] + @override def __iter__(self) -> Iterator[AgentKernelRegistryKey]: return iter(self._registry) + @override def __len__(self) -> int: return len(self._registry) @@ -722,6 +737,7 @@ def __init__( self.watchdog_task = None self._closed = False + @override async def __ainit__(self) -> None: await self._get_socket_pair() await self._create_tasks() @@ -743,6 +759,7 @@ async def refresh_sockets(self) -> None: loop = current_loop() self.read_task = loop.create_task(self.read_output()) + @override def __getstate__(self) -> Mapping[str, Any]: props = self.__dict__.copy() del props["zctx"] diff --git a/src/ai/backend/agent/kubernetes/agent.py b/src/ai/backend/agent/kubernetes/agent.py index 5cf9bb13b6d..7c0eaa14389 100644 --- a/src/ai/backend/agent/kubernetes/agent.py +++ b/src/ai/backend/agent/kubernetes/agent.py @@ -885,6 +885,7 @@ def __init__( ) ) + @override async def __ainit__(self) -> None: await super().__ainit__() ipc_base_path = self.local_config.agent.ipc_base_path @@ -996,6 +997,7 @@ async def fetch_workers(self) -> None: if addr.type == "ExternalIP": self.workers[node.metadata.name]["ExternalIP"] = addr.address + @override async def shutdown(self, stop_signal: signal.Signals) -> None: # Stop agent socket handler task diff --git a/src/ai/backend/agent/kubernetes/intrinsic.py b/src/ai/backend/agent/kubernetes/intrinsic.py index 32c49525678..aaa00f83704 100644 --- a/src/ai/backend/agent/kubernetes/intrinsic.py +++ b/src/ai/backend/agent/kubernetes/intrinsic.py @@ -4,7 +4,7 @@ from collections.abc import Collection, Mapping, Sequence from decimal import Decimal from pathlib import Path -from typing import Any +from typing import Any, override import aiohttp from aiodocker.docker import Docker, DockerContainer @@ -102,15 +102,19 @@ class CPUPlugin(AbstractComputePlugin): (SlotName("cpu"), SlotTypes.COUNT), ] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override async def list_devices(self) -> Collection[CPUDevice]: await K8sConfig.load_kube_config() core_api = K8sClient.CoreV1Api() @@ -134,6 +138,7 @@ async def list_devices(self) -> Collection[CPUDevice]: # if 'node-role.kubernetes.io/master' not in node['metadata']['labels'].keys() ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() log.debug("available_slots: {}", devices) @@ -141,9 +146,11 @@ async def available_slots(self) -> Mapping[SlotName, Decimal]: SlotName("cpu"): Decimal(sum(dev.processing_units for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return { "agent_version": __version__, @@ -151,11 +158,13 @@ async def extra_info(self) -> Mapping[str, str]: "os_type": platform.system(), } + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: # TODO: Create our own k8s metric collector return [] + @override async def gather_container_measures( self, ctx: StatContext, @@ -165,11 +174,13 @@ async def gather_container_measures( return [] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: return [] + @override async def create_alloc_map(self) -> AbstractAllocMap: devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -181,10 +192,12 @@ async def create_alloc_map(self) -> AbstractAllocMap: }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: # TODO: move the sysconf hook in libbaihook.so here return [] + @override async def generate_docker_args( self, docker: Docker, @@ -194,6 +207,7 @@ async def generate_docker_args( # K8s resource allocation return {} + @override async def restore_from_container( self, container: Container, @@ -212,6 +226,7 @@ async def restore_from_container( SlotName("cpu"): resource_spec.allocations[DeviceName("cpu")][SlotName("cpu")], }) + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -228,16 +243,19 @@ async def get_attached_devices( }) return attached_devices + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: return [] + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "cpu", @@ -268,15 +286,19 @@ class MemoryPlugin(AbstractComputePlugin): (SlotName("mem"), SlotTypes.BYTES), ] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: pass + @override async def list_devices(self) -> Collection[MemoryDevice]: await K8sConfig.load_kube_config() core_api = K8sClient.CoreV1Api() @@ -302,33 +324,40 @@ async def list_devices(self) -> Collection[MemoryDevice]: ), ] + @override async def available_slots(self) -> Mapping[SlotName, Decimal]: devices = await self.list_devices() return { SlotName("mem"): Decimal(sum(dev.memory_size for dev in devices)), } + @override def get_version(self) -> str: return __version__ + @override async def extra_info(self) -> Mapping[str, str]: return {} + @override async def gather_node_measures(self, ctx: StatContext) -> Sequence[NodeMeasurement]: # TODO: Create our own k8s metric collector return [] + @override async def gather_container_measures( self, ctx: StatContext, container_ids: Sequence[str] ) -> Sequence[ContainerMeasurement]: # TODO: Implement Kubernetes-specific container metric collection return [] + @override async def gather_process_measures( self, ctx: StatContext, pid_map: Mapping[int, str] ) -> Sequence[ProcessMeasurement]: return [] + @override async def create_alloc_map(self) -> AbstractAllocMap: devices = await self.list_devices() return DiscretePropertyAllocMap( @@ -341,9 +370,11 @@ async def create_alloc_map(self) -> AbstractAllocMap: }, ) + @override async def get_hooks(self, distro: str, arch: str) -> Sequence[Path]: return [] + @override async def generate_docker_args( self, docker: Docker, @@ -353,6 +384,7 @@ async def generate_docker_args( # K8s resource allocation return {} + @override async def restore_from_container( self, container: Container, @@ -367,6 +399,7 @@ async def restore_from_container( SlotName("mem"): {DeviceId("root"): memory_limit}, }) + @override async def get_attached_devices( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]], @@ -383,16 +416,19 @@ async def get_attached_devices( }) return attached_devices + @override async def generate_mounts( self, source_path: Path, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[MountInfo]: return [] + @override async def get_docker_networks( self, device_alloc: Mapping[SlotName, Mapping[DeviceId, Decimal]] ) -> list[str]: return [] + @override def get_metadata(self) -> AcceleratorMetadata: return { "slot_name": "ram", diff --git a/src/ai/backend/agent/resources.py b/src/ai/backend/agent/resources.py index 79860150b0f..622f33c4285 100644 --- a/src/ai/backend/agent/resources.py +++ b/src/ai/backend/agent/resources.py @@ -25,6 +25,7 @@ Self, TextIO, cast, + override, ) import aiodocker @@ -331,9 +332,11 @@ def device_name(self) -> DeviceName: return self._device_name return DeviceName(self.__class__.__name__.removesuffix("Device").lower()) + @override def __hash__(self) -> int: return hash(f"{self.device_name}-{self.device_id}") + @override def __eq__(self, __o: object) -> bool: return hash(self) == hash(__o) @@ -542,6 +545,7 @@ def __init__(self, local_config: AgentUnifiedConfig, etcd: AsyncEtcd) -> None: self.etcd = etcd self.agent_configs = local_config.get_agent_configs() + @override async def __ainit__(self) -> None: alloc_map_mod.log_alloc_map = self.local_config.debug.log_alloc_map computers = await self._load_resources() @@ -818,6 +822,7 @@ class ComputePluginContext(BasePluginContext[AbstractComputePlugin]): plugin_group = "backendai_accelerator_v21" @classmethod + @override def discover_plugins( cls, plugin_group: str, @@ -847,6 +852,7 @@ class Mount: permission: MountPermission = MountPermission.READ_ONLY opts: Mapping[str, Any] | None = None + @override def __str__(self) -> str: return f"{self.source}:{self.target}:{self.permission.value}" diff --git a/src/ai/backend/agent/server.py b/src/ai/backend/agent/server.py index a7b44e3e071..fd2f7071a4d 100644 --- a/src/ai/backend/agent/server.py +++ b/src/ai/backend/agent/server.py @@ -30,6 +30,7 @@ ClassVar, Literal, cast, + override, ) from uuid import UUID @@ -304,6 +305,7 @@ def __init__( self.local_config = local_config self.skip_detect_manager = skip_detect_manager + @override async def __ainit__(self) -> None: if not self.skip_detect_manager: await self.detect_manager() diff --git a/src/ai/backend/agent/stage/kernel_lifecycle/docker/image.py b/src/ai/backend/agent/stage/kernel_lifecycle/docker/image.py index 8f7e4e0b091..af88439c153 100644 --- a/src/ai/backend/agent/stage/kernel_lifecycle/docker/image.py +++ b/src/ai/backend/agent/stage/kernel_lifecycle/docker/image.py @@ -36,9 +36,11 @@ class ImagePullSpec: class ImagePullSpecProvisioner(Provisioner[ImagePullCheckSpec, ImagePullSpec]): @property + @override def name(self) -> str: return "docker-image-pull-check" + @override async def setup(self, spec: ImagePullCheckSpec) -> ImagePullSpec: do_pull = (not spec.image_ref.is_local) and await self._check_image_exist(spec) return ImagePullSpec( @@ -66,6 +68,7 @@ async def _check_image_exist(self, spec: ImagePullCheckSpec) -> bool: raise return False + @override async def teardown(self, resource: ImagePullSpec) -> None: pass diff --git a/src/ai/backend/agent/stats.py b/src/ai/backend/agent/stats.py index 179b763a2cf..b6d83010a59 100644 --- a/src/ai/backend/agent/stats.py +++ b/src/ai/backend/agent/stats.py @@ -20,6 +20,7 @@ Any, Final, cast, + override, ) import aiodocker @@ -294,6 +295,7 @@ def to_serializable_dict(self) -> MovingStatValue: "version": 2, } + @override def __str__(self) -> str: return str({ "min": self.min, diff --git a/src/ai/backend/agent/tasks/cleanup_reported_kernels.py b/src/ai/backend/agent/tasks/cleanup_reported_kernels.py index 132c70931d7..774cf6a0f1b 100644 --- a/src/ai/backend/agent/tasks/cleanup_reported_kernels.py +++ b/src/ai/backend/agent/tasks/cleanup_reported_kernels.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, agent: AbstractAgent[Any, Any]) -> None: self._agent = agent @property + @override def name(self) -> str: return "cleanup_reported_kernels" @property + @override def interval(self) -> float: return _CLEANUP_REPORTED_KERNELS_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.cleanup_reported_kernels() diff --git a/src/ai/backend/agent/tasks/collect_container_stat.py b/src/ai/backend/agent/tasks/collect_container_stat.py index bea2536ef58..8b1787d9913 100644 --- a/src/ai/backend/agent/tasks/collect_container_stat.py +++ b/src/ai/backend/agent/tasks/collect_container_stat.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -22,16 +22,20 @@ def __init__(self, agent: AbstractAgent[Any, Any], local_config: AgentUnifiedCon self._local_config = local_config @property + @override def name(self) -> str: return "collect_container_stat" @property + @override def interval(self) -> float: return self._local_config.agent.utilization_metric.container.interval @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.collect_container_stat() diff --git a/src/ai/backend/agent/tasks/collect_node_stat.py b/src/ai/backend/agent/tasks/collect_node_stat.py index 16be80b08a4..20590acc8bb 100644 --- a/src/ai/backend/agent/tasks/collect_node_stat.py +++ b/src/ai/backend/agent/tasks/collect_node_stat.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from ai.backend.common.cron import PeriodicTask @@ -22,16 +22,20 @@ def __init__(self, runtime: AgentRuntime, local_config: AgentUnifiedConfig) -> N self._local_config = local_config @property + @override def name(self) -> str: return "collect_node_stat" @property + @override def interval(self) -> float: return self._local_config.agent.utilization_metric.node.interval @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._runtime.collect_node_stats() diff --git a/src/ai/backend/agent/tasks/collect_process_stat.py b/src/ai/backend/agent/tasks/collect_process_stat.py index df9ebb56fde..3d0f645f2ff 100644 --- a/src/ai/backend/agent/tasks/collect_process_stat.py +++ b/src/ai/backend/agent/tasks/collect_process_stat.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -22,16 +22,20 @@ def __init__(self, agent: AbstractAgent[Any, Any], local_config: AgentUnifiedCon self._local_config = local_config @property + @override def name(self) -> str: return "collect_process_stat" @property + @override def interval(self) -> float: return self._local_config.agent.utilization_metric.process.interval @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.collect_process_stat() diff --git a/src/ai/backend/agent/tasks/report_kernel_commit_status.py b/src/ai/backend/agent/tasks/report_kernel_commit_status.py index 93ed6f5476f..3e0fe3f2056 100644 --- a/src/ai/backend/agent/tasks/report_kernel_commit_status.py +++ b/src/ai/backend/agent/tasks/report_kernel_commit_status.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, agent: AbstractAgent[Any, Any]) -> None: self._agent = agent @property + @override def name(self) -> str: return "report_kernel_commit_status" @property + @override def interval(self) -> float: return _REPORT_COMMIT_STATUS_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.report_all_kernel_commit_status_map() diff --git a/src/ai/backend/agent/tasks/scan_images.py b/src/ai/backend/agent/tasks/scan_images.py index 344543a9df4..8f34e787109 100644 --- a/src/ai/backend/agent/tasks/scan_images.py +++ b/src/ai/backend/agent/tasks/scan_images.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, agent: AbstractAgent[Any, Any]) -> None: self._agent = agent @property + @override def name(self) -> str: return "scan_images" @property + @override def interval(self) -> float: return _SCAN_IMAGES_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.scan_images_periodically() diff --git a/src/ai/backend/agent/tasks/sync_container_lifecycles.py b/src/ai/backend/agent/tasks/sync_container_lifecycles.py index 7fac65accd1..e54bceead78 100644 --- a/src/ai/backend/agent/tasks/sync_container_lifecycles.py +++ b/src/ai/backend/agent/tasks/sync_container_lifecycles.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from ai.backend.common.cron import PeriodicTask @@ -19,16 +19,20 @@ def __init__(self, agent: AbstractAgent[Any, Any]) -> None: self._agent = agent @property + @override def name(self) -> str: return "sync_container_lifecycles" @property + @override def interval(self) -> float: return self._agent.local_config.agent.sync_container_lifecycles.interval @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._agent.sync_container_lifecycles() diff --git a/src/ai/backend/agent/tasks/update_slots.py b/src/ai/backend/agent/tasks/update_slots.py index a24617350f7..443dfc8c302 100644 --- a/src/ai/backend/agent/tasks/update_slots.py +++ b/src/ai/backend/agent/tasks/update_slots.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, runtime: AgentRuntime) -> None: self._runtime = runtime @property + @override def name(self) -> str: return "update_slots" @property + @override def interval(self) -> float: return _UPDATE_SLOTS_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._runtime.update_agent_slots() diff --git a/src/ai/backend/agent/types.py b/src/ai/backend/agent/types.py index b78e267447a..3f65febfe1e 100644 --- a/src/ai/backend/agent/types.py +++ b/src/ai/backend/agent/types.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from decimal import Decimal from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, cast, override import attrs from aiohttp.typedefs import Middleware @@ -177,6 +177,7 @@ class ContainerLifecycleEvent: exit_code: int | None = None suppress_events: bool = False + @override def __str__(self) -> str: if self.container_id: cid = self.container_id[:13] diff --git a/src/ai/backend/agent/utils.py b/src/ai/backend/agent/utils.py index b13b3ee79d4..29cc0914ae7 100644 --- a/src/ai/backend/agent/utils.py +++ b/src/ai/backend/agent/utils.py @@ -20,6 +20,7 @@ TypedDict, TypeVar, overload, + override, ) from uuid import UUID @@ -65,9 +66,11 @@ class closing_async(AbstractAsyncContextManager[_SupportsAsyncCloseT]): def __init__(self, obj: _SupportsAsyncCloseT) -> None: self.obj = obj + @override async def __aenter__(self) -> _SupportsAsyncCloseT: return self.obj + @override async def __aexit__(self, *exc_info: Any) -> None: await self.obj.close() diff --git a/src/ai/backend/appproxy/common/config.py b/src/ai/backend/appproxy/common/config.py index 20e655adc6b..dbf591c114e 100644 --- a/src/ai/backend/appproxy/common/config.py +++ b/src/ai/backend/appproxy/common/config.py @@ -3,7 +3,7 @@ import types import typing from pathlib import Path -from typing import Annotated, Any +from typing import Annotated, Any, override from pydantic import ( BaseModel, @@ -88,9 +88,11 @@ def host_set_with_protocol(self) -> bool: return True return False + @override def __repr__(self) -> str: return f"{self.host}:{self.port}" + @override def __str__(self) -> str: return self.__repr__() diff --git a/src/ai/backend/appproxy/common/errors/config.py b/src/ai/backend/appproxy/common/errors/config.py index bc7822a02de..9be5a848429 100644 --- a/src/ai/backend/appproxy/common/errors/config.py +++ b/src/ai/backend/appproxy/common/errors/config.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -17,6 +19,7 @@ class ConfigValidationError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/config-validation-error" error_title = "Configuration validation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -31,6 +34,7 @@ class InvalidUIDTypeError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-uid-type" error_title = "UID must be an integer or string." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -45,6 +49,7 @@ class UserNotFoundError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/user-not-found" error_title = "Specified user not found in system." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -59,6 +64,7 @@ class InvalidGIDTypeError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-gid-type" error_title = "GID must be an integer or string." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -73,6 +79,7 @@ class GroupNotFoundError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/group-not-found" error_title = "Specified group not found in system." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -87,6 +94,7 @@ class MissingAnnotationError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-annotation" error_title = "Field annotation is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/common/errors/http.py b/src/ai/backend/appproxy/common/errors/http.py index 83c0da18928..b23566e5687 100644 --- a/src/ai/backend/appproxy/common/errors/http.py +++ b/src/ai/backend/appproxy/common/errors/http.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, override from aiohttp import web @@ -22,6 +22,7 @@ class URLNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/url-not-found" error_title = "Unknown URL path." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -47,6 +48,7 @@ def __init__( self.error_title = f"E00002: No such {object_name}." super().__init__(message, extra_data=extra_data) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -61,6 +63,7 @@ class GenericBadRequest(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -75,6 +78,7 @@ class RejectedByHook(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/rejected-by-hook" error_title = "Operation rejected by a hook plugin." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -98,6 +102,7 @@ class InvalidCredentials(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/invalid-credentials" error_title = "Authentication credentials not valid." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -112,6 +117,7 @@ class GenericForbidden(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/generic-forbidden" error_title = "Forbidden operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -127,6 +133,7 @@ class ClientIPNotAllowed(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/client-ip-not-allowed" error_title = "Client address not allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.APPPROXY, @@ -141,6 +148,7 @@ class InsufficientPrivilege(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/insufficient-privilege" error_title = "Insufficient privilege." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -161,6 +169,7 @@ class MethodNotAllowed(BackendAIError): error_type = "https://api.backend.ai/probs/method-not-allowed" error_title = "HTTP Method Not Allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -175,6 +184,7 @@ class InternalServerError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/internal-server-error" error_title = "Internal server error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -189,6 +199,7 @@ class ServerMisconfiguredError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/server-misconfigured" error_title = "E00001: Service misconfigured." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -203,6 +214,7 @@ class ServiceUnavailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/service-unavailable" error_title = "Service unavailable." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -217,6 +229,7 @@ class QueryNotImplemented(BackendAIError, web.HTTPNotImplemented): error_type = "https://api.backend.ai/probs/not-implemented" error_title = "This API query is not implemented." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -231,6 +244,7 @@ class InvalidAuthParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-auth-params" error_title = "Missing or invalid authorization parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -245,6 +259,7 @@ class AuthorizationFailed(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/auth-failed" error_title = "Credential/signature mismatch." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -259,6 +274,7 @@ class PasswordExpired(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/password-expired" error_title = "Password has expired." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -273,6 +289,7 @@ class InvalidAPIParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-api-params" error_title = "Missing or invalid API parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -287,6 +304,7 @@ class GraphQLError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/graphql-error" error_title = "GraphQL-generated error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/common/errors/service.py b/src/ai/backend/appproxy/common/errors/service.py index 8eae8194b6f..1337393ac1d 100644 --- a/src/ai/backend/appproxy/common/errors/service.py +++ b/src/ai/backend/appproxy/common/errors/service.py @@ -10,6 +10,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -27,6 +29,7 @@ class WorkerNotAvailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/appproxy/worker-not-available" error_title = "Worker not available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -41,6 +44,7 @@ class PortNotAvailable(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/appproxy/port-not-available" error_title = "Designated port already occupied." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -55,6 +59,7 @@ class UnsupportedProtocol(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/appproxy/unsupported-protocol" error_title = "Unsupported protocol." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -69,6 +74,7 @@ class DatabaseError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/appproxy/database-error" error_title = "Error while communicating with database." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -83,6 +89,7 @@ class ContainerConnectionRefused(BackendAIError, web.HTTPBadGateway): error_type = "https://api.backend.ai/probs/appproxy/container-connection-refused" error_title = "Cannot connect to Backend.AI kernel." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -97,6 +104,7 @@ class WorkerRegistrationError(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/appproxy/worker-registration-error" error_title = "E20013: Failed to register worker." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -111,6 +119,7 @@ class CoordinatorConnectionError(BackendAIError, web.HTTPBadGateway): error_type = "https://api.backend.ai/probs/appproxy/coordinator-connection-error" error_title = "E20014: Failed to communicate with coordinator." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/common/etcd.py b/src/ai/backend/appproxy/common/etcd.py index 83bbb47010f..24885a1cddc 100644 --- a/src/ai/backend/appproxy/common/etcd.py +++ b/src/ai/backend/appproxy/common/etcd.py @@ -1,13 +1,15 @@ -from typing import Any, cast +from typing import Any, cast, override from ai.backend.common.etcd import AsyncEtcd class TraefikEtcd(AsyncEtcd): + @override def _mangle_key(self, k: str) -> str: k = k.removeprefix("/") return f"{self.ns}/{k}" + @override def _demangle_key(self, k: bytes | str) -> str: if isinstance(k, bytes): k = k.decode(self.encoding) diff --git a/src/ai/backend/appproxy/common/events.py b/src/ai/backend/appproxy/common/events.py index ef5ea5e2c8d..46d390b0400 100644 --- a/src/ai/backend/appproxy/common/events.py +++ b/src/ai/backend/appproxy/common/events.py @@ -18,6 +18,7 @@ class AppProxyCircuitEvent(AbstractBroadcastEvent): target_worker_authority: str circuits: list[Circuit] + @override def serialize(self) -> tuple[Any, ...]: return ( self.target_worker_authority, @@ -25,6 +26,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( target_worker_authority=value[0], @@ -51,6 +53,7 @@ class AppProxyCircuitRouteUpdatedEvent(AbstractBroadcastEvent): circuit: Circuit routes: list[RouteInfo] + @override def serialize(self) -> tuple[Any, ...]: return ( self.target_worker_authority, @@ -59,6 +62,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( target_worker_authority=value[0], @@ -90,6 +94,7 @@ class GenericWorkerEvent(AbstractAnycastEvent): worker_id: str reason: str + @override def serialize(self) -> tuple[Any, ...]: return ( self.worker_id, @@ -97,6 +102,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls(value[0], value[1]) @@ -150,11 +156,13 @@ def event_name(cls) -> str: class DoCheckWorkerLostEvent(AbstractAnycastEvent): + @override def serialize(self) -> tuple[Any, ...]: return tuple() @classmethod - def deserialize(cls, value: tuple[Any, ...]) -> Self: # noqa: ARG003 + @override + def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls() @classmethod @@ -177,11 +185,13 @@ def user_event(self) -> UserEvent | None: class DoCheckUnusedPortEvent(AbstractAnycastEvent): + @override def serialize(self) -> tuple[Any, ...]: return tuple() @classmethod - def deserialize(cls, value: tuple[Any, ...]) -> Self: # noqa: ARG003 + @override + def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls() @classmethod @@ -211,11 +221,13 @@ class DoReconcileTraefikRoutesEvent(AbstractAnycastEvent): source of truth. """ + @override def serialize(self) -> tuple[Any, ...]: return tuple() @classmethod - def deserialize(cls, value: tuple[Any, ...]) -> Self: # noqa: ARG003 + @override + def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls() @classmethod diff --git a/src/ai/backend/appproxy/common/types.py b/src/ai/backend/appproxy/common/types.py index 39d42d581c8..932fd33cf9b 100644 --- a/src/ai/backend/appproxy/common/types.py +++ b/src/ai/backend/appproxy/common/types.py @@ -10,6 +10,7 @@ Annotated, Any, TypeVar, + override, ) from uuid import UUID @@ -151,6 +152,7 @@ class RouteInfo(BackendAISchema): protocol: ProxyProtocol traffic_ratio: Annotated[float, Field(default=1.0)] + @override def __hash__(self) -> int: return hash(json.dumps(self.model_dump(mode="json"))) diff --git a/src/ai/backend/appproxy/common/utils.py b/src/ai/backend/appproxy/common/utils.py index 8cd0d7fbfd8..5ce1311273e 100644 --- a/src/ai/backend/appproxy/common/utils.py +++ b/src/ai/backend/appproxy/common/utils.py @@ -15,6 +15,7 @@ Any, TypeVar, cast, + override, ) from uuid import UUID @@ -314,6 +315,7 @@ class BackendAIAccessLogger(AccessLogger): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + @override def log(self, request: web.BaseRequest, response: web.StreamResponse, time: float) -> None: if request.get("do_not_print_access_log"): return diff --git a/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/composer.py b/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/composer.py index fd5236364a2..42d7923e5a6 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/composer.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.appproxy.coordinator.config import ServerConfig from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -29,10 +30,12 @@ class BootstrapComposer(DependencyComposer[BootstrapInput, BootstrapResources]): """Composer for bootstrap layer dependencies.""" @property + @override def stage_name(self) -> str: return "bootstrap" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/config.py b/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/config.py index 718fc3670a1..155e18bfd76 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/config.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/bootstrap/config.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.appproxy.coordinator.config import ServerConfig, load from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -20,10 +21,12 @@ class ConfigProvider(NonMonitorableDependencyProvider[ConfigInput, ServerConfig] """Provider for app proxy coordinator configuration.""" @property + @override def stage_name(self) -> str: return "config" @asynccontextmanager + @override async def provide(self, setup_input: ConfigInput) -> AsyncIterator[ServerConfig]: """Load and provide server configuration. diff --git a/src/ai/backend/appproxy/coordinator/dependencies/composer.py b/src/ai/backend/appproxy/coordinator/dependencies/composer.py index fea72176934..1e6af7a4806 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/composer.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -29,10 +30,12 @@ class CoordinatorDependencyComposer(DependencyComposer[DependencyInput, Dependen """Main composer for all app proxy coordinator dependencies.""" @property + @override def stage_name(self) -> str: return "coordinator" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/composer.py b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/composer.py index 11e00d48900..ca976b87689 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/composer.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.appproxy.common.etcd import TraefikEtcd from ai.backend.appproxy.coordinator.config import ServerConfig @@ -27,10 +28,12 @@ class InfrastructureComposer(DependencyComposer[ServerConfig, InfrastructureReso """Composer for infrastructure layer dependencies.""" @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/database.py b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/database.py index 18659b1c8f3..d1398b504b4 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/database.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/database.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.appproxy.coordinator.config import ServerConfig from ai.backend.appproxy.coordinator.health.database import DatabaseHealthChecker @@ -14,15 +15,18 @@ class DatabaseProvider(DependencyProvider[ServerConfig, ExtendedAsyncSAEngine]): """Provider for PostgreSQL database connection.""" @property + @override def stage_name(self) -> str: return "database" @asynccontextmanager + @override async def provide(self, setup_input: ServerConfig) -> AsyncIterator[ExtendedAsyncSAEngine]: """Create and provide database connection.""" async with connect_database(setup_input.db) as db: yield db + @override def gen_readiness_checker(self, resource: ExtendedAsyncSAEngine) -> ServiceHealthChecker: """Readiness only — DB unreachable should drain traffic, not trigger restart.""" return DatabaseHealthChecker(db=resource) diff --git a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/etcd.py b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/etcd.py index 01aaafbaf70..d14b1e98d95 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/etcd.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/etcd.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.appproxy.common.etcd import TraefikEtcd from ai.backend.appproxy.coordinator.config import ServerConfig @@ -17,10 +18,12 @@ class EtcdProvider(DependencyProvider[ServerConfig, TraefikEtcd | None]): """Provider for Traefik etcd connection (optional).""" @property + @override def stage_name(self) -> str: return "etcd" @asynccontextmanager + @override async def provide(self, setup_input: ServerConfig) -> AsyncIterator[TraefikEtcd | None]: """Create and provide Traefik etcd connection if enabled.""" if setup_input.proxy_coordinator.enable_traefik: @@ -52,6 +55,7 @@ async def provide(self, setup_input: ServerConfig) -> AsyncIterator[TraefikEtcd else: yield None + @override def gen_liveness_checker(self, resource: TraefikEtcd | None) -> ServiceHealthChecker | None: """Liveness — stuck etcd connection observed; restart is the recovery path.""" if resource is None: diff --git a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/redis.py b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/redis.py index 150a56ecb0e..6e71daa74d7 100644 --- a/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/redis.py +++ b/src/ai/backend/appproxy/coordinator/dependencies/infrastructure/redis.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any +from typing import Any, override from redis.asyncio import Redis @@ -37,10 +37,12 @@ class RedisProvider(DependencyProvider[ServerConfig, CoordinatorValkeyClients]): """Provider for Redis configuration and Valkey clients.""" @property + @override def stage_name(self) -> str: return "redis" @asynccontextmanager + @override async def provide(self, setup_input: ServerConfig) -> AsyncIterator[CoordinatorValkeyClients]: """Create and provide Redis/Valkey clients.""" redis_profile_target = RedisProfileTarget.from_dict(setup_input.redis.to_dict()) @@ -91,6 +93,7 @@ async def provide(self, setup_input: ServerConfig) -> AsyncIterator[CoordinatorV await valkey_schedule.close() await redis_lock.close() + @override def gen_liveness_checker(self, resource: CoordinatorValkeyClients) -> ServiceHealthChecker: """Liveness — Valkey connection-stuck observed; restart recovers.""" return ValkeyHealthChecker( diff --git a/src/ai/backend/appproxy/coordinator/errors/circuit.py b/src/ai/backend/appproxy/coordinator/errors/circuit.py index dc1b170488f..7ef8493e490 100644 --- a/src/ai/backend/appproxy/coordinator/errors/circuit.py +++ b/src/ai/backend/appproxy/coordinator/errors/circuit.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class CircuitCreationError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/circuit-creation-failed" error_title = "Failed to create circuit." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -33,6 +36,7 @@ class InvalidCircuitConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-circuit-config" error_title = "Invalid circuit configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -47,6 +51,7 @@ class InvalidCircuitStateError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-circuit-state" error_title = "Invalid circuit state." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, diff --git a/src/ai/backend/appproxy/coordinator/errors/config.py b/src/ai/backend/appproxy/coordinator/errors/config.py index 6ceecb04cfb..320d1d6cbc3 100644 --- a/src/ai/backend/appproxy/coordinator/errors/config.py +++ b/src/ai/backend/appproxy/coordinator/errors/config.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class MissingTraefikConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-traefik-config" error_title = "Traefik configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -33,6 +36,7 @@ class MissingDatabaseURLError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-database-url" error_title = "Database URL is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -47,6 +51,7 @@ class MissingConfigFileError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-config-file" error_title = "Configuration file is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -61,6 +66,7 @@ class MissingFrontendConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-frontend-config" error_title = "Frontend configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -75,6 +81,7 @@ class InvalidURLError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-url" error_title = "Invalid URL." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -89,6 +96,7 @@ class InvalidSessionParameterError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-session-parameter" error_title = "Invalid session parameter." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -103,6 +111,7 @@ class InvalidEnumTypeError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/invalid-enum-type" error_title = "Invalid enum type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -117,6 +126,7 @@ class TransactionResultError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/transaction-result-error" error_title = "Transaction did not produce a result." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -131,6 +141,7 @@ class LockContextNotInitializedError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/lock-context-not-initialized" error_title = "Lock context is not initialized." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -145,6 +156,7 @@ class CleanupContextNotInitializedError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/cleanup-context-not-initialized" error_title = "Cleanup context is not initialized." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -159,6 +171,7 @@ class MissingProfilingConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-profiling-config" error_title = "Profiling configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -173,6 +186,7 @@ class MissingRouteInfoError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-route-info" error_title = "Route connection info is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -187,6 +201,7 @@ class MissingHealthCheckInfoError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy/missing-health-check-info" error_title = "Health check info is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, diff --git a/src/ai/backend/appproxy/coordinator/health/database.py b/src/ai/backend/appproxy/coordinator/health/database.py index bf696e77af0..47f88c50609 100644 --- a/src/ai/backend/appproxy/coordinator/health/database.py +++ b/src/ai/backend/appproxy/coordinator/health/database.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.appproxy.coordinator.models.utils import ExtendedAsyncSAEngine from ai.backend.common.health_checker.abc import StaticServiceHealthChecker @@ -35,10 +36,12 @@ def __init__(self, db: ExtendedAsyncSAEngine, timeout: float = 5.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return DATABASE + @override async def check_service(self) -> ServiceHealth: """ Check database connection health by pinging the server. @@ -65,6 +68,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={CID_POSTGRES: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/appproxy/coordinator/models/base.py b/src/ai/backend/appproxy/coordinator/models/base.py index e6bef0c8846..e5b069a6b26 100644 --- a/src/ai/backend/appproxy/coordinator/models/base.py +++ b/src/ai/backend/appproxy/coordinator/models/base.py @@ -11,6 +11,7 @@ ClassVar, TypeVar, cast, + override, ) import asyncpg.pgproto.pgproto @@ -107,16 +108,20 @@ def __init__(self, enum_cls: type[enum.Enum], **opts: Any) -> None: super().__init__(*enums, **opts) self._enum_cls = enum_cls + @override def process_bind_param(self, value: enum.Enum | None, dialect: sa.Dialect) -> str | None: return value.name if value else None + @override def process_result_value(self, value: Any, dialect: sa.Dialect) -> enum.Enum | None: return self._enum_cls[value] if value else None + @override def copy(self, **kw: Any) -> EnumType: return EnumType(self._enum_cls, **self._opts) @property + @override def python_type(self) -> type[enum.Enum]: return self._enum_cls @@ -135,6 +140,7 @@ def __init__(self, enum_cls: type[T_StrEnum], use_name: bool = False, **opts: An self._use_name = use_name self._enum_cls = enum_cls + @override def process_bind_param( self, value: T_StrEnum | None, @@ -148,6 +154,7 @@ def process_bind_param( val: str = value.value return val + @override def process_result_value( self, value: str | None, @@ -159,10 +166,12 @@ def process_result_value( return self._enum_cls[value] return self._enum_cls(value) + @override def copy(self, **kw: Any) -> StrEnumType[T_StrEnum]: return StrEnumType(self._enum_cls, self._use_name, **self._opts) @property + @override def python_type(self) -> type[T_StrEnum]: return self._enum_cls @@ -180,11 +189,13 @@ def __init__(self, schema: type[BackendAISchema]) -> None: super().__init__() self._schema = schema + @override def load_dialect_impl(self, dialect: sa.Dialect) -> TypeEngine[Any]: if dialect.name == "sqlite": return dialect.type_descriptor(sa.JSON()) return super().load_dialect_impl(dialect) + @override def process_bind_param(self, value: Any, dialect: sa.Dialect) -> BaseModel: if value is None: return self._schema() @@ -197,12 +208,14 @@ def process_bind_param(self, value: Any, dialect: sa.Dialect) -> BaseModel: ) from e return validated + @override def process_result_value(self, value: Any, dialect: sa.Dialect) -> BaseModel: if value is None: return self._schema() result: BaseModel = self._schema(**value) return result + @override def copy(self, **kw: Any) -> StructuredJSONColumn: return StructuredJSONColumn(self._schema) @@ -219,16 +232,19 @@ def __init__(self, schema: type[BackendAISchema]) -> None: super().__init__() self._schema = schema + @override def process_bind_param(self, value: BaseModel | None, dialect: sa.Dialect) -> str | None: if value: return value.model_dump_json() return None + @override def process_result_value(self, value: str | None, dialect: sa.Dialect) -> BaseModel | None: if value: return self._schema(**json.loads(value)) return None + @override def copy(self, **kw: Any) -> StructuredJSONObjectColumn: return StructuredJSONObjectColumn(self._schema) @@ -249,11 +265,13 @@ def __init__(self, schema: type[TBaseModel]) -> None: super().__init__() self._schema = schema + @override def process_bind_param(self, value: list[TBaseModel] | None, dialect: sa.Dialect) -> str | None: if value is not None: return TypeAdapter(list[TBaseModel]).dump_json(value).decode("utf-8") return None + @override def process_result_value( self, value: str | None, dialect: sa.Dialect ) -> list[TBaseModel] | None: @@ -261,6 +279,7 @@ def process_result_value( return [self._schema(**i) for i in json.loads(value)] return None + @override def copy(self, **kw: Any) -> StructuredJSONObjectListColumn[TBaseModel]: return StructuredJSONObjectListColumn(self._schema) @@ -287,9 +306,11 @@ def __init__(self, schema: type[TBaseModel]) -> None: super().__init__() self._schema = schema + @override def coerce_compared_value(self, _op: Any, _value: Any) -> JSONB: return JSONB() + @override def process_bind_param( self, value: list[TBaseModel] | None, _dialect: sa.Dialect ) -> list[dict[str, Any]]: @@ -297,6 +318,7 @@ def process_bind_param( return [item.model_dump(mode="json") for item in value] return [] + @override def process_result_value( self, value: list[dict[str, Any]] | str | None, _dialect: sa.Dialect ) -> list[TBaseModel]: @@ -309,6 +331,7 @@ def process_result_value( return [self._schema.model_validate(item) for item in value] return [] + @override def copy(self, **kw: Any) -> PydanticListColumn[TBaseModel]: return PydanticListColumn(self._schema) @@ -321,11 +344,13 @@ class URLColumn(TypeDecorator[yarl.URL]): impl = sa.types.UnicodeText cache_ok = True + @override def process_bind_param(self, value: yarl.URL | str | None, dialect: sa.Dialect) -> str | None: if isinstance(value, yarl.URL): return str(value) return value + @override def process_result_value(self, value: str | None, dialect: sa.Dialect) -> yarl.URL | None: if value is None: return None @@ -340,6 +365,7 @@ class IPColumn(TypeDecorator[ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6 impl = CIDR cache_ok = True + @override def process_bind_param( self, value: ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6Network] | str | None, @@ -363,6 +389,7 @@ def process_bind_param( return None return str(addr) + @override def process_result_value( self, value: str | None, dialect: sa.Dialect ) -> ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6Network] | None: @@ -384,11 +411,13 @@ class GUID[UUID_SubType: uuid.UUID](TypeDecorator[UUID_SubType]): uuid_subtype_func: ClassVar[Callable[[Any], Any]] = lambda v: v cache_ok = True + @override def load_dialect_impl(self, dialect: sa.Dialect) -> TypeEngine[Any]: if dialect.name == "postgresql": return dialect.type_descriptor(UUID()) return dialect.type_descriptor(CHAR(16)) + @override def process_bind_param(self, value: Any, dialect: sa.Dialect) -> str | bytes | None: # NOTE: EndpointId, SessionId, KernelId are *not* actual types defined as classes, # but a "virtual" type that is an identity function at runtime. @@ -404,6 +433,7 @@ def process_bind_param(self, value: Any, dialect: sa.Dialect) -> str | bytes | N return value.bytes return uuid.UUID(value).bytes + @override def process_result_value(self, value: Any, dialect: sa.Dialect) -> UUID_SubType | None: if value is None: return value diff --git a/src/ai/backend/appproxy/coordinator/models/utils.py b/src/ai/backend/appproxy/coordinator/models/utils.py index dc439120120..813d5a1142b 100644 --- a/src/ai/backend/appproxy/coordinator/models/utils.py +++ b/src/ai/backend/appproxy/coordinator/models/utils.py @@ -14,6 +14,7 @@ ParamSpec, TypeVar, overload, + override, ) import sqlalchemy as sa @@ -131,6 +132,7 @@ async def _begin_readonly( self._readonly_txn_count -= 1 @actxmgr + @override async def begin(self, bind: SAConnection | None = None) -> AsyncIterator[SAConnection]: if bind is None: async with self.connect() as _bind, self._begin(_bind) as conn: diff --git a/src/ai/backend/appproxy/coordinator/pglock.py b/src/ai/backend/appproxy/coordinator/pglock.py index ad7cbd7bc82..99258eddd18 100644 --- a/src/ai/backend/appproxy/coordinator/pglock.py +++ b/src/ai/backend/appproxy/coordinator/pglock.py @@ -2,7 +2,7 @@ from contextlib import AbstractAsyncContextManager from types import TracebackType -from typing import Any +from typing import Any, override from ai.backend.common.lock import AbstractDistributedLock @@ -19,10 +19,12 @@ def __init__(self, db: ExtendedAsyncSAEngine, lock_id: LockID) -> None: self.lock_id = lock_id self._lock_ctx = None + @override async def __aenter__(self) -> Any: self._lock_ctx = self.db.advisory_lock(self.lock_id) await self._lock_ctx.__aenter__() + @override async def __aexit__( self, exc_type: type[BaseException] | None, diff --git a/src/ai/backend/appproxy/worker/cli/context.py b/src/ai/backend/appproxy/worker/cli/context.py index 0f544e1fb7d..8eb1ec2ca35 100644 --- a/src/ai/backend/appproxy/worker/cli/context.py +++ b/src/ai/backend/appproxy/worker/cli/context.py @@ -2,7 +2,7 @@ from contextlib import AbstractContextManager from pathlib import Path -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, Self, override import click @@ -18,6 +18,7 @@ def __init__(self, log_level: LogLevel, config_path: Path | None = None) -> None self.config_path = config_path self.log_level = log_level + @override def __enter__(self) -> Self: from ai.backend.logging import LocalLogger @@ -30,6 +31,7 @@ def __enter__(self) -> Self: self._logger.__enter__() return self + @override def __exit__(self, *exc_info: Any) -> None: click_ctx = click.get_current_context() if click_ctx.invoked_subcommand != "start-server": diff --git a/src/ai/backend/appproxy/worker/dependencies/bootstrap/composer.py b/src/ai/backend/appproxy/worker/dependencies/bootstrap/composer.py index 71fba258056..e22031f09d0 100644 --- a/src/ai/backend/appproxy/worker/dependencies/bootstrap/composer.py +++ b/src/ai/backend/appproxy/worker/dependencies/bootstrap/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.appproxy.worker.config import ServerConfig from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -29,10 +30,12 @@ class BootstrapComposer(DependencyComposer[BootstrapInput, BootstrapResources]): """Composer for bootstrap layer dependencies.""" @property + @override def stage_name(self) -> str: return "bootstrap" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/worker/dependencies/bootstrap/config.py b/src/ai/backend/appproxy/worker/dependencies/bootstrap/config.py index 2162ef0bd66..35b05d94e79 100644 --- a/src/ai/backend/appproxy/worker/dependencies/bootstrap/config.py +++ b/src/ai/backend/appproxy/worker/dependencies/bootstrap/config.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.appproxy.worker.config import ServerConfig, load from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -20,10 +21,12 @@ class ConfigProvider(NonMonitorableDependencyProvider[ConfigInput, ServerConfig] """Provider for app proxy worker configuration.""" @property + @override def stage_name(self) -> str: return "config" @asynccontextmanager + @override async def provide(self, setup_input: ConfigInput) -> AsyncIterator[ServerConfig]: """Load and provide server configuration. diff --git a/src/ai/backend/appproxy/worker/dependencies/composer.py b/src/ai/backend/appproxy/worker/dependencies/composer.py index 1ae2eae2073..0c3c18f30bd 100644 --- a/src/ai/backend/appproxy/worker/dependencies/composer.py +++ b/src/ai/backend/appproxy/worker/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -29,10 +30,12 @@ class WorkerDependencyComposer(DependencyComposer[DependencyInput, DependencyRes """Main composer for all app proxy worker dependencies.""" @property + @override def stage_name(self) -> str: return "worker" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/worker/dependencies/infrastructure/composer.py b/src/ai/backend/appproxy/worker/dependencies/infrastructure/composer.py index b0597cfe6b1..3bc18f6ab74 100644 --- a/src/ai/backend/appproxy/worker/dependencies/infrastructure/composer.py +++ b/src/ai/backend/appproxy/worker/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.appproxy.worker.config import ServerConfig from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -21,10 +22,12 @@ class InfrastructureComposer(DependencyComposer[ServerConfig, InfrastructureReso """Composer for infrastructure layer dependencies.""" @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/appproxy/worker/dependencies/infrastructure/redis.py b/src/ai/backend/appproxy/worker/dependencies/infrastructure/redis.py index 972edcc73b0..3ff002d587b 100644 --- a/src/ai/backend/appproxy/worker/dependencies/infrastructure/redis.py +++ b/src/ai/backend/appproxy/worker/dependencies/infrastructure/redis.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.appproxy.worker.config import ServerConfig from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient @@ -26,10 +27,12 @@ class RedisProvider(DependencyProvider[ServerConfig, WorkerValkeyClients]): """Provider for Redis configuration and Valkey clients.""" @property + @override def stage_name(self) -> str: return "redis" @asynccontextmanager + @override async def provide(self, setup_input: ServerConfig) -> AsyncIterator[WorkerValkeyClients]: """Create and provide Valkey clients.""" redis_profile_target = RedisProfileTarget.from_dict(setup_input.redis.to_dict()) @@ -56,6 +59,7 @@ async def provide(self, setup_input: ServerConfig) -> AsyncIterator[WorkerValkey await valkey_live.close() await valkey_stat.close() + @override def gen_liveness_checker(self, resource: WorkerValkeyClients) -> ServiceHealthChecker: """Liveness — Valkey connection-stuck observed; restart recovers.""" return ValkeyHealthChecker( diff --git a/src/ai/backend/appproxy/worker/errors/circuit.py b/src/ai/backend/appproxy/worker/errors/circuit.py index 6d0e2d76802..4954010cb69 100644 --- a/src/ai/backend/appproxy/worker/errors/circuit.py +++ b/src/ai/backend/appproxy/worker/errors/circuit.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class InvalidCircuitDataError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/invalid-circuit-data" error_title = "Invalid circuit data." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -33,6 +36,7 @@ class InvalidFrontendTypeError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/invalid-frontend-type" error_title = "Invalid frontend type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/worker/errors/config.py b/src/ai/backend/appproxy/worker/errors/config.py index 2577693c673..9340862f12d 100644 --- a/src/ai/backend/appproxy/worker/errors/config.py +++ b/src/ai/backend/appproxy/worker/errors/config.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class MissingTraefikConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/missing-traefik-config" error_title = "Traefik configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -33,6 +36,7 @@ class MissingAnnounceAddressError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/missing-announce-address" error_title = "Announce address is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -47,6 +51,7 @@ class MissingProfilingConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/missing-profiling-config" error_title = "Profiling configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -61,6 +66,7 @@ class MissingPortConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/missing-port-config" error_title = "Port configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -75,6 +81,7 @@ class MissingPortProxyConfigError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/missing-port-proxy-config" error_title = "Port proxy configuration is missing." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -89,6 +96,7 @@ class CleanupContextNotInitializedError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/cleanup-context-not-initialized" error_title = "Cleanup context is not initialized." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/worker/errors/process.py b/src/ai/backend/appproxy/worker/errors/process.py index ba62f735c9f..53dd68ebff6 100644 --- a/src/ai/backend/appproxy/worker/errors/process.py +++ b/src/ai/backend/appproxy/worker/errors/process.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class SubprocessPipeError(BackendAIError): error_type = "https://api.backend.ai/probs/appproxy-worker/subprocess-pipe-error" error_title = "Subprocess pipe is not available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/appproxy/worker/proxy/backend/last_access_marker.py b/src/ai/backend/appproxy/worker/proxy/backend/last_access_marker.py index 184bc10b470..b8d20cbd216 100644 --- a/src/ai/backend/appproxy/worker/proxy/backend/last_access_marker.py +++ b/src/ai/backend/appproxy/worker/proxy/backend/last_access_marker.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from ai.backend.common.cron import PeriodicTask @@ -24,16 +24,20 @@ def __init__(self, backend: BaseBackend, route: RouteInfo) -> None: self._route = route @property + @override def name(self) -> str: return "last_access_marker" @property + @override def interval(self) -> float: return _LAST_ACCESS_MARKER_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._backend.mark_last_used_time(self._route) diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/h2/base.py b/src/ai/backend/appproxy/worker/proxy/frontend/h2/base.py index ccb147a1df8..f7dc55e9e33 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/h2/base.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/h2/base.py @@ -1,7 +1,7 @@ import asyncio import logging from asyncio import subprocess -from typing import Any +from typing import Any, override from ai.backend.appproxy.common.errors import ServerMisconfiguredError from ai.backend.appproxy.worker.proxy.backend.h2 import H2Backend @@ -23,6 +23,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.api_port_pool = set(range(h2_config.api_port_pool[0], h2_config.api_port_pool[1] + 1)) + @override async def list_inactive_circuits(self, threshold: int) -> list[Circuit]: # We can't measure activeness of HTTP/2 circuits return [] diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/h2/port.py b/src/ai/backend/appproxy/worker/proxy/frontend/h2/port.py index b19f25e5c84..66f952783ea 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/h2/port.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/h2/port.py @@ -1,7 +1,7 @@ import asyncio import logging from asyncio import subprocess -from typing import Any +from typing import Any, override from yarl import URL @@ -32,6 +32,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.processes = [] self.api_ports = {} + @override async def start(self) -> None: worker_config = self.root_context.local_config.proxy_worker port_proxy_config = worker_config.port_proxy @@ -73,6 +74,7 @@ async def start(self) -> None: self.api_ports[listen_port] = api_port log.info("started nghttpx server at {}:{}", port_proxy_config.bind_host, listen_port) + @override async def stop(self) -> None: for task in self.proc_monitor_tasks: task.cancel() @@ -87,19 +89,23 @@ async def stop(self) -> None: finally: self.api_ports = {} + @override async def initialize_backend(self, circuit: Circuit, routes: list[RouteInfo]) -> H2Backend: api_port = self.api_ports[self.get_circuit_key(circuit)] backend = H2Backend(URL(f"http://localhost:{api_port}"), self.root_context, circuit) return await self.update_backend(backend, routes) + @override async def update_backend(self, backend: H2Backend, routes: list[RouteInfo]) -> H2Backend: backend_configs = [BackendConfig(r.current_kernel_host, r.kernel_port) for r in routes] await backend.update_config(backend_configs) return backend + @override async def terminate_backend(self, backend: H2Backend) -> None: await backend.update_config([]) + @override def get_circuit_key(self, circuit: Circuit) -> int: if not isinstance(circuit.frontend, PortFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/h2/subdomain.py b/src/ai/backend/appproxy/worker/proxy/frontend/h2/subdomain.py index 47043344393..599e9980b8b 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/h2/subdomain.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/h2/subdomain.py @@ -1,7 +1,7 @@ import asyncio import logging from asyncio import subprocess -from typing import Any +from typing import Any, override from yarl import URL @@ -31,6 +31,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.log_monitor_tasks = [] self.proc_monitor_task = None + @override async def start(self) -> None: self.api_port = self.api_port_pool.pop() worker_config = self.root_context.local_config.proxy_worker @@ -67,6 +68,7 @@ async def start(self) -> None: self.proc_monitor_task = asyncio.create_task(self._proc_monitor_task(proc)) log.info("accepting proxy requests at {}:{}", service_addr.host, service_addr.port) + @override async def stop(self) -> None: if not self.proc_monitor_task: return @@ -84,18 +86,22 @@ async def stop(self) -> None: if self.api_port: self.api_port_pool.add(self.api_port) + @override async def initialize_backend(self, circuit: Circuit, routes: list[RouteInfo]) -> H2Backend: backend = H2Backend(URL(f"http://localhost:{self.api_port}"), self.root_context, circuit) return await self.update_backend(backend, routes) + @override async def update_backend(self, backend: H2Backend, routes: list[RouteInfo]) -> H2Backend: backend_configs = [BackendConfig(r.current_kernel_host, r.kernel_port) for r in routes] await backend.update_config(backend_configs) return backend + @override async def terminate_backend(self, backend: H2Backend) -> None: await backend.update_config([]) + @override def get_circuit_key(self, circuit: Circuit) -> str: if not isinstance(circuit.frontend, SubdomainFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/http/base.py b/src/ai/backend/appproxy/worker/proxy/frontend/http/base.py index 43f9aacfbaf..a795079f0ba 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/http/base.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/http/base.py @@ -1,5 +1,6 @@ import logging import time +from typing import override import aiohttp_jinja2 import jwt @@ -87,16 +88,20 @@ def ensure_credential(self, request: web.Request, circuit: Circuit) -> None: if decoded.get("id") != str(circuit.id): raise InvalidCredentials("E20009: Authorization token mismatch") + @override async def initialize_backend(self, circuit: Circuit, routes: list[RouteInfo]) -> HTTPBackend: return HTTPBackend(routes, self.root_context, circuit) + @override async def update_backend(self, backend: HTTPBackend, routes: list[RouteInfo]) -> HTTPBackend: await backend.update_routes(routes) return backend + @override async def terminate_backend(self, backend: HTTPBackend) -> None: await backend.close() + @override async def list_inactive_circuits(self, threshold: int) -> list[Circuit]: now = time.time() return [ diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/http/port.py b/src/ai/backend/appproxy/worker/proxy/frontend/http/port.py index 14914712c85..c133eda195c 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/http/port.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/http/port.py @@ -1,7 +1,7 @@ import importlib.resources import logging import ssl -from typing import Any +from typing import Any, override import aiohttp_jinja2 import jinja2 @@ -26,6 +26,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.sites = [] + @override async def start(self) -> None: ssl_ctx = None proxy_worker_config = self.root_context.local_config.proxy_worker @@ -86,6 +87,7 @@ async def start(self) -> None: port_end, ) + @override async def stop(self) -> None: for site in self.sites: await site.stop() @@ -105,6 +107,7 @@ async def ensure_slot_middleware( request["backend"] = backend return await handler(request) + @override def get_circuit_key(self, circuit: Circuit) -> int: if not isinstance(circuit.frontend, PortFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/http/subdomain.py b/src/ai/backend/appproxy/worker/proxy/frontend/http/subdomain.py index 75abd9699b3..d969c09e77a 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/http/subdomain.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/http/subdomain.py @@ -1,7 +1,7 @@ import importlib.resources import logging import ssl -from typing import Any +from typing import Any, override import aiohttp_jinja2 import jinja2 @@ -31,6 +31,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: raise ServerMisconfiguredError("worker:proxy-worker.wildcard-domain") self.wildcard_config = self.root_context.local_config.proxy_worker.wildcard_domain + @override async def start(self) -> None: ssl_ctx = None proxy_worker_config = self.root_context.local_config.proxy_worker @@ -75,6 +76,7 @@ async def start(self) -> None: self.site = site log.info("accepting proxy requests at {}", service_addr) + @override async def stop(self) -> None: if self.site: await self.site.stop() @@ -102,6 +104,7 @@ async def _exception_safe_handler(request: web.Request) -> web.StreamResponse: return await self.exception_safe_handler_wrapper(request, _exception_safe_handler) + @override def get_circuit_key(self, circuit: Circuit) -> str: if not isinstance(circuit.frontend, SubdomainFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/tcp.py b/src/ai/backend/appproxy/worker/proxy/frontend/tcp.py index 2a8e458e87b..cb6d5061788 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/tcp.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/tcp.py @@ -3,7 +3,7 @@ import logging import socket import time -from typing import Any +from typing import Any, override from aiohttp import web @@ -37,6 +37,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.servers = [] self.server_tasks = [] + @override async def start(self) -> None: proxy_worker_config = self.root_context.local_config.proxy_worker port_proxy_config = proxy_worker_config.port_proxy @@ -72,6 +73,7 @@ async def _listen_task(self, circuit_key: int, server: asyncio.Server) -> None: log.exception("TCPFrontend._listen_task(c: {}): exception:", circuit_key) raise + @override async def stop(self) -> None: for task in self.server_tasks: task.cancel() @@ -104,16 +106,20 @@ def _is_peer_allowed(self, writer: asyncio.StreamWriter, circuit: SerializableCi return False return True + @override async def initialize_backend(self, circuit: Circuit, routes: list[RouteInfo]) -> TCPBackend: return TCPBackend(routes, self.root_context, circuit) + @override async def update_backend(self, backend: TCPBackend, routes: list[RouteInfo]) -> TCPBackend: await backend.update_routes(routes) return backend + @override async def terminate_backend(self, backend: TCPBackend) -> None: await backend.close() + @override async def list_inactive_circuits(self, threshold: int) -> list[Circuit]: now = time.time() return [ @@ -153,6 +159,7 @@ async def pipe( end = time.monotonic() metrics.proxy.observe_downstream_tcp_end(duration=int(end - start)) + @override def get_circuit_key(self, circuit: Circuit) -> int: if not isinstance(circuit.frontend, PortFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/proxy/frontend/traefik.py b/src/ai/backend/appproxy/worker/proxy/frontend/traefik.py index 659e1b8653d..cab46635cef 100644 --- a/src/ai/backend/appproxy/worker/proxy/frontend/traefik.py +++ b/src/ai/backend/appproxy/worker/proxy/frontend/traefik.py @@ -5,7 +5,7 @@ import time import uuid from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override from aiohttp import web @@ -46,17 +46,21 @@ def __init__(self, frontend: AbstractTraefikFrontend[Any]) -> None: self._frontend = frontend @property + @override def name(self) -> str: return "last_used_time_marker_writer" @property + @override def interval(self) -> float: return _LAST_USED_TIME_MARKER_WRITER_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._frontend.last_used_time_marker_writer() @@ -70,17 +74,21 @@ def __init__(self, frontend: AbstractTraefikFrontend[Any]) -> None: self._frontend = frontend @property + @override def name(self) -> str: return "active_circuit_writer" @property + @override def interval(self) -> float: return _ACTIVE_CIRCUIT_WRITER_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._frontend.active_circuit_writer() @@ -110,6 +118,7 @@ def __init__(self, root_context: RootContext) -> None: app.router.add_head("/{key}/mark-inactive", self.mark_inactive) self.runner = web.AppRunner(app, access_log=None) + @override async def start(self) -> None: if not self.root_context.local_config.proxy_worker.traefik: raise MissingTraefikConfigError("Traefik configuration is required") @@ -130,6 +139,7 @@ async def start(self) -> None: ]) await self._local_cron.start() + @override async def stop(self) -> None: if not self.root_context.local_config.proxy_worker.traefik: raise MissingTraefikConfigError("Traefik configuration is required") @@ -180,24 +190,29 @@ async def active_circuit_writer(self) -> None: async with self.redis_keys_lock: self.redis_keys.update(keys) + @override async def initialize_backend(self, circuit: Circuit, routes: list[RouteInfo]) -> TraefikBackend: return TraefikBackend(self.root_context, circuit, routes) + @override async def update_backend( self, backend: TraefikBackend, routes: list[RouteInfo] ) -> TraefikBackend: backend.routes = routes return backend + @override async def terminate_backend(self, backend: TraefikBackend) -> None: pass + @override async def list_inactive_circuits(self, threshold: int) -> list[Circuit]: # FIXME: implement return [] class TraefikPortFrontend(AbstractTraefikFrontend[int]): + @override def get_circuit_key(self, circuit: Circuit) -> int: if not isinstance(circuit.frontend, PortFrontendInfo): raise InvalidFrontendTypeError( @@ -207,6 +222,7 @@ def get_circuit_key(self, circuit: Circuit) -> int: class TraefikSubdomainFrontend(AbstractTraefikFrontend[str]): + @override def get_circuit_key(self, circuit: Circuit) -> str: if not isinstance(circuit.frontend, SubdomainFrontendInfo): raise InvalidFrontendTypeError( @@ -216,6 +232,7 @@ def get_circuit_key(self, circuit: Circuit) -> str: class TraefikTCPFrontend(AbstractTraefikFrontend[int]): + @override def get_circuit_key(self, circuit: Circuit) -> int: if not isinstance(circuit.frontend, PortFrontendInfo): raise InvalidFrontendTypeError( diff --git a/src/ai/backend/appproxy/worker/tasks/heartbeat.py b/src/ai/backend/appproxy/worker/tasks/heartbeat.py index 2e16d9cc6e8..bb6a09aff2b 100644 --- a/src/ai/backend/appproxy/worker/tasks/heartbeat.py +++ b/src/ai/backend/appproxy/worker/tasks/heartbeat.py @@ -4,7 +4,7 @@ import logging import uuid -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from tenacity import AsyncRetrying, TryAgain, retry_if_exception_type, wait_exponential @@ -28,17 +28,21 @@ def __init__(self, root_ctx: RootContext) -> None: self._root_ctx = root_ctx @property + @override def name(self) -> str: return "worker_heartbeat" @property + @override def interval(self) -> float: return self._root_ctx.local_config.proxy_worker.heartbeat_period @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: try: async for attempt in AsyncRetrying( diff --git a/src/ai/backend/cli/extensions.py b/src/ai/backend/cli/extensions.py index e6f62815efc..c0771c342ea 100644 --- a/src/ai/backend/cli/extensions.py +++ b/src/ai/backend/cli/extensions.py @@ -2,7 +2,7 @@ import signal import sys from collections.abc import Callable -from typing import Any +from typing import Any, override import click from click.exceptions import Abort, ClickException @@ -19,6 +19,7 @@ class InterruptAwareCommandMixin(click.BaseCommand): continuing the shell/batch script. """ + @override def main(self, *args: Any, **kwargs: Any) -> None: # type: ignore[override] _interrupted = False try: @@ -80,6 +81,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._commands: dict[str, click.Command] = {} self._aliases: dict[str, str] = {} + @override def command(self, *args: Any, **kwargs: Any) -> Any: aliases = kwargs.pop("aliases", []) decorator = super().command(*args, **kwargs) @@ -96,6 +98,7 @@ def _decorator(f: Callable[..., Any]) -> click.Command: return _decorator + @override def group(self, *args: Any, **kwargs: Any) -> Any: aliases = kwargs.pop("aliases", []) # keep the same class type unless explicitly specified @@ -115,6 +118,7 @@ def _decorator(f: Callable[..., Any]) -> click.Group: return _decorator + @override def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: if cmd_name in self._aliases: cmd_name = self._aliases[cmd_name] @@ -123,6 +127,7 @@ def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None return command return None + @override def format_commands(self, ctx: click.Context, formatter: click.HelpFormatter) -> None: commands = [] for subcommand in self.list_commands(ctx): diff --git a/src/ai/backend/cli/params.py b/src/ai/backend/cli/params.py index ee97af642b5..e3a54268faf 100644 --- a/src/ai/backend/cli/params.py +++ b/src/ai/backend/cli/params.py @@ -211,6 +211,7 @@ def __init__(self, type_: type[TScalar] | None = None) -> None: super().__init__() self.type_ = type_ if type_ is not None else str + @override def convert( self, value: str | int, param: click.Parameter | None, ctx: click.Context | None ) -> int | list[Any]: @@ -239,6 +240,7 @@ def __init__(self, type_: type[TScalar] | type[click.ParamType] | click.ParamTyp super().__init__() self.type_ = type_ + @override def convert( self, value: str | Undefined, diff --git a/src/ai/backend/client/cli/v2/prometheus_query_preset/commands.py b/src/ai/backend/client/cli/v2/prometheus_query_preset/commands.py index 68fd57ebe49..94ed312a1c2 100644 --- a/src/ai/backend/client/cli/v2/prometheus_query_preset/commands.py +++ b/src/ai/backend/client/cli/v2/prometheus_query_preset/commands.py @@ -11,7 +11,7 @@ import asyncio import json from datetime import UTC, datetime -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, cast, override from uuid import UUID import click @@ -42,6 +42,7 @@ class PrometheusTimestampParamType(click.ParamType): def __init__(self) -> None: self._datetime_parser = click.DateTime(formats=_DATETIME_FORMATS) + @override def convert( self, value: object, diff --git a/src/ai/backend/client/exceptions.py b/src/ai/backend/client/exceptions.py index 33f11ae24d8..8363328fc98 100644 --- a/src/ai/backend/client/exceptions.py +++ b/src/ai/backend/client/exceptions.py @@ -1,5 +1,5 @@ import json -from typing import Any +from typing import Any, override __all__: tuple[str, ...] = ( "APIVersionWarning", @@ -12,6 +12,7 @@ class BackendError(Exception): """Exception type to catch all ai.backend-related errors.""" + @override def __str__(self) -> str: return repr(self) diff --git a/src/ai/backend/client/load_balancing.py b/src/ai/backend/client/load_balancing.py index dd8d38ce631..842bcf76a65 100644 --- a/src/ai/backend/client/load_balancing.py +++ b/src/ai/backend/client/load_balancing.py @@ -2,6 +2,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Mapping +from typing import override import attrs from yarl import URL @@ -35,6 +36,7 @@ class SimpleRRLoadBalancer(LoadBalancer): Rotates the endpoints upon every request. """ + @override def rotate(self, endpoints: list[URL]) -> None: if len(endpoints) == 1: return @@ -47,6 +49,7 @@ class PeriodicRRLoadBalancer(LoadBalancer): Rotates the endpoints upon the specified interval. """ + @override def rotate(self, endpoints: list[URL]) -> None: pass @@ -56,6 +59,7 @@ class LowestLatencyLoadBalancer(LoadBalancer): Change the endpoints with the lowest average latency for last N requests. """ + @override def rotate(self, endpoints: list[URL]) -> None: pass diff --git a/src/ai/backend/client/output/console.py b/src/ai/backend/client/output/console.py index 39bc0e87d4a..8dcbd1bd401 100644 --- a/src/ai/backend/client/output/console.py +++ b/src/ai/backend/client/output/console.py @@ -2,7 +2,7 @@ import sys from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any, TypeVar, cast +from typing import Any, TypeVar, cast, override from tabulate import tabulate @@ -21,6 +21,7 @@ class NoItems(Exception): class ConsoleOutputHandler(BaseOutputHandler): + @override def print_item( self, item: _Item | None, @@ -44,6 +45,7 @@ def print_item( ) ) + @override def print_items( self, items: Sequence[_Item], @@ -67,6 +69,7 @@ def print_items( ) ) + @override def print_result_set( self, result_set: ResultSet, @@ -88,6 +91,7 @@ def print_result_set( ) ) + @override def print_list( self, items: Sequence[_Item], @@ -143,6 +147,7 @@ def infinite_fetch() -> Iterator[_Item]: ): print(line, end="") + @override def print_paginated_list( self, fetch_func: Callable[[int, int], PaginatedResult[T]], @@ -204,6 +209,7 @@ def infinite_fetch(_page_size: int) -> Iterator[_Item]: except NoItems: print("No matching items.") + @override def print_mutation_result( self, item: _Item, @@ -227,6 +233,7 @@ def print_mutation_result( ) ) + @override def print_mutation_error( self, error: Exception | None = None, @@ -251,12 +258,14 @@ def print_mutation_error( if error is not None: print_error(error) + @override def print_error( self, error: Exception, ) -> None: print_error(error) + @override def print_fail( self, message: str, diff --git a/src/ai/backend/client/output/formatters.py b/src/ai/backend/client/output/formatters.py index b2e93ff7fd3..c334d410ecf 100644 --- a/src/ai/backend/client/output/formatters.py +++ b/src/ai/backend/client/output/formatters.py @@ -5,7 +5,7 @@ import textwrap from collections import defaultdict from collections.abc import Callable, Mapping -from typing import Any, cast +from typing import Any, cast, override import humanize @@ -76,6 +76,7 @@ class OutputFormatter(AbstractOutputFormatter): The base implementation of output formats. """ + @override def format_console(self, value: Any, field: FieldSpec) -> str: if value is None: return "(null)" @@ -91,6 +92,7 @@ def format_console(self, value: Any, field: FieldSpec) -> str: return "[" + ", ".join(self.format_console(v, field) for v in value) + "]" return str(value) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: if value is None: return None @@ -104,12 +106,14 @@ def format_json(self, value: Any, field: FieldSpec) -> Any: class NestedDictOutputFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: if value is None: return "(null)" value = json.loads(value) return format_nested_dicts(value) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: if value is None: return None @@ -117,11 +121,13 @@ def format_json(self, value: Any, field: FieldSpec) -> Any: class MiBytesOutputFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: if value is not None: value = round(value / 2**20, 1) return super().format_console(value, field) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: if value is not None: value = round(value / 2**20, 1) @@ -129,11 +135,13 @@ def format_json(self, value: Any, field: FieldSpec) -> Any: class SizeBytesOutputFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: if value is not None: value = humanize.naturalsize(value, binary=True, gnu=True) return super().format_console(value, field) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: if value is not None: value = humanize.naturalsize(value, binary=True, gnu=True) @@ -144,9 +152,11 @@ class SubFieldOutputFormatter(OutputFormatter): def __init__(self, subfield_name: str) -> None: self._subfield_name = subfield_name + @override def format_console(self, value: Any, field: FieldSpec) -> str: return super().format_console(value[self._subfield_name], field) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return super().format_json(value[self._subfield_name], field) @@ -167,14 +177,17 @@ def _get_name(self, labels: Any) -> str: raise ValueError("Expected exactly one owner email label") return f"{customized_name[0]} (Owner: {owner_email[0]})" + @override def format_console(self, value: Any, field: FieldSpec) -> str: return super().format_console(self._get_name(value), field) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return super().format_json(self._get_name(value), field) class ResourceSlotFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: value = json.loads(value) if mem := value.get("mem"): @@ -182,6 +195,7 @@ def format_console(self, value: Any, field: FieldSpec) -> str: return ", ".join(f"{k}:{v}" for k, v in value.items()) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return json.loads(value) @@ -194,6 +208,7 @@ def format_json(self, value: Any, field: FieldSpec) -> Any: class AgentStatFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: try: raw_stats = json.loads(value) @@ -264,20 +279,24 @@ def format_value(metric: MetricValue, binary: bool) -> str: bufs.append("\n".join(dev_metric_bufs)) return "\n".join(bufs) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: # TODO: improve return self.format_console(value, field) class GroupListFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: return ", ".join(g["name"] for g in value) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return value class InlineRoutingFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: count_by_status: defaultdict[int, int] = defaultdict(int) for route in value: @@ -286,19 +305,23 @@ def format_console(self, value: Any, field: FieldSpec) -> str: f"{key} {value}" for key, value in sorted(count_by_status.items(), key=lambda kv: kv[0]) ]) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return value class KernelStatFormatter(OutputFormatter): + @override def format_console(self, value: Any, field: FieldSpec) -> str: return format_stats(value) + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return value class NestedObjectFormatter(OutputFormatter): + @override def format_json(self, value: Any, field: FieldSpec) -> Any: if not isinstance(value, list): raise ValueError("NestedObjectFormatter expects a list value") @@ -318,6 +341,7 @@ def _fit_multiline_in_cell(text: str, indent: str) -> str: class ContainerListFormatter(NestedObjectFormatter): + @override def format_console(self, value: Any, field: FieldSpec, indent: str = "") -> str: if not isinstance(value, list): raise ValueError("ContainerListFormatter expects a list value") @@ -337,6 +361,7 @@ def format_console(self, value: Any, field: FieldSpec, indent: str = "") -> str: class DependencyListFormatter(NestedObjectFormatter): + @override def format_console(self, value: Any, field: FieldSpec, indent: str = "") -> str: if not isinstance(value, list): raise ValueError("DependencyListFormatter expects a list value") @@ -358,10 +383,12 @@ def format_console(self, value: Any, field: FieldSpec, indent: str = "") -> str: class ImageObjectFormatter(OutputFormatter): """Formatter for nested image_object field, extracting the name.""" + @override def format_console(self, value: Any, field: FieldSpec) -> str: if value is None: return "(none)" return value.get("name") or "(unknown)" + @override def format_json(self, value: Any, field: FieldSpec) -> Any: return value diff --git a/src/ai/backend/client/output/json.py b/src/ai/backend/client/output/json.py index 6afcdf1b012..2b33080c907 100644 --- a/src/ai/backend/client/output/json.py +++ b/src/ai/backend/client/output/json.py @@ -2,7 +2,7 @@ import json from collections.abc import Callable, Mapping, Sequence -from typing import Any, TypeVar, cast +from typing import Any, TypeVar, cast, override from ai.backend.client.exceptions import BackendAPIError from ai.backend.common.types import ResultSet @@ -15,6 +15,7 @@ class JsonOutputHandler(BaseOutputHandler): + @override def print_item( self, item: Mapping[str, Any] | None, @@ -49,6 +50,7 @@ def print_item( ) ) + @override def print_items( self, items: Sequence[Mapping[str, Any]], @@ -75,12 +77,14 @@ def print_items( ) ) + @override def print_result_set( self, result_set: ResultSet, ) -> None: print(json.dumps(result_set)) + @override def print_list( self, items: Sequence[Mapping[str, Any]], @@ -118,6 +122,7 @@ def print_list( ) ) + @override def print_paginated_list( self, fetch_func: Callable[[int, int], PaginatedResult[T]], @@ -148,6 +153,7 @@ def print_paginated_list( ) ) + @override def print_mutation_result( self, item: Mapping[str, Any], @@ -172,6 +178,7 @@ def print_mutation_result( ) ) + @override def print_mutation_error( self, error: Exception | None = None, @@ -196,6 +203,7 @@ def print_mutation_error( ) ) + @override def print_error( self, error: Exception, @@ -225,6 +233,7 @@ def print_error( ) ) + @override def print_fail( self, message: str, diff --git a/src/ai/backend/client/request.py b/src/ai/backend/client/request.py index 5f3b4fec679..9bf3304a663 100644 --- a/src/ai/backend/client/request.py +++ b/src/ai/backend/client/request.py @@ -16,6 +16,7 @@ from typing import ( Any, cast, + override, ) import aiohttp @@ -75,6 +76,7 @@ async def _coro_return[T](val: T) -> T: class ExtendedJSONEncoder(modjson.JSONEncoder): + @override def default(self, obj: Any) -> Any: if isinstance(obj, uuid.UUID): return str(obj) @@ -657,14 +659,17 @@ def __init__( self._raw_ws = cast(aiohttp.ClientWebSocketResponse, underlying_response) @property + @override def content_type(self) -> str: raise AttributeError("WebSocketResponse does not have an explicit content type.") @property + @override def content_length(self) -> int | None: raise AttributeError("WebSocketResponse does not have a fixed content length.") @property + @override def content(self) -> aiohttp.StreamReader: raise AttributeError("WebSocketResponse does not support reading the content.") diff --git a/src/ai/backend/client/session.py b/src/ai/backend/client/session.py index 4961bd5216e..a9c665bf9c0 100644 --- a/src/ai/backend/client/session.py +++ b/src/ai/backend/client/session.py @@ -12,6 +12,7 @@ Any, Literal, TypeVar, + override, ) import aiohttp @@ -148,6 +149,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.stream_block = threading.Event() self.agen_shutdown = False + @override def run(self) -> None: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) @@ -444,6 +446,7 @@ async def _create_aiohttp_session() -> aiohttp.ClientSession: self.aiohttp_session = self.worker_thread.execute(_create_aiohttp_session()) + @override def open(self) -> None: self._context_token = api_session.set(self) if not self._proxy_mode: @@ -451,6 +454,7 @@ def open(self) -> None: _negotiate_api_version(self.aiohttp_session, self.config) ) + @override def close(self) -> None: """ Terminates the session. It schedules the ``close()`` coroutine @@ -475,6 +479,7 @@ def worker_thread(self) -> _SyncWorkerThread: """ return self._worker_thread + @override def __enter__(self) -> Session: if self.closed: raise RuntimeError("Cannot reuse closed session") @@ -489,6 +494,7 @@ def __enter__(self) -> Session: pass return self + @override def __exit__(self, *exc_info: Any) -> Literal[False]: self.close() return False # raise up the inner exception @@ -535,6 +541,7 @@ async def _aopen(self) -> None: if not self._proxy_mode: self.api_version = await _negotiate_api_version(self.aiohttp_session, self.config) + @override def open(self) -> Awaitable[None]: return self._aopen() @@ -546,9 +553,11 @@ async def _aclose(self) -> None: await _close_aiohttp_session(self.aiohttp_session) api_session.reset(self._context_token) + @override def close(self) -> Awaitable[None]: return self._aclose() + @override async def __aenter__(self) -> AsyncSession: if self.closed: raise RuntimeError("Cannot reuse closed session") @@ -563,6 +572,7 @@ async def __aenter__(self) -> AsyncSession: pass return self + @override async def __aexit__(self, *exc_info: Any) -> Literal[False]: await self.close() return False # raise up the inner exception diff --git a/src/ai/backend/client/utils.py b/src/ai/backend/client/utils.py index 9b8eff4897f..1a76aad1922 100644 --- a/src/ai/backend/client/utils.py +++ b/src/ai/backend/client/utils.py @@ -3,7 +3,7 @@ import textwrap import uuid from pathlib import Path -from typing import Any +from typing import Any, override from ai.backend.client.output.types import FieldSet, FieldSpec @@ -81,9 +81,11 @@ def __init__(self, file_path: str, *, tqdm_instance: Any = None) -> None: self._owns_tqdm = False self.tqdm = tqdm_instance + @override def __enter__(self) -> "ProgressReportingReader": return self + @override def __exit__( self, exc_type: type[BaseException] | None, @@ -94,24 +96,28 @@ def __exit__( self.tqdm.close() self.close() + @override def read(self, *args: Any, **kwargs: Any) -> bytes: chunk = super().read(*args, **kwargs) self.tqdm.set_postfix(file=self._filename, refresh=False) self.tqdm.update(len(chunk)) return chunk + @override def read1(self, *args: Any, **kwargs: Any) -> bytes: chunk = super().read1(*args, **kwargs) self.tqdm.set_postfix(file=self._filename, refresh=False) self.tqdm.update(len(chunk)) return chunk + @override def readinto(self, *args: Any, **kwargs: Any) -> int: count = super().readinto(*args, **kwargs) self.tqdm.set_postfix(file=self._filename, refresh=False) self.tqdm.update(count) return count + @override def readinto1(self, *args: Any, **kwargs: Any) -> int: count = super().readinto1(*args, **kwargs) self.tqdm.set_postfix(file=self._filename, refresh=False) diff --git a/src/ai/backend/common/asyncio.py b/src/ai/backend/common/asyncio.py index 2f049dfa215..419ef575435 100644 --- a/src/ai/backend/common/asyncio.py +++ b/src/ai/backend/common/asyncio.py @@ -9,6 +9,7 @@ Protocol, TypeVar, cast, + override, ) __all__ = ( @@ -149,8 +150,10 @@ class closing_async(AbstractAsyncContextManager[_SupportsAsyncCloseT]): def __init__(self, obj: _SupportsAsyncCloseT) -> None: self.obj = obj + @override async def __aenter__(self) -> _SupportsAsyncCloseT: return self.obj + @override async def __aexit__(self, *exc_info: Any) -> None: await self.obj.close() diff --git a/src/ai/backend/common/auth/__init__.py b/src/ai/backend/common/auth/__init__.py index 8e84bc11150..db4f4915125 100644 --- a/src/ai/backend/common/auth/__init__.py +++ b/src/ai/backend/common/auth/__init__.py @@ -1,4 +1,4 @@ -from typing import NewType +from typing import NewType, override from callosum.auth import ( AbstractClientAuthenticator, @@ -25,12 +25,15 @@ def __init__( self._manager_public_key = manager_public_key self._manager_secret_key = manager_secret_key + @override async def server_public_key(self) -> bytes: return self._agent_public_key + @override async def client_public_key(self) -> bytes: return self._manager_public_key + @override async def client_identity(self) -> Identity: if self._manager_secret_key is None: raise RuntimeError("Manager secret key is not initialized") @@ -50,14 +53,17 @@ def __init__( self._agent_public_key = agent_public_key self._agent_secret_key = agent_secret_key + @override async def server_public_key(self) -> bytes: return self._agent_public_key + @override async def server_identity(self) -> Identity: if self._agent_secret_key is None: raise RuntimeError("Agent secret key is not initialized") return Identity(self.domain, self._agent_secret_key) + @override async def check_client(self, creds: Credential) -> AuthResult: if creds.domain == self.domain and creds.public_key == self._manager_public_key: return AuthResult(success=True, user_id="manager") diff --git a/src/ai/backend/common/bgtask/bgtask.py b/src/ai/backend/common/bgtask/bgtask.py index 32c1d3adf47..3abc8711bdf 100644 --- a/src/ai/backend/common/bgtask/bgtask.py +++ b/src/ai/backend/common/bgtask/bgtask.py @@ -15,6 +15,7 @@ Final, ParamSpec, Self, + override, ) from ai.backend.common.bgtask.exception import InvalidTaskMetadataError @@ -170,12 +171,15 @@ class LocalBgtask(BackgroundTaskMeta): def __init__(self, task: asyncio.Task[Any]) -> None: self._task = task + @override def total_info(self) -> TaskTotalInfo: raise NotImplementedError("LocalBgtask is not recoverable and cannot be stored") + @override def retriable(self) -> bool: return False + @override def async_tasks(self) -> Sequence[asyncio.Task[Any]]: return [self._task] @@ -192,12 +196,15 @@ def __init__( self._total_info = total_info self._task = task + @override def total_info(self) -> TaskTotalInfo: return self._total_info + @override def retriable(self) -> bool: return True + @override def async_tasks(self) -> Sequence[asyncio.Task[Any]]: return [self._task] @@ -214,12 +221,15 @@ def __init__( self._total_info = total_info self._tasks = tasks + @override def total_info(self) -> TaskTotalInfo: return self._total_info + @override def retriable(self) -> bool: return True + @override def async_tasks(self) -> Sequence[asyncio.Task[Any]]: return self._tasks diff --git a/src/ai/backend/common/bgtask/hooks/composite_hook.py b/src/ai/backend/common/bgtask/hooks/composite_hook.py index ca806fbb349..f6f5472251a 100644 --- a/src/ai/backend/common/bgtask/hooks/composite_hook.py +++ b/src/ai/backend/common/bgtask/hooks/composite_hook.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager +from typing import override from .base import AbstractTaskHook, TaskContext @@ -13,6 +14,7 @@ def __init__(self, hooks: list[AbstractTaskHook]) -> None: self._hooks = hooks @asynccontextmanager + @override async def apply(self, context: TaskContext) -> AsyncIterator[TaskContext]: async with AsyncExitStack() as stack: # Apply all hooks in sequence diff --git a/src/ai/backend/common/bgtask/hooks/event_hook.py b/src/ai/backend/common/bgtask/hooks/event_hook.py index d2121b5f308..b66b31dac43 100644 --- a/src/ai/backend/common/bgtask/hooks/event_hook.py +++ b/src/ai/backend/common/bgtask/hooks/event_hook.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.common.events.event_types.bgtask.broadcast import BgtaskUpdatedEvent @@ -17,6 +18,7 @@ def __init__(self, event_producer: EventProducer) -> None: self._event_producer = event_producer @asynccontextmanager + @override async def apply(self, context: TaskContext) -> AsyncIterator[TaskContext]: # Pre-execution: send task started event cache_id = EventCacheDomain.BGTASK.cache_id(str(context.task_id)) diff --git a/src/ai/backend/common/bgtask/hooks/metric_hook.py b/src/ai/backend/common/bgtask/hooks/metric_hook.py index dd10459bf00..3eea0aea250 100644 --- a/src/ai/backend/common/bgtask/hooks/metric_hook.py +++ b/src/ai/backend/common/bgtask/hooks/metric_hook.py @@ -3,7 +3,7 @@ import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from typing import Protocol +from typing import Protocol, override from ai.backend.common.bgtask.types import BgtaskStatus from ai.backend.common.exception import ErrorCode @@ -37,6 +37,7 @@ def __init__(self, observer: BackgroundTaskObserver) -> None: self._observer = observer @asynccontextmanager + @override async def apply(self, context: TaskContext) -> AsyncIterator[TaskContext]: # Pre-execution: start metric observation self._observer.observe_bgtask_started(task_name=context.task_name.value) diff --git a/src/ai/backend/common/bgtask/hooks/valkey_hook.py b/src/ai/backend/common/bgtask/hooks/valkey_hook.py index 9c9c2983d6c..bde9dae0cca 100644 --- a/src/ai/backend/common/bgtask/hooks/valkey_hook.py +++ b/src/ai/backend/common/bgtask/hooks/valkey_hook.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.clients.valkey_client.valkey_bgtask.client import ( TaskSetKey, @@ -23,6 +24,7 @@ def __init__( self._task_set_key = task_set_key @asynccontextmanager + @override async def apply(self, context: TaskContext) -> AsyncIterator[TaskContext]: try: yield context diff --git a/src/ai/backend/common/bgtask/tasks/heartbeat.py b/src/ai/backend/common/bgtask/tasks/heartbeat.py index 58fdfdc9596..753c8632bcd 100644 --- a/src/ai/backend/common/bgtask/tasks/heartbeat.py +++ b/src/ai/backend/common/bgtask/tasks/heartbeat.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, manager: BackgroundTaskManager) -> None: self._manager = manager @property + @override def name(self) -> str: return "bgtask_heartbeat" @property + @override def interval(self) -> float: return _HEARTBEAT_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._manager.do_heartbeat() diff --git a/src/ai/backend/common/bgtask/tasks/retry.py b/src/ai/backend/common/bgtask/tasks/retry.py index 4d0286d97b4..0acdd4580dc 100644 --- a/src/ai/backend/common/bgtask/tasks/retry.py +++ b/src/ai/backend/common/bgtask/tasks/retry.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from ai.backend.common.cron import PeriodicTask @@ -21,16 +21,20 @@ def __init__(self, manager: BackgroundTaskManager) -> None: self._manager = manager @property + @override def name(self) -> str: return "bgtask_retry" @property + @override def interval(self) -> float: return _RETRY_CHECK_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: await self._manager.do_retry_check() diff --git a/src/ai/backend/common/cli.py b/src/ai/backend/common/cli.py index 6520b567a8e..65743da0a75 100644 --- a/src/ai/backend/common/cli.py +++ b/src/ai/backend/common/cli.py @@ -5,7 +5,7 @@ from enum import Enum from importlib import import_module from types import FunctionType -from typing import Any +from typing import Any, override import click @@ -63,6 +63,7 @@ def __init__(self, enum: type[Enum]) -> None: super().__init__(enum_members) self.enum = enum + @override def convert(self, value: Any, param: click.Parameter | None, ctx: click.Context | None) -> Enum: if isinstance(value, self.enum): # for default value, it is already the enum type. @@ -70,6 +71,7 @@ def convert(self, value: Any, param: click.Parameter | None, ctx: click.Context value = super().convert(value, param, ctx) return self.enum[value] + @override def get_metavar(self, param: click.Parameter) -> str: name = self.enum.__name__ name = re.sub(r"([A-Z\d]+)([A-Z][a-z])", r"\1_\2", name) @@ -80,6 +82,7 @@ def get_metavar(self, param: click.Parameter) -> str: class MinMaxRangeParamType(click.ParamType): name = "min-max decimal range" + @override def convert( self, value: Any, param: click.Parameter | None, ctx: click.Context | None ) -> tuple[Decimal | None, Decimal | None]: @@ -97,6 +100,7 @@ def convert( except (ArithmeticError, ValueError): self.fail(f"{value!r} contains an invalid number", param, ctx) + @override def get_metavar(self, param: click.Parameter) -> str: return "MIN:MAX" diff --git a/src/ai/backend/common/clients/http_client/client_pool.py b/src/ai/backend/common/clients/http_client/client_pool.py index 0873c45a68d..fafd39eb909 100644 --- a/src/ai/backend/common/clients/http_client/client_pool.py +++ b/src/ai/backend/common/clients/http_client/client_pool.py @@ -7,7 +7,7 @@ import warnings from collections.abc import MutableMapping from dataclasses import dataclass -from typing import Any, Protocol +from typing import Any, Protocol, override import aiohttp from aiotools import cancel_and_wait @@ -95,6 +95,7 @@ async def close(self) -> None: await client.session.close() self._clients.clear() + @override def __repr__(self) -> str: return ( f"" diff --git a/src/ai/backend/common/clients/valkey_client/client.py b/src/ai/backend/common/clients/valkey_client/client.py index e2ebc30af51..072508cfd4d 100644 --- a/src/ai/backend/common/clients/valkey_client/client.py +++ b/src/ai/backend/common/clients/valkey_client/client.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterable from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass -from typing import Any, Final, Self +from typing import Any, Final, Self, override from aiotools import cancel_and_wait from glide import ( @@ -192,12 +192,14 @@ def __init__( self._human_readable_name = human_readable_name self._pubsub_channels = pubsub_channels + @override async def connect(self) -> None: if self._valkey_client is not None: return await self._create_valkey_client() + @override async def disconnect(self) -> None: if self._valkey_client: await self._valkey_client.close(err_message="ValkeyStandaloneClient is closed.") @@ -243,6 +245,7 @@ async def _create_valkey_client(self) -> None: self._human_readable_name, ) + @override async def ping(self) -> None: """ Ping the server to check if the connection is alive. @@ -255,9 +258,11 @@ async def ping(self) -> None: raise ClientNotConnectedError("ValkeyStandaloneClient is not connected") await self._valkey_client.ping() + @override async def need_reconnect(self) -> bool: return self._valkey_client is None + @override async def check_health(self) -> None: """ Verify that the standalone connection is alive by pinging the server. @@ -271,6 +276,7 @@ async def check_health(self) -> None: await self._valkey_client.ping() @asynccontextmanager + @override async def client(self) -> AsyncIterator[GlideClient]: if self._valkey_client is None: raise ClientNotConnectedError("ValkeyStandaloneClient is not connected") @@ -319,12 +325,14 @@ def __init__( self._valkey_client = None self._master_address = None + @override async def connect(self) -> None: if self._valkey_client is not None: return await self._create_valkey_client() + @override async def disconnect(self) -> None: if self._valkey_client: await self._valkey_client.close(err_message="ValkeySentinelClient is closed.") @@ -386,6 +394,7 @@ async def _get_master_address(self) -> tuple[str, int] | None: ) return None + @override async def ping(self) -> None: """ Ping the current master to check if the connection is alive. @@ -398,6 +407,7 @@ async def ping(self) -> None: raise ClientNotConnectedError("ValkeySentinelClient is not connected") await self._valkey_client.ping() + @override async def need_reconnect(self) -> bool: """Check if reconnection is needed (client disconnected or master changed).""" if self._valkey_client is None: @@ -412,6 +422,7 @@ async def need_reconnect(self) -> bool: return current_master != self._master_address + @override async def check_health(self) -> None: """ Verify that the connected node is still a master via the ROLE command. @@ -439,6 +450,7 @@ async def check_health(self) -> None: ) @asynccontextmanager + @override async def client(self) -> AsyncIterator[GlideClient]: if self._valkey_client is None: raise ClientNotConnectedError("ValkeySentinelClient is not connected") @@ -559,11 +571,13 @@ def __init__( self._operation_failure_count = 0 self._closed = False + @override async def connect(self) -> None: await self._operation_client.connect() await self._monitor_client.connect() self._monitor_task = asyncio.create_task(self._monitor_connection()) + @override async def disconnect(self) -> None: self._closed = True if self._monitor_task: @@ -573,6 +587,7 @@ async def disconnect(self) -> None: await self._monitor_client.disconnect() await self._operation_client.disconnect() + @override async def ping(self) -> None: """ Ping the server to check if the connection is alive. @@ -584,9 +599,11 @@ async def ping(self) -> None: """ await self._monitor_client.ping() + @override async def need_reconnect(self) -> bool: return await self._monitor_client.need_reconnect() + @override async def check_health(self) -> None: """ Run the unified health check via the monitor client. @@ -599,6 +616,7 @@ async def check_health(self) -> None: await self._monitor_client.check_health() @asynccontextmanager + @override async def client(self) -> AsyncIterator[GlideClient]: """ Acquire the operation client for use, tracking consecutive failures. diff --git a/src/ai/backend/common/clients/valkey_client/valkey_tus/exceptions.py b/src/ai/backend/common/clients/valkey_client/valkey_tus/exceptions.py index 68ac41a5f7d..c45699b28f3 100644 --- a/src/ai/backend/common/clients/valkey_client/valkey_tus/exceptions.py +++ b/src/ai/backend/common/clients/valkey_client/valkey_tus/exceptions.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -17,6 +19,7 @@ class TusLeaseHeldError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/storage/tus-lease-held" error_title = "TUS session lease is held by another storage-proxy" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -31,6 +34,7 @@ class TusLeaseLostError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/storage/tus-lease-lost" error_title = "TUS session lease lost mid-write" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -45,6 +49,7 @@ class TusSessionNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/tus-session-not-found" error_title = "TUS upload session not found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, diff --git a/src/ai/backend/common/configs/loader/config_overrider.py b/src/ai/backend/common/configs/loader/config_overrider.py index 7febdab20cf..9c5afe3ad12 100644 --- a/src/ai/backend/common/configs/loader/config_overrider.py +++ b/src/ai/backend/common/configs/loader/config_overrider.py @@ -27,5 +27,6 @@ def set_key(path: Sequence[str], value: Any) -> None: return cfg @property + @override def source_name(self) -> str: return "overrides" diff --git a/src/ai/backend/common/configs/loader/env_loader.py b/src/ai/backend/common/configs/loader/env_loader.py index ac6f85bac95..a8fdaffba11 100644 --- a/src/ai/backend/common/configs/loader/env_loader.py +++ b/src/ai/backend/common/configs/loader/env_loader.py @@ -28,5 +28,6 @@ def set_key(path: Sequence[str], value: Any) -> None: return cfg @property + @override def source_name(self) -> str: return "env" diff --git a/src/ai/backend/common/configs/loader/etcd_loader.py b/src/ai/backend/common/configs/loader/etcd_loader.py index 7ec1df5fa53..5c482f43edf 100644 --- a/src/ai/backend/common/configs/loader/etcd_loader.py +++ b/src/ai/backend/common/configs/loader/etcd_loader.py @@ -18,5 +18,6 @@ async def load(self) -> Mapping[str, Any]: return await self._etcd.get_prefix(self._prefix) @property + @override def source_name(self) -> str: return f"etcd:{self._prefix}" diff --git a/src/ai/backend/common/configs/loader/etcd_watcher.py b/src/ai/backend/common/configs/loader/etcd_watcher.py index c5f0aaf6c10..ac5e4c37b96 100644 --- a/src/ai/backend/common/configs/loader/etcd_watcher.py +++ b/src/ai/backend/common/configs/loader/etcd_watcher.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, override from aiotools import aclosing from etcd_client import WatchEventType @@ -16,6 +16,7 @@ def __init__(self, etcd: AsyncEtcd, config_prefix: str = "ai/backend/config") -> self._etcd = etcd self._config_prefix = config_prefix + @override async def watch(self) -> Any: async with aclosing(self._etcd.watch_prefix(self._config_prefix)) as agen: async for ev in agen: diff --git a/src/ai/backend/common/configs/loader/loader_chain.py b/src/ai/backend/common/configs/loader/loader_chain.py index 756af398f50..63daf405846 100644 --- a/src/ai/backend/common/configs/loader/loader_chain.py +++ b/src/ai/backend/common/configs/loader/loader_chain.py @@ -39,6 +39,7 @@ async def load(self) -> Mapping[str, Any]: return merged @property + @override def source_name(self) -> str: names = [loader.source_name for loader in self._loaders] return f"chain:[{', '.join(names)}]" diff --git a/src/ai/backend/common/configs/loader/toml_loader.py b/src/ai/backend/common/configs/loader/toml_loader.py index 3fe1be70671..cea2647a75e 100644 --- a/src/ai/backend/common/configs/loader/toml_loader.py +++ b/src/ai/backend/common/configs/loader/toml_loader.py @@ -20,5 +20,6 @@ async def load(self) -> Mapping[str, Any]: return raw_cfg @property + @override def source_name(self) -> str: return f"toml:{self._config_path}" diff --git a/src/ai/backend/common/cron/local_cron.py b/src/ai/backend/common/cron/local_cron.py index fdefaebaf50..bb7d9da0ae6 100644 --- a/src/ai/backend/common/cron/local_cron.py +++ b/src/ai/backend/common/cron/local_cron.py @@ -4,7 +4,7 @@ import asyncio import logging -from typing import Final +from typing import Final, override from ai.backend.common.cron.base import Cron, PeriodicTask from ai.backend.logging import BraceStyleAdapter @@ -54,6 +54,7 @@ async def _run_task(self, task: PeriodicTask) -> None: log.debug("Task {} cancelled", task.name) raise + @override async def start(self) -> None: """Start the local cron.""" self._stopped = False @@ -62,6 +63,7 @@ async def start(self) -> None: runner.set_name(f"local-cron-task-{task.name}") self._task_runners.append(runner) + @override async def stop(self) -> None: """Stop the local cron and cancel all of its task runners.""" self._stopped = True diff --git a/src/ai/backend/common/data/endpoint/types.py b/src/ai/backend/common/data/endpoint/types.py index cf500b14fd3..ad7da261d9a 100644 --- a/src/ai/backend/common/data/endpoint/types.py +++ b/src/ai/backend/common/data/endpoint/types.py @@ -1,6 +1,6 @@ from enum import Enum, StrEnum from functools import lru_cache -from typing import Any, Self +from typing import Any, Self, override class EndpointStatus(StrEnum): @@ -61,6 +61,7 @@ def inactive_states(cls) -> set["EndpointLifecycle"]: return {cls.DESTROYING, cls.DESTROYED} @classmethod + @override def _missing_(cls, value: Any) -> Self | None: # Accept v2 :class:`ModelDeploymentStatus` aliases on the wire so # historical / future callers that hand us the new naming still diff --git a/src/ai/backend/common/data/model_deployment/types.py b/src/ai/backend/common/data/model_deployment/types.py index 11f5b1525b9..b6c0560d690 100644 --- a/src/ai/backend/common/data/model_deployment/types.py +++ b/src/ai/backend/common/data/model_deployment/types.py @@ -1,4 +1,4 @@ -from typing import Any, Self +from typing import Any, Self, override from ai.backend.common.types import CIStrEnum @@ -41,6 +41,7 @@ class ModelDeploymentStatus(CIStrEnum): STOPPED = "STOPPED" @classmethod + @override def _missing_(cls, value: Any) -> Self | None: if isinstance(value, str): alias = _LIFECYCLE_TO_DEPLOYMENT_STATUS_ALIASES.get(value.lower()) diff --git a/src/ai/backend/common/data/notification/messages.py b/src/ai/backend/common/data/notification/messages.py index 9972e2e7a2a..70b8819c261 100644 --- a/src/ai/backend/common/data/notification/messages.py +++ b/src/ai/backend/common/data/notification/messages.py @@ -4,7 +4,7 @@ from abc import abstractmethod from collections.abc import Mapping -from typing import Any +from typing import Any, override from pydantic import Field @@ -33,6 +33,7 @@ class NotifiableMessage(BackendAISchema): model_config = {"extra": "forbid"} # Strict validation - reject unknown fields + @override def __init_subclass__(cls) -> None: """Automatically register subclasses by their rule type.""" try: @@ -96,6 +97,7 @@ class SessionStartedMessage(NotifiableMessage): """ @classmethod + @override def rule_type(cls) -> NotificationRuleType: """Return the notification rule type for this message class.""" return NotificationRuleType.SESSION_STARTED @@ -121,6 +123,7 @@ class SessionTerminatedMessage(NotifiableMessage): """ @classmethod + @override def rule_type(cls) -> NotificationRuleType: """Return the notification rule type for this message class.""" return NotificationRuleType.SESSION_TERMINATED @@ -151,6 +154,7 @@ class ArtifactDownloadCompletedMessage(NotifiableMessage): """ @classmethod + @override def rule_type(cls) -> NotificationRuleType: """Return the notification rule type for this message class.""" return NotificationRuleType.ARTIFACT_DOWNLOAD_COMPLETED @@ -181,6 +185,7 @@ class EndpointLifecycleChangedMessage(NotifiableMessage): """ @classmethod + @override def rule_type(cls) -> NotificationRuleType: """Return the notification rule type for this message class.""" return NotificationRuleType.ENDPOINT_LIFECYCLE_CHANGED diff --git a/src/ai/backend/common/dependencies/base.py b/src/ai/backend/common/dependencies/base.py index 912ba74b47a..de8b04bb616 100644 --- a/src/ai/backend/common/dependencies/base.py +++ b/src/ai/backend/common/dependencies/base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from contextlib import AbstractAsyncContextManager -from typing import TYPE_CHECKING, TypeVar, final +from typing import TYPE_CHECKING, TypeVar, final, override if TYPE_CHECKING: from ai.backend.common.health_checker import ServiceHealthChecker @@ -88,10 +88,12 @@ class NonMonitorableDependencyProvider(DependencyProvider[SetupInputT, ResourceT """ @final + @override def gen_liveness_checker(self, resource: ResourceT) -> None: return None @final + @override def gen_readiness_checker(self, resource: ResourceT) -> None: return None diff --git a/src/ai/backend/common/dependencies/stacks/builder.py b/src/ai/backend/common/dependencies/stacks/builder.py index d29085364c9..b880c32c769 100644 --- a/src/ai/backend/common/dependencies/stacks/builder.py +++ b/src/ai/backend/common/dependencies/stacks/builder.py @@ -1,7 +1,7 @@ from __future__ import annotations from contextlib import AsyncExitStack -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from ai.backend.common.dependencies.base import DependencyStack @@ -35,6 +35,7 @@ def __init__(self) -> None: self._liveness_checkers = {} self._readiness_checkers = {} + @override async def enter_dependency( self, provider: DependencyProvider[SetupInputT, ResourceT], @@ -55,6 +56,7 @@ async def enter_dependency( return resource + @override async def enter_composer( self, composer: DependencyComposer[SetupInputT, ResourcesT], diff --git a/src/ai/backend/common/dependencies/stacks/visualizing.py b/src/ai/backend/common/dependencies/stacks/visualizing.py index c008a5ba804..249d1a0a22c 100644 --- a/src/ai/backend/common/dependencies/stacks/visualizing.py +++ b/src/ai/backend/common/dependencies/stacks/visualizing.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from datetime import UTC, datetime from enum import Enum -from typing import Any, TextIO +from typing import Any, TextIO, override from ai.backend.common.dependencies.base import ( DependencyComposer, @@ -129,6 +129,7 @@ def _record_event( self._events.append(event) self._print_event(event) + @override async def enter_dependency( self, provider: DependencyProvider[SetupInputT, ResourceT], @@ -147,6 +148,7 @@ async def enter_dependency( self._record_event(stage_name, DependencyStatus.FAILED, error=e) raise + @override async def enter_composer( self, composer: DependencyComposer[SetupInputT, ResourcesT], diff --git a/src/ai/backend/common/docker.py b/src/ai/backend/common/docker.py index 3b8672966a4..207c5b3cac0 100644 --- a/src/ai/backend/common/docker.py +++ b/src/ai/backend/common/docker.py @@ -19,6 +19,7 @@ NamedTuple, Self, cast, + override, ) import aiohttp @@ -384,9 +385,11 @@ def __init__(self, tags: Iterable[str], value: str | None = None) -> None: value = "" self._data[key] = value + @override def __repr__(self) -> str: return self.__str__() + @override def __str__(self) -> str: return f"PlatformTagSet({self._data!s})" @@ -396,15 +399,19 @@ def has(self, key: str, version: str | None = None) -> bool: _v = self._data.get(key, None) return _v == version + @override def __getitem__(self, key: str) -> str: return self._data[key] + @override def __iter__(self) -> Iterator[str]: return iter(self._data) + @override def __len__(self) -> int: return len(self._data) + @override def __eq__(self, other: Any) -> bool: if isinstance(other, (set, frozenset)): return set(self._data.keys()) == other @@ -429,6 +436,7 @@ def tag_set(self) -> tuple[str, PlatformTagSet]: tags = self.tag.split("-") return (tags[0], PlatformTagSet(tags[1:], self.project_and_image_name)) + @override def __str__(self) -> str: return self.canonical @@ -665,12 +673,15 @@ def short(self) -> str: join = functools.partial(join_non_empty, sep="/") return f"{join(self.project, self.name)}:{self.tag}" + @override def __str__(self) -> str: return self.canonical + @override def __repr__(self) -> str: return f'' + @override def __hash__(self) -> int: return hash((self.project, self.name, self.tag, self.registry, self.architecture)) diff --git a/src/ai/backend/common/dto/manager/context.py b/src/ai/backend/common/dto/manager/context.py index bd6283b39aa..a0dc61e7a16 100644 --- a/src/ai/backend/common/dto/manager/context.py +++ b/src/ai/backend/common/dto/manager/context.py @@ -1,6 +1,6 @@ import uuid from collections.abc import Mapping -from typing import Any, Self +from typing import Any, Self, override from aiohttp import web @@ -14,6 +14,7 @@ class UserIdentityCtx(MiddlewareParam): domain_name: str @classmethod + @override async def from_request(cls, request: web.Request) -> Self: return cls( user_uuid=request["user"]["uuid"], @@ -28,6 +29,7 @@ class KeypairCtx(MiddlewareParam): resource_policy: Mapping[str, Any] @classmethod + @override async def from_request(cls, request: web.Request) -> Self: return cls( access_key=request["keypair"]["access_key"], diff --git a/src/ai/backend/common/dto/manager/query.py b/src/ai/backend/common/dto/manager/query.py index eb0ab4f9f58..d04582b407c 100644 --- a/src/ai/backend/common/dto/manager/query.py +++ b/src/ai/backend/common/dto/manager/query.py @@ -3,7 +3,7 @@ import uuid from collections.abc import Callable from datetime import date, datetime -from typing import TypeVar +from typing import TypeVar, override from pydantic import Field @@ -173,6 +173,7 @@ class NullableDateTimeFilter(DateTimeFilter): description="Filter by null status: true = IS NULL, false = IS NOT NULL", ) + @override def build_query_condition( self, before_factory: Callable[[datetime], _QC], diff --git a/src/ai/backend/common/enum_extension.py b/src/ai/backend/common/enum_extension.py index 7e6a9a34d3e..23b41018b86 100644 --- a/src/ai/backend/common/enum_extension.py +++ b/src/ai/backend/common/enum_extension.py @@ -1,16 +1,18 @@ from __future__ import annotations import enum -from typing import Any +from typing import Any, override __all__ = ("StringSetFlag",) class StringSetFlag(enum.StrEnum): + @override def __eq__(self, other: Any) -> bool: result: bool = self.value == other return result + @override def __hash__(self) -> int: return hash(self.value) @@ -51,5 +53,6 @@ def __rxor__(self, other: Any) -> set[str] | str: return other raise TypeError + @override def __str__(self) -> str: return self.value diff --git a/src/ai/backend/common/etcd.py b/src/ai/backend/common/etcd.py index 949add544ec..2757014d969 100644 --- a/src/ai/backend/common/etcd.py +++ b/src/ai/backend/common/etcd.py @@ -30,6 +30,7 @@ Self, TypeVar, cast, + override, ) from urllib.parse import quote as _quote from urllib.parse import unquote @@ -402,6 +403,7 @@ def _merge_scope_prefix_map( """ return ChainMap(cast(MutableMapping[str, Any], override) or {}, self.scope_prefix_map) + @override async def put( self, key: str, @@ -426,6 +428,7 @@ async def put( mangled_key.encode(self.encoding), str(val).encode(self.encoding) ) + @override async def put_prefix( self, key: str, @@ -459,6 +462,7 @@ async def put_prefix( async with self.etcd.connect() as communicator: await communicator.txn(EtcdTransactionAction().and_then(actions).or_else([])) + @override async def put_dict( self, flattened_dict_obj: Mapping[str, str], @@ -538,6 +542,7 @@ async def atomic_replace_prefixes( return await communicator.txn(EtcdTransactionAction().and_then(actions).or_else([])) + @override async def get( self, key: str, @@ -584,6 +589,7 @@ async def get( return bytes(value).decode(self.encoding) return None + @override async def get_prefix( self, key_prefix: str, @@ -668,6 +674,7 @@ async def get_prefix( # for legacy get_prefix_dict = get_prefix + @override async def replace( self, key: str, @@ -698,6 +705,7 @@ async def replace( return cast(bool, result.succeeded()) + @override async def delete( self, key: str, @@ -710,6 +718,7 @@ async def delete( async with self.etcd.connect() as communicator: await communicator.delete(mangled_key.encode(self.encoding)) + @override async def delete_multi( self, keys: Iterable[str], @@ -728,6 +737,7 @@ async def delete_multi( ) await communicator.txn(EtcdTransactionAction().and_then(actions).or_else([])) + @override async def delete_prefix( self, key_prefix: str, @@ -769,6 +779,7 @@ async def _watch_impl( if cleanup_event: await cleanup_event.notify_waiters() + @override async def watch( self, key: str, @@ -826,6 +837,7 @@ async def watch( else: raise + @override async def watch_prefix( self, key_prefix: str, diff --git a/src/ai/backend/common/events/event_types/artifact/anycast.py b/src/ai/backend/common/events/event_types/artifact/anycast.py index c8a433207a6..25529326196 100644 --- a/src/ai/backend/common/events/event_types/artifact/anycast.py +++ b/src/ai/backend/common/events/event_types/artifact/anycast.py @@ -49,10 +49,12 @@ class ModelVerifyingEvent(BaseArtifactEvent): def event_name(cls) -> str: return "model_verifying" + @override def serialize(self) -> tuple[Any, ...]: return (self.model_id, self.revision, self.registry_type, self.registry_name) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( model_id=value[0], @@ -85,6 +87,7 @@ class ModelImportDoneEvent(BaseArtifactEvent): def event_name(cls) -> str: return "model_import_done" + @override def serialize(self) -> tuple[Any, ...]: verification_result = None if self.verification_result is not None: @@ -101,6 +104,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: verification_result = None if value[6] is not None: @@ -134,6 +138,7 @@ class ModelMetadataFetchDoneEvent(BaseArtifactEvent): def event_name(cls) -> str: return "models_metadata_fetch_done" + @override def serialize(self) -> tuple[Any, ...]: return ( self.model.model_id, @@ -145,6 +150,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( model=ModelMetadataInfo( diff --git a/src/ai/backend/common/events/event_types/bgtask/broadcast.py b/src/ai/backend/common/events/event_types/bgtask/broadcast.py index eabaaadba55..0ebeb9e6c3f 100644 --- a/src/ai/backend/common/events/event_types/bgtask/broadcast.py +++ b/src/ai/backend/common/events/event_types/bgtask/broadcast.py @@ -45,6 +45,7 @@ class BgtaskUpdatedEvent(BaseBgtaskEvent): total_progress: float message: str | None = None + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.task_id), @@ -54,6 +55,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( uuid.UUID(value[0]), diff --git a/src/ai/backend/common/events/event_types/kernel/anycast.py b/src/ai/backend/common/events/event_types/kernel/anycast.py index b439fccb1fd..dd6ffcb5758 100644 --- a/src/ai/backend/common/events/event_types/kernel/anycast.py +++ b/src/ai/backend/common/events/event_types/kernel/anycast.py @@ -171,6 +171,7 @@ def event_name(cls) -> str: class DoSyncKernelLogsEvent(BaseKernelEvent): container_id: str + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.kernel_id), @@ -178,6 +179,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( KernelId(uuid.UUID(value[0])), @@ -189,5 +191,6 @@ def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "do_sync_kernel_logs" diff --git a/src/ai/backend/common/events/event_types/log/anycast.py b/src/ai/backend/common/events/event_types/log/anycast.py index 6a324594122..6210c60a445 100644 --- a/src/ai/backend/common/events/event_types/log/anycast.py +++ b/src/ai/backend/common/events/event_types/log/anycast.py @@ -15,19 +15,24 @@ def serialize(self) -> tuple[Any, ...]: return tuple() @classmethod - def deserialize(cls, value: tuple[Any, ...]) -> Self: # noqa: ARG003 + @override + def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls() @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.LOG + @override def domain_id(self) -> str | None: return None + @override def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "do_log_cleanup" diff --git a/src/ai/backend/common/events/event_types/notification/anycast.py b/src/ai/backend/common/events/event_types/notification/anycast.py index f7f9fd29311..23374adb626 100644 --- a/src/ai/backend/common/events/event_types/notification/anycast.py +++ b/src/ai/backend/common/events/event_types/notification/anycast.py @@ -55,6 +55,7 @@ def domain_id(self) -> str | None: def user_event(self) -> UserEvent | None: return None + @override def serialize(self) -> tuple[bytes, ...]: return ( self.rule_type.encode(), @@ -63,6 +64,7 @@ def serialize(self) -> tuple[bytes, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[bytes, ...]) -> NotificationTriggeredEvent: return cls( rule_type=value[0].decode(), diff --git a/src/ai/backend/common/events/event_types/session/anycast.py b/src/ai/backend/common/events/event_types/session/anycast.py index 735da5a9581..60782595134 100644 --- a/src/ai/backend/common/events/event_types/session/anycast.py +++ b/src/ai/backend/common/events/event_types/session/anycast.py @@ -57,6 +57,7 @@ def user_event(self) -> UserEvent | None: class DoTerminateSessionEvent(BaseSessionEvent): reason: KernelLifecycleEventReason + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.session_id), @@ -181,6 +182,7 @@ class SessionResultEvent(BaseSessionEvent): reason: KernelLifecycleEventReason = KernelLifecycleEventReason.UNKNOWN exit_code: int = -1 + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.session_id), @@ -189,6 +191,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( SessionId(uuid.UUID(value[0])), @@ -248,6 +251,7 @@ def event_name(cls) -> str: return "execution_started" @classmethod + @override def execution_status(cls) -> SessionExecutionStatus: return SessionExecutionStatus.STARTED @@ -260,6 +264,7 @@ def event_name(cls) -> str: return "execution_finished" @classmethod + @override def execution_status(cls) -> SessionExecutionStatus: return SessionExecutionStatus.FINISHED @@ -272,6 +277,7 @@ def event_name(cls) -> str: return "execution_timeout" @classmethod + @override def execution_status(cls) -> SessionExecutionStatus: return SessionExecutionStatus.TIMEOUT @@ -284,5 +290,6 @@ def event_name(cls) -> str: return "execution_cancelled" @classmethod + @override def execution_status(cls) -> SessionExecutionStatus: return SessionExecutionStatus.CANCELED diff --git a/src/ai/backend/common/events/event_types/session/broadcast.py b/src/ai/backend/common/events/event_types/session/broadcast.py index 72e91584734..50a7d61e9b4 100644 --- a/src/ai/backend/common/events/event_types/session/broadcast.py +++ b/src/ai/backend/common/events/event_types/session/broadcast.py @@ -36,6 +36,7 @@ def user_event(self) -> UserEvent | None: class DoTerminateSessionEvent(BaseSessionEvent): reason: KernelLifecycleEventReason + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.session_id), @@ -152,6 +153,7 @@ class SessionResultEvent(BaseSessionEvent): reason: KernelLifecycleEventReason = KernelLifecycleEventReason.UNKNOWN exit_code: int = -1 + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.session_id), @@ -160,6 +162,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( SessionId(uuid.UUID(value[0])), diff --git a/src/ai/backend/common/events/event_types/vfolder/anycast.py b/src/ai/backend/common/events/event_types/vfolder/anycast.py index 02fd1d60a53..8bdcdf56da3 100644 --- a/src/ai/backend/common/events/event_types/vfolder/anycast.py +++ b/src/ai/backend/common/events/event_types/vfolder/anycast.py @@ -26,16 +26,19 @@ def user_event(self) -> UserEvent | None: @dataclass class VFolderDeletionSuccessEvent(VFolderEvent): + @override def serialize(self) -> tuple[Any, ...]: return (str(self.vfid),) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( VFolderID.from_str(value[0]), ) @classmethod + @override def event_name(cls) -> str: return "vfolder_deletion_success" @@ -44,6 +47,7 @@ def event_name(cls) -> str: class VFolderDeletionFailureEvent(VFolderEvent): message: str + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.vfid), @@ -51,6 +55,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( VFolderID.from_str(value[0]), @@ -58,6 +63,7 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: ) @classmethod + @override def event_name(cls) -> str: return "vfolder_deletion_failure" @@ -66,6 +72,7 @@ def event_name(cls) -> str: class VFolderCloneSuccessEvent(VFolderEvent): dst_vfid: VFolderID + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.vfid), @@ -73,6 +80,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( VFolderID.from_str(value[0]), @@ -80,6 +88,7 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: ) @classmethod + @override def event_name(cls) -> str: return "vfolder_clone_success" @@ -89,6 +98,7 @@ class VFolderCloneFailureEvent(VFolderEvent): dst_vfid: VFolderID message: str + @override def serialize(self) -> tuple[Any, ...]: return ( str(self.vfid), @@ -97,6 +107,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( VFolderID.from_str(value[0]), @@ -105,5 +116,6 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: ) @classmethod + @override def event_name(cls) -> str: return "vfolder_clone_failure" diff --git a/src/ai/backend/common/events/event_types/volume/broadcast.py b/src/ai/backend/common/events/event_types/volume/broadcast.py index 5595db47f1e..2463b340fe5 100644 --- a/src/ai/backend/common/events/event_types/volume/broadcast.py +++ b/src/ai/backend/common/events/event_types/volume/broadcast.py @@ -41,6 +41,7 @@ class DoVolumeMountEvent(BaseVolumeEvent): edit_fstab: bool = False fstab_path: str = "/etc/fstab" + @override def serialize(self) -> tuple[Any, ...]: return ( self.dir_name, @@ -55,6 +56,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( dir_name=value[0], @@ -69,6 +71,7 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: ) @classmethod + @override def event_name(cls) -> str: return "do_volume_mount" @@ -87,6 +90,7 @@ class DoVolumeUnmountEvent(BaseVolumeEvent): edit_fstab: bool = False fstab_path: str | None = None + @override def serialize(self) -> tuple[Any, ...]: return ( self.dir_name, @@ -98,6 +102,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( dir_name=value[0], @@ -109,6 +114,7 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: ) @classmethod + @override def event_name(cls) -> str: return "do_volume_unmount" @@ -121,6 +127,7 @@ class BaseAgentVolumeMountEvent(BaseVolumeEvent): quota_scope_id: QuotaScopeID err_msg: str | None = None + @override def serialize(self) -> tuple[Any, ...]: return ( self.node_id, @@ -131,6 +138,7 @@ def serialize(self) -> tuple[Any, ...]: ) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> Self: return cls( value[0], @@ -143,11 +151,13 @@ def deserialize(cls, value: tuple[Any, ...]) -> Self: class VolumeMounted(BaseAgentVolumeMountEvent): @classmethod + @override def event_name(cls) -> str: return "volume_mounted" class VolumeUnmounted(BaseAgentVolumeMountEvent): @classmethod + @override def event_name(cls) -> str: return "volume_unmounted" diff --git a/src/ai/backend/common/events/hub/propagators/bypass.py b/src/ai/backend/common/events/hub/propagators/bypass.py index 86f877aa2e2..50cc693243d 100644 --- a/src/ai/backend/common/events/hub/propagators/bypass.py +++ b/src/ai/backend/common/events/hub/propagators/bypass.py @@ -2,6 +2,7 @@ import logging import uuid from collections.abc import AsyncIterator +from typing import override from ai.backend.common.events.hub.hub import EventPropagator from ai.backend.common.events.types import AbstractEvent @@ -24,6 +25,7 @@ def __init__(self) -> None: self._queue = asyncio.Queue() self._closed = False + @override def id(self) -> uuid.UUID: """ Get the unique identifier for the propagator. @@ -44,9 +46,11 @@ async def receive(self) -> AsyncIterator[AbstractEvent]: except Exception as e: log.error("Error propagating event: {}", e) + @override async def propagate_event(self, event: AbstractEvent) -> None: await self._queue.put(event) + @override async def close(self) -> None: if self._closed: return diff --git a/src/ai/backend/common/events/hub/propagators/cache.py b/src/ai/backend/common/events/hub/propagators/cache.py index e89c451545c..e862ec743e6 100644 --- a/src/ai/backend/common/events/hub/propagators/cache.py +++ b/src/ai/backend/common/events/hub/propagators/cache.py @@ -2,6 +2,7 @@ import logging import uuid from collections.abc import AsyncIterator +from typing import override from ai.backend.common.events.fetcher import EventFetcher from ai.backend.common.events.hub.hub import EventPropagator @@ -30,6 +31,7 @@ def __init__(self, event_fetcher: EventFetcher) -> None: self._queue = asyncio.Queue() self._closed = False + @override def id(self) -> uuid.UUID: """ Get the unique identifier for the propagator. @@ -74,9 +76,11 @@ async def receive(self, cache_id: str) -> AsyncIterator[AbstractEvent]: except Exception as e: log.error("Error propagating event: {}", e) + @override async def propagate_event(self, event: AbstractEvent) -> None: await self._queue.put(event) + @override async def close(self) -> None: if self._closed: return diff --git a/src/ai/backend/common/events/types.py b/src/ai/backend/common/events/types.py index f98b9df7c48..c3319af15d6 100644 --- a/src/ai/backend/common/events/types.py +++ b/src/ai/backend/common/events/types.py @@ -132,6 +132,7 @@ class AbstractBroadcastEvent(AbstractEvent): _register_dict: dict[str, type["AbstractBroadcastEvent"]] = {} + @override def __init_subclass__(cls) -> None: try: name = cls.event_name() diff --git a/src/ai/backend/common/exception.py b/src/ai/backend/common/exception.py index 1d97300c0dc..45e3100a0b5 100644 --- a/src/ai/backend/common/exception.py +++ b/src/ai/backend/common/exception.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Self +from typing import Any, Self, override from aiohttp import web from pydantic_core import ErrorDetails @@ -30,6 +30,7 @@ class UnknownImageReference(ValueError): The first argument of this exception should be the reference given by the user. """ + @override def __str__(self) -> str: return f"Unknown image reference: {self.args[0]}" @@ -40,6 +41,7 @@ class ImageNotAvailable(ValueError): The first argument of this exception should be the reference given by the user. """ + @override def __str__(self) -> str: return f"Unavailable image in the agent: {self.args[0]}" @@ -50,6 +52,7 @@ class UnknownImageRegistry(ValueError): The first argument of this exception should be the registry given by the user. """ + @override def __str__(self) -> str: return f"Unknown image registry: {self.args[0]}" @@ -59,6 +62,7 @@ class InvalidImageName(ValueError): Represents an invalid string for image name. """ + @override def __str__(self) -> str: return f"Invalid image name: {self.args[0]}" @@ -78,6 +82,7 @@ def __init__(self, tag: str, full_name: str | None = None) -> None: self._tag = tag self._full_name = full_name + @override def __str__(self) -> str: if self._full_name is not None: return f"Invalid or duplicate image name tag: {self._tag}, full image name: {self._full_name}" @@ -94,6 +99,7 @@ def __init__(self, project: str, canonical: str) -> None: self._project = project self._canonical = canonical + @override def __str__(self) -> str: return f'Project "{self._project}" mismatch with the image canonical: {self._canonical}' @@ -104,6 +110,7 @@ class AliasResolutionFailed(ValueError): The first argument of this exception should be the alias given by the user. """ + @override def __str__(self) -> str: return f"Failed to resolve alias: {self.args[0]}" @@ -309,6 +316,7 @@ def default(cls) -> Self: error_detail=ErrorDetail.INTERNAL_ERROR, ) + @override def __str__(self) -> str: return f"{self.domain}_{self.operation}_{self.error_detail}" @@ -366,6 +374,7 @@ def __init__( self.body_dict = body self.body = dump_json(body) + @override def __str__(self) -> str: lines = [] if self.extra_msg: @@ -376,6 +385,7 @@ def __str__(self) -> str: lines.append(" -> extra_data: " + repr(self.extra_data)) return "\n".join(lines) + @override def __repr__(self) -> str: lines = [] if self.extra_msg: @@ -388,6 +398,7 @@ def __repr__(self) -> str: lines.append(" -> extra_data: " + repr(self.extra_data)) return "\n".join(lines) + @override def __reduce__(self) -> tuple[type[BackendAIError], tuple[Any, ...], dict[str, Any]]: return ( type(self), @@ -412,6 +423,7 @@ class InvalidErrorCode(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/invalid-error-code" error_title = "Invalid error code in the raised exception." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -424,6 +436,7 @@ class MalformedRequestBody(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Malformed request body." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -436,6 +449,7 @@ class InvalidAPIParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Invalid or Missing API Parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -454,6 +468,7 @@ class BackendAISchemaValidationFailed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/schema-validation-failed" error_title = "Schema validation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -481,6 +496,7 @@ class ModelDefinitionValidationError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/model-definition-validation-failed" error_title = "Model definition validation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -493,6 +509,7 @@ class DeprecatedAPI(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/deprecated" error_title = "This API is deprecated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -505,6 +522,7 @@ class InvalidResourceSlotQuantity(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-resource-slot-quantity" error_title = "Invalid resource slot quantity." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -517,6 +535,7 @@ class UnknownResourceSlotType(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/unknown-resource-slot-type" error_title = "Unknown resource slot type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -529,6 +548,7 @@ class ResourcePresetConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-resource" error_title = "Duplicate Resource Preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RESOURCE_PRESET, @@ -541,6 +561,7 @@ class MiddlewareParamParsingFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/internal-server-error" error_title = "Middleware parameter parsing failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -553,6 +574,7 @@ class ParameterNotParsedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/internal-server-error" error_title = "Parameter Not Parsed Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -565,6 +587,7 @@ class BgtaskNotRegisteredError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/bgtask-not-registered" error_title = "Background Task Not Registered" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BGTASK, @@ -577,6 +600,7 @@ class BgtaskNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/bgtask-not-found" error_title = "Background Task Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BGTASK, @@ -589,6 +613,7 @@ class BgtaskFailedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/bgtask-failed" error_title = "Background Task Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BGTASK, @@ -601,6 +626,7 @@ class BgtaskCancelledError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/bgtask-cancelled" error_title = "Background Task Cancelled" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BGTASK, @@ -613,6 +639,7 @@ class UnreachableError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/unreachable" error_title = "Unreachable" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -625,6 +652,7 @@ class PermissionDeniedError(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/permission-denied" error_title = "Permission Denied." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -637,6 +665,7 @@ class SessionWithInvalidStateError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/session-invalid-state" error_title = "Session with Invalid State" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -649,6 +678,7 @@ class StorageNamespaceNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage-namespace-not-found" error_title = "Storage Namespace Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_NAMESPACE, @@ -661,6 +691,7 @@ class InvalidCursorTypeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-cursor-type" error_title = "Invalid Cursor Type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -673,6 +704,7 @@ class RelationNotLoadedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/relation-not-loaded" error_title = "Relation Not Loaded" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -685,6 +717,7 @@ class ArtifactDefaultRevisionResolveError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/artifact-revision-resolve-failed" error_title = "Cannot Resolve Artifact Default Revision" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -700,6 +733,7 @@ class RuntimeVariantNotSupportedError(BackendAIError, web.HTTPBadRequest): def __init__(self, runtime_variant: str) -> None: super().__init__(extra_msg=f"Runtime variant '{runtime_variant}' is not supported.") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -712,6 +746,7 @@ class GenericNotImplementedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/not-implemented" error_title = "Not Implemented" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -724,6 +759,7 @@ class ObjectStorageBucketNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/object-storage-bucket-not-found" error_title = "Object Storage Bucket Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -736,6 +772,7 @@ class InvalidConfigError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/invalid-configuration" error_title = "Invalid Configuration" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -748,6 +785,7 @@ class ProcessorNotReadyError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/processor-not-ready" error_title = "Processor Not Ready" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -760,6 +798,7 @@ class AgentNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/agent-not-found" error_title = "Agent Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -772,6 +811,7 @@ class ScalingGroupNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/scaling-group-not-found" error_title = "Scaling Group Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -784,6 +824,7 @@ class ScalingGroupConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-scaling-group" error_title = "Duplicate Scaling Group" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -796,6 +837,7 @@ class VFolderNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/vfolder-not-found" error_title = "Virtual Folder Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -808,6 +850,7 @@ class UserNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/user-not-found" error_title = "User Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -820,6 +863,7 @@ class GroupNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/group-not-found" error_title = "Project Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -832,6 +876,7 @@ class DomainNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/domain-not-found" error_title = "Domain Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -844,6 +889,7 @@ class ModelDeploymentNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/model-deployment-not-found" error_title = "Model Deployment Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -856,6 +902,7 @@ class ModelDeploymentUnavailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/model-deployment-unavailable" error_title = "Model Deployment Unavailable" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -868,6 +915,7 @@ class ModelRevisionNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/model-revision-not-found" error_title = "Model Revision Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -936,6 +984,7 @@ def from_http_status( error_message=error_message, ) + @override def error_code(self) -> ErrorCode: return self._error_code @@ -948,6 +997,7 @@ class ClientNotConnectedError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/client-not-connected" error_title = "Client Not Connected" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -964,6 +1014,7 @@ class ValkeySentinelMasterNotFound(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/valkey-sentinel-master-not-found" error_title = "Valkey Sentinel Master Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -981,6 +1032,7 @@ class ValkeyRoleMismatchError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/valkey-role-mismatch" error_title = "Valkey Role Mismatch" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -997,6 +1049,7 @@ class DatabaseError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/database-error" error_title = "Database Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -1009,6 +1062,7 @@ class UserResourcePolicyNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/user-resource-policy-not-found" error_title = "User Resource Policy Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER_RESOURCE_POLICY, @@ -1021,6 +1075,7 @@ class KeypairResourcePolicyNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/keypair-resource-policy-not-found" error_title = "Keypair Resource Policy Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR_RESOURCE_POLICY, @@ -1037,6 +1092,7 @@ class ShowmountFailed(BaseNFSMountCheckFailed, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/showmount-failed" error_title = "showmount command failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, @@ -1049,6 +1105,7 @@ class ShowmountNotFound(BaseNFSMountCheckFailed, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/showmount-not-found" error_title = "showmount command not found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, @@ -1061,6 +1118,7 @@ class ExportPathNotFound(BaseNFSMountCheckFailed, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/nfs-export-path-not-found" error_title = "NFS export path not found on the server" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, @@ -1073,6 +1131,7 @@ class NFSTimeoutError(BaseNFSMountCheckFailed, web.HTTPRequestTimeout): error_type = "https://api.backend.ai/probs/nfs-timeout" error_title = "NFS server is not reachable (timeout)" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, @@ -1085,6 +1144,7 @@ class NFSUnexpectedError(BaseNFSMountCheckFailed, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/nfs-unexpected-error" error_title = "Unexpected NFS error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, @@ -1116,6 +1176,7 @@ def __init__( ) super().__init__(extra_msg=extra_msg) + @override def error_code(self) -> ErrorCode: return self._error_code @@ -1124,6 +1185,7 @@ class ContainerRegistryGroupsAlreadyAssociated(BackendAIError, web.HTTPConflict) error_type = "https://api.backend.ai/probs/container-registry/groups-already-associated" error_title = "Container registry groups already associated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -1136,6 +1198,7 @@ class InvalidNotificationChannelSpec(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-channel-spec" error_title = "Invalid Notification Channel Specification." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -1150,6 +1213,7 @@ class PrometheusConnectionError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/prometheus-connection-error" error_title = "Failed to connect to Prometheus." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.METRIC, @@ -1164,6 +1228,7 @@ class FailedToGetMetric(BackendAIError, web.HTTPBadGateway): error_type = "https://api.backend.ai/probs/failed-to-get-metric" error_title = "Failed to get metric." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.METRIC, @@ -1178,6 +1243,7 @@ class InvalidMetricPresetTemplate(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-metric-preset-template" error_title = "Invalid metric preset template." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.METRIC, @@ -1197,6 +1263,7 @@ class PrometheusQueryEvaluationFailed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/prometheus-query-evaluation-failed" error_title = "Prometheus rejected the query." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.METRIC, @@ -1217,6 +1284,7 @@ class JWTError(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/jwt-error" error_title = "JWT authentication error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -1236,6 +1304,7 @@ class JWTExpiredError(JWTError): error_type = "https://api.backend.ai/probs/jwt-expired" error_title = "JWT token has expired." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -1271,6 +1340,7 @@ class JWTInvalidClaimsError(JWTError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/jwt-invalid-claims" error_title = "JWT claims are invalid." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -1291,6 +1361,7 @@ class JWTDecodeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/jwt-decode-error" error_title = "Failed to decode JWT token." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -1311,6 +1382,7 @@ class JWTPayloadValidationError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/jwt-payload-validation-error" error_title = "JWT payload validation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -1325,6 +1397,7 @@ class RBACTypeConversionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/rbac-type-conversion-error" error_title = "RBAC type conversion failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, @@ -1337,6 +1410,7 @@ class PrometheusQueryPresetNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/prometheus-query-preset-not-found" error_title = "The prometheus query preset does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PROMETHEUS_QUERY_PRESET, @@ -1349,6 +1423,7 @@ class PrometheusQueryPresetInvalidLabel(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/prometheus-query-preset-invalid-label" error_title = "Invalid label specified for prometheus query preset execution." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PROMETHEUS_QUERY_PRESET, @@ -1361,6 +1436,7 @@ class PrometheusQueryPresetCategoryNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/prometheus-query-preset-category-not-found" error_title = "The prometheus query preset category does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PROMETHEUS_QUERY_PRESET_CATEGORY, @@ -1373,6 +1449,7 @@ class CloudDetectionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/cloud-detection-failed" error_title = "Cloud Provider Detection Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.EXTERNAL_SYSTEM, diff --git a/src/ai/backend/common/health_checker/checkers/etcd.py b/src/ai/backend/common/health_checker/checkers/etcd.py index def84dcc2f1..0d1d553e379 100644 --- a/src/ai/backend/common/health_checker/checkers/etcd.py +++ b/src/ai/backend/common/health_checker/checkers/etcd.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.health_checker.abc import StaticServiceHealthChecker @@ -35,10 +36,12 @@ def __init__(self, etcd: AsyncEtcd, timeout: float = 5.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return ETCD + @override async def check_service(self) -> ServiceHealth: """ Check etcd connection health by pinging the server. @@ -65,6 +68,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={CID_ETCD: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/common/health_checker/checkers/http.py b/src/ai/backend/common/health_checker/checkers/http.py index 47e6a276e99..67088a0adc6 100644 --- a/src/ai/backend/common/health_checker/checkers/http.py +++ b/src/ai/backend/common/health_checker/checkers/http.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from datetime import UTC, datetime from http import HTTPMethod +from typing import override import aiohttp @@ -60,10 +61,12 @@ def __init__( self._method = method @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return API + @override async def check_service(self) -> ServiceHealth: """ Check HTTP endpoint health by performing an HTTP request. @@ -110,6 +113,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={self._component_id: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/common/health_checker/checkers/prometheus.py b/src/ai/backend/common/health_checker/checkers/prometheus.py index 992fefc6ef2..342ef3d6071 100644 --- a/src/ai/backend/common/health_checker/checkers/prometheus.py +++ b/src/ai/backend/common/health_checker/checkers/prometheus.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.common.health_checker.abc import StaticServiceHealthChecker from ai.backend.common.health_checker.types import ( @@ -25,9 +26,11 @@ def __init__(self, timeout: float = 1.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: return PROMETHEUS + @override async def check_service(self) -> ServiceHealth: from ai.backend.common.metrics.safe import is_metrics_disabled, metrics_trip_info @@ -44,5 +47,6 @@ async def check_service(self) -> ServiceHealth: ) @property + @override def timeout(self) -> float: return self._timeout diff --git a/src/ai/backend/common/health_checker/checkers/valkey.py b/src/ai/backend/common/health_checker/checkers/valkey.py index 7d3a6869df5..739db788e11 100644 --- a/src/ai/backend/common/health_checker/checkers/valkey.py +++ b/src/ai/backend/common/health_checker/checkers/valkey.py @@ -2,7 +2,7 @@ import asyncio from datetime import UTC, datetime -from typing import Protocol, runtime_checkable +from typing import Protocol, override, runtime_checkable from ai.backend.common.health_checker.abc import StaticServiceHealthChecker from ai.backend.common.health_checker.types import ( @@ -55,6 +55,7 @@ def __init__( self._component_timeout = component_timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return REDIS @@ -99,6 +100,7 @@ async def _check_component( return (component_id, status) + @override async def check_service(self) -> ServiceHealth: """ Check health of all Valkey/Redis clients concurrently. @@ -135,6 +137,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results=results) @property + @override def timeout(self) -> float: """The timeout for the entire health check in seconds.""" return self._timeout diff --git a/src/ai/backend/common/health_checker/exceptions.py b/src/ai/backend/common/health_checker/exceptions.py index 5917f19e8aa..3ee62be5e37 100644 --- a/src/ai/backend/common/health_checker/exceptions.py +++ b/src/ai/backend/common/health_checker/exceptions.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -28,6 +30,7 @@ class HttpHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/http-health-check-failed" error_title = "HTTP health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -42,6 +45,7 @@ class ValkeyHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/valkey-health-check-failed" error_title = "Valkey health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -56,6 +60,7 @@ class EtcdHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/etcd-health-check-failed" error_title = "Etcd health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -70,6 +75,7 @@ class DatabaseHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/database-health-check-failed" error_title = "Database health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -84,6 +90,7 @@ class DockerHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/docker-health-check-failed" error_title = "Docker health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -100,6 +107,7 @@ class HealthCheckerAlreadyRegistered(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/health-checker-already-registered" error_title = "Health checker already registered." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -116,6 +124,7 @@ class HealthCheckerNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/health-checker-not-found" error_title = "Health checker not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, diff --git a/src/ai/backend/common/json.py b/src/ai/backend/common/json.py index 5846a698c44..a0266565ec5 100644 --- a/src/ai/backend/common/json.py +++ b/src/ai/backend/common/json.py @@ -1,12 +1,13 @@ import datetime import json import uuid -from typing import Any, Protocol +from typing import Any, Protocol, override import orjson class ExtendedJSONEncoder(json.JSONEncoder): + @override def default(self, o: Any) -> Any: if isinstance(o, uuid.UUID): return str(o) diff --git a/src/ai/backend/common/jwt/exceptions.py b/src/ai/backend/common/jwt/exceptions.py index e81082115b0..e0609fbd96e 100644 --- a/src/ai/backend/common/jwt/exceptions.py +++ b/src/ai/backend/common/jwt/exceptions.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -20,6 +22,7 @@ class JWTError(BackendAIError): All JWT-specific exceptions inherit from this base class. """ + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -38,6 +41,7 @@ class JWTExpiredError(JWTError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/jwt-expired" error_title = "JWT token has expired." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, diff --git a/src/ai/backend/common/leader/exceptions.py b/src/ai/backend/common/leader/exceptions.py index 4f977d3bc2c..8d90187f884 100644 --- a/src/ai/backend/common/leader/exceptions.py +++ b/src/ai/backend/common/leader/exceptions.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -17,6 +19,7 @@ class LeaderElectionError(BackendAIError): error_type = "https://api.backend.ai/probs/leader-election-error" error_title = "Leader Election Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.LEADER_ELECTION, @@ -31,6 +34,7 @@ class AlreadyStartedError(LeaderElectionError): error_type = "https://api.backend.ai/probs/leader-already-started" error_title = "Leader Election Already Started" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.LEADER_ELECTION, diff --git a/src/ai/backend/common/leader/tasks/event_task.py b/src/ai/backend/common/leader/tasks/event_task.py index 33edeb21d98..ad0e3b416c9 100644 --- a/src/ai/backend/common/leader/tasks/event_task.py +++ b/src/ai/backend/common/leader/tasks/event_task.py @@ -5,7 +5,7 @@ import logging from collections.abc import Callable from dataclasses import dataclass -from typing import Final +from typing import Final, override from ai.backend.common.cron.base import PeriodicTask from ai.backend.common.events.dispatcher import EventProducer @@ -46,6 +46,7 @@ def __init__( self._spec = spec self._event_producer = event_producer + @override async def run(self) -> None: """Execute the task - produce an event.""" try: @@ -56,16 +57,19 @@ async def run(self) -> None: log.exception("Failed to produce event for task {}", self._spec.name) @property + @override def name(self) -> str: """Task name.""" return self._spec.name @property + @override def interval(self) -> float: """Interval between task executions.""" return self._spec.interval @property + @override def initial_delay(self) -> float: """Initial delay before first execution.""" return self._spec.initial_delay diff --git a/src/ai/backend/common/leader/tasks/leader_cron.py b/src/ai/backend/common/leader/tasks/leader_cron.py index 66b6586aaea..cfe80f9f025 100644 --- a/src/ai/backend/common/leader/tasks/leader_cron.py +++ b/src/ai/backend/common/leader/tasks/leader_cron.py @@ -4,7 +4,7 @@ import asyncio import logging -from typing import Final +from typing import Final, override from ai.backend.common.cron.base import PeriodicTask from ai.backend.common.leader.base import LeadershipChecker, LeaderTask @@ -64,6 +64,7 @@ async def _run_task(self, task: PeriodicTask) -> None: log.exception("Unexpected error in task {}", task.name) raise + @override async def start(self, leadership_checker: LeadershipChecker) -> None: """ Start the leader cron. @@ -84,6 +85,7 @@ async def start(self, leadership_checker: LeadershipChecker) -> None: log.info("Leader cron started with {} tasks", len(self._tasks)) + @override async def stop(self) -> None: """ Stop the leader cron. diff --git a/src/ai/backend/common/leader/valkey_leader_election.py b/src/ai/backend/common/leader/valkey_leader_election.py index af04d86215b..de0bfc71f67 100644 --- a/src/ai/backend/common/leader/valkey_leader_election.py +++ b/src/ai/backend/common/leader/valkey_leader_election.py @@ -5,7 +5,7 @@ import asyncio import logging from dataclasses import dataclass -from typing import Final +from typing import Final, override from ai.backend.common.clients.valkey_client.valkey_leader.client import ValkeyLeaderClient from ai.backend.common.leader.base import AbstractLeaderElection, LeaderTask @@ -65,15 +65,18 @@ def __init__( self._leader_tasks = [] @property + @override def is_leader(self) -> bool: """Check if this instance is currently the leader.""" return self._is_leader @property + @override def server_id(self) -> str: """Get the server ID.""" return self._config.server_id + @override def register_task(self, task: LeaderTask) -> None: """ Register a task to run when this instance is the leader. @@ -146,6 +149,7 @@ async def _election_loop(self) -> None: await asyncio.sleep(self._config.renewal_interval) + @override async def start(self) -> None: """ Start the leader election renewal loop and all registered tasks. @@ -169,6 +173,7 @@ async def start(self) -> None: log.info("Valkey leader election started") + @override async def stop(self) -> None: """ Stop the leader election, all tasks, and release leadership if held. diff --git a/src/ai/backend/common/lock.py b/src/ai/backend/common/lock.py index 055305079de..8598b71183e 100644 --- a/src/ai/backend/common/lock.py +++ b/src/ai/backend/common/lock.py @@ -7,7 +7,7 @@ from collections.abc import Mapping from io import IOBase from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, override import trafaret as t from etcd_client import Client as EtcdClient @@ -131,10 +131,12 @@ def release(self) -> None: pass self._file = None + @override async def __aenter__(self) -> FileLock: await self.acquire() return self + @override async def __aexit__(self, *exc_info: Any) -> bool | None: self.release() return None @@ -180,6 +182,7 @@ def __init__( self._debug = debug self._etcd_client = None + @override async def __aenter__(self) -> EtcdCommunicator: self._etcd_client = self.etcd.etcd.with_lock( EtcdLockOption( @@ -196,6 +199,7 @@ async def __aenter__(self) -> EtcdCommunicator: return etcd_communicator + @override async def __aexit__(self, *exc_info: Any) -> bool | None: if self._etcd_client is None: raise RuntimeError("Etcd client is not initialized") @@ -245,6 +249,7 @@ def __init__( else self.default_lock_retry_interval ) + @override async def __aenter__(self) -> None: self._lock = AsyncRedisLock( self._redis, @@ -261,6 +266,7 @@ async def __aenter__(self) -> None: if self._debug: log.debug("RedisLock.__aenter__(): lock acquired") + @override async def __aexit__(self, *exc_info: Any) -> bool | None: if self._lock is None: raise RuntimeError("Lock is not initialized") diff --git a/src/ai/backend/common/message_queue/hiredis_queue.py b/src/ai/backend/common/message_queue/hiredis_queue.py index e2643a2763d..e4eafbb26dc 100644 --- a/src/ai/backend/common/message_queue/hiredis_queue.py +++ b/src/ai/backend/common/message_queue/hiredis_queue.py @@ -3,7 +3,7 @@ import logging import socket from collections.abc import AsyncGenerator, Mapping -from typing import Any, cast +from typing import Any, cast, override import hiredis from aiotools.server import process_index @@ -72,6 +72,7 @@ def __init__(self, target: RedisTarget, args: RedisMQArgs) -> None: asyncio.create_task(self._read_broadcast_messages_loop(args.subscribe_channels)) ) + @override async def send(self, payload: dict[bytes, bytes]) -> None: async with RedisConnection(self._target, db=self._db) as client: pieces = _make_pieces(payload) @@ -84,6 +85,7 @@ async def send(self, payload: dict[bytes, bytes]) -> None: *pieces, ]) + @override async def broadcast(self, payload: Mapping[str, str]) -> None: async with RedisConnection(self._target, db=self._db) as client: payload_bytes = dump_json(payload) @@ -93,6 +95,7 @@ async def broadcast(self, payload: Mapping[str, str]) -> None: payload_bytes, ]) + @override async def broadcast_with_cache(self, cache_id: str, payload: Mapping[str, str]) -> None: async with RedisConnection(self._target, db=self._db) as client: payload_bytes = dump_json(payload) @@ -114,6 +117,7 @@ async def broadcast_with_cache(self, cache_id: str, payload: Mapping[str, str]) ], ]) + @override async def fetch_cached_broadcast_message(self, cache_id: str) -> Mapping[str, str] | None: if self._closed: raise RuntimeError("Queue is closed") @@ -123,6 +127,7 @@ async def fetch_cached_broadcast_message(self, cache_id: str) -> Mapping[str, st return None return cast(Mapping[str, str] | None, load_json(reply)) + @override async def broadcast_batch(self, events: list[BroadcastPayload]) -> None: """ Broadcast multiple messages in a batch with optional caching. @@ -184,6 +189,7 @@ async def subscribe_queue(self) -> AsyncGenerator[BroadcastMessage, None]: # ty except asyncio.CancelledError: break + @override async def done(self, msg_id: MessageId) -> None: await self._done(self._anycast_stream_key, msg_id) @@ -196,6 +202,7 @@ async def _done(self, stream_key: str, msg_id: MessageId) -> None: msg_id, ]) + @override async def close(self) -> None: if self._closed: return diff --git a/src/ai/backend/common/message_queue/redis_queue/exceptions.py b/src/ai/backend/common/message_queue/redis_queue/exceptions.py index d387586cc81..901eb1a2bb3 100644 --- a/src/ai/backend/common/message_queue/redis_queue/exceptions.py +++ b/src/ai/backend/common/message_queue/redis_queue/exceptions.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -11,6 +13,7 @@ class MessageQueueClosedError(BackendAIError): error_type = "https://api.backend.ai/probs/message-queue-closed" error_title = "Message Queue Closed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MESSAGE_QUEUE, diff --git a/src/ai/backend/common/metrics/safe.py b/src/ai/backend/common/metrics/safe.py index 20e8e4572c5..9a5020aadce 100644 --- a/src/ai/backend/common/metrics/safe.py +++ b/src/ai/backend/common/metrics/safe.py @@ -12,7 +12,7 @@ import logging import threading from datetime import UTC, datetime -from typing import Any +from typing import Any, override from prometheus_client import Counter, Gauge, Histogram @@ -138,6 +138,7 @@ def observe(self, amount: float) -> None: class SafeCounter(Counter): """Counter that becomes no-op on any recording error.""" + @override def labels(self, *args: Any, **kwargs: Any) -> Any: if _metrics_disabled: return _NOOP_LABELED @@ -147,6 +148,7 @@ def labels(self, *args: Any, **kwargs: Any) -> Any: _trip(e) return _NOOP_LABELED + @override def inc(self, amount: float = 1, exemplar: dict[str, str] | None = None) -> None: if _metrics_disabled: return @@ -159,6 +161,7 @@ def inc(self, amount: float = 1, exemplar: dict[str, str] | None = None) -> None class SafeGauge(Gauge): """Gauge that becomes no-op on any recording error.""" + @override def labels(self, *args: Any, **kwargs: Any) -> Any: if _metrics_disabled: return _NOOP_LABELED @@ -168,6 +171,7 @@ def labels(self, *args: Any, **kwargs: Any) -> Any: _trip(e) return _NOOP_LABELED + @override def inc(self, amount: float = 1) -> None: if _metrics_disabled: return @@ -176,6 +180,7 @@ def inc(self, amount: float = 1) -> None: except Exception as e: _trip(e) + @override def dec(self, amount: float = 1) -> None: if _metrics_disabled: return @@ -184,6 +189,7 @@ def dec(self, amount: float = 1) -> None: except Exception as e: _trip(e) + @override def set(self, value: float) -> None: if _metrics_disabled: return @@ -196,6 +202,7 @@ def set(self, value: float) -> None: class SafeHistogram(Histogram): """Histogram that becomes no-op on any recording error.""" + @override def labels(self, *args: Any, **kwargs: Any) -> Any: if _metrics_disabled: return _NOOP_LABELED @@ -205,6 +212,7 @@ def labels(self, *args: Any, **kwargs: Any) -> Any: _trip(e) return _NOOP_LABELED + @override def observe(self, amount: float, exemplar: dict[str, str] | None = None) -> None: if _metrics_disabled: return diff --git a/src/ai/backend/common/plugin/event.py b/src/ai/backend/common/plugin/event.py index 210c9021385..af5121e5e4e 100644 --- a/src/ai/backend/common/plugin/event.py +++ b/src/ai/backend/common/plugin/event.py @@ -2,7 +2,7 @@ import logging from abc import ABCMeta, abstractmethod -from typing import Any +from typing import Any, override from ai.backend.common.events.types import AbstractEvent from ai.backend.common.types import AgentId @@ -14,9 +14,11 @@ class AbstractEventDispatcherPlugin(AbstractPlugin, metaclass=ABCMeta): + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass diff --git a/src/ai/backend/common/plugin/monitor.py b/src/ai/backend/common/plugin/monitor.py index d5d7b371a2f..476475387f9 100644 --- a/src/ai/backend/common/plugin/monitor.py +++ b/src/ai/backend/common/plugin/monitor.py @@ -3,7 +3,7 @@ import enum from abc import ABCMeta, abstractmethod from collections.abc import Mapping -from typing import Any +from typing import Any, override from . import AbstractPlugin, BasePluginContext @@ -27,9 +27,11 @@ class StatMetricTypes(enum.Enum): class AbstractStatReporterPlugin(AbstractPlugin, metaclass=ABCMeta): + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass @@ -44,9 +46,11 @@ async def report_metric( class AbstractErrorReporterPlugin(AbstractPlugin, metaclass=ABCMeta): + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass diff --git a/src/ai/backend/common/redis_client.py b/src/ai/backend/common/redis_client.py index 57ae79a1298..c9f14c68a3f 100644 --- a/src/ai/backend/common/redis_client.py +++ b/src/ai/backend/common/redis_client.py @@ -3,7 +3,7 @@ import socket from collections.abc import AsyncIterator, Sequence from contextlib import AbstractAsyncContextManager -from typing import Any, Final +from typing import Any, Final, override import hiredis @@ -274,6 +274,7 @@ async def connect(self) -> RedisClient: return client + @override async def __aenter__(self) -> RedisClient: return await self.connect() @@ -286,5 +287,6 @@ async def disconnect(self) -> None: pass # there's no more room we can do anything raise e + @override async def __aexit__(self, *exc_info: Any) -> None: return await self.disconnect() diff --git a/src/ai/backend/common/resilience/policies/metrics.py b/src/ai/backend/common/resilience/policies/metrics.py index a5d6873746c..11b75eafd78 100644 --- a/src/ai/backend/common/resilience/policies/metrics.py +++ b/src/ai/backend/common/resilience/policies/metrics.py @@ -4,7 +4,7 @@ import time from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import ParamSpec, TypeVar +from typing import ParamSpec, TypeVar, override from ai.backend.common.metrics.metric import DomainType, LayerMetricObserver, LayerType from ai.backend.common.resilience.policy import Policy @@ -57,6 +57,7 @@ def __init__(self, args: MetricArgs) -> None: self._layer = args.layer self._observer = LayerMetricObserver.instance() + @override async def execute( self, next_call: Callable[P, Awaitable[R]], diff --git a/src/ai/backend/common/resilience/policies/retry.py b/src/ai/backend/common/resilience/policies/retry.py index b4d1e7c2a07..620f053b574 100644 --- a/src/ai/backend/common/resilience/policies/retry.py +++ b/src/ai/backend/common/resilience/policies/retry.py @@ -5,7 +5,7 @@ import logging from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import Final, ParamSpec, TypeVar, cast +from typing import Final, ParamSpec, TypeVar, cast, override from ai.backend.common.exception import ( BackendAIError, @@ -30,6 +30,7 @@ class ResilienceRetryError(BackendAIError): error_type = "https://api.backend.ai/probs/resilience-retry-error" error_title = "Resilience retry error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -98,6 +99,7 @@ def __init__(self, args: RetryArgs) -> None: *args.non_retryable_exceptions, ) + @override async def execute( self, next_call: Callable[P, Awaitable[R]], diff --git a/src/ai/backend/common/resilience/policies/timeout.py b/src/ai/backend/common/resilience/policies/timeout.py index 52064e39027..10de9990281 100644 --- a/src/ai/backend/common/resilience/policies/timeout.py +++ b/src/ai/backend/common/resilience/policies/timeout.py @@ -4,7 +4,7 @@ import logging from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import ParamSpec, TypeVar +from typing import ParamSpec, TypeVar, override from ai.backend.common.exception import ( BackendAIError, @@ -28,6 +28,7 @@ class ResilienceTimeoutError(BackendAIError): error_type = "https://api.backend.ai/probs/resilience-timeout" error_title = "Resilience timeout exceeded." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -61,6 +62,7 @@ def __init__(self, args: TimeoutArgs) -> None: """ self._timeout = args.timeout + @override async def execute( self, next_call: Callable[P, Awaitable[R]], diff --git a/src/ai/backend/common/service_discovery/etcd_discovery/service_discovery.py b/src/ai/backend/common/service_discovery/etcd_discovery/service_discovery.py index 6d0664f34b5..8ceb6ec4e3b 100644 --- a/src/ai/backend/common/service_discovery/etcd_discovery/service_discovery.py +++ b/src/ai/backend/common/service_discovery/etcd_discovery/service_discovery.py @@ -1,7 +1,7 @@ import uuid from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, cast +from typing import Any, cast, override from ai.backend.common.etcd import AsyncEtcd, ConfigScopes from ai.backend.common.service_discovery.service_discovery import ( @@ -28,20 +28,24 @@ def __init__(self, args: ETCDServiceDiscoveryArgs) -> None: self._etcd = args.etcd self._prefix = args.prefix + @override async def register(self, service_meta: ServiceMetadata) -> None: prefix = self._service_prefix(service_meta.service_group, service_meta.id) await self._etcd.put_prefix(prefix, service_meta.to_dict(), scope=ConfigScopes.GLOBAL) + @override async def unregister(self, service_group: str, service_id: uuid.UUID) -> None: await self._etcd.delete_prefix( self._service_prefix(service_group, service_id), scope=ConfigScopes.GLOBAL ) + @override async def heartbeat(self, service_meta: ServiceMetadata) -> None: service_meta.health_status.update_heartbeat() prefix = self._service_prefix(service_meta.service_group, service_meta.id) await self._etcd.put_prefix(prefix, service_meta.to_dict(), scope=ConfigScopes.GLOBAL) + @override async def discover(self) -> Sequence[ServiceMetadata]: raw_service_groups = await self._etcd.get_prefix(self._prefix, scope=ConfigScopes.GLOBAL) services = [] @@ -53,6 +57,7 @@ async def discover(self) -> Sequence[ServiceMetadata]: services.append(ServiceMetadata.from_dict(service_config)) return services + @override async def get_service_group(self, service_group: str) -> Sequence[ServiceMetadata]: service_group_prefix = self._service_group_prefix(service_group) raw_service_configs = await self._etcd.get_prefix( @@ -66,6 +71,7 @@ async def get_service_group(self, service_group: str) -> Sequence[ServiceMetadat services.append(ServiceMetadata.from_dict(service_config)) return services + @override async def get_service(self, service_group: str, service_id: uuid.UUID) -> ServiceMetadata: service_prefix = self._service_prefix(service_group, service_id) raw_service_config = await self._etcd.get_prefix(service_prefix, scope=ConfigScopes.GLOBAL) @@ -74,6 +80,7 @@ async def get_service(self, service_group: str, service_id: uuid.UUID) -> Servic service_config = cast(dict[str, Any], raw_service_config) return ServiceMetadata.from_dict(service_config) + @override async def sync_model_service_routes( self, routes: Sequence[ModelServiceMetadata], diff --git a/src/ai/backend/common/service_discovery/event_publisher.py b/src/ai/backend/common/service_discovery/event_publisher.py index 2308efcc5fb..b74855047f6 100644 --- a/src/ai/backend/common/service_discovery/event_publisher.py +++ b/src/ai/backend/common/service_discovery/event_publisher.py @@ -8,7 +8,7 @@ import logging from datetime import datetime -from typing import Final +from typing import Final, override from ai.backend.common.configs.service_discovery import ServiceDiscoveryConfig from ai.backend.common.cron import LocalCron, PeriodicTask @@ -36,19 +36,23 @@ def __init__(self, publisher: ServiceDiscoveryEventPublisher, interval: float) - self._interval = interval @property + @override def name(self) -> str: return "service_discovery_heartbeat" @property + @override def interval(self) -> float: return self._interval @property + @override def initial_delay(self) -> float: # The original loop slept for `interval` before publishing the first # heartbeat (the initial registration is published by start() directly). return self._interval + @override async def run(self) -> None: await self._publisher.publish_registered() diff --git a/src/ai/backend/common/service_discovery/redis_discovery/service_discovery.py b/src/ai/backend/common/service_discovery/redis_discovery/service_discovery.py index d7e828bf38d..cffe4d0399f 100644 --- a/src/ai/backend/common/service_discovery/redis_discovery/service_discovery.py +++ b/src/ai/backend/common/service_discovery/redis_discovery/service_discovery.py @@ -2,7 +2,7 @@ import uuid from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, Self +from typing import Any, Self, override from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient from ai.backend.common.json import dump_json_str @@ -74,20 +74,24 @@ async def _hget_json(self, name: str) -> ServiceMetadata: return ServiceMetadata.from_dict(result) + @override async def register(self, service_meta: ServiceMetadata) -> None: key = self._service_prefix(service_meta.service_group, service_meta.id) await self._hsetex(key, service_meta.to_dict()) + @override async def unregister(self, service_group: str, service_id: uuid.UUID) -> None: """Unregister a service using ValkeyLiveClient.""" key = self._service_prefix(service_group, service_id) await self._valkey_client.delete_key(key) + @override async def heartbeat(self, service_meta: ServiceMetadata) -> None: service_meta.health_status.update_heartbeat() key = self._service_prefix(service_meta.service_group, service_meta.id) await self._hsetex(key, service_meta.to_dict()) + @override async def discover(self) -> Sequence[ServiceMetadata]: pattern = f"{self._prefix}.*.*" keys = await self._scan_keys(pattern) @@ -95,6 +99,7 @@ async def discover(self) -> Sequence[ServiceMetadata]: raise ValueError("No service groups found.") return [await self._hget_json(key) for key in keys] + @override async def get_service_group(self, service_group: str) -> Sequence[ServiceMetadata]: pattern = f"{self._service_group_prefix(service_group)}.*" keys = await self._scan_keys(pattern) @@ -102,10 +107,12 @@ async def get_service_group(self, service_group: str) -> Sequence[ServiceMetadat raise ValueError(f"No services found in group {service_group}.") return [await self._hget_json(key) for key in keys] + @override async def get_service(self, service_group: str, service_id: uuid.UUID) -> ServiceMetadata: key = self._service_prefix(service_group, service_id) return await self._hget_json(key) + @override async def sync_model_service_routes( self, routes: Sequence[ModelServiceMetadata], diff --git a/src/ai/backend/common/service_discovery/service_discovery.py b/src/ai/backend/common/service_discovery/service_discovery.py index 282213720ea..e506a49e69f 100644 --- a/src/ai/backend/common/service_discovery/service_discovery.py +++ b/src/ai/backend/common/service_discovery/service_discovery.py @@ -5,7 +5,7 @@ import warnings from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import Any, Final, Self +from typing import Any, Final, Self, override from pydantic import Field @@ -102,6 +102,7 @@ def update_heartbeat(self) -> None: """ self.last_heartbeat = time.time() + @override def __eq__(self, value: object) -> bool: if not isinstance(value, HealthStatus): return False diff --git a/src/ai/backend/common/stage/types.py b/src/ai/backend/common/stage/types.py index 4467924f118..28f7582545f 100644 --- a/src/ai/backend/common/stage/types.py +++ b/src/ai/backend/common/stage/types.py @@ -98,6 +98,7 @@ def __init__(self, provisioner: Provisioner) -> None: # type: ignore[type-arg] self._resource = None self._setup_completed = asyncio.Event() + @override async def setup(self, spec_generator: SpecGenerator[TSpec]) -> None: """ Sets up the lifecycle stage. @@ -111,12 +112,14 @@ async def setup(self, spec_generator: SpecGenerator[TSpec]) -> None: finally: self._setup_completed.set() + @override async def wait_for_resource(self) -> TResource: await self._setup_completed.wait() if self._resource is None: raise RuntimeError("Resource setup failed") return self._resource + @override async def teardown(self) -> None: """ Tears down the lifecycle stage. diff --git a/src/ai/backend/common/typed_validators.py b/src/ai/backend/common/typed_validators.py index b786ca536a4..5cf6a0c806a 100644 --- a/src/ai/backend/common/typed_validators.py +++ b/src/ai/backend/common/typed_validators.py @@ -5,7 +5,7 @@ from collections.abc import Mapping, Sequence from datetime import tzinfo from pathlib import Path -from typing import Annotated, Any, ClassVar, Final, TypeVar +from typing import Annotated, Any, ClassVar, Final, TypeVar, override import jwt from dateutil import tz @@ -300,9 +300,11 @@ def __getitem__(self, *args: Any) -> int | str: def to_legacy(self) -> LegacyHostPortPair: return LegacyHostPortPair(host=self.host, port=self.port) + @override def __str__(self) -> str: return self.address + @override def __repr__(self) -> str: return self.address diff --git a/src/ai/backend/common/types.py b/src/ai/backend/common/types.py index 84dde7a5e0c..2aff80390b3 100644 --- a/src/ai/backend/common/types.py +++ b/src/ai/backend/common/types.py @@ -243,6 +243,7 @@ def _validation_failure_info(cls, exc: ValidationError) -> SchemaValidationFailu return SchemaValidationFailureInfo(summary=str(exc), errors=sanitized) @classmethod + @override def model_validate(cls, *args: Any, **kwargs: Any) -> Self: try: return super().model_validate(*args, **kwargs) @@ -250,6 +251,7 @@ def model_validate(cls, *args: Any, **kwargs: Any) -> Self: raise cls.build_validation_error(cls._validation_failure_info(e)) from e @classmethod + @override def model_validate_json(cls, *args: Any, **kwargs: Any) -> Self: try: return super().model_validate_json(*args, **kwargs) @@ -844,9 +846,11 @@ def __init__(self, expression: str, *, escape_map: Mapping[str, str] | None = No self.escape_map.update(escape_map) # self.unescape_map = {v: k for k, v in self.escape_map.items()} + @override def __str__(self) -> str: return self.expression + @override def __repr__(self) -> str: return self.__str__() @@ -865,6 +869,7 @@ class HostPortPair(namedtuple("HostPortPair", "host port")): def as_sockaddr(self) -> tuple[str, int]: return str(self.host), self.port + @override def __str__(self) -> str: if isinstance(self.host, ipaddress.IPv6Address): return f"[{self.host}]:{self.port}" @@ -910,9 +915,11 @@ def _to_ip_network(val: str) -> Address: def address(self) -> Address | None: return self._address + @override def __str__(self) -> str: return str(self._address) + @override def __eq__(self, other: object) -> bool: if other is self: return True @@ -1038,6 +1045,7 @@ def _quantize(val: int | Decimal, multiplier: int) -> Decimal: value = d.quantize(Decimal(".00")).normalize() return value + @override def __str__(self) -> str: suffix_idx = self._preformat() if suffix_idx == 0: @@ -1049,6 +1057,7 @@ def __str__(self) -> str: value = self._quantize(self, multiplier) return f"{value:f} {suffix.upper()}iB" + @override def __format__(self, format_spec: str) -> str: if len(format_spec) != 1: raise ValueError("format-string for BinarySize can be only one character.") @@ -1124,6 +1133,7 @@ def _process_raw_value(cls, key: str, value: RawResourceValue) -> Decimal: v = Decimal(value) return v + @override def __setitem__(self, key: str | SlotName, value: RawResourceValue | None) -> None: normalized_key = str(key) if value is None: @@ -1131,10 +1141,12 @@ def __setitem__(self, key: str | SlotName, value: RawResourceValue | None) -> No return self.data[normalized_key] = self._process_raw_value(normalized_key, value) + @override def __getitem__(self, key: str | SlotName) -> Decimal: normalized_key = str(key) return self.data[normalized_key] + @override def copy(self) -> Self: return type(self)(self.data.copy()) @@ -1163,6 +1175,7 @@ def __sub__(self, other: ResourceSlot) -> ResourceSlot: def __neg__(self) -> ResourceSlot: return type(self)({k: -v for k, v in self.data.items()}) + @override def __eq__(self, other: object) -> bool: if other is self: return True @@ -1173,6 +1186,7 @@ def __eq__(self, other: object) -> bool: other_values = [other.data[k] for k in sorted(other.data.keys())] return self_values == other_values + @override def __ne__(self, other: object) -> bool: if not isinstance(other, ResourceSlot): return NotImplemented @@ -1486,6 +1500,7 @@ def validate_quota_scope_id(v: Any) -> QuotaScopeID: ), ) + @override def __str__(self) -> str: match self.scope_id: case UUID(): @@ -1493,6 +1508,7 @@ def __str__(self) -> str: case _: raise ValueError(f"Invalid quota scope ID: {self.scope_id!r}") + @override def __repr__(self) -> str: return self.__str__() @@ -1559,17 +1575,20 @@ def __init__(self, quota_scope_id: QuotaScopeID | str | None, folder_id: UUID) - case None: self.quota_scope_id = None + @override def __str__(self) -> str: if self.quota_scope_id is None: return self.folder_id.hex return f"{self.quota_scope_id}/{self.folder_id.hex}" + @override def __eq__(self, other: Any) -> bool: result: bool = ( self.quota_scope_id == other.quota_scope_id and self.folder_id == other.folder_id ) return result + @override def __hash__(self) -> int: qsid = str(self.quota_scope_id) if self.quota_scope_id is not None else None return hash((qsid, self.folder_id)) @@ -1599,6 +1618,7 @@ class VFolderMount(JSONSerializableMixin): mount_perm: MountPermission usage_mode: VFolderUsageMode + @override def to_json(self) -> dict[str, Any]: return { "name": self.name, @@ -1611,6 +1631,7 @@ def to_json(self) -> dict[str, Any]: } @classmethod + @override def from_json(cls, obj: Mapping[str, Any]) -> Self: base = cls.as_trafaret().check(obj) return cls(**base) @@ -1641,6 +1662,7 @@ def to_dataclass(self) -> VFolderMountData: ) @classmethod + @override def as_trafaret(cls) -> t.Trafaret: from . import validators as tx @@ -1667,6 +1689,7 @@ def __or__[T1, T2]( self, value: dict[T1, T2], / ) -> dict[str | T1, set[VFolderHostPermission] | T2]: ... + @override def __or__( self, value: dict[Any, Any], / ) -> dict[str, set[VFolderHostPermission]] | dict[str | Any, set[VFolderHostPermission] | Any]: @@ -1683,16 +1706,19 @@ def __or__( union_map[host] |= set(perm_list) return VFolderHostPermissionMap(union_map) + @override def to_json(self) -> dict[str, Any]: return {host: [perm.value for perm in perms] for host, perms in self.items()} @classmethod + @override def from_json(cls, obj: Mapping[str, Any]) -> Self: return cls({ host: {VFolderHostPermission(perm) for perm in perms} for host, perms in obj.items() }) @classmethod + @override def as_trafaret(cls) -> t.Trafaret: from . import validators as tx diff --git a/src/ai/backend/common/utils.py b/src/ai/backend/common/utils.py index ae013ffd1b9..43b8a2f0d7e 100644 --- a/src/ai/backend/common/utils.py +++ b/src/ai/backend/common/utils.py @@ -18,6 +18,7 @@ TYPE_CHECKING, Any, TypeVar, + override, ) import aiofiles @@ -278,9 +279,11 @@ def __init__( self.d = d self.p = p + @override def __eq__(self, o: Any) -> bool: return str(self) == str(o) + @override def __str__(self) -> str: return f"{self.device} {self.mountpoint} {self.fstype} {self.options} {self.d} {self.p}" diff --git a/src/ai/backend/common/validators.py b/src/ai/backend/common/validators.py index 73eec6993da..93a2bc2abec 100644 --- a/src/ai/backend/common/validators.py +++ b/src/ai/backend/common/validators.py @@ -22,6 +22,7 @@ Literal, TypeVar, cast, + override, ) import dateutil.tz @@ -274,6 +275,7 @@ def __init__( self._allow_devnull = allow_devnull self._resolve = resolve + @override def check_and_return(self, value: Any) -> _Path: try: p = _Path(value).resolve() if self._resolve else _Path(value) diff --git a/src/ai/backend/common/web/session/__init__.py b/src/ai/backend/common/web/session/__init__.py index 9766a48fab3..751befef581 100644 --- a/src/ai/backend/common/web/session/__init__.py +++ b/src/ai/backend/common/web/session/__init__.py @@ -102,6 +102,7 @@ def __init__( if session_data is not None: self._mapping.update(session_data) + @override def __repr__(self) -> str: return f"<{self.__class__.__name__} [new:{self.new}, changed:{self._changed}, created:{self.created} expiration:{self.expiration_dt}] {self._mapping!r}>" @@ -155,23 +156,29 @@ def set_new_identity(self, identity: Any | None) -> None: self._identity = identity + @override def __len__(self) -> int: return len(self._mapping) + @override def __iter__(self) -> Iterator[str]: return iter(self._mapping) + @override def __contains__(self, key: object) -> bool: return key in self._mapping + @override def __getitem__(self, key: str) -> Any: return self._mapping[key] + @override def __setitem__(self, key: str, value: Any) -> None: self._mapping[key] = value self._changed = True self._created = get_time() + @override def __delitem__(self, key: str) -> None: del self._mapping[key] self._changed = True @@ -381,6 +388,7 @@ def __init__( decoder=decoder, ) + @override async def load_session(self, request: web.Request) -> Session: cookie = self.load_cookie(request) if cookie is None: diff --git a/src/ai/backend/common/web/session/redis_storage.py b/src/ai/backend/common/web/session/redis_storage.py index 8dfd120b801..3f1d32ab0f2 100644 --- a/src/ai/backend/common/web/session/redis_storage.py +++ b/src/ai/backend/common/web/session/redis_storage.py @@ -51,6 +51,7 @@ def __init__( async def get_redis_time(self) -> int: return await self._valkey_client.get_server_time_second() + @override async def load_session(self, request: web.Request) -> Session: # If X-BackendAI-SessionID exists in request, login will use the value as SessionID, # instead of Cookie value. diff --git a/src/ai/backend/install/app.py b/src/ai/backend/install/app.py index dbb83d923da..8fae60d6839 100644 --- a/src/ai/backend/install/app.py +++ b/src/ai/backend/install/app.py @@ -5,7 +5,7 @@ import shutil import textwrap from pathlib import Path -from typing import Any, cast +from typing import Any, cast, override from weakref import WeakSet from rich.text import Text @@ -56,6 +56,7 @@ def __init__(self, *, non_interactive: bool = False, **kwargs: Any) -> None: self._non_interactive = non_interactive self._task = None + @override def compose(self) -> ComposeResult: yield Label("Development Setup", classes="mode-title") with TabbedContent(): @@ -119,6 +120,7 @@ def __init__(self, *, non_interactive: bool = False, **kwargs: Any) -> None: self._non_interactive = non_interactive self._task = None + @override def compose(self) -> ComposeResult: yield Label("Package Setup", classes="mode-title") with TabbedContent(): @@ -215,6 +217,7 @@ def __init__( self._install_variable = install_variable self._non_interactive = non_interactive + @override def compose(self) -> ComposeResult: yield Label("Choose deployment type:", id="pkg-type-heading") with ListView(id="pkg-type-list") as lv: @@ -265,6 +268,7 @@ def feed_variables(self, install_variable: InstallVariable) -> None: address_input = self.app.query_one("#public-ip", Input) address_input.value = install_variable.public_facing_address + @override def compose(self) -> ComposeResult: yield Label("Configure", classes="mode-title") with Vertical(id="configure-form"): @@ -307,6 +311,7 @@ def __init__(self, install_info: InstallInfo, **kwargs: Any) -> None: super().__init__(**kwargs) self.install_info = install_info + @override def compose(self) -> ComposeResult: service = self.install_info.service_config yield Markdown( @@ -573,6 +578,7 @@ def __init__( enable_telemetry=args.enable_telemetry, ) + @override def compose(self) -> ComposeResult: yield Label(id="heading") if self._dist_info_path is None: @@ -701,6 +707,7 @@ def __init__(self, args: CliArgs | None = None) -> None: ) self._args = args + @override def compose(self) -> ComposeResult: yield Header(show_clock=True) logo_text = textwrap.dedent( diff --git a/src/ai/backend/install/context.py b/src/ai/backend/install/context.py index 39b60de1860..ce348237e1d 100644 --- a/src/ai/backend/install/context.py +++ b/src/ai/backend/install/context.py @@ -21,7 +21,7 @@ from datetime import datetime from functools import partial from pathlib import Path -from typing import Any, Final +from typing import Any, Final, override import aiofiles import aiotools @@ -1855,6 +1855,7 @@ async def populate_images(self) -> None: class DevContext(Context): + @override def hydrate_install_info(self) -> InstallInfo: # TODO: customize addr/user/password options # TODO: multi-node setup @@ -1938,6 +1939,7 @@ def hydrate_install_info(self) -> InstallInfo: accelerator=self.install_variable.accelerator, ) + @override def copy_config(self, template_name: str) -> Path: with self.resource_path("ai.backend.install.configs", template_name) as src_path: dst_path = Path.cwd() / template_name @@ -1948,6 +1950,7 @@ def copy_config(self, template_name: str) -> Path: shutil.copy(src_path, dst_path) return dst_path + @override async def check_prerequisites(self) -> None: await super().check_prerequisites() await install_git_lfs(self) @@ -1961,6 +1964,7 @@ async def install(self) -> None: await install_editable_webui(self) await self.install_halfstack() + @override async def _configure_mock_accelerator(self, accelerator: Accelerator) -> None: """ cp "configs/accelerator/mock-accelerator.toml" mock-accelerator.toml @@ -2026,6 +2030,7 @@ async def configure(self) -> None: class PackageContext(Context): + @override def hydrate_install_info(self) -> InstallInfo: # TODO: customize addr/user/password options # TODO: multi-node setup @@ -2099,6 +2104,7 @@ def hydrate_install_info(self) -> InstallInfo: accelerator=self.install_variable.accelerator, ) + @override def copy_config(self, template_name: str) -> Path: with self.resource_path("ai.backend.install.configs", template_name) as src_path: dst_path = self.dist_info.target_path / template_name @@ -2109,6 +2115,7 @@ def copy_config(self, template_name: str) -> Path: shutil.copy(src_path, dst_path) return dst_path + @override async def check_prerequisites(self) -> None: await super().check_prerequisites() if self.install_variable.with_harbor: @@ -2275,6 +2282,7 @@ async def install(self) -> None: self.log_header("Installing databases (halfstack)...") await self.install_halfstack() + @override async def _configure_mock_accelerator(self, accelerator: Accelerator) -> None: """ cp "configs/accelerator/mock-accelerator.toml" mock-accelerator.toml diff --git a/src/ai/backend/install/types.py b/src/ai/backend/install/types.py index a956885d122..2cf2bd020b9 100644 --- a/src/ai/backend/install/types.py +++ b/src/ai/backend/install/types.py @@ -4,6 +4,7 @@ import enum from datetime import datetime from pathlib import Path +from typing import override from pydantic import Field from rich.console import ConsoleRenderable, RichCast @@ -92,6 +93,7 @@ def __init__(self, msg: str, *, instruction: str | None = None) -> None: self.msg = msg self.instruction = instruction + @override def __rich__(self) -> ConsoleRenderable: text = f"[bold red]:warning: [bold white]{self.msg}[/]" if self.instruction: @@ -138,6 +140,7 @@ class OSInfo(RichCast): distro: str distro_variants: set[str] + @override def __rich__(self) -> ConsoleRenderable: variants = [self.distro, *sorted(self.distro_variants)] variant = ", ".join(variants) diff --git a/src/ai/backend/install/widgets.py b/src/ai/backend/install/widgets.py index 9147618c890..f2e7b6a0382 100644 --- a/src/ai/backend/install/widgets.py +++ b/src/ai/backend/install/widgets.py @@ -3,7 +3,7 @@ import asyncio import io from pathlib import Path -from typing import Any +from typing import Any, override from rich.console import Console, ConsoleRenderable from rich.text import Text @@ -26,6 +26,7 @@ class DirectoryPathValidator(Validator): + @override def validate(self, value: str) -> ValidationResult: if Path(value).is_dir(): return self.success() @@ -38,6 +39,7 @@ def __init__(self, label: str, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._label = label + @override def compose(self) -> ComposeResult: yield Label(Text.from_markup(self._label), classes="progress-name") yield ProgressBar(classes="progress-download") @@ -73,6 +75,7 @@ def __init__( width=200, ) + @override def on_mount(self) -> None: if self.app.is_headless: self._stdout_console = Console(force_terminal=True) @@ -102,6 +105,7 @@ def _to_plain(self, content: object) -> str: return f"{content}\n" return self._render_io.getvalue() + @override def write( self, data: object, @@ -197,6 +201,7 @@ def __init__( self._focus_save = None self._validator = validator + @override def compose(self) -> ComposeResult: yield Label(self._label) yield Input( diff --git a/src/ai/backend/kernel/app/__init__.py b/src/ai/backend/kernel/app/__init__.py index a71e6c978b1..7734af5dc68 100644 --- a/src/ai/backend/kernel/app/__init__.py +++ b/src/ai/backend/kernel/app/__init__.py @@ -5,7 +5,7 @@ import logging from collections.abc import Mapping -from typing import Any +from typing import Any, override from ai.backend.kernel import BaseRunner @@ -21,17 +21,21 @@ class Runner(BaseRunner): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + @override async def init_with_loop(self) -> None: pass + @override async def build_heuristic(self) -> int: log.warning("batch-mode execution is not supported") return 0 + @override async def execute_heuristic(self) -> int: log.warning("batch-mode execution is not supported") return 0 + @override async def start_service(self, service_info: Mapping[str, Any]) -> tuple[None, dict[str, Any]]: # app kernels use service-definition templates. return None, {} diff --git a/src/ai/backend/kernel/logging.py b/src/ai/backend/kernel/logging.py index 176fc325829..fa7dc51df53 100644 --- a/src/ai/backend/kernel/logging.py +++ b/src/ai/backend/kernel/logging.py @@ -4,13 +4,14 @@ from contextlib import closing from io import StringIO from logging.handlers import QueueHandler -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override if TYPE_CHECKING: from janus import _SyncQueueProxy class RelativeCreatedFormatter(logging.Formatter): + @override def format(self, record: logging.LogRecord) -> str: record.relative_seconds = record.relativeCreated / 1000 return super().format(record) @@ -23,11 +24,13 @@ def __init__(self, fmt: str, args: tuple[Any, ...]) -> None: self.fmt = fmt self.args = args + @override def __str__(self) -> str: return self.fmt.format(*self.args) class BraceStyleAdapter(logging.LoggerAdapter[logging.Logger]): + @override def log(self, level: int, msg: object, *args: object, **kwargs: object) -> None: if self.isEnabledFor(level): _msg, _kwargs = self.process(msg, kwargs) @@ -50,6 +53,7 @@ def setup_logger_basic(log_prefix: str, debug: bool) -> None: class LogQHandler(QueueHandler): + @override def enqueue(self, record: logging.LogRecord) -> None: if self.formatter is None: raise RuntimeError("Formatter is not initialized") diff --git a/src/ai/backend/kernel/python/__init__.py b/src/ai/backend/kernel/python/__init__.py index afaa1948607..ae7ae9f8057 100644 --- a/src/ai/backend/kernel/python/__init__.py +++ b/src/ai/backend/kernel/python/__init__.py @@ -5,7 +5,7 @@ import tempfile from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import janus @@ -34,6 +34,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.input_queue: janus.Queue[tuple[str, str]] | None = None self.output_queue: janus.Queue[bytes] | None = None + @override async def init_with_loop(self) -> None: self.input_queue = janus.Queue[tuple[str, str]]() self.output_queue = janus.Queue[bytes]() @@ -61,6 +62,7 @@ async def init_with_loop(self) -> None: user_site_path.mkdir(parents=True, exist_ok=True) shutil.copy(str(sitecustomize_path), str(user_site_path / "sitecustomize.py")) + @override async def build_heuristic(self) -> int: if Path("setup.py").is_file(): cmd = [ @@ -77,6 +79,7 @@ async def build_heuristic(self) -> int: log.warning('skipping the build phase due to missing "setup.py" file') return 0 + @override async def execute_heuristic(self) -> int: if Path("main.py").is_file(): cmd = [ @@ -88,6 +91,7 @@ async def execute_heuristic(self) -> int: log.error('cannot find the main script ("main.py").') return 127 + @override async def start_service( self, service_info: Mapping[str, Any] ) -> ( diff --git a/src/ai/backend/kernel/python/drawing/turtle.py b/src/ai/backend/kernel/python/drawing/turtle.py index 9fe4f7eec6c..66281654261 100644 --- a/src/ai/backend/kernel/python/drawing/turtle.py +++ b/src/ai/backend/kernel/python/drawing/turtle.py @@ -1,7 +1,7 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from .color import Colors @@ -15,14 +15,17 @@ class Vec2D(tuple[float, float]): def __new__(cls, x: float | int, y: float | int) -> Vec2D: return tuple.__new__(cls, (x, y)) + @override def __add__(self, other: Vec2D) -> Vec2D: # type: ignore[override] return Vec2D(self[0] + other[0], self[1] + other[1]) + @override def __mul__(self, other: Vec2D | float | int) -> Vec2D | float: # type: ignore[override] if isinstance(other, Vec2D): return self[0] * other[0] + self[1] * other[1] return Vec2D(self[0] * other, self[1] * other) + @override def __rmul__(self, other: float | int) -> Vec2D: # type: ignore[override] return Vec2D(self[0] * other, self[1] * other) @@ -46,6 +49,7 @@ def rotate(self, angle: float | int) -> Vec2D: def __getnewargs__(self) -> tuple[float, float]: return (self[0], self[1]) + @override def __repr__(self) -> str: return "({:.2f},{:.2f})".format(*self) diff --git a/src/ai/backend/kernel/utils.py b/src/ai/backend/kernel/utils.py index 44820fa95ea..88e65d8b569 100644 --- a/src/ai/backend/kernel/utils.py +++ b/src/ai/backend/kernel/utils.py @@ -4,7 +4,7 @@ import site import traceback from pathlib import Path -from typing import Any, Final +from typing import Any, Final, override __all__ = ( "current_loop", @@ -45,6 +45,7 @@ def __init__(self, path_prefix: str) -> None: self.path_prefix = path_prefix self.site_prefix = site.getsitepackages()[0] + @override def filter(self, record: logging.LogRecord) -> bool: if record.exc_info: _, _, exc_tb = record.exc_info diff --git a/src/ai/backend/logging/config.py b/src/ai/backend/logging/config.py index 5381087c72a..177ce5469d6 100644 --- a/src/ai/backend/logging/config.py +++ b/src/ai/backend/logging/config.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import StrEnum -from typing import Annotated, Any, Self +from typing import Annotated, Any, Self, override from pydantic import ( AliasChoices, @@ -63,9 +63,11 @@ class HostPortPair(BaseConfigModel): host: str = Field(examples=["127.0.0.1"]) port: int = Field(gt=0, lt=65536, examples=[8201]) + @override def __repr__(self) -> str: return f"{self.host}:{self.port}" + @override def __str__(self) -> str: return self.__repr__() diff --git a/src/ai/backend/logging/formatter.py b/src/ai/backend/logging/formatter.py index 99cf7d0f42f..a428248ad93 100644 --- a/src/ai/backend/logging/formatter.py +++ b/src/ai/backend/logging/formatter.py @@ -7,7 +7,7 @@ from collections.abc import Sequence from datetime import UTC, datetime from types import TracebackType -from typing import Any, cast +from typing import Any, cast, override import coloredlogs from pythonjsonlogger.json import JsonFormatter @@ -34,14 +34,17 @@ def format_exception( class SerializedExceptionFormatter(logging.Formatter): + @override def formatException(self, ei: Sequence[str] | _SysExcInfoType) -> str: return format_exception(self, ei) class ConsoleFormatter(logging.Formatter): + @override def formatException(self, ei: Sequence[str] | _SysExcInfoType) -> str: return format_exception(self, ei) + @override def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str: ct = self.converter(record.created) if datefmt: @@ -52,9 +55,11 @@ def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> s class CustomJsonFormatter(JsonFormatter): + @override def formatException(self, ei: Sequence[str] | _SysExcInfoType) -> str: return format_exception(self, ei) + @override def add_fields( self, log_record: dict[str, Any], # the manipulated entry object @@ -84,5 +89,6 @@ class pretty: def __init__(self, obj: Any) -> None: self.obj = obj + @override def __repr__(self) -> str: return pprint.pformat(self.obj) diff --git a/src/ai/backend/logging/handler/intrinsic.py b/src/ai/backend/logging/handler/intrinsic.py index fd124d468b9..7ff8a73d33d 100644 --- a/src/ai/backend/logging/handler/intrinsic.py +++ b/src/ai/backend/logging/handler/intrinsic.py @@ -35,6 +35,7 @@ def __init__(self, *, endpoint: str, msgpack_options: MsgpackOptions) -> None: else: self._sock = None + @override def close(self) -> None: if self._sock is not None: self._sock.close() diff --git a/src/ai/backend/logging/utils.py b/src/ai/backend/logging/utils.py index 27c6cbf7e14..64a83c30783 100644 --- a/src/ai/backend/logging/utils.py +++ b/src/ai/backend/logging/utils.py @@ -52,6 +52,7 @@ def __init__(self, fmt: str, args: tuple[Any, ...], kwargs: Mapping[str, Any]) - self.args = args self.kwargs = kwargs + @override def __str__(self) -> str: return self.fmt.format(*self.args, **self.kwargs) diff --git a/src/ai/backend/manager/actions/action/base.py b/src/ai/backend/manager/actions/action/base.py index c69a890cef3..9fe718edf1b 100644 --- a/src/ai/backend/manager/actions/action/base.py +++ b/src/ai/backend/manager/actions/action/base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import datetime, timedelta -from typing import TypeVar +from typing import TypeVar, override from ai.backend.common.data.permission.types import EntityType from ai.backend.common.exception import ErrorCode @@ -52,6 +52,7 @@ def entity_id(self) -> str | None: class SearchActionResult[T](BaseActionResult): result: SearchResult[T] + @override def entity_id(self) -> str | None: return None diff --git a/src/ai/backend/manager/actions/validators/auth_validator.py b/src/ai/backend/manager/actions/validators/auth_validator.py index c507bc0af2b..8a89ceada68 100644 --- a/src/ai/backend/manager/actions/validators/auth_validator.py +++ b/src/ai/backend/manager/actions/validators/auth_validator.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.common.contexts.user import current_user from ai.backend.manager.actions.action import BaseAction, BaseActionTriggerMeta from ai.backend.manager.actions.validator.base import ActionValidator @@ -6,6 +8,7 @@ class AuthorizationValidator(ActionValidator): + @override async def validate(self, action: BaseAction, meta: BaseActionTriggerMeta) -> None: user = current_user() if user is None: diff --git a/src/ai/backend/manager/api/gql/agent/types.py b/src/ai/backend/manager/api/gql/agent/types.py index b173bce4c8b..b7ad997f2b2 100644 --- a/src/ai/backend/manager/api/gql/agent/types.py +++ b/src/ai/backend/manager/api/gql/agent/types.py @@ -3,7 +3,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override import strawberry from strawberry import Info @@ -536,6 +536,7 @@ async def resource_slots( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/app_config_allow_list/types.py b/src/ai/backend/manager/api/gql/app_config_allow_list/types.py index 3bddd112901..765fc75ec73 100644 --- a/src/ai/backend/manager/api/gql/app_config_allow_list/types.py +++ b/src/ai/backend/manager/api/gql/app_config_allow_list/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import Any, Self, cast +from typing import Any, Self, cast, override from uuid import UUID from strawberry import Info @@ -92,6 +92,7 @@ class AppConfigAllowListGQL(PydanticNodeMixin[AppConfigAllowListNode]): updated_at: datetime = gql_field(description="Last update timestamp (UTC).") @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/app_config_definition/types.py b/src/ai/backend/manager/api/gql/app_config_definition/types.py index 4dac5e6196a..a9622767f0c 100644 --- a/src/ai/backend/manager/api/gql/app_config_definition/types.py +++ b/src/ai/backend/manager/api/gql/app_config_definition/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import Any, Self, cast +from typing import Any, Self, cast, override from uuid import UUID from strawberry import Info @@ -84,6 +84,7 @@ class AppConfigDefinitionGQL(PydanticNodeMixin[AppConfigDefinitionNode]): updated_at: datetime = gql_field(description="Last update timestamp (UTC).") @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/artifact/types.py b/src/ai/backend/manager/api/gql/artifact/types.py index 6ec092645bc..7178d17ec0b 100644 --- a/src/ai/backend/manager/api/gql/artifact/types.py +++ b/src/ai/backend/manager/api/gql/artifact/types.py @@ -3,7 +3,7 @@ import uuid from collections.abc import Iterable from datetime import datetime -from typing import Any, Self, cast +from typing import Any, Self, cast, override import strawberry from strawberry import ID, UNSET, Info @@ -747,6 +747,7 @@ class ArtifactRevision(PydanticNodeMixin[ArtifactRevisionNode]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/audit_log/types/node.py b/src/ai/backend/manager/api/gql/audit_log/types/node.py index 719b444327f..850439f58d3 100644 --- a/src/ai/backend/manager/api/gql/audit_log/types/node.py +++ b/src/ai/backend/manager/api/gql/audit_log/types/node.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -88,6 +88,7 @@ async def user( return user_data @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, diff --git a/src/ai/backend/manager/api/gql/container_registry/types.py b/src/ai/backend/manager/api/gql/container_registry/types.py index e8e82383f83..6c4f74c6657 100644 --- a/src/ai/backend/manager/api/gql/container_registry/types.py +++ b/src/ai/backend/manager/api/gql/container_registry/types.py @@ -4,7 +4,7 @@ from collections.abc import Iterable from enum import StrEnum -from typing import Any, Self, cast +from typing import Any, Self, cast, override from uuid import UUID from strawberry import Info @@ -72,6 +72,7 @@ class ContainerRegistryGQL(PydanticNodeMixin[Any]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/access_token.py b/src/ai/backend/manager/api/gql/deployment/types/access_token.py index 7fe17af9d21..492e38f5b74 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/access_token.py +++ b/src/ai/backend/manager/api/gql/deployment/types/access_token.py @@ -4,7 +4,7 @@ from collections.abc import Iterable from datetime import datetime -from typing import Any, Self, cast +from typing import Any, Self, cast, override from uuid import UUID from strawberry import ID, Info @@ -84,6 +84,7 @@ class AccessToken(PydanticNodeMixin[AccessTokenNodeDTO]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/auto_scaling.py b/src/ai/backend/manager/api/gql/deployment/types/auto_scaling.py index 8ccef12dd9d..4566819b0d2 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/auto_scaling.py +++ b/src/ai/backend/manager/api/gql/deployment/types/auto_scaling.py @@ -6,7 +6,7 @@ from datetime import datetime from decimal import Decimal from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -153,6 +153,7 @@ async def query_preset( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/deployment.py b/src/ai/backend/manager/api/gql/deployment/types/deployment.py index 2b716561f22..6c28f3af16a 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/deployment.py +++ b/src/ai/backend/manager/api/gql/deployment/types/deployment.py @@ -4,7 +4,7 @@ from collections.abc import Iterable from datetime import datetime -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -573,6 +573,7 @@ async def access_tokens( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/replica.py b/src/ai/backend/manager/api/gql/deployment/types/replica.py index e5a0d3e44c2..278399ad29d 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/replica.py +++ b/src/ai/backend/manager/api/gql/deployment/types/replica.py @@ -4,7 +4,7 @@ from collections.abc import Iterable from datetime import datetime -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -285,6 +285,7 @@ async def deployment( return await info.context.data_loaders.deployment_loader.load(UUID(str(self.deployment_id))) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/revision.py b/src/ai/backend/manager/api/gql/deployment/types/revision.py index bc5883ae25c..8613a7cf97e 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/revision.py +++ b/src/ai/backend/manager/api/gql/deployment/types/revision.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -665,6 +665,7 @@ async def resource_slots( return [AllocatedResourceSlotGQL.from_pydantic(item) for item in payload.items] @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/deployment/types/route.py b/src/ai/backend/manager/api/gql/deployment/types/route.py index 518d5bfe469..13e6b8cab80 100644 --- a/src/ai/backend/manager/api/gql/deployment/types/route.py +++ b/src/ai/backend/manager/api/gql/deployment/types/route.py @@ -6,7 +6,7 @@ from datetime import datetime from enum import StrEnum from functools import lru_cache -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -175,6 +175,7 @@ async def revision( return await info.context.data_loaders.revision_loader.load(UUID(str(self.revision_id))) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/domain.py b/src/ai/backend/manager/api/gql/domain.py index 064af880f97..a21eb478965 100644 --- a/src/ai/backend/manager/api/gql/domain.py +++ b/src/ai/backend/manager/api/gql/domain.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -21,6 +22,7 @@ class Domain(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/domain_v2/types/node.py b/src/ai/backend/manager/api/gql/domain_v2/types/node.py index 30cea53f257..ffd764a28ff 100644 --- a/src/ai/backend/manager/api/gql/domain_v2/types/node.py +++ b/src/ai/backend/manager/api/gql/domain_v2/types/node.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, override import strawberry from strawberry import Info @@ -320,6 +320,7 @@ async def users( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/extensions/exception_handler.py b/src/ai/backend/manager/api/gql/extensions/exception_handler.py index 4ed63a716fa..7f2ad91b9f1 100644 --- a/src/ai/backend/manager/api/gql/extensions/exception_handler.py +++ b/src/ai/backend/manager/api/gql/extensions/exception_handler.py @@ -3,7 +3,7 @@ import asyncio import logging from collections.abc import Awaitable, Callable -from typing import Any, Final +from typing import Any, Final, override from graphql import GraphQLError, GraphQLResolveInfo from strawberry.extensions.base_extension import SchemaExtension @@ -18,6 +18,7 @@ class GQLExceptionHandlerExtension(SchemaExtension): """Transforms internal exceptions into client-safe GraphQL errors with error codes.""" + @override def resolve( self, _next: Callable[..., Any], diff --git a/src/ai/backend/manager/api/gql/extensions/logging_ext.py b/src/ai/backend/manager/api/gql/extensions/logging_ext.py index a0198a1e9d4..1fed88ae103 100644 --- a/src/ai/backend/manager/api/gql/extensions/logging_ext.py +++ b/src/ai/backend/manager/api/gql/extensions/logging_ext.py @@ -2,7 +2,7 @@ import logging from collections.abc import Callable -from typing import Any, Final +from typing import Any, Final, override from graphql import GraphQLResolveInfo from strawberry.extensions.base_extension import SchemaExtension @@ -16,6 +16,7 @@ class GQLLoggingExtension(SchemaExtension): """Logs GraphQL operation details for audit purposes.""" + @override def resolve( self, _next: Callable[..., Any], diff --git a/src/ai/backend/manager/api/gql/extensions/metric.py b/src/ai/backend/manager/api/gql/extensions/metric.py index f226b5dd9c0..d3da9b303ca 100644 --- a/src/ai/backend/manager/api/gql/extensions/metric.py +++ b/src/ai/backend/manager/api/gql/extensions/metric.py @@ -3,7 +3,7 @@ import asyncio import time from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from graphql import get_named_type, is_leaf_type from opentelemetry import trace @@ -20,6 +20,7 @@ class GQLMetricExtension(SchemaExtension): """Records per-field GraphQL metrics with OpenTelemetry tracing.""" + @override def resolve( self, _next: Callable[..., AwaitableOrValue[object]], diff --git a/src/ai/backend/manager/api/gql/extensions/validation.py b/src/ai/backend/manager/api/gql/extensions/validation.py index 8d9c32a59d7..7cbafe4551f 100644 --- a/src/ai/backend/manager/api/gql/extensions/validation.py +++ b/src/ai/backend/manager/api/gql/extensions/validation.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator +from typing import override from graphql import ValidationRule from strawberry.extensions.base_extension import SchemaExtension @@ -17,6 +18,7 @@ class GQLValidationExtension(SchemaExtension): which validation rules to add at request time. """ + @override def on_validate(self) -> Iterator[None]: ctx: StrawberryGQLContext = self.execution_context.context config = ctx.config_provider.config diff --git a/src/ai/backend/manager/api/gql/huggingface_registry.py b/src/ai/backend/manager/api/gql/huggingface_registry.py index 1b1c2e2dba4..72437440bf1 100644 --- a/src/ai/backend/manager/api/gql/huggingface_registry.py +++ b/src/ai/backend/manager/api/gql/huggingface_registry.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterable, Sequence -from typing import Self, cast +from typing import Self, cast, override from uuid import UUID import strawberry @@ -62,6 +62,7 @@ class HuggingFaceRegistry(PydanticNodeMixin[HuggingFaceRegistryNode]): token: str | None @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/image/types.py b/src/ai/backend/manager/api/gql/image/types.py index 526e8d16998..2ef4c0f69ac 100644 --- a/src/ai/backend/manager/api/gql/image/types.py +++ b/src/ai/backend/manager/api/gql/image/types.py @@ -11,7 +11,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import Any, Self, cast +from typing import Any, Self, cast, override import strawberry from strawberry import Info @@ -155,6 +155,7 @@ class ImageV2AliasGQL(PydanticNodeMixin[ImageAliasNode]): alias: str = gql_field(description="The alias string for the image.") @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, @@ -335,6 +336,7 @@ async def aliases( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/image_federation.py b/src/ai/backend/manager/api/gql/image_federation.py index a2fd141d589..fe65d8fa3f0 100644 --- a/src/ai/backend/manager/api/gql/image_federation.py +++ b/src/ai/backend/manager/api/gql/image_federation.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -29,6 +30,7 @@ class Image(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/kernel/types.py b/src/ai/backend/manager/api/gql/kernel/types.py index fb82af22e2a..a4e672c56da 100644 --- a/src/ai/backend/manager/api/gql/kernel/types.py +++ b/src/ai/backend/manager/api/gql/kernel/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -529,6 +529,7 @@ async def resource_allocations( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/legacy_node_stubs.py b/src/ai/backend/manager/api/gql/legacy_node_stubs.py index fcde1a6b8ce..f4fc3d54c10 100644 --- a/src/ai/backend/manager/api/gql/legacy_node_stubs.py +++ b/src/ai/backend/manager/api/gql/legacy_node_stubs.py @@ -15,6 +15,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -36,6 +37,7 @@ class AgentNodeStub(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, info: Info, node_ids: Iterable[str], required: bool = False ) -> list[AgentNodeStub]: @@ -59,6 +61,7 @@ class NetworkNodeStub(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, info: Info, node_ids: Iterable[str], required: bool = False ) -> list[NetworkNodeStub]: @@ -82,6 +85,7 @@ class ModelCardStub(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, info: Info, node_ids: Iterable[str], required: bool = False ) -> list[ModelCardStub]: @@ -105,6 +109,7 @@ class ContainerRegistryNodeStub(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, info: Info, node_ids: Iterable[str], required: bool = False ) -> list[ContainerRegistryNodeStub]: diff --git a/src/ai/backend/manager/api/gql/notification/types.py b/src/ai/backend/manager/api/gql/notification/types.py index 825a9ba7e19..d8765d4fa4f 100644 --- a/src/ai/backend/manager/api/gql/notification/types.py +++ b/src/ai/backend/manager/api/gql/notification/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import Self +from typing import Self, override from uuid import UUID import strawberry @@ -215,6 +215,7 @@ class NotificationChannel(PydanticNodeMixin[NotificationChannelNode]): created_at: datetime @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, @@ -239,6 +240,7 @@ class NotificationRule(PydanticNodeMixin[NotificationRuleNode]): created_at: datetime @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/object_storage.py b/src/ai/backend/manager/api/gql/object_storage.py index f55c57d6ca6..1dc35287e87 100644 --- a/src/ai/backend/manager/api/gql/object_storage.py +++ b/src/ai/backend/manager/api/gql/object_storage.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Self, cast +from typing import Self, cast, override from uuid import UUID import strawberry @@ -81,6 +81,7 @@ class ObjectStorage(PydanticNodeMixin[ObjectStorageNode]): region: str @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/project.py b/src/ai/backend/manager/api/gql/project.py index b710c294527..446b7d5e8f8 100644 --- a/src/ai/backend/manager/api/gql/project.py +++ b/src/ai/backend/manager/api/gql/project.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -21,6 +22,7 @@ class Project(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/project_v2/types/node.py b/src/ai/backend/manager/api/gql/project_v2/types/node.py index 3ab37518157..e942fc7d6ef 100644 --- a/src/ai/backend/manager/api/gql/project_v2/types/node.py +++ b/src/ai/backend/manager/api/gql/project_v2/types/node.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -281,6 +281,7 @@ async def users( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/pydantic_compat.py b/src/ai/backend/manager/api/gql/pydantic_compat.py index 727427e38dd..6f699d10797 100644 --- a/src/ai/backend/manager/api/gql/pydantic_compat.py +++ b/src/ai/backend/manager/api/gql/pydantic_compat.py @@ -40,7 +40,18 @@ class CreateFooInputGQL(PydanticInputMixin[CreateFooInputDTO]): import types from decimal import Decimal from enum import Enum -from typing import Any, ClassVar, Self, Union, cast, final, get_args, get_origin, get_type_hints +from typing import ( + Any, + ClassVar, + Self, + Union, + cast, + final, + get_args, + get_origin, + get_type_hints, + override, +) from uuid import UUID from pydantic import BaseModel @@ -230,6 +241,7 @@ class CreateFooInputGQL(PydanticInputMixin[CreateFooInputDTO]): __dto_type__: ClassVar[type] + @override def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) for base in getattr(cls, "__orig_bases__", ()): diff --git a/src/ai/backend/manager/api/gql/rbac/types/entity.py b/src/ai/backend/manager/api/gql/rbac/types/entity.py index 7039f0db823..bb2229f6450 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/entity.py +++ b/src/ai/backend/manager/api/gql/rbac/types/entity.py @@ -6,7 +6,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import Any, Self, cast +from typing import Any, Self, cast, override from strawberry import Info from strawberry.relay import Connection, Edge, NodeID @@ -128,6 +128,7 @@ async def scope( return await _load_rbac_element(info, element_type, self.scope_id) @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, diff --git a/src/ai/backend/manager/api/gql/rbac/types/permission.py b/src/ai/backend/manager/api/gql/rbac/types/permission.py index 5cb2f458758..786ebe063e0 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/permission.py +++ b/src/ai/backend/manager/api/gql/rbac/types/permission.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -159,6 +159,7 @@ class PermissionGQL(PydanticNodeMixin[PermissionNodeDTO]): created_at: datetime @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, diff --git a/src/ai/backend/manager/api/gql/rbac/types/role.py b/src/ai/backend/manager/api/gql/rbac/types/role.py index ffa404f76f8..244f9fc2687 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/role.py +++ b/src/ai/backend/manager/api/gql/rbac/types/role.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -167,6 +167,7 @@ class RoleGQL(PydanticNodeMixin[Any]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, @@ -415,6 +416,7 @@ class RoleAssignmentGQL(PydanticNodeMixin[RoleAssignmentNode]): granted_at: datetime @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, diff --git a/src/ai/backend/manager/api/gql/reservoir_registry.py b/src/ai/backend/manager/api/gql/reservoir_registry.py index 4d565ac087c..42ef41f6197 100644 --- a/src/ai/backend/manager/api/gql/reservoir_registry.py +++ b/src/ai/backend/manager/api/gql/reservoir_registry.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Sequence -from typing import Self, cast +from typing import Self, cast, override from uuid import UUID import strawberry @@ -62,6 +62,7 @@ class ReservoirRegistry(PydanticNodeMixin[ReservoirRegistryNode]): api_version: str @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/resource_group/federation.py b/src/ai/backend/manager/api/gql/resource_group/federation.py index b3144fa6a60..3b8ccafd825 100644 --- a/src/ai/backend/manager/api/gql/resource_group/federation.py +++ b/src/ai/backend/manager/api/gql/resource_group/federation.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -21,6 +22,7 @@ class ResourceGroup(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/resource_group/types.py b/src/ai/backend/manager/api/gql/resource_group/types.py index 103bbd350fd..650ad18f553 100644 --- a/src/ai/backend/manager/api/gql/resource_group/types.py +++ b/src/ai/backend/manager/api/gql/resource_group/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from decimal import Decimal from enum import StrEnum -from typing import Self +from typing import Self, override from uuid import UUID from strawberry import Info @@ -364,6 +364,7 @@ class ResourceGroupGQL(PydanticNodeMixin[ResourceGroupDetailNode]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/resource_slot/types.py b/src/ai/backend/manager/api/gql/resource_slot/types.py index f1a987ced56..a4237aabad1 100644 --- a/src/ai/backend/manager/api/gql/resource_slot/types.py +++ b/src/ai/backend/manager/api/gql/resource_slot/types.py @@ -12,7 +12,7 @@ from collections.abc import Iterable from decimal import Decimal from enum import StrEnum -from typing import Any, Self +from typing import Any, Self, override from strawberry import Info from strawberry.relay import Connection, Edge, NodeID @@ -155,6 +155,7 @@ class ResourceSlotTypeGQL(PydanticNodeMixin[Any]): rank: int = gql_field(description="Display ordering rank. Lower values appear first.") @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, @@ -266,6 +267,7 @@ class AgentResourceSlotGQL(PydanticNodeMixin[AgentResourceNodeDTO]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, @@ -378,6 +380,7 @@ class KernelResourceAllocationGQL(PydanticNodeMixin[Any]): ) @classmethod + @override async def resolve_nodes( # type: ignore[override] cls, *, diff --git a/src/ai/backend/manager/api/gql/scheduling_history/types.py b/src/ai/backend/manager/api/gql/scheduling_history/types.py index a31d1fefe6e..4138dab71be 100644 --- a/src/ai/backend/manager/api/gql/scheduling_history/types.py +++ b/src/ai/backend/manager/api/gql/scheduling_history/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Self, cast +from typing import TYPE_CHECKING, Annotated, Self, cast, override from uuid import UUID import strawberry @@ -210,6 +210,7 @@ async def session( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, @@ -268,6 +269,7 @@ async def deployment( return await info.context.data_loaders.deployment_loader.load(UUID(str(self.deployment_id))) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, @@ -336,6 +338,7 @@ async def deployment( return await info.context.data_loaders.deployment_loader.load(UUID(str(self.deployment_id))) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/schema.py b/src/ai/backend/manager/api/gql/schema.py index b85c951d41c..a89d5663324 100644 --- a/src/ai/backend/manager/api/gql/schema.py +++ b/src/ai/backend/manager/api/gql/schema.py @@ -1,4 +1,5 @@ import re +from typing import override import strawberry from graphql.pyutils.undefined import Undefined as GraphQLUndefined @@ -955,6 +956,7 @@ class Subscription: class CustomizedSchema(Schema): + @override def as_str(self) -> str: # Strawberry picks up pydantic field defaults (including SENTINEL) as GraphQL # schema field default_values. SENTINEL is not a valid GraphQL scalar value, so diff --git a/src/ai/backend/manager/api/gql/session/types.py b/src/ai/backend/manager/api/gql/session/types.py index 79d44da7b15..41e44784d91 100644 --- a/src/ai/backend/manager/api/gql/session/types.py +++ b/src/ai/backend/manager/api/gql/session/types.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from datetime import datetime from enum import StrEnum -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -468,6 +468,7 @@ async def replica( # TODO: Add `vfolder_mounts` dynamic field (VFolder connection type needed) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/session_federation.py b/src/ai/backend/manager/api/gql/session_federation.py index 7fc6c587523..d4543ff8ff2 100644 --- a/src/ai/backend/manager/api/gql/session_federation.py +++ b/src/ai/backend/manager/api/gql/session_federation.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -21,6 +22,7 @@ class Session(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/storage_namespace.py b/src/ai/backend/manager/api/gql/storage_namespace.py index a6f258ddd98..c5885dbb9cc 100644 --- a/src/ai/backend/manager/api/gql/storage_namespace.py +++ b/src/ai/backend/manager/api/gql/storage_namespace.py @@ -2,7 +2,7 @@ import uuid from collections.abc import Iterable -from typing import Any, Self, cast +from typing import Any, Self, cast, override import strawberry from strawberry import ID, Info @@ -54,6 +54,7 @@ class StorageNamespace(PydanticNodeMixin[Any]): namespace: str @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/user/types/node.py b/src/ai/backend/manager/api/gql/user/types/node.py index e95ce3d408e..b445003e6e8 100644 --- a/src/ai/backend/manager/api/gql/user/types/node.py +++ b/src/ai/backend/manager/api/gql/user/types/node.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import TYPE_CHECKING, Annotated, Any, Self, cast +from typing import TYPE_CHECKING, Annotated, Any, Self, cast, override from uuid import UUID import strawberry @@ -279,6 +279,7 @@ async def projects( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/user_federation.py b/src/ai/backend/manager/api/gql/user_federation.py index 221a3c75c13..09fc40aa90d 100644 --- a/src/ai/backend/manager/api/gql/user_federation.py +++ b/src/ai/backend/manager/api/gql/user_federation.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -21,6 +22,7 @@ class User(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/vfolder.py b/src/ai/backend/manager/api/gql/vfolder.py index b79e9d4bfbc..3c3d3f553c9 100644 --- a/src/ai/backend/manager/api/gql/vfolder.py +++ b/src/ai/backend/manager/api/gql/vfolder.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterable +from typing import override from strawberry import ID, Info, relay @@ -24,6 +25,7 @@ class VFolder(relay.Node): id: relay.NodeID[str] @classmethod + @override def resolve_nodes( cls, *, diff --git a/src/ai/backend/manager/api/gql/vfolder_v2/types/node.py b/src/ai/backend/manager/api/gql/vfolder_v2/types/node.py index 5e01aba9076..431909151b5 100644 --- a/src/ai/backend/manager/api/gql/vfolder_v2/types/node.py +++ b/src/ai/backend/manager/api/gql/vfolder_v2/types/node.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import Any, override from uuid import UUID from strawberry import Info @@ -159,6 +159,7 @@ async def model_cards( ) @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/gql/vfs_storage.py b/src/ai/backend/manager/api/gql/vfs_storage.py index a3ce1d851d5..8c81dc7f3c2 100644 --- a/src/ai/backend/manager/api/gql/vfs_storage.py +++ b/src/ai/backend/manager/api/gql/vfs_storage.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Self, cast +from typing import Self, cast, override from uuid import UUID import strawberry @@ -62,6 +62,7 @@ class VFSStorage(PydanticNodeMixin[VFSStorageNode]): base_path: str @classmethod + @override async def resolve_nodes( # type: ignore[override] # Strawberry Node uses AwaitableOrValue overloads incompatible with async def cls, *, diff --git a/src/ai/backend/manager/api/rest/manager/handler.py b/src/ai/backend/manager/api/rest/manager/handler.py index 5adc7ffd088..74e3d338d22 100644 --- a/src/ai/backend/manager/api/rest/manager/handler.py +++ b/src/ai/backend/manager/api/rest/manager/handler.py @@ -14,7 +14,7 @@ import logging import textwrap from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, override import trafaret as t from aiohttp import web @@ -67,48 +67,58 @@ def __init__(self, node_id: str, pid: int, val: int) -> None: self.mgr_id = f"{node_id}:{pid}" self.val = val + @override def metric_value_string(self, metric_name: str, _primitive: PromMetricPrimitive) -> str: return f"""{metric_name}{{mgr_id="{self.mgr_id}"}} {self.val}""" class SQLAlchemyTotalConnectionMetricGroup(PromMetricGroup[SQLAlchemyConnectionMetric]): @property + @override def metric_name(self) -> str: return "sqlalchemy_total_connection" @property + @override def description(self) -> str: return "The number of total connections in SQLAlchemy connection pool." @property + @override def metric_primitive(self) -> PromMetricPrimitive: return PromMetricPrimitive.gauge class SQLAlchemyOpenConnectionMetricGroup(PromMetricGroup[SQLAlchemyConnectionMetric]): @property + @override def metric_name(self) -> str: return "sqlalchemy_open_connection" @property + @override def description(self) -> str: return "The number of open connections in SQLAlchemy connection pool." @property + @override def metric_primitive(self) -> PromMetricPrimitive: return PromMetricPrimitive.gauge class SQLAlchemyClosedConnectionMetricGroup(PromMetricGroup[SQLAlchemyConnectionMetric]): @property + @override def metric_name(self) -> str: return "sqlalchemy_closed_connection" @property + @override def description(self) -> str: return "The number of closed connections in SQLAlchemy connection pool." @property + @override def metric_primitive(self) -> PromMetricPrimitive: return PromMetricPrimitive.gauge @@ -119,6 +129,7 @@ def __init__(self, node_id: str, pid: int, redis_obj_name: str, val: int) -> Non self.mgr_id = f"{node_id}:{pid}" self.val = val + @override def metric_value_string(self, metric_name: str, _primitive: PromMetricPrimitive) -> str: return ( f"""{metric_name}{{mgr_id="{self.mgr_id}",name="{self.redis_obj_name}"}} {self.val}""" @@ -127,14 +138,17 @@ def metric_value_string(self, metric_name: str, _primitive: PromMetricPrimitive) class RedisConnectionMetricGroup(PromMetricGroup[RedisConnectionMetric]): @property + @override def metric_name(self) -> str: return "redis_connection" @property + @override def description(self) -> str: return "The number of connections in Redis Client's connection pool." @property + @override def metric_primitive(self) -> PromMetricPrimitive: return PromMetricPrimitive.gauge diff --git a/src/ai/backend/manager/api/utils.py b/src/ai/backend/manager/api/utils.py index c427776b42a..96b6573d3f1 100644 --- a/src/ai/backend/manager/api/utils.py +++ b/src/ai/backend/manager/api/utils.py @@ -30,6 +30,7 @@ TypeAlias, TypeVar, cast, + override, ) import trafaret as t @@ -291,6 +292,7 @@ def __float__(self) -> float: def __int__(self) -> int: return 0xFFFF_FFFF_FFFF_FFFF # a practical 64-bit maximum + @override def __hash__(self) -> int: return hash(self) diff --git a/src/ai/backend/manager/api/wsproxy.py b/src/ai/backend/manager/api/wsproxy.py index 206f95eedfb..bdf33a3d037 100644 --- a/src/ai/backend/manager/api/wsproxy.py +++ b/src/ai/backend/manager/api/wsproxy.py @@ -8,7 +8,7 @@ import logging from abc import ABCMeta, abstractmethod from collections.abc import Callable, Coroutine -from typing import Any, TypeVar +from typing import Any, TypeVar, override import aiohttp import aiotools @@ -70,6 +70,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.down_task: asyncio.Task[Any] | None = None + @override async def proxy(self) -> web.WebSocketResponse: reader: asyncio.StreamReader writer: asyncio.StreamWriter diff --git a/src/ai/backend/manager/clients/agent/client.py b/src/ai/backend/manager/clients/agent/client.py index 0cc004c2cd1..41ef8e6e086 100644 --- a/src/ai/backend/manager/clients/agent/client.py +++ b/src/ai/backend/manager/clients/agent/client.py @@ -2,7 +2,7 @@ import logging from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast, override from pydantic import ConfigDict @@ -82,6 +82,7 @@ def __init__( def agent_id(self) -> AgentId: return self._agent_id + @override async def connect(self) -> None: """Establish connection to the agent. @@ -109,6 +110,7 @@ async def connect(self) -> None: ) raise + @override async def close(self) -> None: """Close connection to the agent. diff --git a/src/ai/backend/manager/clients/prometheus/querier.py b/src/ai/backend/manager/clients/prometheus/querier.py index bb616099694..8df1fedb0ac 100644 --- a/src/ai/backend/manager/clients/prometheus/querier.py +++ b/src/ai/backend/manager/clients/prometheus/querier.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass +from typing import override from uuid import UUID from ai.backend.manager.clients.prometheus.preset import LabelMatcher @@ -35,6 +36,7 @@ class ContainerMetricQuerier(MetricQuerier): user_id: UUID | None = None project_id: UUID | None = None + @override def labels(self) -> Mapping[str, LabelMatcher]: """Return the labels for the container metric query.""" result: dict[str, LabelMatcher] = { diff --git a/src/ai/backend/manager/config/bootstrap.py b/src/ai/backend/manager/config/bootstrap.py index 98e5eaf61f2..e123d4842cc 100644 --- a/src/ai/backend/manager/config/bootstrap.py +++ b/src/ai/backend/manager/config/bootstrap.py @@ -1,6 +1,6 @@ from pathlib import Path from pprint import pformat -from typing import Any, Self +from typing import Any, Self, override from pydantic import Field @@ -75,6 +75,7 @@ class BootstrapConfig(BackendAISchema): """, ) + @override def __repr__(self) -> str: return pformat(self.model_dump()) diff --git a/src/ai/backend/manager/config/loader/legacy_etcd_loader.py b/src/ai/backend/manager/config/loader/legacy_etcd_loader.py index d0451cbdc95..8ef11cc0d88 100644 --- a/src/ai/backend/manager/config/loader/legacy_etcd_loader.py +++ b/src/ai/backend/manager/config/loader/legacy_etcd_loader.py @@ -42,6 +42,7 @@ def __init__(self, etcd: AsyncEtcd, config_prefix: str | None = None) -> None: async def load(self) -> Mapping[str, Any]: return await self._etcd.get_prefix(self._config_prefix) + @override def __hash__(self) -> int: # When used as a key in dicts, we don't care our contents. # Just treat it like an opaque object. diff --git a/src/ai/backend/manager/config/unified.py b/src/ai/backend/manager/config/unified.py index 9cf77b3832a..ed28435844b 100644 --- a/src/ai/backend/manager/config/unified.py +++ b/src/ai/backend/manager/config/unified.py @@ -182,7 +182,7 @@ from ipaddress import IPv4Network from pathlib import Path from pprint import pformat -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, override from pydantic import ( AliasChoices, @@ -3725,5 +3725,6 @@ class ManagerUnifiedConfig(BaseConfigSchema): extra="allow", ) + @override def __repr__(self) -> str: return pformat(self.model_dump()) diff --git a/src/ai/backend/manager/config_legacy.py b/src/ai/backend/manager/config_legacy.py index ca3036a9940..ab15ecb52bb 100644 --- a/src/ai/backend/manager/config_legacy.py +++ b/src/ai/backend/manager/config_legacy.py @@ -13,6 +13,7 @@ from typing import ( Any, Final, + override, ) import click @@ -400,6 +401,7 @@ async def dispatch_watch_callbacks(self, updated_keys: Sequence[str]) -> None: class LocalConfig(AbstractConfig): + @override async def reload(self) -> None: raise NotImplementedError diff --git a/src/ai/backend/manager/container_registry/local.py b/src/ai/backend/manager/container_registry/local.py index 0588e39c5cc..d472296561d 100644 --- a/src/ai/backend/manager/container_registry/local.py +++ b/src/ai/backend/manager/container_registry/local.py @@ -26,6 +26,7 @@ class LocalRegistry(BaseContainerRegistry): @actxmgr + @override async def prepare_client_session(self) -> AsyncIterator[tuple[yarl.URL, aiohttp.ClientSession]]: connector = get_docker_connector() async with aiohttp.ClientSession(connector=connector.connector) as sess: diff --git a/src/ai/backend/manager/data/common/sentinel.py b/src/ai/backend/manager/data/common/sentinel.py index ed4c1aea2f2..5bce96a1b33 100644 --- a/src/ai/backend/manager/data/common/sentinel.py +++ b/src/ai/backend/manager/data/common/sentinel.py @@ -1,12 +1,13 @@ from __future__ import annotations from collections.abc import MutableMapping -from typing import Any +from typing import Any, override class Singleton(type): _instances: MutableMapping[Any, Any] = {} + @override def __call__(cls, *args: Any, **kwargs: Any) -> Any: if cls not in cls._instances: cls._instances[cls] = super().__call__(*args, **kwargs) diff --git a/src/ai/backend/manager/data/group/types.py b/src/ai/backend/manager/data/group/types.py index 48635bdce96..1639024bbae 100644 --- a/src/ai/backend/manager/data/group/types.py +++ b/src/ai/backend/manager/data/group/types.py @@ -25,6 +25,7 @@ class ProjectType(enum.StrEnum): MODEL_STORE = "model-store" @classmethod + @override def _missing_(cls, value: Any) -> ProjectType | None: if not isinstance(value, str): raise DataTransformationFailed( diff --git a/src/ai/backend/manager/dependencies/agents/agent_client_pool.py b/src/ai/backend/manager/dependencies/agents/agent_client_pool.py index 66f5c6f5ffc..81aad30caf2 100644 --- a/src/ai/backend/manager/dependencies/agents/agent_client_pool.py +++ b/src/ai/backend/manager/dependencies/agents/agent_client_pool.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.manager.agent_cache import AgentRPCCache @@ -22,10 +23,12 @@ class AgentClientPoolDependency( """Provides AgentClientPool lifecycle management.""" @property + @override def stage_name(self) -> str: return "agent-client-pool" @asynccontextmanager + @override async def provide(self, setup_input: AgentClientPoolInput) -> AsyncIterator[AgentClientPool]: """Initialize and provide an agent client pool. diff --git a/src/ai/backend/manager/dependencies/agents/appproxy_client_pool.py b/src/ai/backend/manager/dependencies/agents/appproxy_client_pool.py index 79dbfdecdb3..b796b47ff55 100644 --- a/src/ai/backend/manager/dependencies/agents/appproxy_client_pool.py +++ b/src/ai/backend/manager/dependencies/agents/appproxy_client_pool.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.manager.clients.appproxy.client import AppProxyClientPool @@ -13,10 +14,12 @@ class AppProxyClientPoolDependency( """Provides AppProxyClientPool lifecycle management.""" @property + @override def stage_name(self) -> str: return "appproxy-client-pool" @asynccontextmanager + @override async def provide(self, setup_input: None) -> AsyncIterator[AppProxyClientPool]: """Initialize and provide an app proxy client pool. diff --git a/src/ai/backend/manager/dependencies/agents/composer.py b/src/ai/backend/manager/dependencies/agents/composer.py index 6fdf42bb4b1..169af286f89 100644 --- a/src/ai/backend/manager/dependencies/agents/composer.py +++ b/src/ai/backend/manager/dependencies/agents/composer.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.events.dispatcher import EventProducer @@ -86,10 +86,12 @@ class AgentsComposer(DependencyComposer[AgentsInput, AgentsResources]): """Composes all agent-layer dependencies (Layer 4).""" @property + @override def stage_name(self) -> str: return "agents" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/agents/deployment_controller.py b/src/ai/backend/manager/dependencies/agents/deployment_controller.py index a5ce37d2a7e..16123440a63 100644 --- a/src/ai/backend/manager/dependencies/agents/deployment_controller.py +++ b/src/ai/backend/manager/dependencies/agents/deployment_controller.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.clients.valkey_client.valkey_schedule.client import ValkeyScheduleClient from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -46,10 +46,12 @@ class DeploymentControllerDependency( """Provides DeploymentController lifecycle management.""" @property + @override def stage_name(self) -> str: return "deployment-controller" @asynccontextmanager + @override async def provide( self, setup_input: DeploymentControllerInput ) -> AsyncIterator[DeploymentController]: diff --git a/src/ai/backend/manager/dependencies/agents/registry.py b/src/ai/backend/manager/dependencies/agents/registry.py index 74751541b53..54b5b92c2a0 100644 --- a/src/ai/backend/manager/dependencies/agents/registry.py +++ b/src/ai/backend/manager/dependencies/agents/registry.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.auth import PublicKey, SecretKey from ai.backend.common.clients.valkey_client.valkey_image.client import ValkeyImageClient @@ -54,10 +55,12 @@ class AgentRegistryDependency( """Provides AgentRegistry lifecycle management.""" @property + @override def stage_name(self) -> str: return "agent-registry" @asynccontextmanager + @override async def provide(self, setup_input: AgentRegistryInput) -> AsyncIterator[AgentRegistry]: """Initialize and provide an agent registry. diff --git a/src/ai/backend/manager/dependencies/agents/route_controller.py b/src/ai/backend/manager/dependencies/agents/route_controller.py index 0beb4fd6fdf..9ab1128e131 100644 --- a/src/ai/backend/manager/dependencies/agents/route_controller.py +++ b/src/ai/backend/manager/dependencies/agents/route_controller.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_schedule.client import ValkeyScheduleClient from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -25,10 +26,12 @@ class RouteControllerDependency( """Provides RouteController lifecycle management.""" @property + @override def stage_name(self) -> str: return "route-controller" @asynccontextmanager + @override async def provide(self, setup_input: RouteControllerInput) -> AsyncIterator[RouteController]: """Initialize and provide a route controller. diff --git a/src/ai/backend/manager/dependencies/agents/scheduling_controller.py b/src/ai/backend/manager/dependencies/agents/scheduling_controller.py index 1e5d6de0df9..8103bbcbe01 100644 --- a/src/ai/backend/manager/dependencies/agents/scheduling_controller.py +++ b/src/ai/backend/manager/dependencies/agents/scheduling_controller.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_schedule.client import ValkeyScheduleClient from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -37,10 +38,12 @@ class SchedulingControllerDependency( """Provides SchedulingController lifecycle management.""" @property + @override def stage_name(self) -> str: return "scheduling-controller" @asynccontextmanager + @override async def provide( self, setup_input: SchedulingControllerInput ) -> AsyncIterator[SchedulingController]: diff --git a/src/ai/backend/manager/dependencies/bootstrap/composer.py b/src/ai/backend/manager/dependencies/bootstrap/composer.py index 9a6ec493c92..d27685e6c0e 100644 --- a/src/ai/backend/manager/dependencies/bootstrap/composer.py +++ b/src/ai/backend/manager/dependencies/bootstrap/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.etcd import AsyncEtcd @@ -47,10 +48,12 @@ class BootstrapComposer(DependencyComposer[BootstrapInput, BootstrapResources]): """ @property + @override def stage_name(self) -> str: return "bootstrap" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/bootstrap/config.py b/src/ai/backend/manager/dependencies/bootstrap/config.py index f53930e66c6..6df370c7479 100644 --- a/src/ai/backend/manager/dependencies/bootstrap/config.py +++ b/src/ai/backend/manager/dependencies/bootstrap/config.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.logging.types import LogLevel @@ -30,10 +31,12 @@ class BootstrapConfigDependency( """ @property + @override def stage_name(self) -> str: return "bootstrap-config" @asynccontextmanager + @override async def provide(self, setup_input: BootstrapConfigInput) -> AsyncIterator[BootstrapConfig]: """Load and provide bootstrap configuration. diff --git a/src/ai/backend/manager/dependencies/bootstrap/etcd.py b/src/ai/backend/manager/dependencies/bootstrap/etcd.py index 441d676783d..124020f3e6a 100644 --- a/src/ai/backend/manager/dependencies/bootstrap/etcd.py +++ b/src/ai/backend/manager/dependencies/bootstrap/etcd.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.health_checker import ServiceHealthChecker @@ -15,10 +16,12 @@ class EtcdDependency(BootstrapDependency[AsyncEtcd]): """Provides etcd client lifecycle management for bootstrap stage.""" @property + @override def stage_name(self) -> str: return "etcd" @asynccontextmanager + @override async def provide(self, setup_input: BootstrapConfig) -> AsyncIterator[AsyncEtcd]: """Initialize and provide an etcd client. @@ -31,6 +34,7 @@ async def provide(self, setup_input: BootstrapConfig) -> AsyncIterator[AsyncEtcd async with AsyncEtcd.create_from_config(setup_input.etcd.to_dataclass()) as etcd: yield etcd + @override def gen_liveness_checker(self, resource: AsyncEtcd) -> ServiceHealthChecker: """Liveness — stuck etcd connection observed; restart is the recovery path.""" return EtcdHealthChecker(etcd=resource) diff --git a/src/ai/backend/manager/dependencies/components/agent_cache.py b/src/ai/backend/manager/dependencies/components/agent_cache.py index 34aed5b2c89..eb4fc292761 100644 --- a/src/ai/backend/manager/dependencies/components/agent_cache.py +++ b/src/ai/backend/manager/dependencies/components/agent_cache.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from zmq.auth.certs import load_certificate @@ -29,10 +30,12 @@ class AgentCacheDependency(NonMonitorableDependencyProvider[AgentCacheInput, Age """Provides AgentRPCCache lifecycle management.""" @property + @override def stage_name(self) -> str: return "agent-cache" @asynccontextmanager + @override async def provide(self, setup_input: AgentCacheInput) -> AsyncIterator[AgentRPCCache]: """Initialize and provide agent RPC cache. diff --git a/src/ai/backend/manager/dependencies/components/composer.py b/src/ai/backend/manager/dependencies/components/composer.py index 13cba3ea18b..4c9ca935348 100644 --- a/src/ai/backend/manager/dependencies/components/composer.py +++ b/src/ai/backend/manager/dependencies/components/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.etcd import AsyncEtcd @@ -47,10 +48,12 @@ class ComponentsComposer(DependencyComposer[ComponentsInput, ComponentsResources """ @property + @override def stage_name(self) -> str: return "components" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/components/storage.py b/src/ai/backend/manager/dependencies/components/storage.py index 27d90bb6557..d03f008ea14 100644 --- a/src/ai/backend/manager/dependencies/components/storage.py +++ b/src/ai/backend/manager/dependencies/components/storage.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.manager.config.unified import ManagerUnifiedConfig from ai.backend.manager.models.storage import StorageSessionManager @@ -13,10 +14,12 @@ class StorageManagerDependency(ComponentDependency[StorageSessionManager]): """Provides StorageSessionManager lifecycle management.""" @property + @override def stage_name(self) -> str: return "storage-manager" @asynccontextmanager + @override async def provide( self, setup_input: ManagerUnifiedConfig ) -> AsyncIterator[StorageSessionManager]: diff --git a/src/ai/backend/manager/dependencies/composer.py b/src/ai/backend/manager/dependencies/composer.py index 2db38b5b565..4f13cbb6a31 100644 --- a/src/ai/backend/manager/dependencies/composer.py +++ b/src/ai/backend/manager/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.logging.types import LogLevel @@ -101,10 +102,12 @@ class ManagerDependencyComposer(DependencyComposer[DependencyInput, DependencyRe """ @property + @override def stage_name(self) -> str: return "manager-dependencies" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/config/provider.py b/src/ai/backend/manager/dependencies/config/provider.py index 5d855c22250..f6a4d71a3e8 100644 --- a/src/ai/backend/manager/dependencies/config/provider.py +++ b/src/ai/backend/manager/dependencies/config/provider.py @@ -4,7 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, override from ai.backend.common.configs.loader import ( ConfigOverrider, @@ -44,10 +44,12 @@ class ConfigProviderDependency( """ @property + @override def stage_name(self) -> str: return "config-provider" @asynccontextmanager + @override async def provide( self, setup_input: ConfigProviderInput ) -> AsyncIterator[ManagerConfigProvider]: diff --git a/src/ai/backend/manager/dependencies/domain/composer.py b/src/ai/backend/manager/dependencies/domain/composer.py index 006fd84450b..a700fb9bf6a 100644 --- a/src/ai/backend/manager/dependencies/domain/composer.py +++ b/src/ai/backend/manager/dependencies/domain/composer.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.manager.notification.notification_center import NotificationCenter @@ -69,10 +69,12 @@ class DomainComposer(DependencyComposer[DomainInput, DomainResources]): """ @property + @override def stage_name(self) -> str: return "domain" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/domain/distributed_lock.py b/src/ai/backend/manager/dependencies/domain/distributed_lock.py index 9bbad4ec416..2b620889e19 100644 --- a/src/ai/backend/manager/dependencies/domain/distributed_lock.py +++ b/src/ai/backend/manager/dependencies/domain/distributed_lock.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common import redis_helper from ai.backend.common.defs import REDIS_STREAM_LOCK, RedisRole @@ -91,10 +91,12 @@ class DistributedLockFactoryDependency( """ @property + @override def stage_name(self) -> str: return "distributed-lock-factory" @asynccontextmanager + @override async def provide( self, setup_input: DistributedLockInput ) -> AsyncIterator[DistributedLockFactory]: diff --git a/src/ai/backend/manager/dependencies/domain/notification.py b/src/ai/backend/manager/dependencies/domain/notification.py index e3c6ddeccdd..df416d261f1 100644 --- a/src/ai/backend/manager/dependencies/domain/notification.py +++ b/src/ai/backend/manager/dependencies/domain/notification.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.manager.notification.notification_center import NotificationCenter @@ -16,10 +17,12 @@ class NotificationCenterDependency(DomainDependency[None, NotificationCenter]): """ @property + @override def stage_name(self) -> str: return "notification-center" @asynccontextmanager + @override async def provide(self, setup_input: None) -> AsyncIterator[NotificationCenter]: """Initialize and provide a NotificationCenter. diff --git a/src/ai/backend/manager/dependencies/domain/repositories.py b/src/ai/backend/manager/dependencies/domain/repositories.py index 5287deb82ce..691d39c65fd 100644 --- a/src/ai/backend/manager/dependencies/domain/repositories.py +++ b/src/ai/backend/manager/dependencies/domain/repositories.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.manager.repositories.ops import DBOpsProvider from ai.backend.manager.repositories.repositories import Repositories @@ -47,10 +47,12 @@ class RepositoriesDependency(DomainDependency[RepositoriesInput, Repositories]): """ @property + @override def stage_name(self) -> str: return "repositories" @asynccontextmanager + @override async def provide(self, setup_input: RepositoriesInput) -> AsyncIterator[Repositories]: """Initialize and provide Repositories. diff --git a/src/ai/backend/manager/dependencies/domain/services.py b/src/ai/backend/manager/dependencies/domain/services.py index 4b6d4ac5c36..8b8fc324f25 100644 --- a/src/ai/backend/manager/dependencies/domain/services.py +++ b/src/ai/backend/manager/dependencies/domain/services.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.manager.repositories.container_registry_quota.repository import ( PerProjectRegistryQuotaRepository, @@ -34,10 +34,12 @@ class ServicesContextDependency(DomainDependency[ServicesInput, ServicesContext] """ @property + @override def stage_name(self) -> str: return "services-context" @asynccontextmanager + @override async def provide(self, setup_input: ServicesInput) -> AsyncIterator[ServicesContext]: """Initialize and provide a ServicesContext. diff --git a/src/ai/backend/manager/dependencies/errors.py b/src/ai/backend/manager/dependencies/errors.py index 0508f3436e3..745ab6164a8 100644 --- a/src/ai/backend/manager/dependencies/errors.py +++ b/src/ai/backend/manager/dependencies/errors.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class InvalidManagerKeypairError(BackendAIError): error_type = "https://api.backend.ai/probs/dependency/invalid-manager-keypair" error_title = "Invalid manager keypair configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -33,6 +36,7 @@ class DependencyInitializationError(BackendAIError): error_type = "https://api.backend.ai/probs/dependency/initialization-failure" error_title = "Dependency initialization failure." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, diff --git a/src/ai/backend/manager/dependencies/infrastructure/composer.py b/src/ai/backend/manager/dependencies/infrastructure/composer.py index 05163cd21d4..994bf1d94fe 100644 --- a/src/ai/backend/manager/dependencies/infrastructure/composer.py +++ b/src/ai/backend/manager/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.etcd import AsyncEtcd @@ -46,10 +47,12 @@ class InfrastructureComposer(DependencyComposer[InfrastructureInput, Infrastruct """ @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/infrastructure/database.py b/src/ai/backend/manager/dependencies/infrastructure/database.py index 60940aa8ac8..577b2f85b43 100644 --- a/src/ai/backend/manager/dependencies/infrastructure/database.py +++ b/src/ai/backend/manager/dependencies/infrastructure/database.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.health_checker import ServiceHealthChecker from ai.backend.manager.config.unified import ManagerUnifiedConfig @@ -15,10 +16,12 @@ class DatabaseDependency(InfrastructureDependency[ExtendedAsyncSAEngine]): """Provides database connection lifecycle management.""" @property + @override def stage_name(self) -> str: return "database" @asynccontextmanager + @override async def provide( self, setup_input: ManagerUnifiedConfig ) -> AsyncIterator[ExtendedAsyncSAEngine]: @@ -33,6 +36,7 @@ async def provide( async with connect_database(setup_input.db) as db: yield db + @override def gen_readiness_checker( self, resource: ExtendedAsyncSAEngine, diff --git a/src/ai/backend/manager/dependencies/infrastructure/redis.py b/src/ai/backend/manager/dependencies/infrastructure/redis.py index 975d1810a84..8f61960f4ad 100644 --- a/src/ai/backend/manager/dependencies/infrastructure/redis.py +++ b/src/ai/backend/manager/dependencies/infrastructure/redis.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_artifact.client import ( ValkeyArtifactDownloadTrackingClient, @@ -79,10 +80,12 @@ class ValkeyDependency(InfrastructureDependency[ValkeyClients]): """Provides lifecycle management for 9 specialized Valkey clients.""" @property + @override def stage_name(self) -> str: return "valkey" @asynccontextmanager + @override async def provide(self, setup_input: ManagerUnifiedConfig) -> AsyncIterator[ValkeyClients]: """Initialize and provide all Valkey clients. @@ -156,6 +159,7 @@ async def provide(self, setup_input: ManagerUnifiedConfig) -> AsyncIterator[Valk finally: await clients.close() + @override def gen_liveness_checker( self, resource: ValkeyClients, diff --git a/src/ai/backend/manager/dependencies/messaging/composer.py b/src/ai/backend/manager/dependencies/messaging/composer.py index 276c3112c43..a04f2b102ac 100644 --- a/src/ai/backend/manager/dependencies/messaging/composer.py +++ b/src/ai/backend/manager/dependencies/messaging/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.events.dispatcher import EventProducer @@ -51,10 +52,12 @@ class MessagingComposer(DependencyComposer[MessagingInput, MessagingResources]): """ @property + @override def stage_name(self) -> str: return "messaging" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/messaging/event_fetcher.py b/src/ai/backend/manager/dependencies/messaging/event_fetcher.py index 6dbb0ff5eae..678937929d1 100644 --- a/src/ai/backend/manager/dependencies/messaging/event_fetcher.py +++ b/src/ai/backend/manager/dependencies/messaging/event_fetcher.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.fetcher import EventFetcher @@ -12,10 +13,12 @@ class EventFetcherDependency(NonMonitorableDependencyProvider[AbstractMessageQue """Provides EventFetcher lifecycle management.""" @property + @override def stage_name(self) -> str: return "event-fetcher" @asynccontextmanager + @override async def provide(self, setup_input: AbstractMessageQueue) -> AsyncIterator[EventFetcher]: """Initialize and provide the event fetcher. diff --git a/src/ai/backend/manager/dependencies/messaging/event_hub.py b/src/ai/backend/manager/dependencies/messaging/event_hub.py index fcca2cbdaf8..dd289cbe13a 100644 --- a/src/ai/backend/manager/dependencies/messaging/event_hub.py +++ b/src/ai/backend/manager/dependencies/messaging/event_hub.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.hub.hub import EventHub @@ -11,10 +12,12 @@ class EventHubDependency(NonMonitorableDependencyProvider[None, EventHub]): """Provides EventHub lifecycle management.""" @property + @override def stage_name(self) -> str: return "event-hub" @asynccontextmanager + @override async def provide(self, setup_input: None) -> AsyncIterator[EventHub]: """Initialize and provide the event hub. diff --git a/src/ai/backend/manager/dependencies/messaging/event_producer.py b/src/ai/backend/manager/dependencies/messaging/event_producer.py index 2e6a8826fec..5cef8e6774c 100644 --- a/src/ai/backend/manager/dependencies/messaging/event_producer.py +++ b/src/ai/backend/manager/dependencies/messaging/event_producer.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.dispatcher import EventProducer @@ -27,10 +28,12 @@ class EventProducerDependency(NonMonitorableDependencyProvider[EventProducerInpu """Provides EventProducer lifecycle management.""" @property + @override def stage_name(self) -> str: return "event-producer" @asynccontextmanager + @override async def provide(self, setup_input: EventProducerInput) -> AsyncIterator[EventProducer]: """Initialize and provide the event producer. diff --git a/src/ai/backend/manager/dependencies/messaging/message_queue.py b/src/ai/backend/manager/dependencies/messaging/message_queue.py index 87a374f46e0..4390e8538e6 100644 --- a/src/ai/backend/manager/dependencies/messaging/message_queue.py +++ b/src/ai/backend/manager/dependencies/messaging/message_queue.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Final +from typing import Final, override from ai.backend.common.defs import REDIS_STREAM_DB, RedisRole from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -31,10 +31,12 @@ class MessageQueueDependency( """Provides AbstractMessageQueue lifecycle management.""" @property + @override def stage_name(self) -> str: return "message-queue" @asynccontextmanager + @override async def provide(self, setup_input: MessageQueueInput) -> AsyncIterator[AbstractMessageQueue]: """Initialize and provide the message queue. diff --git a/src/ai/backend/manager/dependencies/orchestration/composer.py b/src/ai/backend/manager/dependencies/orchestration/composer.py index acfb2f3c342..78599ccc104 100644 --- a/src/ai/backend/manager/dependencies/orchestration/composer.py +++ b/src/ai/backend/manager/dependencies/orchestration/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_schedule import ValkeyScheduleClient from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient @@ -100,10 +101,12 @@ class OrchestrationComposer(DependencyComposer[OrchestrationInput, Orchestration """ @property + @override def stage_name(self) -> str: return "orchestration" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/orchestration/idle_checker.py b/src/ai/backend/manager/dependencies/orchestration/idle_checker.py index b53a4c8808b..e98f2b2583c 100644 --- a/src/ai/backend/manager/dependencies/orchestration/idle_checker.py +++ b/src/ai/backend/manager/dependencies/orchestration/idle_checker.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.dispatcher import EventProducer @@ -32,10 +33,12 @@ class IdleCheckerHostDependency( """ @property + @override def stage_name(self) -> str: return "idle-checker-host" @asynccontextmanager + @override async def provide(self, setup_input: IdleCheckerInput) -> AsyncIterator[IdleCheckerHost]: """Initialize and provide idle checker host. diff --git a/src/ai/backend/manager/dependencies/orchestration/leader_election.py b/src/ai/backend/manager/dependencies/orchestration/leader_election.py index a8ee0db2c01..b950d27d96d 100644 --- a/src/ai/backend/manager/dependencies/orchestration/leader_election.py +++ b/src/ai/backend/manager/dependencies/orchestration/leader_election.py @@ -5,6 +5,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_leader.client import ValkeyLeaderClient from ai.backend.common.defs import REDIS_STREAM_LOCK, RedisRole @@ -47,10 +48,12 @@ class LeaderElectionDependency( """ @property + @override def stage_name(self) -> str: return "leader-election" @asynccontextmanager + @override async def provide( self, setup_input: LeaderElectionInput ) -> AsyncIterator[ValkeyLeaderElection]: diff --git a/src/ai/backend/manager/dependencies/orchestration/sokovan.py b/src/ai/backend/manager/dependencies/orchestration/sokovan.py index 8e49ff3f7c9..e12ccdec612 100644 --- a/src/ai/backend/manager/dependencies/orchestration/sokovan.py +++ b/src/ai/backend/manager/dependencies/orchestration/sokovan.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.http_client.client_pool import ( ClientPool, @@ -97,10 +98,12 @@ class SokovanOrchestratorDependency( """ @property + @override def stage_name(self) -> str: return "sokovan-orchestrator" @asynccontextmanager + @override async def provide( self, setup_input: SokovanOrchestratorInput ) -> AsyncIterator[SokovanOrchestrator]: diff --git a/src/ai/backend/manager/dependencies/plugins/composer.py b/src/ai/backend/manager/dependencies/plugins/composer.py index 7a3c0bf5dcb..22144335b8d 100644 --- a/src/ai/backend/manager/dependencies/plugins/composer.py +++ b/src/ai/backend/manager/dependencies/plugins/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.plugin.event import EventDispatcherPluginContext @@ -38,10 +39,12 @@ class PluginsComposer(DependencyComposer[PluginsInput, PluginsResources]): """ @property + @override def stage_name(self) -> str: return "plugins" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/plugins/event_dispatcher.py b/src/ai/backend/manager/dependencies/plugins/event_dispatcher.py index e6380ec7830..13ba8a7497a 100644 --- a/src/ai/backend/manager/dependencies/plugins/event_dispatcher.py +++ b/src/ai/backend/manager/dependencies/plugins/event_dispatcher.py @@ -3,6 +3,7 @@ import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.plugin.event import EventDispatcherPluginContext from ai.backend.logging import BraceStyleAdapter @@ -16,10 +17,12 @@ class EventDispatcherPluginDependency(PluginDependency[EventDispatcherPluginCont """Provides EventDispatcherPluginContext lifecycle management.""" @property + @override def stage_name(self) -> str: return "event-dispatcher-plugin" @asynccontextmanager + @override async def provide( self, setup_input: PluginsInput ) -> AsyncIterator[EventDispatcherPluginContext]: diff --git a/src/ai/backend/manager/dependencies/plugins/hook.py b/src/ai/backend/manager/dependencies/plugins/hook.py index ab77c464284..bf1f1cd791f 100644 --- a/src/ai/backend/manager/dependencies/plugins/hook.py +++ b/src/ai/backend/manager/dependencies/plugins/hook.py @@ -3,6 +3,7 @@ import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.plugin.hook import ALL_COMPLETED, PASSED, HookPluginContext from ai.backend.logging import BraceStyleAdapter @@ -20,10 +21,12 @@ class HookPluginDependency(PluginDependency[HookPluginContext]): """ @property + @override def stage_name(self) -> str: return "hook-plugin" @asynccontextmanager + @override async def provide(self, setup_input: PluginsInput) -> AsyncIterator[HookPluginContext]: """Initialize and provide a HookPluginContext. diff --git a/src/ai/backend/manager/dependencies/plugins/monitoring.py b/src/ai/backend/manager/dependencies/plugins/monitoring.py index c1aebe6e4fa..70237ebd6ad 100644 --- a/src/ai/backend/manager/dependencies/plugins/monitoring.py +++ b/src/ai/backend/manager/dependencies/plugins/monitoring.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from ai.backend.common.dependencies import NonMonitorableDependencyProvider, ResourceT from ai.backend.common.etcd import AsyncEtcd @@ -45,10 +45,12 @@ class ErrorMonitorDependency(MonitoringDependency[ManagerErrorPluginContext]): """ @property + @override def stage_name(self) -> str: return "error-monitor-plugin" @asynccontextmanager + @override async def provide( self, setup_input: MonitoringInput ) -> AsyncIterator[ManagerErrorPluginContext]: @@ -74,10 +76,12 @@ class StatsMonitorDependency(MonitoringDependency[ManagerStatsPluginContext]): """ @property + @override def stage_name(self) -> str: return "stats-monitor-plugin" @asynccontextmanager + @override async def provide( self, setup_input: MonitoringInput ) -> AsyncIterator[ManagerStatsPluginContext]: diff --git a/src/ai/backend/manager/dependencies/plugins/network.py b/src/ai/backend/manager/dependencies/plugins/network.py index 174691ee2b4..b786b15d032 100644 --- a/src/ai/backend/manager/dependencies/plugins/network.py +++ b/src/ai/backend/manager/dependencies/plugins/network.py @@ -3,6 +3,7 @@ import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.plugin.network import NetworkPluginContext @@ -16,10 +17,12 @@ class NetworkPluginDependency(PluginDependency[NetworkPluginContext]): """Provides NetworkPluginContext lifecycle management.""" @property + @override def stage_name(self) -> str: return "network-plugin" @asynccontextmanager + @override async def provide(self, setup_input: PluginsInput) -> AsyncIterator[NetworkPluginContext]: """Initialize and provide a NetworkPluginContext. diff --git a/src/ai/backend/manager/dependencies/processing/agent_lost_checker.py b/src/ai/backend/manager/dependencies/processing/agent_lost_checker.py index f3294a20d6d..d5e531935a6 100644 --- a/src/ai/backend/manager/dependencies/processing/agent_lost_checker.py +++ b/src/ai/backend/manager/dependencies/processing/agent_lost_checker.py @@ -9,7 +9,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.cron import LocalCron from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -36,10 +36,12 @@ class AgentLostCheckerDependency( """Provides a periodic timer that detects agent heartbeat timeouts.""" @property + @override def stage_name(self) -> str: return "agent-lost-checker" @asynccontextmanager + @override async def provide(self, setup_input: AgentLostCheckerInput) -> AsyncIterator[LocalCron]: cron = LocalCron([ AgentLostCheckerTask( diff --git a/src/ai/backend/manager/dependencies/processing/bgtask_registry.py b/src/ai/backend/manager/dependencies/processing/bgtask_registry.py index b236c437441..3496fb8fdb0 100644 --- a/src/ai/backend/manager/dependencies/processing/bgtask_registry.py +++ b/src/ai/backend/manager/dependencies/processing/bgtask_registry.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.bgtask.bgtask import BackgroundTaskManager from ai.backend.common.bgtask.task.registry import BackgroundTaskHandlerRegistry @@ -42,10 +43,12 @@ class BgtaskRegistryDependency( """ @property + @override def stage_name(self) -> str: return "bgtask-registry" @asynccontextmanager + @override async def provide( self, setup_input: BgtaskRegistryInput ) -> AsyncIterator[BackgroundTaskHandlerRegistry]: diff --git a/src/ai/backend/manager/dependencies/processing/composer.py b/src/ai/backend/manager/dependencies/processing/composer.py index ee9518c9fb5..b4c7c9dee5d 100644 --- a/src/ai/backend/manager/dependencies/processing/composer.py +++ b/src/ai/backend/manager/dependencies/processing/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.bgtask.bgtask import BackgroundTaskManager from ai.backend.common.clients.valkey_client.valkey_artifact.client import ( @@ -199,10 +200,12 @@ class ProcessingComposer(DependencyComposer[ProcessingInput, ProcessingResources """ @property + @override def stage_name(self) -> str: return "processing" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/processing/event_dispatcher.py b/src/ai/backend/manager/dependencies/processing/event_dispatcher.py index 93ab5fedf48..d907ddd04fa 100644 --- a/src/ai/backend/manager/dependencies/processing/event_dispatcher.py +++ b/src/ai/backend/manager/dependencies/processing/event_dispatcher.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.dispatcher import EventDispatcher, EventObserver @@ -28,10 +29,12 @@ class EventDispatcherDependency( """ @property + @override def stage_name(self) -> str: return "event-dispatcher" @asynccontextmanager + @override async def provide(self, setup_input: EventDispatcherInput) -> AsyncIterator[EventDispatcher]: dispatcher = EventDispatcher( setup_input.message_queue, diff --git a/src/ai/backend/manager/dependencies/processing/log_cleanup_timer.py b/src/ai/backend/manager/dependencies/processing/log_cleanup_timer.py index 45436952246..7ef11ce71b7 100644 --- a/src/ai/backend/manager/dependencies/processing/log_cleanup_timer.py +++ b/src/ai/backend/manager/dependencies/processing/log_cleanup_timer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.distributed import GlobalTimer @@ -30,10 +31,12 @@ class LogCleanupTimerDependency( """ @property + @override def stage_name(self) -> str: return "log-cleanup-timer" @asynccontextmanager + @override async def provide(self, setup_input: LogCleanupTimerInput) -> AsyncIterator[GlobalTimer]: timer = GlobalTimer( setup_input.distributed_lock_factory(LockID.LOCKID_LOG_CLEANUP_TIMER, 20.0), diff --git a/src/ai/backend/manager/dependencies/processing/manager_status_watcher.py b/src/ai/backend/manager/dependencies/processing/manager_status_watcher.py index 29f17ce7a40..a2be44c1f9a 100644 --- a/src/ai/backend/manager/dependencies/processing/manager_status_watcher.py +++ b/src/ai/backend/manager/dependencies/processing/manager_status_watcher.py @@ -10,7 +10,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.types import QueueSentinel @@ -98,10 +98,12 @@ class ManagerStatusWatcherDependency( """Provides background tasks that watch and report manager status.""" @property + @override def stage_name(self) -> str: return "manager-status-watcher" @asynccontextmanager + @override async def provide( self, setup_input: ManagerStatusWatcherInput ) -> AsyncIterator[ManagerStatusWatcherResult]: diff --git a/src/ai/backend/manager/dependencies/processing/processors.py b/src/ai/backend/manager/dependencies/processing/processors.py index 86047d96276..f134c259252 100644 --- a/src/ai/backend/manager/dependencies/processing/processors.py +++ b/src/ai/backend/manager/dependencies/processing/processors.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.events.fetcher import EventFetcher @@ -32,10 +33,12 @@ class ProcessorsDependency(NonMonitorableDependencyProvider[ProcessorsProviderIn """ @property + @override def stage_name(self) -> str: return "processors" @asynccontextmanager + @override async def provide(self, setup_input: ProcessorsProviderInput) -> AsyncIterator[Processors]: processors = create_processors( ProcessorArgs( diff --git a/src/ai/backend/manager/dependencies/processing/stats_reporter.py b/src/ai/backend/manager/dependencies/processing/stats_reporter.py index dc00c2b72fc..7a4314df6b4 100644 --- a/src/ai/backend/manager/dependencies/processing/stats_reporter.py +++ b/src/ai/backend/manager/dependencies/processing/stats_reporter.py @@ -9,7 +9,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.cron import LocalCron from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -34,10 +34,12 @@ class StatsReporterDependency(NonMonitorableDependencyProvider[StatsReporterInpu """Provides a periodic timer that reports system metrics.""" @property + @override def stage_name(self) -> str: return "stats-reporter" @asynccontextmanager + @override async def provide(self, setup_input: StatsReporterInput) -> AsyncIterator[LocalCron]: cron = LocalCron([ StatsReporterTask( diff --git a/src/ai/backend/manager/dependencies/system/background_task_manager.py b/src/ai/backend/manager/dependencies/system/background_task_manager.py index ba73f8134bd..9e1167d55b6 100644 --- a/src/ai/backend/manager/dependencies/system/background_task_manager.py +++ b/src/ai/backend/manager/dependencies/system/background_task_manager.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.bgtask.bgtask import BackgroundTaskManager, BackgroundTaskManagerArgs from ai.backend.common.bgtask.hooks.metric_hook import BackgroundTaskObserver @@ -27,10 +28,12 @@ class BackgroundTaskManagerDependency( """Provides BackgroundTaskManager with managed lifecycle.""" @property + @override def stage_name(self) -> str: return "background-task-manager" @asynccontextmanager + @override async def provide( self, setup_input: BackgroundTaskManagerInput ) -> AsyncIterator[BackgroundTaskManager]: diff --git a/src/ai/backend/manager/dependencies/system/composer.py b/src/ai/backend/manager/dependencies/system/composer.py index f2e1e031ea8..ff3bbf7a2a4 100644 --- a/src/ai/backend/manager/dependencies/system/composer.py +++ b/src/ai/backend/manager/dependencies/system/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.bgtask.bgtask import BackgroundTaskManager from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -74,10 +75,12 @@ class SystemComposer(DependencyComposer[SystemInput, SystemResources]): """ @property + @override def stage_name(self) -> str: return "system" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/manager/dependencies/system/cors_options.py b/src/ai/backend/manager/dependencies/system/cors_options.py index 5dea1203add..e28ad458ce4 100644 --- a/src/ai/backend/manager/dependencies/system/cors_options.py +++ b/src/ai/backend/manager/dependencies/system/cors_options.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override import aiohttp_cors @@ -14,10 +15,12 @@ class CORSOptionsDependency(SystemDependency[CORSOptions]): """Provides CORS options configuration.""" @property + @override def stage_name(self) -> str: return "cors-options" @asynccontextmanager + @override async def provide(self, setup_input: object) -> AsyncIterator[CORSOptions]: yield { "*": aiohttp_cors.ResourceOptions( # type: ignore[no-untyped-call] diff --git a/src/ai/backend/manager/dependencies/system/gql_adapter.py b/src/ai/backend/manager/dependencies/system/gql_adapter.py index b2a55a7fe35..f888d3b7c72 100644 --- a/src/ai/backend/manager/dependencies/system/gql_adapter.py +++ b/src/ai/backend/manager/dependencies/system/gql_adapter.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.manager.api.gql.adapter import BaseGQLAdapter @@ -12,9 +13,11 @@ class GQLAdapterDependency(SystemDependency[BaseGQLAdapter]): """Provides BaseGQLAdapter instance.""" @property + @override def stage_name(self) -> str: return "gql-adapter" @asynccontextmanager + @override async def provide(self, setup_input: object) -> AsyncIterator[BaseGQLAdapter]: yield BaseGQLAdapter() diff --git a/src/ai/backend/manager/dependencies/system/health_probe.py b/src/ai/backend/manager/dependencies/system/health_probe.py index 250778cd3c4..a0dac92dec6 100644 --- a/src/ai/backend/manager/dependencies/system/health_probe.py +++ b/src/ai/backend/manager/dependencies/system/health_probe.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyProvider from ai.backend.common.etcd import AsyncEtcd @@ -39,10 +40,12 @@ class HealthProbeDependency(DependencyProvider[HealthProbeInput, HealthProbe]): """Provides HealthProbe with health checker registration and lifecycle management.""" @property + @override def stage_name(self) -> str: return "health-probe" @asynccontextmanager + @override async def provide(self, setup_input: HealthProbeInput) -> AsyncIterator[HealthProbe]: probe = HealthProbe(options=HealthProbeOptions(check_interval=60)) diff --git a/src/ai/backend/manager/dependencies/system/jwt_validator.py b/src/ai/backend/manager/dependencies/system/jwt_validator.py index ffefca6cc69..698fe037ac2 100644 --- a/src/ai/backend/manager/dependencies/system/jwt_validator.py +++ b/src/ai/backend/manager/dependencies/system/jwt_validator.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.common.jwt.validator import JWTValidator @@ -12,10 +13,12 @@ class JWTValidatorDependency(NonMonitorableDependencyProvider[ManagerUnifiedConf """Provides JWTValidator instance from manager configuration.""" @property + @override def stage_name(self) -> str: return "jwt-validator" @asynccontextmanager + @override async def provide(self, setup_input: ManagerUnifiedConfig) -> AsyncIterator[JWTValidator]: jwt_config = setup_input.jwt.to_jwt_config() yield JWTValidator(jwt_config) diff --git a/src/ai/backend/manager/dependencies/system/metrics.py b/src/ai/backend/manager/dependencies/system/metrics.py index e0eb4d097ec..d0a0abe833c 100644 --- a/src/ai/backend/manager/dependencies/system/metrics.py +++ b/src/ai/backend/manager/dependencies/system/metrics.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.metrics.metric import CommonMetricRegistry @@ -12,9 +13,11 @@ class MetricsDependency(SystemDependency[CommonMetricRegistry]): """Provides CommonMetricRegistry singleton.""" @property + @override def stage_name(self) -> str: return "metrics" @asynccontextmanager + @override async def provide(self, setup_input: object) -> AsyncIterator[CommonMetricRegistry]: yield CommonMetricRegistry.instance() diff --git a/src/ai/backend/manager/dependencies/system/prometheus_client.py b/src/ai/backend/manager/dependencies/system/prometheus_client.py index 4a925d21259..47ea2c86012 100644 --- a/src/ai/backend/manager/dependencies/system/prometheus_client.py +++ b/src/ai/backend/manager/dependencies/system/prometheus_client.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.clients.http_client import ClientPool, tcp_client_session_factory from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -19,10 +20,12 @@ class PrometheusClientDependency( """Provides PrometheusClient with managed ClientPool lifecycle.""" @property + @override def stage_name(self) -> str: return "prometheus-client" @asynccontextmanager + @override async def provide(self, setup_input: ManagerUnifiedConfig) -> AsyncIterator[PrometheusClient]: client_pool = ClientPool(tcp_client_session_factory) client = PrometheusClient( diff --git a/src/ai/backend/manager/dependencies/system/service_discovery.py b/src/ai/backend/manager/dependencies/system/service_discovery.py index c189de0c063..26cd28419e0 100644 --- a/src/ai/backend/manager/dependencies/system/service_discovery.py +++ b/src/ai/backend/manager/dependencies/system/service_discovery.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.defs import RedisRole from ai.backend.common.dependencies import NonMonitorableDependencyProvider @@ -49,10 +50,12 @@ class ServiceDiscoveryDependency( """Provides ServiceDiscovery and ServiceDiscoveryLoop.""" @property + @override def stage_name(self) -> str: return "service-discovery" @asynccontextmanager + @override async def provide( self, setup_input: ServiceDiscoveryInput ) -> AsyncIterator[ServiceDiscoveryResources]: diff --git a/src/ai/backend/manager/errors/agent.py b/src/ai/backend/manager/errors/agent.py index 37509a65356..cf1b573129b 100644 --- a/src/ai/backend/manager/errors/agent.py +++ b/src/ai/backend/manager/errors/agent.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -27,6 +29,7 @@ def __init__(self, agent_id: AgentId, failure_reason: str) -> None: self.failure_reason = failure_reason super().__init__(f"Agent {agent_id} connection unavailable: {failure_reason}") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/manager/errors/api.py b/src/ai/backend/manager/errors/api.py index e180b07703f..1de66b163ec 100644 --- a/src/ai/backend/manager/errors/api.py +++ b/src/ai/backend/manager/errors/api.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -19,6 +21,7 @@ class UnsupportedOperation(BackendAIError): error_type = "https://api.backend.ai/probs/unsupported-operation" error_title = "This operation is not supported." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -31,6 +34,7 @@ class NotImplementedAPI(BackendAIError, web.HTTPNotImplemented): error_type = "https://api.backend.ai/probs/not-implemented" error_title = "This API is not implemented." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -43,6 +47,7 @@ class InvalidAPIParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-api-params" error_title = "Missing or invalid API parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -55,6 +60,7 @@ class GraphQLError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/graphql-error" error_title = "GraphQL-generated error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -67,6 +73,7 @@ class RateLimitExceeded(BackendAIError, web.HTTPTooManyRequests): error_type = "https://api.backend.ai/probs/rate-limit-exceeded" error_title = "You have reached your API query rate limit." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -79,6 +86,7 @@ class InvalidGraphQLParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-graphql-params" error_title = "Invalid GraphQL parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -91,6 +99,7 @@ class InvalidCursor(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-cursor" error_title = "Invalid cursor format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, diff --git a/src/ai/backend/manager/errors/appproxy.py b/src/ai/backend/manager/errors/appproxy.py index 594f98e1e33..bb93def92c7 100644 --- a/src/ai/backend/manager/errors/appproxy.py +++ b/src/ai/backend/manager/errors/appproxy.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -17,6 +19,7 @@ class AppProxyConnectionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/appproxy-connection-error" error_title = "Failed to connect to AppProxy." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.APPPROXY, @@ -31,6 +34,7 @@ class AppProxyResponseError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/appproxy-response-error" error_title = "Invalid response from AppProxy." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.APPPROXY, diff --git a/src/ai/backend/manager/errors/artifact.py b/src/ai/backend/manager/errors/artifact.py index dc1ff60b62e..2f6fd2a5946 100644 --- a/src/ai/backend/manager/errors/artifact.py +++ b/src/ai/backend/manager/errors/artifact.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class ArtifactNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-not-found" error_title = "Artifact Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -25,6 +28,7 @@ class ArtifactNotVerified(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/artifact-not-verified" error_title = "Artifact Not Verified" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -37,6 +41,7 @@ class ArtifactUpdateError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/artifact-update-failed" error_title = "Artifact Update Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -49,6 +54,7 @@ class ArtifactDeletionBadRequestError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/artifact-deletion-failed" error_title = "Artifact Deletion Bad Request" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -61,6 +67,7 @@ class ArtifactDeletionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/artifact-deletion-failed" error_title = "Artifact Deletion Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -73,6 +80,7 @@ class ArtifactAssociationCreationError(BackendAIError, web.HTTPInternalServerErr error_type = "https://api.backend.ai/probs/artifact-association-creation-failed" error_title = "Artifact Association Creation Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_ASSOCIATION, @@ -85,6 +93,7 @@ class ArtifactAssociationDeletionError(BackendAIError, web.HTTPInternalServerErr error_type = "https://api.backend.ai/probs/artifact-association-deletion-failed" error_title = "Artifact Association Deletion Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_ASSOCIATION, @@ -97,6 +106,7 @@ class ArtifactAssociationNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-association-not-found" error_title = "Artifact Association Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_ASSOCIATION, @@ -109,6 +119,7 @@ class ArtifactNotApproved(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/artifact-not-approved" error_title = "Artifact Not Approved" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -121,6 +132,7 @@ class ArtifactReadonly(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/artifact-readonly" error_title = "You cannot upload files to readonly artifact storage" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -133,6 +145,7 @@ class InvalidArtifactModifierTypeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-modifier-type" error_title = "Invalid Artifact Modifier Type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -145,6 +158,7 @@ class ArtifactRevisionNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-revision-not-found" error_title = "Artifact Revision Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -157,6 +171,7 @@ class ArtifactScanLimitExceededError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/artifact-scan-limit-exceeded" error_title = "Artifact Scan Limit Exceeded" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -169,6 +184,7 @@ class ArtifactImportBadRequestError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-bad-import-request" error_title = "Artifact Bad Import Request" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -181,6 +197,7 @@ class ArtifactImportDelegationError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/artifact-import-delegation-failed" error_title = "Artifact Import Delegation Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -193,6 +210,7 @@ class RemoteReservoirArtifactImportError(BackendAIError, web.HTTPInternalServerE error_type = "https://api.backend.ai/probs/remote-reservoir-artifact-import-error" error_title = "Remote Reservoir Artifact Import Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, diff --git a/src/ai/backend/manager/errors/artifact_registry.py b/src/ai/backend/manager/errors/artifact_registry.py index ec8a1c957fe..47c26d0e702 100644 --- a/src/ai/backend/manager/errors/artifact_registry.py +++ b/src/ai/backend/manager/errors/artifact_registry.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class ArtifactRegistryNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-registry-not-found" error_title = "Artifact Registry Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -25,6 +28,7 @@ class ArtifactRegistryBadScanRequestError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/artifact-registry-bad-scan-request" error_title = "Artifact Registry Bad Scan Request" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -37,6 +41,7 @@ class InvalidArtifactRegistryTypeError(BackendAIError, web.HTTPInternalServerErr error_type = "https://api.backend.ai/probs/invalid-artifact-registry-type" error_title = "Invalid Artifact Registry Type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -49,6 +54,7 @@ class ReservoirConnectionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/reservoir-connection-error" error_title = "Reservoir Connection Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -61,6 +67,7 @@ class RemoteReservoirScanError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/remote-reservoir-scan-error" error_title = "Remote Reservoir Scan Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, diff --git a/src/ai/backend/manager/errors/auth.py b/src/ai/backend/manager/errors/auth.py index e314809c196..84fa1c6bfc2 100644 --- a/src/ai/backend/manager/errors/auth.py +++ b/src/ai/backend/manager/errors/auth.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class InvalidCredentials(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-credentials" error_title = "Invalid credentials for authentication." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -33,6 +36,7 @@ class InsufficientPrivilege(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/insufficient-privilege" error_title = "Insufficient privilege." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -45,6 +49,7 @@ class InvalidAuthParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-auth-params" error_title = "Missing or invalid authorization parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -57,6 +62,7 @@ class AuthorizationFailed(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/auth-failed" error_title = "Credential/signature mismatch." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -69,6 +75,7 @@ class PasswordExpired(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/password-expired" error_title = "Password has expired." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -81,6 +88,7 @@ class EmailAlreadyExistsError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/email-already-exists" error_title = "Email already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -93,6 +101,7 @@ class UserCreationError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/user-creation-failed" error_title = "Failed to create user account." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -104,6 +113,7 @@ def error_code(self) -> ErrorCode: class UserNotFound(ObjectNotFound): object_name = "user" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -115,6 +125,7 @@ def error_code(self) -> ErrorCode: class AccessKeyNotFound(ObjectNotFound): object_name = "access key" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -127,6 +138,7 @@ class GroupMembershipNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/group-membership-not-found" error_title = "User is not a member of the specified group." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -139,6 +151,7 @@ class InvalidClientIPConfig(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/invalid-client-ip-config" error_title = "Invalid client IP configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -150,6 +163,7 @@ def error_code(self) -> ErrorCode: class LoginSessionNotFoundError(ObjectNotFound): object_name = "login_session" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -162,6 +176,7 @@ class LoginSessionExpiredError(BackendAIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/login-session-expired" error_title = "Login session has expired." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -174,6 +189,7 @@ class LoginBlockedError(BackendAIError, web.HTTPTooManyRequests): error_type = "https://api.backend.ai/probs/login-blocked" error_title = "Too many failed login attempts." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -185,6 +201,7 @@ def error_code(self) -> ErrorCode: class LoginClientTypeNotFound(ObjectNotFound): object_name = "login_client_type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -197,6 +214,7 @@ class LoginClientTypeConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/login-client-type-conflict" error_title = "A login client type with the same name already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, @@ -209,6 +227,7 @@ class TooManyConcurrentLoginSessions(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/active-login-session-exists" error_title = "Too many concurrent login sessions for this user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AUTH, diff --git a/src/ai/backend/manager/errors/bgtask.py b/src/ai/backend/manager/errors/bgtask.py index e2c361e601c..716afcde2bc 100644 --- a/src/ai/backend/manager/errors/bgtask.py +++ b/src/ai/backend/manager/errors/bgtask.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -19,6 +21,7 @@ class InvalidBgtaskId(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-bgtask-id" error_title = "Invalid background task ID format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BGTASK, diff --git a/src/ai/backend/manager/errors/common.py b/src/ai/backend/manager/errors/common.py index 9ad9933392d..a2d7357965b 100644 --- a/src/ai/backend/manager/errors/common.py +++ b/src/ai/backend/manager/errors/common.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, override from aiohttp import web @@ -22,6 +22,7 @@ class URLNotFound(BackendAIError, web.HTTPNotFound): # TODO: Misused now. error_type = "https://api.backend.ai/probs/url-not-found" error_title = "Unknown URL path." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -47,6 +48,7 @@ def __init__( self.error_title = f"No such {self.object_name}." super().__init__(extra_msg, extra_data, **kwargs) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -59,6 +61,7 @@ class GenericBadRequest(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -71,6 +74,7 @@ class GenericForbidden(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/generic-forbidden" error_title = "Forbidden operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -83,6 +87,7 @@ class RejectedByHook(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/rejected-by-hook" error_title = "Operation rejected by a hook plugin." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PLUGIN, @@ -104,6 +109,7 @@ class MethodNotAllowed(BackendAIError, web.HTTPMethodNotAllowed): error_type = "https://api.backend.ai/probs/method-not-allowed" error_title = "HTTP Method Not Allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -116,6 +122,7 @@ class InternalServerError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/internal-server-error" error_title = "Internal server error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -128,6 +135,7 @@ class ServerMisconfiguredError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/server-misconfigured" error_title = "Service misconfigured." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -140,6 +148,7 @@ class ServiceUnavailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/service-unavailable" error_title = "Serivce unavailable." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -152,6 +161,7 @@ class ServerFrozen(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/server-frozen" error_title = "The server is frozen due to maintenance. Please try again later." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -164,6 +174,7 @@ class Forbidden(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/forbidden" error_title = "Forbidden operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, diff --git a/src/ai/backend/manager/errors/container_registry.py b/src/ai/backend/manager/errors/container_registry.py index 5b47c4e913a..2df493f95f0 100644 --- a/src/ai/backend/manager/errors/container_registry.py +++ b/src/ai/backend/manager/errors/container_registry.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -10,6 +12,7 @@ class InvalidContainerRegistryProject(BackendAIError, web.HTTPBadRequest): + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -19,6 +22,7 @@ def error_code(self) -> ErrorCode: class InvalidContainerRegistryURL(BackendAIError, web.HTTPBadRequest): + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, diff --git a/src/ai/backend/manager/errors/deployment.py b/src/ai/backend/manager/errors/deployment.py index 6ff9b0af07c..eb936746338 100644 --- a/src/ai/backend/manager/errors/deployment.py +++ b/src/ai/backend/manager/errors/deployment.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class DefinitionFileNotFound(ObjectNotFound): object_name = "definition-file" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -24,6 +27,7 @@ def error_code(self) -> ErrorCode: class EndpointNotFound(ObjectNotFound): object_name = "endpoint" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -35,6 +39,7 @@ def error_code(self) -> ErrorCode: class DeploymentRevisionNotFound(ObjectNotFound): object_name = "deployment-revision" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -46,6 +51,7 @@ def error_code(self) -> ErrorCode: class UserNotFoundInDeployment(ObjectNotFound): object_name = "user in deployment" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -57,6 +63,7 @@ def error_code(self) -> ErrorCode: class NoActiveKeypairForDeployment(ObjectNotFound): object_name = "active keypair for deployment user" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -69,6 +76,7 @@ class DeploymentHasNoTargetRevision(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/deployment-has-no-target-revision" error_title = "Deployment has no target revision." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -92,6 +100,7 @@ class RevisionMissingModelVFolder(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/revision-missing-model-vfolder" error_title = "Deployment revision has no model vfolder." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -114,6 +123,7 @@ class RevisionNotDeployable(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/revision-not-deployable" error_title = "Deployment revision references deleted resources." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -126,6 +136,7 @@ class InvalidDeploymentStrategy(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-deployment-strategy" error_title = "Unknown or invalid deployment strategy." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -138,6 +149,7 @@ class RouteSessionNotFound(BackendAIError): error_type = "https://api.backend.ai/probs/route-session-not-found" error_title = "No session associated with route." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -154,6 +166,7 @@ def __init__(self, session_status: str) -> None: super().__init__(f"Session status: {session_status}") self.session_status = session_status + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -166,6 +179,7 @@ class RouteUnhealthy(BackendAIError): error_type = "https://api.backend.ai/probs/route-unhealthy" error_title = "Route health check failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -178,6 +192,7 @@ class IncompleteRevisionData(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/incomplete-revision-data" error_title = "Revision data is missing required fields." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, diff --git a/src/ai/backend/manager/errors/export.py b/src/ai/backend/manager/errors/export.py index ddf5af98938..2e8c333c5ef 100644 --- a/src/ai/backend/manager/errors/export.py +++ b/src/ai/backend/manager/errors/export.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -22,6 +24,7 @@ class ExportReportNotFound(BackendAIError, web.HTTPNotFound): def __init__(self, report_key: str) -> None: super().__init__(extra_msg=f"Unknown report: {report_key}") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -39,6 +42,7 @@ class InvalidExportFieldKeys(BackendAIError, web.HTTPBadRequest): def __init__(self, invalid_keys: list[str]) -> None: super().__init__(extra_msg=f"Invalid field keys: {invalid_keys}") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -56,6 +60,7 @@ class TooManyConcurrentExports(BackendAIError, web.HTTPServiceUnavailable): def __init__(self) -> None: super().__init__(extra_msg="Too many concurrent export requests. Please try again later.") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, diff --git a/src/ai/backend/manager/errors/fair_share.py b/src/ai/backend/manager/errors/fair_share.py index 9b93bf371ad..96b93e4c919 100644 --- a/src/ai/backend/manager/errors/fair_share.py +++ b/src/ai/backend/manager/errors/fair_share.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class FairShareError(BackendAIError): error_type = "https://api.backend.ai/probs/fair-share-error" error_title = "Fair share operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -33,6 +36,7 @@ class FairShareNotFoundError(FairShareError): error_type = "https://api.backend.ai/probs/fair-share-not-found" error_title = "Fair share entity not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -51,6 +55,7 @@ def __init__(self, invalid_types: list[str]) -> None: self.invalid_types = invalid_types super().__init__(f"Resource types not available in capacity: {', '.join(invalid_types)}") + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, diff --git a/src/ai/backend/manager/errors/image.py b/src/ai/backend/manager/errors/image.py index 08f00ab31e9..e9212dfe696 100644 --- a/src/ai/backend/manager/errors/image.py +++ b/src/ai/backend/manager/errors/image.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -20,6 +22,7 @@ class ImageNotFound(ObjectNotFound): object_name = "environment image" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -31,6 +34,7 @@ def error_code(self) -> ErrorCode: class ImageAliasNotFound(ObjectNotFound): object_name = "image alias" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE_ALIAS, @@ -42,6 +46,7 @@ def error_code(self) -> ErrorCode: class ContainerRegistryNotFound(ObjectNotFound): object_name = "container_registry" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -53,6 +58,7 @@ def error_code(self) -> ErrorCode: class ContainerRegistryGroupsAssociationNotFound(ObjectNotFound): object_name = "association of container_registry and group" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -65,6 +71,7 @@ class ContainerRegistryWebhookAuthorizationFailed(BackendAIError, web.HTTPUnauth error_type = "https://api.backend.ai/probs/webhook/auth-failed" error_title = "Container Registry Webhook authorization failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -77,6 +84,7 @@ class HarborWebhookContainerRegistryRowNotFound(InternalServerError): error_type = "https://api.backend.ai/probs/webhook/harbor/container-registry-not-found" error_title = "Container registry row not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, @@ -88,6 +96,7 @@ def error_code(self) -> ErrorCode: class UnknownImageReferenceError(ObjectNotFound): object_name = "image reference" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -102,6 +111,7 @@ class ImageAccessForbiddenError(BackendAIError): error_type = "https://api.backend.ai/probs/generic-forbidden" error_title = "User does not have permission to access this image." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -114,6 +124,7 @@ class ForgetImageNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/generic-not-found" error_title = "The image you are trying to delete does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -126,6 +137,7 @@ class AliasImageActionValueError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-parameters" error_title = "Invalid parameters for image alias." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE_ALIAS, @@ -142,6 +154,7 @@ class AliasImageActionDBError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/image-db-error" error_title = "Database error while managing image alias." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE_ALIAS, @@ -154,6 +167,7 @@ class ModifyImageActionValueError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-parameters" error_title = "Invalid parameters for image modification." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -170,6 +184,7 @@ class PurgeImageActionByIdObjectDBError(BackendAIError, web.HTTPInternalServerEr error_type = "https://api.backend.ai/probs/image-db-error" error_title = "Database error while purging image." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, @@ -181,6 +196,7 @@ def error_code(self) -> ErrorCode: class RegistryNotFoundForImage(ObjectNotFound): object_name = "registry for image" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.CONTAINER_REGISTRY, diff --git a/src/ai/backend/manager/errors/kernel.py b/src/ai/backend/manager/errors/kernel.py index bd7437566f1..7770116ef3c 100644 --- a/src/ai/backend/manager/errors/kernel.py +++ b/src/ai/backend/manager/errors/kernel.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from aiohttp import web @@ -29,6 +29,7 @@ class KernelNotReady(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/kernel-not-ready" error_title = "Kernel not ready." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -41,6 +42,7 @@ class InvalidSessionId(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-session-id" error_title = "Invalid session ID format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -52,6 +54,7 @@ def error_code(self) -> ErrorCode: class SessionNotFound(ObjectNotFound): object_name = "session" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -63,6 +66,7 @@ def error_code(self) -> ErrorCode: class MainKernelNotFound(ObjectNotFound): object_name = "main kernel" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -74,6 +78,7 @@ def error_code(self) -> ErrorCode: class KernelNotFound(ObjectNotFound): object_name = "kernel" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -86,6 +91,7 @@ class TooManySessionsMatched(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/too-many-sessions-matched" error_title = "Too many sessions matched." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -117,6 +123,7 @@ class TooManyKernelsFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/too-many-kernels" error_title = "There are two or more matching kernels." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -129,6 +136,7 @@ class SessionAlreadyExists(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/session-already-exists" error_title = "The session already exists but you requested not to reuse existing one." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -141,6 +149,7 @@ class QuotaExceeded(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/quota-exceeded" error_title = "You have reached your resource limit." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -160,6 +169,7 @@ class BackendAgentError(BackendAIError): "FAILURE": "https://api.backend.ai/probs/agent-failure", } + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -217,16 +227,19 @@ def __init__( "agent-details": agent_details, }) + @override def __str__(self) -> str: if self.agent_exception: return f"{self.agent_error_title} ({self.agent_exception})" return f"{self.agent_error_title}" + @override def __repr__(self) -> str: if self.agent_exception: return f"<{type(self).__name__}: {self.agent_error_title} ({self.agent_exception})>" return f"<{type(self).__name__}: {self.agent_error_title}>" + @override def __reduce__(self) -> tuple[type[BackendAgentError], tuple[Any, ...]]: # type: ignore[override] return (type(self), (self.agent_error_type, self.agent_details)) @@ -235,6 +248,7 @@ class KernelCreationFailed(BackendAgentError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/kernel-creation-failed" error_title = "Kernel creation has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -247,6 +261,7 @@ class KernelDestructionFailed(BackendAgentError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/kernel-destruction-failed" error_title = "Kernel destruction has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -259,6 +274,7 @@ class KernelRestartFailed(BackendAgentError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/kernel-restart-failed" error_title = "Kernel restart has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -271,6 +287,7 @@ class KernelExecutionFailed(BackendAgentError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/kernel-execution-failed" error_title = "Executing user code in the kernel has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -283,6 +300,7 @@ class InvalidStreamMode(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-stream-mode" error_title = "Invalid or missing stream mode parameter." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -295,6 +313,7 @@ class InvalidSessionData(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/invalid-session-data" error_title = "Session data has an invalid type or format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -307,6 +326,7 @@ class InvalidKernelData(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/invalid-kernel-data" error_title = "Kernel data has an invalid type or format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -319,6 +339,7 @@ class InvalidKernelConfig(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-kernel-config" error_title = "Invalid kernel configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -331,6 +352,7 @@ class IncompleteSessionSpec(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/incomplete-session-spec" error_title = "Session spec has unresolved required fields." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -342,6 +364,7 @@ def error_code(self) -> ErrorCode: class IdlePolicyNotFound(ObjectNotFound): object_name = "idle policy" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -354,6 +377,7 @@ class InvalidKernelStatus(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/invalid-kernel-status" error_title = "Invalid kernel status for this operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, diff --git a/src/ai/backend/manager/errors/keypair.py b/src/ai/backend/manager/errors/keypair.py index e62680664ff..469007c9645 100644 --- a/src/ai/backend/manager/errors/keypair.py +++ b/src/ai/backend/manager/errors/keypair.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -19,6 +21,7 @@ class InvalidSSHPrivateKey(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-ssh-private-key" error_title = "The SSH private key is invalid or in an unsupported format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, @@ -31,6 +34,7 @@ class InvalidSSHPublicKey(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-ssh-public-key" error_title = "The SSH public key is invalid or in an unsupported format." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, @@ -43,6 +47,7 @@ class SSHKeypairMismatch(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/ssh-keypair-mismatch" error_title = "The SSH public key does not match the private key." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, diff --git a/src/ai/backend/manager/errors/notification.py b/src/ai/backend/manager/errors/notification.py index 7434a5a3520..0357ba1e31d 100644 --- a/src/ai/backend/manager/errors/notification.py +++ b/src/ai/backend/manager/errors/notification.py @@ -1,5 +1,7 @@ """Notification-related error definitions.""" +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -26,6 +28,7 @@ class NotificationChannelNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/notification-channel-not-found" error_title = "The notification channel does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -38,6 +41,7 @@ class NotificationRuleNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/notification-rule-not-found" error_title = "The notification rule does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -50,6 +54,7 @@ class NotificationChannelConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/notification-channel-conflict" error_title = "The notification channel already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -62,6 +67,7 @@ class NotificationRuleConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/notification-rule-conflict" error_title = "The notification rule already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -74,6 +80,7 @@ class NotificationProcessingFailure(BackendAIError, web.HTTPInternalServerError) error_type = "https://api.backend.ai/probs/notification-processing-failure" error_title = "Failed to process notification." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -86,6 +93,7 @@ class NotificationTemplateRenderingFailure(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/notification-template-rendering-failure" error_title = "Failed to render notification template." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -98,6 +106,7 @@ class InvalidNotificationChannelType(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-notification-channel-type" error_title = "Invalid notification channel type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, @@ -110,6 +119,7 @@ class InvalidNotificationSpec(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-notification-spec" error_title = "Invalid notification specification." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.NOTIFICATION, diff --git a/src/ai/backend/manager/errors/object_storage.py b/src/ai/backend/manager/errors/object_storage.py index 834d313c56b..a6bea418cab 100644 --- a/src/ai/backend/manager/errors/object_storage.py +++ b/src/ai/backend/manager/errors/object_storage.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class ObjectStorageNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/object-storage-not-found" error_title = "Object Storage Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -25,6 +28,7 @@ class ObjectStorageOperationNotSupported(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/object-storage-operation-not-supported" error_title = "Object Storage Operation Not Supported" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, diff --git a/src/ai/backend/manager/errors/permission.py b/src/ai/backend/manager/errors/permission.py index ecaf69ed3b4..ef78e569ebe 100644 --- a/src/ai/backend/manager/errors/permission.py +++ b/src/ai/backend/manager/errors/permission.py @@ -1,5 +1,7 @@ """Permission and RBAC-related error definitions.""" +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -26,6 +28,7 @@ class RoleNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/role-not-found" error_title = "The role does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -45,6 +48,7 @@ class UserSystemRoleNotProvisioned(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/user-system-role-not-provisioned" error_title = "The user's SYSTEM role is not provisioned." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -57,6 +61,7 @@ class RoleAlreadyAssigned(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/role-already-assigned" error_title = "The role is already assigned to the user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -69,6 +74,7 @@ class RoleNotAssigned(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/role-not-assigned" error_title = "The role is not assigned to the user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -81,6 +87,7 @@ class NotEnoughPermission(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/not-enough-permission" error_title = "Insufficient permission to perform this operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -93,6 +100,7 @@ class PermissionNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/permission-not-found" error_title = "The permission does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, @@ -105,6 +113,7 @@ class ObjectPermissionNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/object-permission-not-found" error_title = "The object permission does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, @@ -117,6 +126,7 @@ class ReplaceRolePermissionRoleIdMismatch(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/replace-role-permission-role-id-mismatch" error_title = "Permission entry role_id does not match the request role_id." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, diff --git a/src/ai/backend/manager/errors/repository.py b/src/ai/backend/manager/errors/repository.py index 741aebef58f..a5ed8c44e34 100644 --- a/src/ai/backend/manager/errors/repository.py +++ b/src/ai/backend/manager/errors/repository.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, override from aiohttp import web @@ -23,6 +23,7 @@ class RepositoryError(BackendAIError): error_type = "https://api.backend.ai/probs/repository-error" error_title = "Repository operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -37,6 +38,7 @@ class UpsertEmptyResultError(RepositoryError): error_type = "https://api.backend.ai/probs/upsert-empty-result" error_title = "Upsert operation did not return any row." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -51,6 +53,7 @@ class UnsupportedCompositePrimaryKeyError(RepositoryError): error_type = "https://api.backend.ai/probs/unsupported-composite-pk" error_title = "Unsupported composite primary key." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -70,6 +73,7 @@ class EmptySearchScopeError(RepositoryError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/empty-search-scope" error_title = "Search scope must not be empty." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -112,6 +116,7 @@ def __init__( self.detail = detail self.pgcode = pgcode + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -126,6 +131,7 @@ class UniqueConstraintViolationError(RepositoryIntegrityError): error_type = "https://api.backend.ai/probs/unique-constraint-violation" error_title = "Unique constraint violated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -140,6 +146,7 @@ class ForeignKeyViolationError(RepositoryIntegrityError): error_type = "https://api.backend.ai/probs/foreign-key-violation" error_title = "Foreign key constraint violated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -154,6 +161,7 @@ class CheckConstraintViolationError(RepositoryIntegrityError): error_type = "https://api.backend.ai/probs/check-constraint-violation" error_title = "Check constraint violated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -172,6 +180,7 @@ class NotNullViolationError(RepositoryIntegrityError): error_title = "Not-null constraint violated." status_code = 400 + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -186,6 +195,7 @@ class ExclusionViolationError(RepositoryIntegrityError): error_type = "https://api.backend.ai/probs/exclusion-violation" error_title = "Exclusion constraint violated." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, diff --git a/src/ai/backend/manager/errors/resource.py b/src/ai/backend/manager/errors/resource.py index a993983ec8d..758848801bc 100644 --- a/src/ai/backend/manager/errors/resource.py +++ b/src/ai/backend/manager/errors/resource.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -20,6 +22,7 @@ class DomainNotFound(ObjectNotFound): object_name = "domain" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -32,6 +35,7 @@ class ProjectHasActiveKernelsError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/project-has-active-kernels" error_title = "Project has active kernels." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -44,6 +48,7 @@ class ProjectHasVFoldersMountedError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/project-has-vfolders-mounted" error_title = "Project has vfolders mounted to active kernels." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -56,6 +61,7 @@ class ProjectHasActiveEndpointsError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/project-has-active-endpoints" error_title = "Project has active endpoints." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -67,6 +73,7 @@ def error_code(self) -> ErrorCode: class ScalingGroupNotFound(ObjectNotFound): object_name = "scaling group" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -79,6 +86,7 @@ class ScalingGroupSessionTypeNotAllowed(BackendAIError, web.HTTPUnprocessableEnt error_type = "https://api.backend.ai/probs/scaling-group-session-type-not-allowed" error_title = "Scaling group does not allow this session type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -91,6 +99,7 @@ class ScalingGroupDeletionFailure(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/scaling-group-deletion-failure" error_title = "Failed to delete scaling group." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -102,6 +111,7 @@ def error_code(self) -> ErrorCode: class InstanceNotFound(ObjectNotFound): object_name = "agent instance" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.INSTANCE, @@ -114,6 +124,7 @@ class InstanceNotAvailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/instance-not-available" error_title = "There is no available instance." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.INSTANCE, @@ -126,6 +137,7 @@ class ProjectNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/project-not-found" error_title = "Project not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -137,6 +149,7 @@ def error_code(self) -> ErrorCode: class TaskTemplateNotFound(ObjectNotFound): object_name = "task template" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.TEMPLATE, @@ -148,6 +161,7 @@ def error_code(self) -> ErrorCode: class AppNotFound(ObjectNotFound): object_name = "app service" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -165,6 +179,7 @@ class DomainDataProcessingError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/domain-data-processing-error" error_title = "Failed to process domain data." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -176,6 +191,7 @@ def error_code(self) -> ErrorCode: class ScalingGroupProxyTargetNotFound(ObjectNotFound): object_name = "scaling group proxy target" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -187,6 +203,7 @@ def error_code(self) -> ErrorCode: class ResourcePresetNotFound(ObjectNotFound): object_name = "resource preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RESOURCE_PRESET, @@ -198,6 +215,7 @@ def error_code(self) -> ErrorCode: class RuntimeVariantNotFound(ObjectNotFound): object_name = "runtime variant" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RUNTIME_VARIANT, @@ -210,6 +228,7 @@ class RuntimeVariantConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-runtime-variant" error_title = "Duplicate Runtime Variant" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RUNTIME_VARIANT, @@ -221,6 +240,7 @@ def error_code(self) -> ErrorCode: class RuntimeVariantPresetNotFound(ObjectNotFound): object_name = "runtime variant preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RUNTIME_VARIANT, @@ -233,6 +253,7 @@ class RuntimeVariantPresetConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-runtime-variant-preset" error_title = "Duplicate Runtime Variant Preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RUNTIME_VARIANT, @@ -244,6 +265,7 @@ def error_code(self) -> ErrorCode: class ModelCardNotFound(ObjectNotFound): object_name = "model card" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_CARD, @@ -256,6 +278,7 @@ class ModelCardConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-model-card" error_title = "Duplicate Model Card" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_CARD, @@ -268,6 +291,7 @@ class InvalidProjectTypeForModelCard(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-project-type-for-model-card" error_title = "Project is not a MODEL_STORE type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_CARD, @@ -280,6 +304,7 @@ class NotAModelVFolder(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/not-a-model-vfolder" error_title = "VFolder usage_mode is not MODEL." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_CARD, @@ -291,6 +316,7 @@ def error_code(self) -> ErrorCode: class DeploymentRevisionPresetNotFound(ObjectNotFound): object_name = "deployment revision preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -303,6 +329,7 @@ class DeploymentRevisionPresetConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-deployment-revision-preset" error_title = "Duplicate Deployment Revision Preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -314,6 +341,7 @@ def error_code(self) -> ErrorCode: class AgentNotFound(ObjectNotFound): object_name = "agent" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -326,6 +354,7 @@ class DomainCreationFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/domain-creation-failed" error_title = "Failed to create domain." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -338,6 +367,7 @@ class DomainNodeCreationFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/domain-node-creation-failed" error_title = "Failed to create domain node." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -350,6 +380,7 @@ class DomainHasActiveKernels(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/domain-has-active-kernels" error_title = "Domain has active kernels." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -362,6 +393,7 @@ class DomainHasUsers(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/domain-has-users" error_title = "Domain has associated users." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -374,6 +406,7 @@ class DomainHasGroups(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/domain-has-groups" error_title = "Domain has associated groups." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -386,6 +419,7 @@ class DomainDeletionFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/domain-deletion-failed" error_title = "Failed to delete domain." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -398,6 +432,7 @@ class DomainUpdateNotAllowed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/domain-update-not-allowed" error_title = "Domain update not allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -410,6 +445,7 @@ class InvalidDomainConfiguration(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-domain-configuration" error_title = "Invalid domain configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -422,6 +458,7 @@ class AllocationFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/allocation-failed" error_title = "Failed to allocate resources." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -434,6 +471,7 @@ class InvalidUserUpdateMode(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-user-update-mode" error_title = "Invalid user update mode." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -446,6 +484,7 @@ class InvalidPresetQuery(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-preset-query" error_title = "Invalid resource preset query parameters." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RESOURCE_PRESET, @@ -458,6 +497,7 @@ class NoAvailableScalingGroup(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/no-available-scaling-group" error_title = "No scaling groups available for this session." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -470,6 +510,7 @@ class AgentNotAllocated(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/agent-not-allocated" error_title = "Agent ID has not been allocated for the session." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -482,6 +523,7 @@ class SessionNotAllocated(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/session-not-allocated" error_title = "Session ID is not available during allocation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -494,6 +536,7 @@ class NoCurrentTaskContext(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/no-current-task-context" error_title = "No current asyncio task context available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -506,6 +549,7 @@ class DatabaseConnectionUnavailable(BackendAIError, web.HTTPInternalServerError) error_type = "https://api.backend.ai/probs/database-connection-unavailable" error_title = "Database connection is not available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -518,6 +562,7 @@ class InvalidSchedulerState(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/invalid-scheduler-state" error_title = "Scheduler is in an invalid state." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SCALING_GROUP, @@ -530,6 +575,7 @@ class ConfigurationLoadFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/configuration-load-failed" error_title = "Failed to load configuration." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -542,6 +588,7 @@ class DataTransformationFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/data-transformation-failed" error_title = "Failed to transform data." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, @@ -554,6 +601,7 @@ class DBOperationFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/db-operation-failed" error_title = "Database operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DATABASE, diff --git a/src/ai/backend/manager/errors/resource_slot.py b/src/ai/backend/manager/errors/resource_slot.py index 3bd35f2ede1..7db25716e27 100644 --- a/src/ai/backend/manager/errors/resource_slot.py +++ b/src/ai/backend/manager/errors/resource_slot.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -19,6 +21,7 @@ class ResourceSlotTypeNotFound(BackendAIError): error_type = "https://api.backend.ai/probs/resource-slot-type-not-found" error_title = "Resource slot type not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.RESOURCE_PRESET, @@ -33,6 +36,7 @@ class AgentResourceNotFound(BackendAIError): error_type = "https://api.backend.ai/probs/agent-resource-not-found" error_title = "Agent resource not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, @@ -47,6 +51,7 @@ class ResourceAllocationNotFound(BackendAIError): error_type = "https://api.backend.ai/probs/resource-allocation-not-found" error_title = "Resource allocation not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KERNEL, @@ -61,6 +66,7 @@ class AgentResourceCapacityExceeded(BackendAIError): error_type = "https://api.backend.ai/probs/agent-resource-capacity-exceeded" error_title = "Agent resource capacity exceeded." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.AGENT, diff --git a/src/ai/backend/manager/errors/role_invitation.py b/src/ai/backend/manager/errors/role_invitation.py index 848b0aed59f..f04f7866c4f 100644 --- a/src/ai/backend/manager/errors/role_invitation.py +++ b/src/ai/backend/manager/errors/role_invitation.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class RoleInvitationNotFound(ObjectNotFound): object_name = "role-invitation" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE_INVITATION, @@ -25,6 +28,7 @@ class DuplicateRoleInvitationError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-role-invitation" error_title = "Duplicate role invitation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE_INVITATION, @@ -37,6 +41,7 @@ class RoleInvitationInvalidState(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/role-invitation-invalid-state" error_title = "Invalid role invitation state transition." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE_INVITATION, diff --git a/src/ai/backend/manager/errors/role_preset.py b/src/ai/backend/manager/errors/role_preset.py index 2e70cb48787..cf7c94a4f7e 100644 --- a/src/ai/backend/manager/errors/role_preset.py +++ b/src/ai/backend/manager/errors/role_preset.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -18,6 +20,7 @@ class RolePresetNotFound(ObjectNotFound): object_name = "role_preset" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, @@ -30,6 +33,7 @@ class RolePermissionPresetConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/duplicate-role-permission-preset" error_title = "Duplicate role permission preset entry." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROLE, diff --git a/src/ai/backend/manager/errors/service.py b/src/ai/backend/manager/errors/service.py index e7c5f07cd23..00a22b50aa0 100644 --- a/src/ai/backend/manager/errors/service.py +++ b/src/ai/backend/manager/errors/service.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -23,6 +25,7 @@ class NoUpdatesToApply(BackendAIError): def __init__(self, message: str = "No updates to apply") -> None: super().__init__(message) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -34,6 +37,7 @@ def error_code(self) -> ErrorCode: class EndpointNotFound(ObjectNotFound): object_name = "endpoint" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -45,6 +49,7 @@ def error_code(self) -> ErrorCode: class ModelDefinitionNotFound(ObjectNotFound): object_name = "model_definition" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -58,6 +63,7 @@ class ScalingImpossible(BackendAIError, web.HTTPBadRequest): "Scaling operation cannot be performed due to insufficient resources or constraints." ) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -69,6 +75,7 @@ def error_code(self) -> ErrorCode: class AutoScalingRuleNotFound(ObjectNotFound): object_name = "auto_scaling_rule" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -80,6 +87,7 @@ def error_code(self) -> ErrorCode: class AutoScalingPolicyNotFound(ObjectNotFound): object_name = "auto_scaling_policy" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -91,6 +99,7 @@ def error_code(self) -> ErrorCode: class DeploymentPolicyNotFound(ObjectNotFound): object_name = "deployment_policy" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -102,6 +111,7 @@ def error_code(self) -> ErrorCode: class RoutingNotFound(ObjectNotFound): object_name = "routing" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROUTE, @@ -113,6 +123,7 @@ def error_code(self) -> ErrorCode: class EndpointTokenNotFound(ObjectNotFound): object_name = "endpoint_token" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -124,6 +135,7 @@ def error_code(self) -> ErrorCode: class ModelServiceNotFound(ObjectNotFound): object_name = "model service" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -135,6 +147,7 @@ def error_code(self) -> ErrorCode: class RouteNotFound(ObjectNotFound): object_name = "route" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ROUTE, @@ -146,6 +159,7 @@ def error_code(self) -> ErrorCode: class ModelServiceDependencyNotCleared(BackendAIError, web.HTTPBadRequest): error_title = "Cannot delete model VFolders bound to alive model services." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_SERVICE, @@ -158,6 +172,7 @@ class AppServiceStartFailed(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/app-service-start-failed" error_title = "Failed to start the application service." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -171,6 +186,7 @@ class EndpointAccessForbiddenError(GenericForbidden): error_title = "Access to this endpoint is forbidden." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT, @@ -182,6 +198,7 @@ def error_code(self) -> ErrorCode: class EndpointAutoScalingRuleNotFound(ObjectNotFound): object_name = "endpoint auto scaling rule" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ENDPOINT_AUTO_SCALING, diff --git a/src/ai/backend/manager/errors/storage.py b/src/ai/backend/manager/errors/storage.py index e77f63f1606..5b3e8f0195f 100644 --- a/src/ai/backend/manager/errors/storage.py +++ b/src/ai/backend/manager/errors/storage.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, override from aiohttp import web @@ -24,6 +24,7 @@ class TooManyVFoldersFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/too-many-vfolders" error_title = "Multiple vfolders found for the operation for a single vfolder." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -49,6 +50,7 @@ def __init__(self, matched_rows: Sequence[Mapping[str, Any]]) -> None: class VFolderNotFound(ObjectNotFound): object_name = "virtual folder" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -60,6 +62,7 @@ def error_code(self) -> ErrorCode: class QuotaScopeNotFoundError(ObjectNotFound): object_name = "quota scope" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -72,6 +75,7 @@ class ModelCardParseError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/model-card-parse-error" error_title = "Model Card Parse Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_CARD, @@ -84,6 +88,7 @@ class VFolderAlreadyExists(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/vfolder-already-exists" error_title = "The virtual folder already exists with the same name." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -96,6 +101,7 @@ class VFolderGone(BackendAIError, web.HTTPGone): error_type = "https://api.backend.ai/probs/vfolder-gone" error_title = "The virtual folder is gone." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -108,6 +114,7 @@ class VFolderBadRequest(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-operation-failed" error_title = "Virtual folder operation has failed due to bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -120,6 +127,7 @@ class VFolderOperationFailed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-operation-failed" error_title = "Virtual folder operation has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -132,6 +140,7 @@ class VFolderFilterStatusFailed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-filter-status-failed" error_title = "Virtual folder status filtering has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -144,6 +153,7 @@ class VFolderFilterStatusNotAvailable(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-filter-status-not-available" error_title = "There is no available virtual folder to filter its status." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -156,6 +166,7 @@ class VFolderPermissionError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-permission-error" error_title = "The virtual folder does not permit the specified permission." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -168,6 +179,7 @@ class VFolderInvitationNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/vfolder-invitation-not-found" error_title = "Virtual folder invitation not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER_INVITATION, @@ -180,6 +192,7 @@ class VFolderCreationFailure(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-creation-failed" error_title = "Virtual folder creation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -192,6 +205,7 @@ class VFolderGrantAlreadyExists(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/vfolder-grant-already-exists" error_title = "Virtual folder grant already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -204,6 +218,7 @@ class VFolderDeletionNotAllowed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-deletion-not-allowed" error_title = "Virtual folder deletion is not allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -216,6 +231,7 @@ class VFolderHasLinkedModelCard(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-has-linked-model-card" error_title = "Virtual folder has linked model card(s)." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -228,6 +244,7 @@ class InsufficientStoragePermission(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage-permission-not-allowed" error_title = "The specified storage permission is not allowed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -240,6 +257,7 @@ class VFolderInvalidParameter(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/vfolder-invalid-parameter" error_title = "Invalid parameter for virtual folder operation." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -252,6 +270,7 @@ class DotfileCreationFailed(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Dotfile creation has failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOTFILE, @@ -264,6 +283,7 @@ class DotfileAlreadyExists(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/generic-bad-request" error_title = "Dotfile already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOTFILE, @@ -275,6 +295,7 @@ def error_code(self) -> ErrorCode: class DotfileNotFound(ObjectNotFound): object_name = "dotfile" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOTFILE, @@ -287,6 +308,7 @@ class DotfileVFolderPathConflict(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/dotfile-vfolder-path-conflict" error_title = "The dotfile path conflicts with a virtual folder path." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOTFILE, @@ -299,6 +321,7 @@ class StorageProxyNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage-proxy-not-found" error_title = "Storage proxy not found." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -311,6 +334,7 @@ class StorageProxyConnectionError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/storage-proxy-connection-error" error_title = "Failed to connect to storage proxy." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -323,6 +347,7 @@ class StorageProxyTimeoutError(BackendAIError, web.HTTPGatewayTimeout): error_type = "https://api.backend.ai/probs/storage-proxy-timeout" error_title = "Request to storage proxy timed out." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -335,6 +360,7 @@ class UnexpectedStorageProxyResponseError(BackendAIError, web.HTTPInternalServer error_type = "https://api.backend.ai/probs/unexpected-storage-proxy-response" error_title = "Unexpected response from storage proxy." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -347,6 +373,7 @@ class UnsupportedStorageTypeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/unsupported-storage-type" error_title = "Unsupported storage type." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, diff --git a/src/ai/backend/manager/errors/user.py b/src/ai/backend/manager/errors/user.py index b56c9f6fcdd..6b82d83e58f 100644 --- a/src/ai/backend/manager/errors/user.py +++ b/src/ai/backend/manager/errors/user.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class UserNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/user-not-found" error_title = "The user does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -25,6 +28,7 @@ class UserConflict(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/user-conflict" error_title = "The user already exists." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -37,6 +41,7 @@ class UserModificationBadRequest(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/user-modification-bad-request" error_title = "Failed to modify user due to bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -49,6 +54,7 @@ class UserCreationBadRequest(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/user-creation-bad-request" error_title = "Failed to create user due to bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -61,6 +67,7 @@ class UserCreationFailure(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/user-creation-failure" error_title = "Failed to create user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -73,6 +80,7 @@ class UserModificationFailure(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/user-modification-failure" error_title = "Failed to modify user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -85,6 +93,7 @@ class UserPurgeFailure(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/user-purge-failure" error_title = "Failed to purge user." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -97,6 +106,7 @@ class KeyPairNotFound(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/keypair-not-found" error_title = "The key pair does not exist." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, @@ -109,6 +119,7 @@ class KeyPairForbidden(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/keypair-forbidden" error_title = "The key pair is not allowed to be used." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, diff --git a/src/ai/backend/manager/errors/vfs_storage.py b/src/ai/backend/manager/errors/vfs_storage.py index d1ccbfe1751..5c1a74c1848 100644 --- a/src/ai/backend/manager/errors/vfs_storage.py +++ b/src/ai/backend/manager/errors/vfs_storage.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class VFSStorageNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/vfs-storage-not-found" error_title = "VFS Storage Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFS_STORAGE, diff --git a/src/ai/backend/manager/event_dispatcher/reporters.py b/src/ai/backend/manager/event_dispatcher/reporters.py index 1f5a5e87a66..3e00c6d7557 100644 --- a/src/ai/backend/manager/event_dispatcher/reporters.py +++ b/src/ai/backend/manager/event_dispatcher/reporters.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.common.events.reporter import ( AbstractEventReporter, CompleteEventReportArgs, @@ -12,6 +14,7 @@ class EventLogger(AbstractEventReporter): def __init__(self, db: ExtendedAsyncSAEngine) -> None: self._db = db + @override async def prepare_event_report( self, event: AbstractEvent, @@ -22,6 +25,7 @@ async def prepare_event_report( session.add(event_log) await session.flush() + @override async def complete_event_report( self, event: AbstractEvent, diff --git a/src/ai/backend/manager/events/hub/propagators/session.py b/src/ai/backend/manager/events/hub/propagators/session.py index cc96e457d92..7aeb4cbd599 100644 --- a/src/ai/backend/manager/events/hub/propagators/session.py +++ b/src/ai/backend/manager/events/hub/propagators/session.py @@ -4,7 +4,7 @@ import logging import uuid from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import sqlalchemy as sa from sqlalchemy.orm import load_only @@ -88,10 +88,12 @@ def __init__( self._filters = filters self._closed_lock = asyncio.Lock() + @override def id(self) -> uuid.UUID: """Get the unique identifier for the propagator.""" return self._id + @override async def propagate_event(self, event: AbstractEvent) -> None: """Propagate an event to the SSE stream if it passes filtering.""" if self._closed: @@ -283,6 +285,7 @@ async def _should_send_event(self, event_data: Mapping[str, Any]) -> bool: return True + @override async def close(self) -> None: """Close the propagator and stop event streaming.""" async with self._closed_lock: diff --git a/src/ai/backend/manager/health/database.py b/src/ai/backend/manager/health/database.py index 250a37e005c..3a030f4d9d3 100644 --- a/src/ai/backend/manager/health/database.py +++ b/src/ai/backend/manager/health/database.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.common.health_checker.abc import StaticServiceHealthChecker from ai.backend.common.health_checker.types import ( @@ -35,10 +36,12 @@ def __init__(self, db: ExtendedAsyncSAEngine, timeout: float = 5.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return DATABASE + @override async def check_service(self) -> ServiceHealth: """ Check database connection health by pinging the server. @@ -65,6 +68,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={CID_POSTGRES: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/manager/health/rpc.py b/src/ai/backend/manager/health/rpc.py index 699c414a7d2..3fbb0b0860d 100644 --- a/src/ai/backend/manager/health/rpc.py +++ b/src/ai/backend/manager/health/rpc.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from aiohttp import web @@ -23,6 +24,7 @@ class RpcHealthCheckError(HealthCheckError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/rpc-health-check-failed" error_title = "RPC health check failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.HEALTH_CHECK, @@ -53,10 +55,12 @@ def __init__(self, agent_client: AgentClient, timeout: float = 5.0) -> None: self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: """The service group this checker monitors.""" return AGENT + @override async def check_service(self) -> ServiceHealth: """ Check agent RPC health by calling the health RPC method. @@ -86,6 +90,7 @@ async def check_service(self) -> ServiceHealth: return ServiceHealth(results={component_id: status}) @property + @override def timeout(self) -> float: """The timeout for the health check in seconds.""" return self._timeout diff --git a/src/ai/backend/manager/idle.py b/src/ai/backend/manager/idle.py index ab630089f65..5f6ac1b649a 100644 --- a/src/ai/backend/manager/idle.py +++ b/src/ai/backend/manager/idle.py @@ -502,6 +502,7 @@ class NewUserGracePeriodChecker(AbstractIdleCheckReporter): }, ).allow_extra("*") + @override async def populate_config(self, raw_config: Mapping[str, Any]) -> None: config = self._config_iv.check(raw_config) self.user_initial_grace_period = config["user_initial_grace_period"] @@ -516,6 +517,7 @@ async def populate_config(self, raw_config: Mapping[str, Any]) -> None: _grace_period, ) + @override async def get_extra_info( self, _redis_obj: ValkeyLiveClient, _session_id: SessionId ) -> dict[str, Any] | None: @@ -548,6 +550,7 @@ def grace_period_const(self) -> float: else 0 ) + @override async def get_checker_result( self, redis_obj: ValkeyLiveClient, @@ -738,6 +741,7 @@ class NetworkTimeoutIdleChecker(BaseIdleChecker): def terminate_reason(self) -> KernelLifecycleEventReason: return KernelLifecycleEventReason.IDLE_TIMEOUT + @override async def populate_config(self, raw_config: Mapping[str, Any]) -> None: config = self._config_iv.check(raw_config) self.idle_timeout = config["threshold"] or DEFAULT_NETWORK_CHECKER_IDLE_TIMEOUT @@ -746,6 +750,7 @@ async def populate_config(self, raw_config: Mapping[str, Any]) -> None: self.idle_timeout.total_seconds(), ) + @override async def get_extra_info( self, _redis_obj: ValkeyLiveClient, _session_id: SessionId ) -> dict[str, Any] | None: @@ -798,6 +803,7 @@ async def check_idleness( ) return remaining >= 0 + @override async def get_checker_result( self, redis_obj: ValkeyLiveClient, @@ -814,9 +820,11 @@ class SessionLifetimeChecker(BaseIdleChecker): report_key: ClassVar[str] = "session_lifetime" extra_info_key: ClassVar[str] = "session_lifetime_extra" + @override async def populate_config(self, raw_config: Mapping[str, Any]) -> None: pass + @override async def get_extra_info( self, _redis_obj: ValkeyLiveClient, _session_id: SessionId ) -> dict[str, Any] | None: @@ -856,6 +864,7 @@ async def check_idleness( return remaining > 0 return True + @override async def get_checker_result( self, redis_obj: ValkeyLiveClient, @@ -969,6 +978,7 @@ class UtilizationIdleChecker(BaseIdleChecker): def terminate_reason(self) -> KernelLifecycleEventReason: return KernelLifecycleEventReason.IDLE_UTILIZATION + @override async def populate_config(self, raw_config: Mapping[str, Any]) -> None: config = UtilizationConfig(**config_key_to_snake_case(raw_config)) self.resource_thresholds = config.resource_thresholds @@ -996,9 +1006,11 @@ def _to_timedelta(val: tv.TVariousDelta) -> timedelta: ) @classmethod + @override def get_extra_info_key(cls, session_id: SessionId) -> str | None: return f"session.{session_id}.{cls.extra_info_key}" + @override async def get_extra_info( self, redis_obj: ValkeyLiveClient, session_id: SessionId ) -> dict[str, Any] | None: @@ -1275,6 +1287,7 @@ async def get_current_utilization( log.warning(_msg, exc_info=e) return None + @override async def get_checker_result( self, redis_obj: ValkeyLiveClient, diff --git a/src/ai/backend/manager/models/agent/row.py b/src/ai/backend/manager/models/agent/row.py index a5b2c57e60a..b70697bceeb 100644 --- a/src/ai/backend/manager/models/agent/row.py +++ b/src/ai/backend/manager/models/agent/row.py @@ -318,12 +318,14 @@ def apply_sgroup_permission_ctx( ) -> None: self.sgroup_permission_ctx = sgroup_permission_ctx + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(AgentRow).where(cond) + @override async def calculate_final_permission(self, rbac_obj: AgentRow) -> frozenset[AgentPermission]: agent_row = rbac_obj agent_id = cast(AgentId, agent_row.id) diff --git a/src/ai/backend/manager/models/artifact/row.py b/src/ai/backend/manager/models/artifact/row.py index de811de486f..834b1ed356c 100644 --- a/src/ai/backend/manager/models/artifact/row.py +++ b/src/ai/backend/manager/models/artifact/row.py @@ -3,7 +3,7 @@ import logging import uuid from datetime import datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -102,6 +102,7 @@ class ArtifactRow(Base): # type: ignore[misc] primaryjoin=_get_artifact_revision_join_cond, ) + @override def __str__(self) -> str: return ( f"ArtifactRow(" @@ -120,6 +121,7 @@ def __str__(self) -> str: f"extra={self.extra})" ) + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/artifact_registries/row.py b/src/ai/backend/manager/models/artifact_registries/row.py index aea833faba9..a5177b77760 100644 --- a/src/ai/backend/manager/models/artifact_registries/row.py +++ b/src/ai/backend/manager/models/artifact_registries/row.py @@ -2,7 +2,7 @@ import logging import uuid -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -65,6 +65,7 @@ class ArtifactRegistryRow(Base): # type: ignore[misc] viewonly=True, ) + @override def __str__(self) -> str: return f"ArtifactRegistryRow(id={self.id}, registry_id={self.registry_id}, type={self.type}, name={self.name})" diff --git a/src/ai/backend/manager/models/artifact_revision/row.py b/src/ai/backend/manager/models/artifact_revision/row.py index 9f7b6d2e4f2..64e1740ce2f 100644 --- a/src/ai/backend/manager/models/artifact_revision/row.py +++ b/src/ai/backend/manager/models/artifact_revision/row.py @@ -3,7 +3,7 @@ import logging import uuid from datetime import datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -115,6 +115,7 @@ class ArtifactRevisionRow(Base): # type: ignore[misc] ) ) + @override def __str__(self) -> str: readme_display = self.readme[:15] if self.readme else None created_at_str = self.created_at.isoformat() if self.created_at else None diff --git a/src/ai/backend/manager/models/audit_log/row.py b/src/ai/backend/manager/models/audit_log/row.py index 29f8463f23d..30865210c6f 100644 --- a/src/ai/backend/manager/models/audit_log/row.py +++ b/src/ai/backend/manager/models/audit_log/row.py @@ -3,6 +3,7 @@ import logging import uuid from datetime import datetime, timedelta +from typing import override import sqlalchemy as sa from sqlalchemy.orm import Mapped, mapped_column @@ -82,6 +83,7 @@ def __init__( self.status = status self.created_at = created_at + @override def __str__(self) -> str: return ( f"AuditLogRow(" @@ -98,6 +100,7 @@ def __str__(self) -> str: f")" ) + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/base.py b/src/ai/backend/manager/models/base.py index 987bb929faf..139b669b3b2 100644 --- a/src/ai/backend/manager/models/base.py +++ b/src/ai/backend/manager/models/base.py @@ -21,6 +21,7 @@ Self, TypeVar, cast, + override, ) import sqlalchemy as sa @@ -195,6 +196,7 @@ def __init__(self, enum_cls: type[T_Enum], **opts: Any) -> None: super().__init__(*enums, **opts) self._enum_cls = enum_cls + @override def process_bind_param( self, value: T_Enum | None, @@ -202,6 +204,7 @@ def process_bind_param( ) -> str | None: return value.name if value else None + @override def process_result_value( self, value: Any | None, @@ -209,10 +212,12 @@ def process_result_value( ) -> T_Enum | None: return self._enum_cls[value] if value else None + @override def copy(self, **_kw: Any) -> Self: return EnumType(self._enum_cls, **self._opts) # type: ignore[return-value] @property + @override def python_type(self) -> type[T_Enum]: return self._enum_cls @@ -249,6 +254,7 @@ def __init__(self, enum_cls: type[T_Enum], **opts: Any) -> None: super().__init__(*enums, **opts) self._enum_cls = enum_cls + @override def process_bind_param( self, value: T_Enum | None, @@ -256,6 +262,7 @@ def process_bind_param( ) -> str | None: return value.value if value else None + @override def process_result_value( self, value: Any | None, @@ -263,10 +270,12 @@ def process_result_value( ) -> T_Enum | None: return self._enum_cls(value) if value else None + @override def copy(self, **_kw: Any) -> Self: return EnumValueType(self._enum_cls, **self._opts) # type: ignore[return-value] @property + @override def python_type(self) -> type[T_Enum]: return self._enum_cls @@ -287,6 +296,7 @@ def __init__( self._use_name = use_name self._enum_cls = enum_cls + @override def process_bind_param( self, value: T_StrEnum | None, @@ -298,6 +308,7 @@ def process_bind_param( return value.name return cast(str, value.value) + @override def process_result_value( self, value: str | None, @@ -309,10 +320,12 @@ def process_result_value( return self._enum_cls[value] return self._enum_cls(value) + @override def copy(self, **_kw: Any) -> Self: return StrEnumType(self._enum_cls, self._use_name, **self._opts) # type: ignore[return-value] @property + @override def python_type(self) -> type[T_StrEnum]: return self._enum_cls @@ -336,6 +349,7 @@ def __init__(self, enum_cls: type[T_IntFlag], **opts: Any) -> None: super().__init__(**opts) self._enum_cls = enum_cls + @override def process_bind_param( self, value: T_IntFlag | None, @@ -345,6 +359,7 @@ def process_bind_param( return None return int(value) + @override def process_result_value( self, value: int | None, @@ -354,10 +369,12 @@ def process_result_value( return None return self._enum_cls(value) + @override def copy(self, **_kw: Any) -> Self: return IntFlagType(self._enum_cls, **self._opts) # type: ignore[return-value] @property + @override def python_type(self) -> type[T_IntFlag]: return self._enum_cls @@ -377,9 +394,11 @@ class CurvePublicKeyColumn(TypeDecorator[PublicKey]): impl = sa.String cache_ok = True + @override def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: return dialect.type_descriptor(sa.String(40)) + @override def process_bind_param( self, value: PublicKey | None, @@ -387,6 +406,7 @@ def process_bind_param( ) -> str | None: return value.decode("ascii") if value else None + @override def process_result_value( self, value: str | None, @@ -405,9 +425,11 @@ class QuotaScopeIDType(TypeDecorator[QuotaScopeID]): impl = sa.String cache_ok = True + @override def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: return dialect.type_descriptor(sa.String(64)) + @override def process_bind_param( self, value: QuotaScopeID | None, @@ -415,6 +437,7 @@ def process_bind_param( ) -> str | None: return str(value) if value else None + @override def process_result_value( self, value: str | None, @@ -431,6 +454,7 @@ class ResourceSlotColumn(TypeDecorator[ResourceSlot]): impl = JSONB cache_ok = True + @override def process_bind_param( self, value: ResourceSlot | None, @@ -440,6 +464,7 @@ def process_bind_param( return None return value.to_json() + @override def process_result_value( self, value: dict[str, str] | None, @@ -466,11 +491,13 @@ def __init__(self, schema: t.Trafaret) -> None: super().__init__() self._schema = schema + @override def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: if dialect.name == "sqlite": return dialect.type_descriptor(sa.JSON()) return super().load_dialect_impl(dialect) + @override def process_bind_param( self, value: Any | None, @@ -487,6 +514,7 @@ def process_bind_param( ) from e return value + @override def process_result_value( self, value: Any | None, @@ -496,6 +524,7 @@ def process_result_value( return self._schema.check({}) return self._schema.check(value) + @override def copy(self, **_kw: Any) -> Self: return StructuredJSONColumn(self._schema) # type: ignore[return-value] @@ -512,6 +541,7 @@ def __init__(self, schema: type[JSONSerializableMixin]) -> None: super().__init__() self._schema = schema + @override def process_bind_param( self, value: JSONSerializableMixin | None, _dialect: Dialect ) -> dict[str, Any] | None: @@ -519,6 +549,7 @@ def process_bind_param( return None return self._schema.to_json(value) + @override def process_result_value( self, value: dict[str, Any] | None, dialect: Dialect ) -> JSONSerializableMixin | None: @@ -526,6 +557,7 @@ def process_result_value( return None return self._schema.from_json(value) + @override def copy(self, **_kw: Any) -> Self: return StructuredJSONObjectColumn(self._schema) # type: ignore[return-value] @@ -543,14 +575,17 @@ def __init__(self, schema: type[JSONSerializableMixin]) -> None: super().__init__() self._schema = schema + @override def coerce_compared_value(self, _op: Any, _value: Any) -> JSONB: return JSONB() + @override def process_bind_param( self, value: list[JSONSerializableMixin] | None, _dialect: Dialect ) -> list[dict[str, Any]]: return [self._schema.to_json(item) for item in value] if value is not None else [] + @override def process_result_value( self, value: list[Any] | None, _dialect: Dialect ) -> list[JSONSerializableMixin]: @@ -558,6 +593,7 @@ def process_result_value( return [] return [self._schema.from_json(item) for item in value] + @override def copy(self, **_kw: Any) -> Self: return StructuredJSONObjectListColumn(self._schema) # type: ignore[return-value] @@ -575,6 +611,7 @@ def __init__(self, schema: type[TBaseModel]) -> None: super().__init__() self._schema = schema + @override def process_bind_param( self, value: TBaseModel | None, @@ -585,6 +622,7 @@ def process_bind_param( return value.model_dump(mode="json") return None + @override def process_result_value( self, value: dict[str, Any] | None, @@ -595,6 +633,7 @@ def process_result_value( return self._schema.model_validate(value) return None + @override def copy(self, **_kw: Any) -> Self: return PydanticColumn(self._schema) # type: ignore[return-value] @@ -612,9 +651,11 @@ def __init__(self, schema: type[TBaseModel]) -> None: super().__init__() self._schema = schema + @override def coerce_compared_value(self, _op: Any, _value: Any) -> JSONB: return JSONB() + @override def process_bind_param( self, value: list[TBaseModel] | None, _dialect: Dialect ) -> list[dict[str, Any]]: @@ -623,6 +664,7 @@ def process_bind_param( return [item.model_dump(mode="json") for item in value] return [] + @override def process_result_value( self, value: list[dict[str, Any]] | str | None, _dialect: Dialect ) -> list[TBaseModel]: @@ -634,6 +676,7 @@ def process_result_value( return [self._schema.model_validate(item) for item in value] return [] + @override def copy(self, **_kw: Any) -> Self: return PydanticListColumn(self._schema) # type: ignore[return-value] @@ -668,6 +711,7 @@ def __init__(self, schema: type[ABCColumnPayload]) -> None: super().__init__() self._schema = schema + @override def process_bind_param( self, value: ABCColumnPayload | None, @@ -678,6 +722,7 @@ def process_bind_param( return None return value.serialize() + @override def process_result_value( self, value: dict[str, Any] | None, @@ -688,6 +733,7 @@ def process_result_value( return None return self._schema.load(value) + @override def copy(self, **_kw: Any) -> Self: return ABCColumn(self._schema) # type: ignore[return-value] @@ -700,9 +746,11 @@ class URLColumn(TypeDecorator[yarl.URL]): impl = UnicodeText cache_ok = True + @override def process_bind_param(self, value: yarl.URL | None, _dialect: Dialect) -> str | None: return str(value) + @override def process_result_value(self, value: str | None, _dialect: Dialect) -> yarl.URL | None: if value is None: return None @@ -717,6 +765,7 @@ class IPColumn(TypeDecorator[ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6 impl = CIDR cache_ok = True + @override def process_bind_param( self, value: str | ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6Network] | None, @@ -733,6 +782,7 @@ def process_bind_param( raise InvalidAPIParameters(f"{value} is invalid IP address value") from e return str(cidr) + @override def process_result_value( self, value: str | None, _dialect: Dialect ) -> ReadableCIDR[ipaddress.IPv4Network | ipaddress.IPv6Network] | None: @@ -753,6 +803,7 @@ def __init__(self, perm_type: type[AbstractPermission]) -> None: super().__init__(sa.String) self._perm_type = perm_type + @override def process_bind_param( self, value: set[AbstractPermission] | None, @@ -765,6 +816,7 @@ def process_bind_param( except ValueError as e: raise InvalidAPIParameters(f"Invalid value for binding to {self._perm_type}") from e + @override def process_result_value( self, value: Sequence[str] | None, @@ -784,6 +836,7 @@ class VFolderHostPermissionColumn(TypeDecorator[VFolderHostPermissionMap]): cache_ok = True perm_col = PermissionListColumn(VFolderHostPermission) + @override def process_bind_param( self, value: Mapping[str, Any] | None, @@ -795,6 +848,7 @@ def process_bind_param( host: self.perm_col.process_bind_param(perms, dialect) for host, perms in value.items() } + @override def process_result_value( self, value: Mapping[str, Any] | None, @@ -840,11 +894,13 @@ def __init__(self, subtype: Callable[[uuid.UUID], TUUIDSubType] | None = None) - super().__init__() self._subtype_func: Callable[[uuid.UUID], Any] | None = subtype + @override def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: if dialect.name == "postgresql": return dialect.type_descriptor(UUID()) return dialect.type_descriptor(CHAR(16)) + @override def process_bind_param(self, value: Any | None, dialect: Dialect) -> str | bytes | None: # NOTE: DeploymentID, SessionId, KernelId are *not* actual types defined as classes, # but a "virtual" type that is an identity function at runtime. @@ -867,6 +923,7 @@ def _apply_subtype(self, raw: uuid.UUID) -> Any: return self._subtype_func(raw) return type(self).uuid_subtype_func(raw) + @override def process_result_value(self, value: Any, _dialect: Dialect) -> TUUIDSubType | None: if value is None: return value @@ -878,6 +935,7 @@ def process_result_value(self, value: Any, _dialect: Dialect) -> TUUIDSubType | return cast(TUUIDSubType, self._apply_subtype(uuid.UUID(bytes=value.bytes))) return cast(TUUIDSubType, self._apply_subtype(uuid.UUID(value))) + @override def copy(self, **_kw: Any) -> Self: return type(self)(self._subtype_func) @@ -906,9 +964,11 @@ def __init__( allow_unicode=allow_unicode, ) + @override def coerce_compared_value(self, _op: Any, _value: Any) -> Unicode: return Unicode() + @override def process_bind_param(self, value: Any | None, _dialect: Dialect) -> str | None: if value is None: return value @@ -1186,6 +1246,7 @@ class DecimalType(TypeDecorator[Decimal], Decimal): impl = sa.VARCHAR cache_ok = True + @override def process_bind_param( self, value: Decimal | None, @@ -1193,6 +1254,7 @@ def process_bind_param( ) -> str | None: return f"{value:f}" if value is not None else None + @override def process_result_value( self, value: Any | None, @@ -1201,5 +1263,6 @@ def process_result_value( return Decimal(value) if value is not None else None @property + @override def python_type(self) -> type[Decimal]: return Decimal diff --git a/src/ai/backend/manager/models/domain/row.py b/src/ai/backend/manager/models/domain/row.py index c06a52f3d08..ad73dfafafc 100644 --- a/src/ai/backend/manager/models/domain/row.py +++ b/src/ai/backend/manager/models/domain/row.py @@ -176,6 +176,7 @@ class DomainModel(RBACModel[DomainPermission]): _permissions: frozenset[DomainPermission] = field(default_factory=frozenset) @property + @override def permissions(self) -> Container[DomainPermission]: return self._permissions @@ -282,12 +283,14 @@ def _OR_coalesce( ) return cond + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(DomainRow).where(cond) + @override async def calculate_final_permission(self, rbac_obj: DomainRow) -> frozenset[DomainPermission]: domain_row = rbac_obj domain_name = domain_row.name diff --git a/src/ai/backend/manager/models/group/row.py b/src/ai/backend/manager/models/group/row.py index f512fbd4089..7efeef5f04b 100644 --- a/src/ai/backend/manager/models/group/row.py +++ b/src/ai/backend/manager/models/group/row.py @@ -391,6 +391,7 @@ class ProjectModel(RBACModel[ProjectPermission]): _permissions: frozenset[ProjectPermission] = field(default_factory=frozenset) @property + @override def permissions(self) -> Container[ProjectPermission]: return self._permissions @@ -605,12 +606,14 @@ def _OR_coalesce( ) return cond + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(GroupRow).where(cond) + @override async def calculate_final_permission(self, rbac_obj: GroupRow) -> frozenset[ProjectPermission]: project_row = rbac_obj project_id = project_row.id diff --git a/src/ai/backend/manager/models/hasher/bcrypt.py b/src/ai/backend/manager/models/hasher/bcrypt.py index b024ff95a4a..46487235d86 100644 --- a/src/ai/backend/manager/models/hasher/bcrypt.py +++ b/src/ai/backend/manager/models/hasher/bcrypt.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + import bcrypt from .base import PasswordHasher @@ -13,6 +15,7 @@ class BcryptHasher(PasswordHasher): """Bcrypt password hashing implementation.""" + @override def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: # Note: bcrypt generates its own salt, so salt_size is ignored hash_value = bcrypt.hashpw( @@ -26,6 +29,7 @@ def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: hash_value=hash_value, ) + @override def verify(self, password: str, hash_info: HashInfo) -> bool: try: # For bcrypt, the full hash is stored in hash_value diff --git a/src/ai/backend/manager/models/hasher/pbkdf2.py b/src/ai/backend/manager/models/hasher/pbkdf2.py index f272cf7734d..a2af15eef42 100644 --- a/src/ai/backend/manager/models/hasher/pbkdf2.py +++ b/src/ai/backend/manager/models/hasher/pbkdf2.py @@ -6,6 +6,7 @@ import hashlib import secrets +from typing import override from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm @@ -16,6 +17,7 @@ class PBKDF2_SHA256Hasher(PasswordHasher): """PBKDF2-SHA256 password hashing implementation.""" + @override def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: salt = secrets.token_bytes(password_info.salt_size) # Use hashlib.pbkdf2_hmac for standard PBKDF2 @@ -30,6 +32,7 @@ def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: hash_value=hashed.hex(), ) + @override def verify(self, password: str, hash_info: HashInfo) -> bool: if hash_info.algorithm != PasswordHashAlgorithm.PBKDF2_SHA256: return False @@ -50,6 +53,7 @@ def verify(self, password: str, hash_info: HashInfo) -> bool: class PBKDF2_SHA3_256Hasher(PasswordHasher): """PBKDF2-SHA3-256 password hashing implementation.""" + @override def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: salt = secrets.token_bytes(password_info.salt_size) # Use hashlib.pbkdf2_hmac with SHA3-256 @@ -64,6 +68,7 @@ def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: hash_value=hashed.hex(), ) + @override def verify(self, password: str, hash_info: HashInfo) -> bool: if hash_info.algorithm != PasswordHashAlgorithm.PBKDF2_SHA3_256: return False diff --git a/src/ai/backend/manager/models/hasher/sha.py b/src/ai/backend/manager/models/hasher/sha.py index 014e2343170..76232c9b264 100644 --- a/src/ai/backend/manager/models/hasher/sha.py +++ b/src/ai/backend/manager/models/hasher/sha.py @@ -6,6 +6,7 @@ import hashlib import secrets +from typing import override from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm @@ -16,6 +17,7 @@ class SHA256Hasher(PasswordHasher): """SHA-256 password hashing implementation with PBKDF2-like iterations.""" + @override def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: salt = secrets.token_bytes(password_info.salt_size).hex() hashed = password_info.password @@ -29,6 +31,7 @@ def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: hash_value=hashed, ) + @override def verify(self, password: str, hash_info: HashInfo) -> bool: if hash_info.algorithm != PasswordHashAlgorithm.SHA256: return False @@ -46,6 +49,7 @@ def verify(self, password: str, hash_info: HashInfo) -> bool: class SHA3_256Hasher(PasswordHasher): """SHA3-256 password hashing implementation with PBKDF2-like iterations.""" + @override def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: salt = secrets.token_bytes(password_info.salt_size).hex() hashed = password_info.password @@ -59,6 +63,7 @@ def generate_new_hash(self, password_info: PasswordInfo) -> HashInfo: hash_value=hashed, ) + @override def verify(self, password: str, hash_info: HashInfo) -> bool: if hash_info.algorithm != PasswordHashAlgorithm.SHA3_256: return False diff --git a/src/ai/backend/manager/models/hasher/types.py b/src/ai/backend/manager/models/hasher/types.py index db495236bc3..0afc82f5b16 100644 --- a/src/ai/backend/manager/models/hasher/types.py +++ b/src/ai/backend/manager/models/hasher/types.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass -from typing import Any, Final +from typing import Any, Final, override from sqlalchemy.types import VARCHAR, TypeDecorator @@ -138,6 +138,7 @@ class PasswordColumn(TypeDecorator[str]): impl = VARCHAR cache_ok = True + @override def process_bind_param(self, value: Any, _dialect: Any) -> str | None: if value is None: return None diff --git a/src/ai/backend/manager/models/huggingface_registry/row.py b/src/ai/backend/manager/models/huggingface_registry/row.py index 21265676d5b..c551c828cb2 100644 --- a/src/ai/backend/manager/models/huggingface_registry/row.py +++ b/src/ai/backend/manager/models/huggingface_registry/row.py @@ -2,7 +2,7 @@ import logging import uuid -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -59,9 +59,11 @@ class HuggingFaceRegistryRow(Base): # type: ignore[misc] viewonly=True, ) + @override def __str__(self) -> str: return f"HuggingFaceRegistryRow(id={self.id}, url={self.url}, token={self.token})" + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/image/row.py b/src/ai/backend/manager/models/image/row.py index 807f6bcdb52..08a4fdf1a4d 100644 --- a/src/ai/backend/manager/models/image/row.py +++ b/src/ai/backend/manager/models/image/row.py @@ -577,9 +577,11 @@ async def list( result = await session.execute(query) return list(result.scalars().all()) + @override def __str__(self) -> str: return self.image_ref.canonical + f" ({self.image_ref.architecture})" + @override def __repr__(self) -> str: return self.__str__() @@ -931,12 +933,14 @@ def _OR_coalesce( ) return cond + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(ImageRow).where(cond) + @override async def calculate_final_permission(self, rbac_obj: ImageRow) -> frozenset[ImagePermission]: image_row = rbac_obj image_id = image_row.id diff --git a/src/ai/backend/manager/models/object_storage/row.py b/src/ai/backend/manager/models/object_storage/row.py index 2a1d241b079..f8ee48ca5e0 100644 --- a/src/ai/backend/manager/models/object_storage/row.py +++ b/src/ai/backend/manager/models/object_storage/row.py @@ -2,7 +2,7 @@ import logging import uuid -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -88,6 +88,7 @@ class ObjectStorageRow(Base): # type: ignore[misc] primaryjoin=_get_object_storage_namespace_join_cond, ) + @override def __str__(self) -> str: return ( f"ObjectStorageRow(" @@ -100,6 +101,7 @@ def __str__(self) -> str: f"region={self.region})" ) + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/rbac/__init__.py b/src/ai/backend/manager/models/rbac/__init__.py index ab4f800e74c..e374fbc4680 100644 --- a/src/ai/backend/manager/models/rbac/__init__.py +++ b/src/ai/backend/manager/models/rbac/__init__.py @@ -5,7 +5,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Callable, Container, Iterable, Mapping, Sequence from dataclasses import dataclass, field -from typing import Any, Self, TypeVar, cast +from typing import Any, Self, TypeVar, cast, override import sqlalchemy as sa from sqlalchemy.ext.asyncio import AsyncSession @@ -267,6 +267,7 @@ async def _calculate_role_in_scope_for_user( class BaseScope(metaclass=ABCMeta): @abstractmethod + @override def __str__(self) -> str: pass @@ -282,13 +283,16 @@ def deserialize(cls, val: str) -> Self: @dataclass(frozen=True) class SystemScope(BaseScope): + @override def __str__(self) -> str: return "system scope()" + @override def serialize(self) -> str: return "system:" @classmethod + @override def deserialize(cls, val: str) -> Self: type_, _, _ = val.partition(":") if type_ != "system": @@ -300,13 +304,16 @@ def deserialize(cls, val: str) -> Self: class DomainScope(BaseScope): domain_name: str + @override def __str__(self) -> str: return f"Domain(name: {self.domain_name})" + @override def serialize(self) -> str: return f"domain:{self.domain_name}" @classmethod + @override def deserialize(cls, val: str) -> Self: type_, _, domain_name = val.partition(":") if type_ != "domain": @@ -319,15 +326,18 @@ class ProjectScope(BaseScope): project_id: uuid.UUID domain_name: str | None = None + @override def __str__(self) -> str: return f"Project(id: {self.project_id}, domain: {self.domain_name}])" + @override def serialize(self) -> str: if self.domain_name is not None: return f"project:{self.project_id}:{self.domain_name}" return f"project:{self.project_id}" @classmethod + @override def deserialize(cls, val: str) -> Self: type_, _, values = val.partition(":") if type_ != "project": @@ -343,15 +353,18 @@ class UserScope(BaseScope): user_id: uuid.UUID domain_name: str | None = None + @override def __str__(self) -> str: return f"User(id: {self.user_id}, domain: {self.domain_name})" + @override def serialize(self) -> str: if self.domain_name is not None: return f"user:{self.user_id}:{self.domain_name}" return f"user:{self.user_id}" @classmethod + @override def deserialize(cls, val: str) -> Self: type_, _, values = val.partition(":") if type_ != "user": @@ -438,9 +451,11 @@ class ScalingGroup(ExtraScope): class ContainerRegistryScope(ExtraScope): registry_id: uuid.UUID + @override def __str__(self) -> str: return f"container_registry:{self.registry_id!s}" + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/rbac/exceptions.py b/src/ai/backend/manager/models/rbac/exceptions.py index 9acf56ebd96..c3d9654eeeb 100644 --- a/src/ai/backend/manager/models/rbac/exceptions.py +++ b/src/ai/backend/manager/models/rbac/exceptions.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class InvalidScope(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/invalid-scope" error_title = "Invalid scope specified." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, @@ -25,6 +28,7 @@ class ScopeTypeMismatch(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/scope-type-mismatch" error_title = "Scope type mismatch." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, @@ -37,6 +41,7 @@ class NotEnoughPermission(BackendAIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/not-enough-permission" error_title = "Not enough permission." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.PERMISSION, diff --git a/src/ai/backend/manager/models/reservoir_registry/row.py b/src/ai/backend/manager/models/reservoir_registry/row.py index 7b55e9a6fe9..feb7ade315f 100644 --- a/src/ai/backend/manager/models/reservoir_registry/row.py +++ b/src/ai/backend/manager/models/reservoir_registry/row.py @@ -2,7 +2,7 @@ import logging import uuid -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -61,9 +61,11 @@ class ReservoirRegistryRow(Base): # type: ignore[misc] viewonly=True, ) + @override def __str__(self) -> str: return f"ReservoirRegistryRow(id={self.id}, endpoint={self.endpoint}, api_version={self.api_version})" + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/models/scaling_group/row.py b/src/ai/backend/manager/models/scaling_group/row.py index d514714825d..24443026642 100644 --- a/src/ai/backend/manager/models/scaling_group/row.py +++ b/src/ai/backend/manager/models/scaling_group/row.py @@ -452,6 +452,7 @@ class ScalingGroupModel(RBACModel[ScalingGroupPermission]): _permissions: frozenset[ScalingGroupPermission] = field(default_factory=frozenset) @property + @override def permissions(self) -> Container[ScalingGroupPermission]: return self._permissions @@ -623,12 +624,14 @@ def _OR_coalesce( ) return cond + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(ScalingGroupRow).where(cond) + @override async def calculate_final_permission(self, rbac_obj: str) -> frozenset[ScalingGroupPermission]: host_name = rbac_obj return self.object_id_to_additional_permission_map.get(host_name, frozenset()) diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index ead1cb36450..e80e534134a 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -1461,12 +1461,14 @@ def _OR_coalesce( ) return cond + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(SessionRow).where(cond) + @override async def calculate_final_permission( self, rbac_obj: SessionRow ) -> frozenset[ComputeSessionPermission]: diff --git a/src/ai/backend/manager/models/storage.py b/src/ai/backend/manager/models/storage.py index 1a898c80c1f..bbbf799a9b7 100644 --- a/src/ai/backend/manager/models/storage.py +++ b/src/ai/backend/manager/models/storage.py @@ -114,9 +114,11 @@ class PermissionContext(AbstractPermissionContext[StorageHostPermission, str, st def host_to_permissions_map(self) -> StorageHostToPermissionMap: return self.object_id_to_additional_permission_map + @override async def build_query(self) -> sa.sql.Select[Any] | None: return None + @override async def calculate_final_permission(self, rbac_obj: str) -> frozenset[StorageHostPermission]: host_name = rbac_obj return self.object_id_to_additional_permission_map.get(host_name, frozenset()) diff --git a/src/ai/backend/manager/models/utils.py b/src/ai/backend/manager/models/utils.py index 5fb1bf05229..98090c6967a 100644 --- a/src/ai/backend/manager/models/utils.py +++ b/src/ai/backend/manager/models/utils.py @@ -14,6 +14,7 @@ ParamSpec, TypeVar, overload, + override, ) import sqlalchemy as sa @@ -131,6 +132,7 @@ async def _begin_readonly( self._readonly_txn_count -= 1 @actxmgr + @override async def begin(self, bind: SAConnection | None = None) -> AsyncIterator[SAConnection]: if bind is None: async with self.connect() as _bind, self._begin(_bind) as conn: diff --git a/src/ai/backend/manager/models/vfolder/row.py b/src/ai/backend/manager/models/vfolder/row.py index 5f3c622b5e7..81c127d2ae3 100644 --- a/src/ai/backend/manager/models/vfolder/row.py +++ b/src/ai/backend/manager/models/vfolder/row.py @@ -1555,12 +1555,14 @@ def _OR_coalesce( def apply_host_permission_ctx(self, host_permission_ctx: StorageHostPermissionContext) -> None: self.host_permission_ctx = host_permission_ctx + @override async def build_query(self) -> sa.sql.Select[Any] | None: cond = self.query_condition if cond is None: return None return sa.select(VFolderRow).where(cond) + @override async def calculate_final_permission( self, rbac_obj: VFolderRow ) -> frozenset[VFolderRBACPermission]: diff --git a/src/ai/backend/manager/models/vfs_storage/row.py b/src/ai/backend/manager/models/vfs_storage/row.py index 11ca46f4041..247e3d53c35 100644 --- a/src/ai/backend/manager/models/vfs_storage/row.py +++ b/src/ai/backend/manager/models/vfs_storage/row.py @@ -3,7 +3,7 @@ import logging import uuid from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import sqlalchemy as sa from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship @@ -58,6 +58,7 @@ class VFSStorageRow(Base): # type: ignore[misc] ) ) + @override def __str__(self) -> str: return ( f"VFSStorageRow(" @@ -67,6 +68,7 @@ def __str__(self) -> str: f"base_path={self.base_path})" ) + @override def __repr__(self) -> str: return self.__str__() diff --git a/src/ai/backend/manager/network/overlay.py b/src/ai/backend/manager/network/overlay.py index ae2c9fec6d1..f5c7a9815c2 100644 --- a/src/ai/backend/manager/network/overlay.py +++ b/src/ai/backend/manager/network/overlay.py @@ -4,7 +4,7 @@ import uuid from collections.abc import Mapping from http import HTTPStatus -from typing import Any +from typing import Any, override import aiodocker import trafaret as t @@ -29,6 +29,7 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, self.plugin_config = plugin_config_iv.check(plugin_config) + @override async def init(self, context: Any = None) -> None: self.docker = aiodocker.Docker() @@ -40,12 +41,15 @@ async def init(self, context: Any = None) -> None: "'docker swarm init' or 'docker swarm init --advertise-addr '" ) + @override async def cleanup(self) -> None: await self.docker.close() + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: return await super().update_plugin_config(plugin_config) + @override async def create_network( self, *, identifier: str | None = None, options: dict[str, Any] | None = None ) -> NetworkInfo: @@ -96,6 +100,7 @@ async def create_network( }, ) + @override async def destroy_network(self, network_id: str) -> None: try: await asyncio.sleep(2.0) diff --git a/src/ai/backend/manager/notification/channels/email/channel.py b/src/ai/backend/manager/notification/channels/email/channel.py index 30804eb9937..1664be5dbdd 100644 --- a/src/ai/backend/manager/notification/channels/email/channel.py +++ b/src/ai/backend/manager/notification/channels/email/channel.py @@ -7,6 +7,7 @@ import smtplib from email.mime.text import MIMEText from functools import partial +from typing import override from ai.backend.common.data.notification.types import EmailSpec from ai.backend.logging import BraceStyleAdapter @@ -28,6 +29,7 @@ def __init__(self, email_spec: EmailSpec) -> None: """Initialize email channel with configuration.""" self._spec = email_spec + @override async def send(self, message: NotificationMessage) -> SendResult: """ Send notification via email using SMTP. diff --git a/src/ai/backend/manager/notification/channels/webhook/channel.py b/src/ai/backend/manager/notification/channels/webhook/channel.py index caff64633f7..c8af8e0dcec 100644 --- a/src/ai/backend/manager/notification/channels/webhook/channel.py +++ b/src/ai/backend/manager/notification/channels/webhook/channel.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from typing import override from urllib.parse import urlparse import aiohttp @@ -37,6 +38,7 @@ def __init__(self, http_client_pool: ClientPool, webhook_spec: WebhookSpec) -> N self._http_client_pool = http_client_pool self._webhook_spec = webhook_spec + @override async def send( self, message: NotificationMessage, diff --git a/src/ai/backend/manager/pglock.py b/src/ai/backend/manager/pglock.py index 6a3cd963422..833a34a8e32 100644 --- a/src/ai/backend/manager/pglock.py +++ b/src/ai/backend/manager/pglock.py @@ -1,7 +1,7 @@ from __future__ import annotations from contextlib import AbstractAsyncContextManager -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from ai.backend.common.lock import AbstractDistributedLock @@ -20,10 +20,12 @@ def __init__(self, db: ExtendedAsyncSAEngine, lock_id: LockID) -> None: self.lock_id = lock_id self._lock_ctx = None + @override async def __aenter__(self) -> Any: self._lock_ctx = self.db.advisory_lock(self.lock_id) await self._lock_ctx.__aenter__() + @override async def __aexit__(self, *exc_info: Any) -> bool | None: if self._lock_ctx is None: raise DBOperationFailed("Lock context is not initialized") diff --git a/src/ai/backend/manager/plugin/error_monitor.py b/src/ai/backend/manager/plugin/error_monitor.py index 94ec04d0820..99d08f53e4b 100644 --- a/src/ai/backend/manager/plugin/error_monitor.py +++ b/src/ai/backend/manager/plugin/error_monitor.py @@ -25,6 +25,7 @@ class ErrorMonitor(AbstractErrorReporterPlugin): _error_log_repository: ErrorLogRepository + @override async def init(self, context: Any | None = None) -> None: if context is None: log.warning( @@ -36,12 +37,15 @@ async def init(self, context: Any | None = None) -> None: self.enabled = True self._error_log_repository = context["error_log_repository"] + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: pass + @override async def capture_message(self, message: str) -> None: pass + @override async def capture_exception( self, exc_instance: Exception | None = None, @@ -95,6 +99,7 @@ async def capture_exception( class ErrorEventDispatcher(AbstractEventDispatcherPlugin): _error_log_repository: ErrorLogRepository + @override async def init(self, context: Any | None = None) -> None: if context is None: log.warning( diff --git a/src/ai/backend/manager/plugin/keypair/hook.py b/src/ai/backend/manager/plugin/keypair/hook.py index 69c476f2e27..722f40db646 100644 --- a/src/ai/backend/manager/plugin/keypair/hook.py +++ b/src/ai/backend/manager/plugin/keypair/hook.py @@ -5,7 +5,7 @@ import logging import re from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, override import sqlalchemy as sa import trafaret as t @@ -36,17 +36,21 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, super().__init__(plugin_config, local_config) self.plugin_config = plugin_config_checker.check(self.plugin_config) + @override def get_handlers(self) -> Sequence[tuple[str, HookHandler]]: return [ ("AUTHORIZE", self.authorize), ] + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: self.plugin_config = plugin_config + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass diff --git a/src/ai/backend/manager/plugin/keypair/webapp.py b/src/ai/backend/manager/plugin/keypair/webapp.py index 3e931431b13..0ee1a1e99b6 100644 --- a/src/ai/backend/manager/plugin/keypair/webapp.py +++ b/src/ai/backend/manager/plugin/keypair/webapp.py @@ -1,7 +1,7 @@ import json import logging from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, override import aiohttp_cors import yarl @@ -65,15 +65,19 @@ async def _webapp_shutdown(app: web.Application) -> None: class KeypairAuthWebAppPlugin(WebappPlugin): + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: self.plugin_config = new_plugin_config + @override async def create_app( self, cors_options: CORSOptions, diff --git a/src/ai/backend/manager/plugin/openid/exceptions.py b/src/ai/backend/manager/plugin/openid/exceptions.py index 7b7d3723e40..2def4dee6d9 100644 --- a/src/ai/backend/manager/plugin/openid/exceptions.py +++ b/src/ai/backend/manager/plugin/openid/exceptions.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -14,6 +16,7 @@ class InvalidSession(BackendAIError, web.HTTPUnauthorized): error_title = "The OpenID session is invalid or has expired." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, diff --git a/src/ai/backend/manager/plugin/openid/hook.py b/src/ai/backend/manager/plugin/openid/hook.py index 82308844ac5..90628c6d1a2 100644 --- a/src/ai/backend/manager/plugin/openid/hook.py +++ b/src/ai/backend/manager/plugin/openid/hook.py @@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence from typing import ( Any, + override, ) import jwt @@ -34,16 +35,20 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, super().__init__(plugin_config, local_config) self._config = OIDCHookConfig(**plugin_config) + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: self.plugin_config = plugin_config self._config = OIDCHookConfig(**plugin_config) + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override def get_handlers(self) -> Sequence[tuple[str, HookHandler]]: return [ ("AUTHORIZE", self.pre_auth_hook), diff --git a/src/ai/backend/manager/plugin/openid/webapp.py b/src/ai/backend/manager/plugin/openid/webapp.py index fee1dc39220..1aba48eab46 100644 --- a/src/ai/backend/manager/plugin/openid/webapp.py +++ b/src/ai/backend/manager/plugin/openid/webapp.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, + override, ) import aiohttp @@ -257,17 +258,21 @@ def __init__(self, app: web.Application) -> None: self._app = app @property + @override def name(self) -> str: return "openid.jwks_refresh" @property + @override def interval(self) -> float: return _JWKS_REFRESH_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: async with aiohttp.ClientSession() as sess: async with sess.get(self._app["openid.jwks_uri"]) as resp: @@ -284,12 +289,15 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, super().__init__(plugin_config, local_config) self._config = OIDCWebAppConfig(**plugin_config) + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_etcd_config: Mapping[str, Any]) -> None: self.plugin_config = new_etcd_config self._config = OIDCWebAppConfig(**new_etcd_config) @@ -426,6 +434,7 @@ async def redirect(self, request: web.Request) -> web.Response: token = encode_jwt_token(token_data, self._config.secret) return web.HTTPFound(redirect_uri.update_query({"sToken": token})) + @override async def create_app( self, cors_options: CORSOptions, diff --git a/src/ai/backend/manager/plugin/totp/exception.py b/src/ai/backend/manager/plugin/totp/exception.py index 6f020edd42f..7a6ec2bfb70 100644 --- a/src/ai/backend/manager/plugin/totp/exception.py +++ b/src/ai/backend/manager/plugin/totp/exception.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -14,6 +16,7 @@ class ExpiredToken(BackendAIError, web.HTTPBadRequest): error_title = "Expired Token" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -27,6 +30,7 @@ class InvalidToken(BackendAIError, web.HTTPBadRequest): error_title = "Invalid Token" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -40,6 +44,7 @@ class AuthorizationFailed(BackendAIError, web.HTTPUnauthorized): error_title = "TOTP authorization failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, diff --git a/src/ai/backend/manager/plugin/totp/hook.py b/src/ai/backend/manager/plugin/totp/hook.py index ccd3d7f38ba..24e39cd41b3 100644 --- a/src/ai/backend/manager/plugin/totp/hook.py +++ b/src/ai/backend/manager/plugin/totp/hook.py @@ -1,6 +1,6 @@ import logging from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, override import pyotp from aiohttp import web @@ -39,17 +39,21 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, self._plugin_config.token_secret, self._plugin_config.token_lifetime ) + @override def get_handlers(self) -> Sequence[tuple[str, HookHandler]]: return [ ("POST_AUTHORIZE", self.validate_otp), ] + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: self.plugin_config = plugin_config self._plugin_config = TOTPConfig(**plugin_config) diff --git a/src/ai/backend/manager/plugin/totp/webapp.py b/src/ai/backend/manager/plugin/totp/webapp.py index 20b56e33813..f90453e6f95 100644 --- a/src/ai/backend/manager/plugin/totp/webapp.py +++ b/src/ai/backend/manager/plugin/totp/webapp.py @@ -2,7 +2,7 @@ import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any, cast +from typing import Any, cast, override import aiohttp_cors import pyotp @@ -252,18 +252,22 @@ def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, self._plugin_config.token_secret, self._plugin_config.token_lifetime ) + @override async def init(self, context: Any = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: self.plugin_config = new_plugin_config self._plugin_config = TOTPConfig(**self.plugin_config) self._token_parser.set_secret(self._plugin_config.token_secret) self._token_parser.set_lifetime(self._plugin_config.token_lifetime) + @override async def create_app( self, cors_options: CORSOptions, diff --git a/src/ai/backend/manager/repositories/auth/types.py b/src/ai/backend/manager/repositories/auth/types.py index a4837ed8b98..474229e12dc 100644 --- a/src/ai/backend/manager/repositories/auth/types.py +++ b/src/ai/backend/manager/repositories/auth/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -21,6 +22,7 @@ class MyLoginSessionSearchScope(SearchScope): user_id: UUID + @override def to_condition(self) -> QueryCondition: user_id = self.user_id @@ -30,6 +32,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: return [ ExistenceCheck( @@ -46,6 +49,7 @@ class MyLoginHistorySearchScope(SearchScope): user_id: UUID + @override def to_condition(self) -> QueryCondition: user_id = self.user_id @@ -55,6 +59,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: return [ ExistenceCheck( diff --git a/src/ai/backend/manager/repositories/base/pagination.py b/src/ai/backend/manager/repositories/base/pagination.py index 7736c5338a5..3a6a2e49fae 100644 --- a/src/ai/backend/manager/repositories/base/pagination.py +++ b/src/ai/backend/manager/repositories/base/pagination.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import sqlalchemy as sa @@ -75,13 +75,16 @@ class NoPagination(QueryPagination): """ @property + @override def uses_window_function(self) -> bool: return False + @override def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: """No pagination applied - returns query unchanged.""" return query + @override def compute_page_info( self, rows: list[Row[Any]], _total_count: int ) -> PageInfoResult[Row[Any]]: @@ -110,9 +113,11 @@ class OffsetPagination(QueryPagination): """Number of items to skip from the beginning (must be non-negative).""" @property + @override def uses_window_function(self) -> bool: return True + @override def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: """Apply offset-based pagination to query.""" @@ -121,6 +126,7 @@ def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: query = query.offset(self.offset) return query + @override def compute_page_info(self, rows: list[Row[Any]], total_count: int) -> PageInfoResult[Row[Any]]: """Compute pagination info for offset-based pagination.""" @@ -156,9 +162,11 @@ class CursorForwardPagination(QueryPagination): """Optional QueryCondition for cursor position. If None, starts from the beginning.""" @property + @override def uses_window_function(self) -> bool: return False + @override def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: """ Apply cursor-based forward pagination to query. @@ -170,6 +178,7 @@ def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: query = query.order_by(self.cursor_order) return query.limit(self.first + 1) + @override def compute_page_info( self, rows: list[Row[Any]], _total_count: int ) -> PageInfoResult[Row[Any]]: @@ -210,9 +219,11 @@ class CursorBackwardPagination(QueryPagination): """Optional QueryCondition for cursor position. If None, starts from the end.""" @property + @override def uses_window_function(self) -> bool: return False + @override def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: """ Apply cursor-based backward pagination to query. @@ -224,6 +235,7 @@ def apply(self, query: sa.sql.Select[Any]) -> sa.sql.Select[Any]: query = query.order_by(self.cursor_order) return query.limit(self.last + 1) + @override def compute_page_info( self, rows: list[Row[Any]], _total_count: int ) -> PageInfoResult[Row[Any]]: diff --git a/src/ai/backend/manager/repositories/deployment/types/endpoint.py b/src/ai/backend/manager/repositories/deployment/types/endpoint.py index 85142329ea9..2ea41ce2381 100644 --- a/src/ai/backend/manager/repositories/deployment/types/endpoint.py +++ b/src/ai/backend/manager/repositories/deployment/types/endpoint.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field from datetime import datetime -from typing import Any +from typing import Any, override from uuid import UUID import sqlalchemy as sa @@ -156,6 +156,7 @@ class ProjectDeploymentSearchScope(SearchScope): project_id: UUID + @override def to_condition(self) -> QueryCondition: project_id = self.project_id @@ -165,6 +166,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: return [ ExistenceCheck( diff --git a/src/ai/backend/manager/repositories/domain/types.py b/src/ai/backend/manager/repositories/domain/types.py index 074fb6d5573..d0ff684b98b 100644 --- a/src/ai/backend/manager/repositories/domain/types.py +++ b/src/ai/backend/manager/repositories/domain/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override import sqlalchemy as sa @@ -30,6 +31,7 @@ class DomainSearchScope(SearchScope): resource_group: str """Required. The scaling group (resource group) to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for ScalingGroupForDomainRow.""" resource_group = self.resource_group @@ -40,6 +42,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ diff --git a/src/ai/backend/manager/repositories/fair_share/types.py b/src/ai/backend/manager/repositories/fair_share/types.py index 9beea047146..ec7481fe313 100644 --- a/src/ai/backend/manager/repositories/fair_share/types.py +++ b/src/ai/backend/manager/repositories/fair_share/types.py @@ -8,6 +8,7 @@ import uuid from collections.abc import Sequence from dataclasses import dataclass +from typing import override import sqlalchemy as sa @@ -52,6 +53,7 @@ class DomainFairShareSearchScope(SearchScope): resource_group: str """Required. The scaling group to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for DomainRow. @@ -65,6 +67,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ @@ -89,6 +92,7 @@ class ProjectFairShareSearchScope(SearchScope): domain_name: str """Required. The domain to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for GroupRow filtered by domain. @@ -103,6 +107,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ @@ -135,6 +140,7 @@ class UserFairShareSearchScope(SearchScope): project_id: uuid.UUID """Required. The project to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for GroupRow filtered by domain and project. @@ -153,6 +159,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks( self, ) -> Sequence[ExistenceCheck[str] | ExistenceCheck[uuid.UUID]]: diff --git a/src/ai/backend/manager/repositories/group/types.py b/src/ai/backend/manager/repositories/group/types.py index dfa4b948fca..30df4cf6707 100644 --- a/src/ai/backend/manager/repositories/group/types.py +++ b/src/ai/backend/manager/repositories/group/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -46,6 +47,7 @@ class DomainProjectSearchScope(SearchScope): domain_name: str """Required. The domain to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for GroupRow.""" domain_name = self.domain_name @@ -56,6 +58,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ @@ -79,6 +82,7 @@ class UserProjectSearchScope(SearchScope): user_uuid: UUID """Required. The user UUID to search projects for.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition on AssociationScopesEntitiesRow. @@ -99,6 +103,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation. diff --git a/src/ai/backend/manager/repositories/image/creators.py b/src/ai/backend/manager/repositories/image/creators.py index 9f8f4d6f06f..75294420b1b 100644 --- a/src/ai/backend/manager/repositories/image/creators.py +++ b/src/ai/backend/manager/repositories/image/creators.py @@ -16,6 +16,7 @@ class ImageAliasCreatorSpec(CreatorSpec[ImageAliasRow]): alias: str image_id: UUID + @override def build_row(self) -> ImageAliasRow: return ImageAliasRow( alias=self.alias, diff --git a/src/ai/backend/manager/repositories/keypair/types.py b/src/ai/backend/manager/repositories/keypair/types.py index a3bc66b4e37..6215b3b4e20 100644 --- a/src/ai/backend/manager/repositories/keypair/types.py +++ b/src/ai/backend/manager/repositories/keypair/types.py @@ -7,6 +7,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -35,6 +36,7 @@ class UserKeypairSearchScope(SearchScope): user_uuid: UUID """Required. The user whose keypairs to search.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for KeyPairRow.""" user_uuid = self.user_uuid @@ -45,6 +47,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ @@ -69,6 +72,7 @@ class KeypairResourcePolicyKeypairSearchScope(SearchScope): resource_policy_name: str """Required. The keypair resource policy whose keypairs to search.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for KeyPairRow.""" resource_policy_name = self.resource_policy_name @@ -79,6 +83,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ diff --git a/src/ai/backend/manager/repositories/model_card/types.py b/src/ai/backend/manager/repositories/model_card/types.py index 4e6dc0816e2..ede0bafcec5 100644 --- a/src/ai/backend/manager/repositories/model_card/types.py +++ b/src/ai/backend/manager/repositories/model_card/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -60,6 +61,7 @@ class ProjectModelCardSearchScope(SearchScope): project_id: UUID user_id: UUID + @override def to_condition(self) -> QueryCondition: project_id = self.project_id @@ -69,6 +71,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: return [ ExistenceCheck( @@ -101,6 +104,7 @@ class VFolderModelCardSearchScope(SearchScope): vfolder_id: VFolderUUID + @override def to_condition(self) -> QueryCondition: vfolder_id = self.vfolder_id @@ -110,5 +114,6 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: return () diff --git a/src/ai/backend/manager/repositories/ops/rbac/provider.py b/src/ai/backend/manager/repositories/ops/rbac/provider.py index 872aa92501c..214ca483c5d 100644 --- a/src/ai/backend/manager/repositories/ops/rbac/provider.py +++ b/src/ai/backend/manager/repositories/ops/rbac/provider.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterator, Collection, Sequence from contextlib import asynccontextmanager +from typing import override from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession as SASession @@ -76,6 +77,7 @@ class RBACOpsProvider(DBOpsProvider): """Hands out :class:`RBACWriteOps` for the read-write surface.""" @asynccontextmanager + @override async def write_ops(self) -> AsyncIterator[RBACWriteOps]: async with self._db.begin_session_read_committed() as sess: yield RBACWriteOps(sess) diff --git a/src/ai/backend/manager/repositories/ops/sokovan/provider.py b/src/ai/backend/manager/repositories/ops/sokovan/provider.py index fc8663b2143..3a8038e274a 100644 --- a/src/ai/backend/manager/repositories/ops/sokovan/provider.py +++ b/src/ai/backend/manager/repositories/ops/sokovan/provider.py @@ -12,7 +12,7 @@ from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import final +from typing import final, override import sqlalchemy as sa @@ -70,6 +70,7 @@ class SokovanOpsProvider(RBACOpsProvider): """Hands out :class:`SokovanWriteOps` for the read-write surface.""" @asynccontextmanager + @override async def write_ops(self) -> AsyncIterator[SokovanWriteOps]: async with self._db.begin_session_read_committed() as sess: yield SokovanWriteOps(sess) diff --git a/src/ai/backend/manager/repositories/permission_controller/types.py b/src/ai/backend/manager/repositories/permission_controller/types.py index 352ed37d0e3..dbcad12ab74 100644 --- a/src/ai/backend/manager/repositories/permission_controller/types.py +++ b/src/ai/backend/manager/repositories/permission_controller/types.py @@ -3,7 +3,7 @@ import uuid from collections.abc import Sequence from dataclasses import dataclass -from typing import Any +from typing import Any, override import sqlalchemy as sa @@ -35,6 +35,7 @@ class PermissionSearchScope(SearchScope): role_id: uuid.UUID + @override def to_condition(self) -> QueryCondition: role_id = self.role_id @@ -44,6 +45,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [ ExistenceCheck( @@ -61,6 +63,7 @@ class ScopedRoleSearchScope(SearchScope): element_type: RBACElementType scope_id: str + @override def to_condition(self) -> QueryCondition: element_type = self.element_type scope_id = self.scope_id @@ -76,6 +79,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [] @@ -86,6 +90,7 @@ class ObjectPermissionSearchScope(SearchScope): role_id: uuid.UUID + @override def to_condition(self) -> QueryCondition: role_id = self.role_id @@ -95,6 +100,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [ ExistenceCheck( diff --git a/src/ai/backend/manager/repositories/resource_preset/db_source/types.py b/src/ai/backend/manager/repositories/resource_preset/db_source/types.py index e9dda6265e1..e8a7eb379fa 100644 --- a/src/ai/backend/manager/repositories/resource_preset/db_source/types.py +++ b/src/ai/backend/manager/repositories/resource_preset/db_source/types.py @@ -5,7 +5,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass -from typing import Any +from typing import Any, override from uuid import UUID from ai.backend.common.types import ( @@ -136,6 +136,7 @@ class AccessKeyFilter(ResourceOccupancyFilter): def __init__(self, access_key: AccessKey) -> None: self.access_key = access_key + @override def get_condition(self) -> Any: return KernelRow.access_key == self.access_key @@ -146,6 +147,7 @@ class GroupIdFilter(ResourceOccupancyFilter): def __init__(self, group_id: UUID) -> None: self.group_id = group_id + @override def get_condition(self) -> Any: return KernelRow.group_id == self.group_id @@ -156,5 +158,6 @@ class DomainNameFilter(ResourceOccupancyFilter): def __init__(self, domain_name: str) -> None: self.domain_name = domain_name + @override def get_condition(self) -> Any: return KernelRow.domain_name == self.domain_name diff --git a/src/ai/backend/manager/repositories/resource_usage_history/types.py b/src/ai/backend/manager/repositories/resource_usage_history/types.py index 54f82324b78..8c4411f722d 100644 --- a/src/ai/backend/manager/repositories/resource_usage_history/types.py +++ b/src/ai/backend/manager/repositories/resource_usage_history/types.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from dataclasses import dataclass from datetime import date, datetime -from typing import Any +from typing import Any, override import sqlalchemy as sa @@ -211,6 +211,7 @@ class DomainUsageBucketSearchScope(SearchScope): resource_group: str domain_name: str + @override def to_condition(self) -> QueryCondition: resource_group = self.resource_group domain_name = self.domain_name @@ -224,6 +225,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [ ExistenceCheck( @@ -247,6 +249,7 @@ class ProjectUsageBucketSearchScope(SearchScope): domain_name: str project_id: uuid.UUID + @override def to_condition(self) -> QueryCondition: resource_group = self.resource_group domain_name = self.domain_name @@ -262,6 +265,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [ ExistenceCheck( @@ -291,6 +295,7 @@ class UserUsageBucketSearchScope(SearchScope): project_id: uuid.UUID user_uuid: uuid.UUID + @override def to_condition(self) -> QueryCondition: resource_group = self.resource_group domain_name = self.domain_name @@ -308,6 +313,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return [ ExistenceCheck( diff --git a/src/ai/backend/manager/repositories/role_invitation/types.py b/src/ai/backend/manager/repositories/role_invitation/types.py index 2089f30fba7..5026e634dd3 100644 --- a/src/ai/backend/manager/repositories/role_invitation/types.py +++ b/src/ai/backend/manager/repositories/role_invitation/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -30,6 +31,7 @@ class InviteeSearchScope(SearchScope): invitee_user_id: UUID + @override def to_condition(self) -> QueryCondition: invitee_user_id = self.invitee_user_id @@ -39,6 +41,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: # Invitee existence is guaranteed by RBAC authentication before reaching here. return [] @@ -50,6 +53,7 @@ class InviterSearchScope(SearchScope): inviter_user_id: UUID + @override def to_condition(self) -> QueryCondition: inviter_user_id = self.inviter_user_id @@ -59,6 +63,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: # Inviter existence is guaranteed by RBAC authentication before reaching here. return [] @@ -70,6 +75,7 @@ class RoleInvitationSearchScope(SearchScope): role_id: UUID + @override def to_condition(self) -> QueryCondition: role_id = self.role_id @@ -79,6 +85,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: # Role existence is validated by the RBAC permission check before reaching here. return [] diff --git a/src/ai/backend/manager/repositories/scheduling_history/types.py b/src/ai/backend/manager/repositories/scheduling_history/types.py index 1e80f7a5727..150834e4043 100644 --- a/src/ai/backend/manager/repositories/scheduling_history/types.py +++ b/src/ai/backend/manager/repositories/scheduling_history/types.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import Any, override from uuid import UUID from ai.backend.common.data.filter_specs import UUIDEqualMatchSpec @@ -42,6 +42,7 @@ class SessionSchedulingHistorySearchScope(SearchScope): session_id: UUID """Required. The session to search history for.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for SessionSchedulingHistoryRow.""" return SessionSchedulingHistoryConditions.by_session_id_filter( @@ -49,6 +50,7 @@ def to_condition(self) -> QueryCondition: ) @property + @override def existence_checks(self) -> list[ExistenceCheck[Any]]: """Check that the session exists.""" return [ @@ -73,6 +75,7 @@ class DeploymentHistorySearchScope(SearchScope): deployment_id: UUID """Required. The deployment to search history for.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for DeploymentHistoryRow.""" return DeploymentHistoryConditions.by_deployment_id_filter( @@ -80,6 +83,7 @@ def to_condition(self) -> QueryCondition: ) @property + @override def existence_checks(self) -> list[ExistenceCheck[Any]]: """Check that the deployment (endpoint) exists.""" return [ @@ -104,6 +108,7 @@ class RouteHistorySearchScope(SearchScope): route_id: ReplicaID """Required. The route to search history for.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for RouteHistoryRow.""" return RouteHistoryConditions.by_route_id_filter( @@ -111,6 +116,7 @@ def to_condition(self) -> QueryCondition: ) @property + @override def existence_checks(self) -> list[ExistenceCheck[Any]]: """Check that the route exists.""" return [ diff --git a/src/ai/backend/manager/repositories/session/types.py b/src/ai/backend/manager/repositories/session/types.py index 584cd95de1e..df697aba5f9 100644 --- a/src/ai/backend/manager/repositories/session/types.py +++ b/src/ai/backend/manager/repositories/session/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -27,6 +28,7 @@ class ProjectSessionSearchScope(SearchScope): project_id: UUID """Required. The project (group) to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for SessionRow.""" project_id = self.project_id @@ -37,6 +39,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ diff --git a/src/ai/backend/manager/repositories/user/types.py b/src/ai/backend/manager/repositories/user/types.py index fd0afcb5acd..73b8a4e10e6 100644 --- a/src/ai/backend/manager/repositories/user/types.py +++ b/src/ai/backend/manager/repositories/user/types.py @@ -7,6 +7,7 @@ from collections.abc import Sequence from dataclasses import dataclass +from typing import override from uuid import UUID import sqlalchemy as sa @@ -41,6 +42,7 @@ class DomainUserSearchScope(SearchScope): domain_name: str """Required. The domain to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for UserRow.""" domain_name = self.domain_name @@ -51,6 +53,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[str]]: """Return existence checks for scope validation.""" return [ @@ -73,6 +76,7 @@ class ProjectUserSearchScope(SearchScope): project_id: UUID """Required. The project (group) to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for AssociationScopesEntitiesRow. @@ -87,6 +91,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ @@ -108,6 +113,7 @@ class RoleUserSearchScope(SearchScope): role_id: UUID """Required. The role to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for UserRoleRow.""" role_id = self.role_id @@ -118,6 +124,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ diff --git a/src/ai/backend/manager/repositories/vfolder/types.py b/src/ai/backend/manager/repositories/vfolder/types.py index 99cc0a6d1a4..3971f2c7061 100644 --- a/src/ai/backend/manager/repositories/vfolder/types.py +++ b/src/ai/backend/manager/repositories/vfolder/types.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field +from typing import override from uuid import UUID import sqlalchemy as sa @@ -52,6 +53,7 @@ class ProjectVFolderSearchScope(SearchScope): project_id: UUID """Required. The project (group) to search within.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for VFolderRow.""" project_id = self.project_id @@ -62,6 +64,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ @@ -83,6 +86,7 @@ class UserVFolderSearchScope(SearchScope): user_id: UUID """Required. The user whose vfolders to search.""" + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition for VFolderRow. @@ -103,6 +107,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[UUID]]: """Return existence checks for scope validation.""" return [ diff --git a/src/ai/backend/manager/services/image/actions/modify_image.py b/src/ai/backend/manager/services/image/actions/modify_image.py index 805498e01e6..5ab34f9ea78 100644 --- a/src/ai/backend/manager/services/image/actions/modify_image.py +++ b/src/ai/backend/manager/services/image/actions/modify_image.py @@ -44,6 +44,7 @@ class ModifyImageActionUnknownImageReferenceError(BackendAIError): error_type = "https://api.backend.ai/probs/image-not-found" error_title = "Unknown image reference." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.IMAGE, diff --git a/src/ai/backend/manager/services/model_serving/actions/force_sync.py b/src/ai/backend/manager/services/model_serving/actions/force_sync.py index 3d992aaa02b..528be728749 100644 --- a/src/ai/backend/manager/services/model_serving/actions/force_sync.py +++ b/src/ai/backend/manager/services/model_serving/actions/force_sync.py @@ -11,10 +11,12 @@ class ForceSyncAction(ModelServiceAction): service_id: uuid.UUID + @override def entity_id(self) -> str | None: return None @classmethod + @override def operation_type(cls) -> ActionOperationType: return ActionOperationType.UPDATE diff --git a/src/ai/backend/manager/services/model_serving/actions/scale_service_replicas.py b/src/ai/backend/manager/services/model_serving/actions/scale_service_replicas.py index 294735a283a..8e75e4c0b7d 100644 --- a/src/ai/backend/manager/services/model_serving/actions/scale_service_replicas.py +++ b/src/ai/backend/manager/services/model_serving/actions/scale_service_replicas.py @@ -41,6 +41,7 @@ def entity_id(self) -> str | None: return None @classmethod + @override def operation_type(cls) -> ActionOperationType: return ActionOperationType.UPDATE diff --git a/src/ai/backend/manager/sokovan/deployment/exceptions.py b/src/ai/backend/manager/sokovan/deployment/exceptions.py index 9aa4c8b60a7..9747255ea78 100644 --- a/src/ai/backend/manager/sokovan/deployment/exceptions.py +++ b/src/ai/backend/manager/sokovan/deployment/exceptions.py @@ -2,6 +2,8 @@ Exceptions for deployment management. """ +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -17,6 +19,7 @@ class DeploymentError(BackendAIError): error_type = "https://api.backend.ai/probs/deployment-failed" error_title = "Deployment operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -31,6 +34,7 @@ class InvalidEndpointState(DeploymentError): error_type = "https://api.backend.ai/probs/invalid-endpoint-state" error_title = "Invalid endpoint state." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -45,6 +49,7 @@ class RouteCreationFailed(DeploymentError): error_type = "https://api.backend.ai/probs/route-creation-failed" error_title = "Failed to create route." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -59,6 +64,7 @@ class ScalingOperationFailed(DeploymentError): error_type = "https://api.backend.ai/probs/scaling-failed" error_title = "Scaling operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, @@ -73,6 +79,7 @@ class ServiceInfoRetrievalFailed(DeploymentError): error_type = "https://api.backend.ai/probs/service-info-retrieval-failed" error_title = "Failed to retrieve service information." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.MODEL_DEPLOYMENT, diff --git a/src/ai/backend/manager/sokovan/deployment/group/lifecycle/source.py b/src/ai/backend/manager/sokovan/deployment/group/lifecycle/source.py index 238e0830e5a..a5846973238 100644 --- a/src/ai/backend/manager/sokovan/deployment/group/lifecycle/source.py +++ b/src/ai/backend/manager/sokovan/deployment/group/lifecycle/source.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.deployment.types import ReplicaGroupHandlerCategory from ai.backend.manager.models.replica_group.conditions import ReplicaGroupConditions from ai.backend.manager.repositories.base import BatchQuerier, NoPagination @@ -26,6 +28,7 @@ class GroupLifecycleSource( def __init__(self, replica_group_repository: ReplicaGroupRepository) -> None: self._replica_group_repository = replica_group_repository + @override async def fetch_reconcile_info( self, category: ReplicaGroupHandlerCategory, @@ -59,6 +62,7 @@ class GroupAutoscaleSource( def __init__(self, replica_group_repository: ReplicaGroupRepository) -> None: self._replica_group_repository = replica_group_repository + @override async def fetch_reconcile_info( self, category: ReplicaGroupHandlerCategory, diff --git a/src/ai/backend/manager/sokovan/deployment/group/lifecycle/types.py b/src/ai/backend/manager/sokovan/deployment/group/lifecycle/types.py index 7c31c7e933c..913ce9f010c 100644 --- a/src/ai/backend/manager/sokovan/deployment/group/lifecycle/types.py +++ b/src/ai/backend/manager/sokovan/deployment/group/lifecycle/types.py @@ -6,6 +6,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field from datetime import datetime +from typing import override from uuid import UUID from ai.backend.common.identifier.deployment import DeploymentID @@ -56,9 +57,11 @@ class GroupLifecycleReconcileInfo(BaseReconcilerInfo): views: list[ReplicaGroupLifecycleReconcileView] current_time: datetime + @override def entity_ids(self) -> Sequence[UUID]: return [view.group_id for view in self.views] + @override def now(self) -> datetime: return self.current_time @@ -70,9 +73,11 @@ class GroupAutoscaleReconcileInfo(BaseReconcilerInfo): views: list[ReplicaGroupAutoscaleReconcileView] current_time: datetime + @override def entity_ids(self) -> Sequence[UUID]: return [view.group_id for view in self.views] + @override def now(self) -> datetime: return self.current_time @@ -97,15 +102,19 @@ class GroupLifecycleDecision(ReconcilerDecision): next_current_revision_id: TriState[DeploymentRevisionID] = field(default_factory=TriState.nop) next_target_revision_id: TriState[DeploymentRevisionID] = field(default_factory=TriState.nop) + @override def entity_id(self) -> UUID: return self.replica_group_id + @override def outcome(self) -> HandlerOutcome: return self.handler_outcome + @override def last_history(self) -> LastHistory | None: return self.prior_history + @override def policy_resolver(self) -> HandlerPolicyResolver: return self.handler_options @@ -118,11 +127,14 @@ class GroupLifecycleReconcileResult(BaseReconcilerResult): processed: int = 0 failed: int = 0 + @override def processed_count(self) -> int: return self.processed + @override def failed_count(self) -> int: return self.failed + @override def decisions(self) -> Sequence[ReconcilerDecision]: return self.lifecycle_decisions diff --git a/src/ai/backend/manager/sokovan/deployment/group/scaling/source.py b/src/ai/backend/manager/sokovan/deployment/group/scaling/source.py index 383e758ef8c..ca313c443da 100644 --- a/src/ai/backend/manager/sokovan/deployment/group/scaling/source.py +++ b/src/ai/backend/manager/sokovan/deployment/group/scaling/source.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.deployment.types import ReplicaGroupHandlerCategory from ai.backend.manager.models.replica_group.conditions import ReplicaGroupConditions from ai.backend.manager.repositories.base import BatchQuerier, NoPagination @@ -25,6 +27,7 @@ class GroupScalingSource( def __init__(self, replica_group_repository: ReplicaGroupRepository) -> None: self._replica_group_repository = replica_group_repository + @override async def fetch_reconcile_info( self, category: ReplicaGroupHandlerCategory, diff --git a/src/ai/backend/manager/sokovan/deployment/group/scaling/types.py b/src/ai/backend/manager/sokovan/deployment/group/scaling/types.py index efba64446a0..16df9c4100e 100644 --- a/src/ai/backend/manager/sokovan/deployment/group/scaling/types.py +++ b/src/ai/backend/manager/sokovan/deployment/group/scaling/types.py @@ -5,6 +5,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field from datetime import datetime +from typing import override from uuid import UUID from ai.backend.common.identifier.deployment import DeploymentID @@ -42,9 +43,11 @@ class GroupScalingReconcileInfo(BaseReconcilerInfo): views: list[ReplicaGroupScalingReconcileView] current_time: datetime + @override def entity_ids(self) -> Sequence[UUID]: return [view.group_id for view in self.views] + @override def now(self) -> datetime: return self.current_time @@ -67,15 +70,19 @@ class GroupScalingDecision(ReconcilerDecision): handler_options: DeploymentHandlerOptions error_code: str | None = None + @override def entity_id(self) -> UUID: return self.replica_group_id + @override def outcome(self) -> HandlerOutcome: return self.handler_outcome + @override def last_history(self) -> LastHistory | None: return self.prior_history + @override def policy_resolver(self) -> HandlerPolicyResolver: return self.handler_options @@ -90,11 +97,14 @@ class GroupScalingReconcileResult(BaseReconcilerResult): processed: int = 0 failed: int = 0 + @override def processed_count(self) -> int: return self.processed + @override def failed_count(self) -> int: return self.failed + @override def decisions(self) -> Sequence[ReconcilerDecision]: return self.scaling_decisions diff --git a/src/ai/backend/manager/sokovan/deployment/handlers/destroying.py b/src/ai/backend/manager/sokovan/deployment/handlers/destroying.py index d156e01637a..8b73293a02e 100644 --- a/src/ai/backend/manager/sokovan/deployment/handlers/destroying.py +++ b/src/ai/backend/manager/sokovan/deployment/handlers/destroying.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.deployment.types import ( @@ -40,25 +41,30 @@ def __init__( self._route_controller = route_controller @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "destroying-deployments" @classmethod + @override def category(cls) -> DeploymentHandlerCategory: return DeploymentHandlerCategory.LIFECYCLE @property + @override def lock_id(self) -> LockID | None: """Lock for destroying deployments.""" return LockID.LOCKID_DEPLOYMENT_DESTROYING @classmethod + @override def target_statuses(cls) -> DeploymentTargetStatuses: """Get the target deployment statuses for this handler.""" return DeploymentTargetStatuses(lifecycle_stages=[EndpointLifecycle.DESTROYING]) @classmethod + @override def status_transitions(cls) -> DeploymentStatusTransitions: """Define state transitions for destroying deployment handler (BEP-1030). @@ -72,6 +78,7 @@ def status_transitions(cls) -> DeploymentStatusTransitions: give_up=DeploymentLifecycleStatus(lifecycle=EndpointLifecycle.DESTROYED), ) + @override async def execute( self, deployments: Sequence[DeploymentWithHistory] ) -> DeploymentExecutionResult: @@ -81,6 +88,7 @@ async def execute( # Execute destruction logic via executor return await self._deployment_executor.destroy_deployment(deployments) + @override async def post_process(self, result: DeploymentExecutionResult) -> None: """Handle post-processing after destroying deployments.""" log.info("Destroyed {} deployments", len(result.successes)) diff --git a/src/ai/backend/manager/sokovan/deployment/handlers/replica.py b/src/ai/backend/manager/sokovan/deployment/handlers/replica.py index c0d83c6d2a1..4c445f25dca 100644 --- a/src/ai/backend/manager/sokovan/deployment/handlers/replica.py +++ b/src/ai/backend/manager/sokovan/deployment/handlers/replica.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.deployment.types import ( @@ -35,20 +36,24 @@ def __init__( self._deployment_controller = deployment_controller @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "check-replica-deployments" @classmethod + @override def category(cls) -> DeploymentHandlerCategory: return DeploymentHandlerCategory.SCALING @property + @override def lock_id(self) -> LockID | None: """Lock for checking replicas.""" return LockID.LOCKID_DEPLOYMENT_CHECK_REPLICA @classmethod + @override def target_statuses(cls) -> DeploymentTargetStatuses: """Target READY **and** DEPLOYING endpoints whose replica count may need adjusting. @@ -69,11 +74,13 @@ def target_statuses(cls) -> DeploymentTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> DeploymentStatusTransitions: """No status transition: this handler only records the deployment's desired replica count. The group autoscale reconcile reads that count and scales the serving group.""" return DeploymentStatusTransitions() + @override async def execute( self, deployments: Sequence[DeploymentWithHistory] ) -> DeploymentExecutionResult: @@ -96,6 +103,7 @@ async def execute( # Calculate desired replicas and adjust return await self._deployment_executor.calculate_desired_replicas(scalable) + @override async def post_process(self, result: DeploymentExecutionResult) -> None: """No follow-up: the group autoscale reconcile picks up the recorded desired count.""" return diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/appproxy_sync.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/appproxy_sync.py index 134145bd992..da4e2338799 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/appproxy_sync.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/appproxy_sync.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -45,20 +46,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: return "appproxy-sync" @property + @override def lock_id(self) -> LockID | None: # Sync is idempotent (set + event); the long cycle keeps state # convergent without an explicit distributed lock. return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.SYNC @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.RUNNING], @@ -66,6 +71,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: # Pure side-effect handler; never mutates lifecycle/health columns. return RouteStatusTransitions( @@ -74,9 +80,11 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: return await self._route_executor.sync_appproxy(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: synced = len(result.successes) failed = len(result.errors) diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/draining.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/draining.py index 8d276b671ff..417e2282ed8 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/draining.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/draining.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -41,20 +42,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "drain-routes" @property + @override def lock_id(self) -> LockID | None: """No lock needed for draining routes.""" return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.TERMINATING], @@ -62,6 +67,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Drained → COOLING_DOWN; the transition history timestamps the start of the termination grace period.""" @@ -73,6 +79,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute traffic draining for routes. @@ -83,6 +90,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: log.debug("Draining {} routes", len(routes)) return await self._route_executor.drain_routes(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after draining routes.""" log.info( diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/health_check.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/health_check.py index e266c18a94f..94cae3d5194 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/health_check.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/health_check.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -35,20 +36,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "health-check-routes" @property + @override def lock_id(self) -> LockID | None: """Lock for health check routes.""" return LockID.LOCKID_DEPLOYMENT_HEALTH_CHECK_ROUTES @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.HEALTH @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.RUNNING], @@ -61,6 +66,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Health check only changes health_status, not lifecycle status.""" return RouteStatusTransitions( @@ -69,6 +75,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=RouteTransitionTarget(health_status=RouteHealthStatus.DEGRADED), ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute health check for routes.""" log.debug("Checking health for {} routes", len(routes)) @@ -76,6 +83,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: # Execute route health check logic via executor return await self._route_executor.check_route_health(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Log health-check results. diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/observer/health_check.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/observer/health_check.py index 14dee1de7b4..5e854b28f64 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/observer/health_check.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/observer/health_check.py @@ -11,6 +11,7 @@ import logging from collections.abc import Sequence from dataclasses import dataclass +from typing import override import aiohttp @@ -63,9 +64,11 @@ def __init__( self._valkey_schedule = valkey_schedule @classmethod + @override def name(cls) -> str: return "route-health-observer" + @override async def observe(self, routes: Sequence[RouteData]) -> RouteObservationResult: """Check health for routes using ReplicaProbeTarget from Valkey.""" if not routes: diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/probe_target_sync.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/probe_target_sync.py index c79ff5963df..b387d0bbb23 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/probe_target_sync.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/probe_target_sync.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.deployment.types import ( @@ -37,18 +38,22 @@ def __init__(self, route_executor: RouteExecutor) -> None: self._route_executor = route_executor @classmethod + @override def name(cls) -> str: return "replica-probe-target-sync" @property + @override def lock_id(self) -> LockID | None: return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.SYNC @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.RUNNING], @@ -60,6 +65,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: return RouteStatusTransitions( success=None, @@ -67,10 +73,12 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Sync probe targets for routes that have health_check and replica info.""" return await self._route_executor.sync_route_probe_targets(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: if result.errors: log.warning( diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/provisioning.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/provisioning.py index b7c28b4870c..a38f9c8b62f 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/provisioning.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/provisioning.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -35,20 +36,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "provision-routes" @property + @override def lock_id(self) -> LockID | None: """Lock for provisioning routes.""" return LockID.LOCKID_DEPLOYMENT_PROVISIONING_ROUTES @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.PROVISIONING], @@ -56,6 +61,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Session enqueued → sub_status=STARTING; enqueue failed → FAILED_TO_START.""" return RouteStatusTransitions( @@ -68,6 +74,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute provisioning for routes.""" log.debug("Provisioning {} routes", len(routes)) @@ -75,6 +82,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: # Execute route provisioning logic via executor return await self._route_executor.provision_routes(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after provisioning routes.""" log.info( diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/route_eviction.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/route_eviction.py index 0c9cbee2e37..6993c1c5d77 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/route_eviction.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/route_eviction.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -47,20 +48,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "evict-routes" @property + @override def lock_id(self) -> LockID | None: """No lock needed for eviction.""" return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.HEALTH @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: # The handler runs over every active route (PROVISIONING / RUNNING) # regardless of health, because the orphan-revision check applies @@ -72,6 +77,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Eviction: success → TERMINATING (draining stage), failure → no change.""" return RouteStatusTransitions( @@ -83,6 +89,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """ Execute eviction for routes flagged by either eviction reason. @@ -95,6 +102,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: # Use executor logic to filter routes by scaling group config return await self._route_executor.cleanup_routes_by_config(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after eviction check.""" if result.successes: diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/running.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/running.py index 88525a9b7b0..e0f12509367 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/running.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/running.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -36,26 +37,31 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "check-running-routes" @property + @override def lock_id(self) -> LockID | None: """Lock for checking running routes.""" return LockID.LOCKID_DEPLOYMENT_RUNNING_ROUTES @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.RUNNING], ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Running check: success=no change, failure=TERMINATING (session died) + reset health.""" return RouteStatusTransitions( @@ -68,6 +74,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute health check for running routes.""" log.debug("Checking health for {} running routes", len(routes)) @@ -75,6 +82,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: # Execute route health check logic via executor return await self._route_executor.check_running_routes(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after checking running routes.""" if result.errors: diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/service_discovery_sync.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/service_discovery_sync.py index 73c2d74381b..977375dbeea 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/service_discovery_sync.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/service_discovery_sync.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -34,20 +35,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "service-discovery-sync" @property + @override def lock_id(self) -> LockID | None: """No lock needed for service discovery sync.""" return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.SYNC @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.RUNNING], @@ -55,6 +60,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Define state transitions for service discovery sync handler (BEP-1030). @@ -67,6 +73,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute service discovery synchronization for healthy routes.""" log.debug("Syncing {} healthy routes to service discovery", len(routes)) @@ -74,6 +81,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: # Execute service discovery sync logic via executor return await self._route_executor.sync_service_discovery(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after service discovery sync.""" synced_count = len(result.successes) diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/starting.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/starting.py index a38ce57e15f..de105b62f61 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/starting.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/starting.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -39,18 +40,22 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: return "check-starting-routes" @property + @override def lock_id(self) -> None: return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.PROVISIONING], @@ -58,6 +63,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Host/port ready → sub_status=WARMING_UP; session dead → FAILED_TO_START.""" return RouteStatusTransitions( @@ -71,10 +77,12 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: log.debug("Checking {} starting routes for host/port readiness", len(routes)) return await self._route_executor.check_starting_routes(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: log.info( "Starting check: {} routes ready (→ warming_up), {} failed", diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/terminating.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/terminating.py index 34043efb981..a6c59013b31 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/terminating.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/terminating.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -41,20 +42,24 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "terminate-routes" @property + @override def lock_id(self) -> LockID | None: """No lock needed for terminating routes.""" return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.TERMINATING], @@ -62,6 +67,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Terminating → TERMINATED on success, reset health to NOT_CHECKED. @@ -77,6 +83,7 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: """Execute termination for routes. @@ -87,6 +94,7 @@ async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: log.debug("Terminating {} routes", len(routes)) return await self._route_executor.terminate_routes(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: """Handle post-processing after terminating routes.""" log.info( diff --git a/src/ai/backend/manager/sokovan/deployment/route/handlers/warming_up.py b/src/ai/backend/manager/sokovan/deployment/route/handlers/warming_up.py index 4f065b7b74c..9a8e33dfa8b 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/handlers/warming_up.py +++ b/src/ai/backend/manager/sokovan/deployment/route/handlers/warming_up.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.events.dispatcher import EventProducer from ai.backend.logging import BraceStyleAdapter @@ -39,18 +40,22 @@ def __init__( self._event_producer = event_producer @classmethod + @override def name(cls) -> str: return "check-warming-up-routes" @property + @override def lock_id(self) -> None: return None @classmethod + @override def category(cls) -> RouteHandlerCategory: return RouteHandlerCategory.LIFECYCLE @classmethod + @override def target_statuses(cls) -> RouteTargetStatuses: return RouteTargetStatuses( lifecycle=[RouteStatus.PROVISIONING], @@ -58,6 +63,7 @@ def target_statuses(cls) -> RouteTargetStatuses: ) @classmethod + @override def status_transitions(cls) -> RouteStatusTransitions: """Health probe passes → RUNNING + HEALTHY + traffic=ACTIVE. Failure stays WARMING_UP (retry); expired/give_up → TERMINATING. @@ -74,10 +80,12 @@ def status_transitions(cls) -> RouteStatusTransitions: stale=None, ) + @override async def execute(self, routes: Sequence[RouteData]) -> RouteExecutionResult: log.debug("Checking {} warming-up routes for initial health", len(routes)) return await self._route_executor.check_warming_up_health(routes) + @override async def post_process(self, result: RouteExecutionResult) -> None: log.info( "Warming-up check: {} routes activated (→ running), {} still probing", diff --git a/src/ai/backend/manager/sokovan/reconciler/base.py b/src/ai/backend/manager/sokovan/reconciler/base.py index f61f347e1f0..8b9b85de3f4 100644 --- a/src/ai/backend/manager/sokovan/reconciler/base.py +++ b/src/ai/backend/manager/sokovan/reconciler/base.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from datetime import datetime from enum import StrEnum +from typing import override from uuid import UUID from ai.backend.common.events.types import AbstractAnycastEvent @@ -219,9 +220,11 @@ def __init__( self._metadata = metadata @property + @override def lock_id(self) -> LockID | None: return self._metadata.lock_id + @override async def run(self) -> None: metrics = ReconcilerMetricObserver.instance() kind = self._metadata.kind diff --git a/src/ai/backend/manager/sokovan/reconciler/flag.py b/src/ai/backend/manager/sokovan/reconciler/flag.py index d634a478655..f653ff204a5 100644 --- a/src/ai/backend/manager/sokovan/reconciler/flag.py +++ b/src/ai/backend/manager/sokovan/reconciler/flag.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.clients.valkey_client.valkey_schedule import ValkeyScheduleClient from ai.backend.manager.sokovan.reconciler.coordinator import ReconcilerFlag @@ -14,5 +16,6 @@ class ValkeyReconcilerFlag(ReconcilerFlag): def __init__(self, valkey_schedule: ValkeyScheduleClient) -> None: self._valkey_schedule = valkey_schedule + @override async def check_mark_needed(self, reconcile_type: str) -> bool: return await self._valkey_schedule.load_and_delete_schedule_mark(reconcile_type) diff --git a/src/ai/backend/manager/sokovan/recorder/exceptions.py b/src/ai/backend/manager/sokovan/recorder/exceptions.py index f89f8976a42..be83075c2fa 100644 --- a/src/ai/backend/manager/sokovan/recorder/exceptions.py +++ b/src/ai/backend/manager/sokovan/recorder/exceptions.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import override + from ai.backend.common.exception import ( BackendAIError, ErrorCode, @@ -15,6 +17,7 @@ class RecorderError(BackendAIError): error_type = "https://api.backend.ai/probs/recorder-error" error_title = "Recorder operation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -38,6 +41,7 @@ def __init__(self, new_phase: str, active_phase: str) -> None: ) ) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -61,6 +65,7 @@ def __init__(self, step_name: str) -> None: ) ) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, diff --git a/src/ai/backend/manager/sokovan/scheduler/exceptions.py b/src/ai/backend/manager/sokovan/scheduler/exceptions.py index c9176d60d3c..33a50291735 100644 --- a/src/ai/backend/manager/sokovan/scheduler/exceptions.py +++ b/src/ai/backend/manager/sokovan/scheduler/exceptions.py @@ -3,6 +3,7 @@ """ from abc import ABC, abstractmethod +from typing import override from ai.backend.common.exception import ( BackendAIError, @@ -24,6 +25,7 @@ class SchedulingError(BackendAIError, ABC): error_type = "https://api.backend.ai/probs/scheduling-failed" error_title = "Scheduling failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/cleanup/force_terminated.py b/src/ai/backend/manager/sokovan/scheduler/handlers/cleanup/force_terminated.py index 9041d45647a..a1ab06167a7 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/cleanup/force_terminated.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/cleanup/force_terminated.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.clients.valkey_client.valkey_schedule.client import ValkeyScheduleClient from ai.backend.common.types import SessionId @@ -41,12 +41,15 @@ def __init__( self._valkey_schedule = valkey_schedule @classmethod + @override def name(cls) -> str: return "cleanup-force-terminated" + @override async def fetch_session_ids(self) -> Sequence[SessionId]: return await self._valkey_schedule.get_force_terminated_sessions() + @override async def execute(self, session_ids: Sequence[SessionId]) -> None: log.info("Processing {} force-terminated sessions for container cleanup", len(session_ids)) diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/kernel/sweep_stale_kernels.py b/src/ai/backend/manager/sokovan/scheduler/handlers/kernel/sweep_stale_kernels.py index ec70cd37e3f..816ae747264 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/kernel/sweep_stale_kernels.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/kernel/sweep_stale_kernels.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.kernel.types import KernelInfo, KernelStatus @@ -41,16 +41,19 @@ def __init__( self._terminator = terminator @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "sweep-stale-kernels" @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus]: """Running kernels that may be stale.""" return [KernelStatus.RUNNING] @classmethod + @override def status_transitions(cls) -> KernelStatusTransitions: """Define state transitions for sweep stale kernels handler. @@ -63,10 +66,12 @@ def status_transitions(cls) -> KernelStatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """Lock for operations targeting TERMINATING sessions.""" return LockID.LOCKID_SOKOVAN_TARGET_TERMINATING + @override async def execute( self, _scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/check_precondition.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/check_precondition.py index 1276febb1a1..47ce9e4129a 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/check_precondition.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/check_precondition.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.types import AccessKey from ai.backend.logging import BraceStyleAdapter @@ -44,21 +44,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "check-precondition" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions in SCHEDULED state.""" return [SessionStatus.SCHEDULED] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """Include sessions where kernels are in SCHEDULED status.""" return [KernelStatus.SCHEDULED] @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for check precondition handler (BEP-1030). @@ -91,10 +95,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """Lock for operations targeting SCHEDULED sessions transitioning to PREPARING.""" return LockID.LOCKID_SOKOVAN_TARGET_PREPARING + @override async def execute( self, _scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/deprioritize_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/deprioritize_sessions.py index 4c309a98283..744c143eb2a 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/deprioritize_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/deprioritize_sessions.py @@ -4,6 +4,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.common.defs.session import SESSION_PRIORITY_MIN from ai.backend.common.types import AccessKey @@ -44,21 +45,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "deprioritize-sessions" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions in DEPRIORITIZING state.""" return [SessionStatus.DEPRIORITIZING] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """No kernel filtering for deprioritize.""" return None @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for deprioritize handler (BEP-1030). @@ -86,10 +91,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """No lock needed for deprioritizing sessions.""" return None + @override async def execute( self, scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py index a8a1f8cc016..d3a02a2bc79 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.types import AccessKey from ai.backend.logging import BraceStyleAdapter @@ -46,21 +46,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "schedule-sessions" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions in PENDING state.""" return [SessionStatus.PENDING] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """No kernel filtering for scheduling.""" return None @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for scheduling handler (BEP-1030). @@ -86,10 +90,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """Lock for operations targeting PENDING sessions.""" return LockID.LOCKID_SOKOVAN_TARGET_PENDING + @override async def execute( self, scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/start_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/start_sessions.py index 36e7b043775..4b728aec275 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/start_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/start_sessions.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.types import AccessKey from ai.backend.logging import BraceStyleAdapter @@ -46,21 +46,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "start-sessions" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions in PREPARED state.""" return [SessionStatus.PREPARED] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """Include sessions where kernels are in PREPARED status.""" return [KernelStatus.PREPARED] @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for start sessions handler (BEP-1030). @@ -90,10 +94,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """Lock for operations targeting PREPARED sessions transitioning to CREATING.""" return LockID.LOCKID_SOKOVAN_TARGET_CREATING + @override async def execute( self, _scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/terminate_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/terminate_sessions.py index 856617aa443..9c7afd3dda2 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/terminate_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/terminate_sessions.py @@ -4,7 +4,7 @@ import logging from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.kernel.types import KernelStatus @@ -43,21 +43,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "terminate-sessions" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions in TERMINATING state.""" return [SessionStatus.TERMINATING] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """No kernel filtering - termination applies to any kernel state.""" return None @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for terminate sessions handler (BEP-1030). @@ -80,10 +84,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """Lock for operations targeting TERMINATING sessions.""" return LockID.LOCKID_SOKOVAN_TARGET_TERMINATING + @override async def execute( self, _scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/maintenance/sweep_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/maintenance/sweep_sessions.py index b5b60233cfd..5a104ea1264 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/maintenance/sweep_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/maintenance/sweep_sessions.py @@ -4,6 +4,7 @@ import logging from collections.abc import Sequence +from typing import override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.kernel.types import KernelStatus @@ -36,21 +37,25 @@ def __init__( self._repository = repository @classmethod + @override def name(cls) -> str: """Get the name of the handler.""" return "sweep-sessions" @classmethod + @override def target_statuses(cls) -> list[SessionStatus]: """Sessions that may need sweeping - PENDING with timeout.""" return [SessionStatus.PENDING] @classmethod + @override def target_kernel_statuses(cls) -> list[KernelStatus] | None: """No kernel filtering for sweep check.""" return None @classmethod + @override def status_transitions(cls) -> StatusTransitions: """Define state transitions for sweep sessions handler (BEP-1030). @@ -75,10 +80,12 @@ def status_transitions(cls) -> StatusTransitions: ) @property + @override def lock_id(self) -> LockID | None: """No lock needed for sweeping stale sessions.""" return None + @override async def execute( self, _scaling_group: str, diff --git a/src/ai/backend/manager/sokovan/scheduler/hooks/status.py b/src/ai/backend/manager/sokovan/scheduler/hooks/status.py index caac853fd68..501c4af4fcc 100644 --- a/src/ai/backend/manager/sokovan/scheduler/hooks/status.py +++ b/src/ai/backend/manager/sokovan/scheduler/hooks/status.py @@ -9,6 +9,7 @@ import logging from abc import ABC, abstractmethod from dataclasses import dataclass +from typing import override from ai.backend.common.types import ( AgentId, @@ -72,6 +73,7 @@ class RunningTransitionHook(StatusTransitionHook): def __init__(self, deps: RunningHookDependencies) -> None: self._deps = deps + @override async def execute(self, session: SessionWithKernels) -> None: """Execute RUNNING transition hook. @@ -149,6 +151,7 @@ class TerminatedTransitionHook(StatusTransitionHook): def __init__(self, deps: TerminatedHookDependencies) -> None: self._deps = deps + @override async def execute(self, session: SessionWithKernels) -> None: """Execute TERMINATED transition hook.""" await self._destroy_network(session) diff --git a/src/ai/backend/manager/sokovan/scheduler/post_processors/cache_invalidation.py b/src/ai/backend/manager/sokovan/scheduler/post_processors/cache_invalidation.py index c09e8a677b0..ee9545ca28d 100644 --- a/src/ai/backend/manager/sokovan/scheduler/post_processors/cache_invalidation.py +++ b/src/ai/backend/manager/sokovan/scheduler/post_processors/cache_invalidation.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.common.types import AccessKey from ai.backend.logging import BraceStyleAdapter @@ -22,6 +22,7 @@ class CacheInvalidationPostProcessor(PostProcessor): def __init__(self, repository: SchedulerRepository) -> None: self._repository = repository + @override async def execute(self, context: PostProcessorContext) -> None: """Invalidate cache for all affected access keys.""" affected_keys: set[AccessKey] = set() diff --git a/src/ai/backend/manager/sokovan/scheduler/post_processors/kernel_schedule_marking.py b/src/ai/backend/manager/sokovan/scheduler/post_processors/kernel_schedule_marking.py index ed1d681e15c..8bbba477652 100644 --- a/src/ai/backend/manager/sokovan/scheduler/post_processors/kernel_schedule_marking.py +++ b/src/ai/backend/manager/sokovan/scheduler/post_processors/kernel_schedule_marking.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.kernel.types import KernelStatus @@ -32,6 +32,7 @@ class KernelScheduleMarkingPostProcessor(KernelPostProcessor): def __init__(self, scheduling_controller: SchedulingController) -> None: self._scheduling_controller = scheduling_controller + @override async def execute(self, context: KernelPostProcessorContext) -> None: """Mark the next schedule types for all target statuses.""" if not context.target_statuses: diff --git a/src/ai/backend/manager/sokovan/scheduler/post_processors/schedule_marking.py b/src/ai/backend/manager/sokovan/scheduler/post_processors/schedule_marking.py index 57634974b81..44b64243875 100644 --- a/src/ai/backend/manager/sokovan/scheduler/post_processors/schedule_marking.py +++ b/src/ai/backend/manager/sokovan/scheduler/post_processors/schedule_marking.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.data.session.types import SessionStatus @@ -37,6 +37,7 @@ class ScheduleMarkingPostProcessor(PostProcessor): def __init__(self, scheduling_controller: SchedulingController) -> None: self._scheduling_controller = scheduling_controller + @override async def execute(self, context: PostProcessorContext) -> None: """Mark the next schedule types for all target statuses.""" if not context.target_statuses: diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/allocators/repository_allocator.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/allocators/repository_allocator.py index 8e5651ca27d..40978caeffe 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/allocators/repository_allocator.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/allocators/repository_allocator.py @@ -5,6 +5,8 @@ to ensure transactional consistency. """ +from typing import override + from ai.backend.common.types import SessionId from ai.backend.manager.data.sokovan import AllocationBatch from ai.backend.manager.repositories.scheduler import SchedulerRepository @@ -22,12 +24,14 @@ class RepositoryAllocator(SchedulingAllocator): _repository: SchedulerRepository + @override def name(self) -> str: """ Return the allocator name for predicates. """ return "RepositoryAllocator" + @override def success_message(self) -> str: """ Return a message describing successful allocation. @@ -43,6 +47,7 @@ def __init__(self, schedule_repository: SchedulerRepository) -> None: """ self._repository = schedule_repository + @override async def allocate(self, batch: AllocationBatch) -> list[SessionId]: """ Reserve and assign the batch's sessions to agents by delegating to the diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/concentrated.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/concentrated.py index 0ee51129b67..075c8617056 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/concentrated.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/concentrated.py @@ -10,6 +10,7 @@ import sys from collections.abc import Sequence from decimal import Decimal +from typing import override from ai.backend.common.types import SessionTypes @@ -42,18 +43,21 @@ def __init__(self, agent_selection_resource_priority: list[str]) -> None: """ self.agent_selection_resource_priority = agent_selection_resource_priority + @override def name(self) -> str: """ Return the selector name for predicates. """ return "ConcentratedAgentSelector" + @override def success_message(self) -> str: """ Return a message describing successful agent selection. """ return "Agent selected using concentrated strategy for maximum resource utilization" + @override def select_tracker_by_strategy( self, trackers: Sequence[AgentStateTracker], diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/dispersed.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/dispersed.py index d483bb4d2a6..eecba7d131f 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/dispersed.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/dispersed.py @@ -8,6 +8,7 @@ import sys from collections.abc import Sequence from decimal import Decimal +from typing import override from .selector import ( AbstractAgentSelector, @@ -37,18 +38,21 @@ def __init__(self, agent_selection_resource_priority: list[str]) -> None: """ self.agent_selection_resource_priority = agent_selection_resource_priority + @override def name(self) -> str: """ Return the selector name for predicates. """ return "DispersedAgentSelector" + @override def success_message(self) -> str: """ Return a message describing successful agent selection. """ return "Agent selected using dispersed strategy for balanced workload distribution" + @override def select_tracker_by_strategy( self, trackers: Sequence[AgentStateTracker], diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/legacy.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/legacy.py index ca90f0f1785..04778b03071 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/legacy.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/legacy.py @@ -8,6 +8,7 @@ import sys from collections.abc import Sequence from decimal import Decimal +from typing import override from .selector import ( AbstractAgentSelector, @@ -31,18 +32,21 @@ class LegacyAgentSelector(AbstractAgentSelector): def __init__(self, agent_selection_resource_priority: list[str]) -> None: self.agent_selection_resource_priority = agent_selection_resource_priority + @override def name(self) -> str: """ Return the selector name for predicates. """ return "LegacyAgentSelector" + @override def success_message(self) -> str: """ Return a message describing successful agent selection. """ return "Agent selected using legacy priority-based strategy" + @override def select_tracker_by_strategy( self, trackers: Sequence[AgentStateTracker], diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/roundrobin.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/roundrobin.py index 518b6de667b..0a408783338 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/roundrobin.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/selectors/roundrobin.py @@ -6,6 +6,7 @@ """ from collections.abc import Sequence +from typing import override from .selector import ( AbstractAgentSelector, @@ -33,18 +34,21 @@ def __init__(self, next_index: int = 0) -> None: """ self.next_index = next_index + @override def name(self) -> str: """ Return the selector name for predicates. """ return "RoundRobinAgentSelector" + @override def success_message(self) -> str: """ Return a message describing successful agent selection. """ return "Agent selected using round-robin strategy for even workload distribution" + @override def select_tracker_by_strategy( self, trackers: Sequence[AgentStateTracker], diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/concurrency.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/concurrency.py index 562064fd826..217a949c931 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/concurrency.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/concurrency.py @@ -1,5 +1,7 @@ """Validator for concurrent session limits.""" +from typing import override + from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot from .exceptions import ConcurrencyLimitExceeded @@ -12,14 +14,17 @@ class ConcurrencyValidator(ValidatorRule): This corresponds to check_concurrency predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "ConcurrencyLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Session count is within the maximum concurrent session limit for the access key" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the keypair's resource policy policy = snapshot.resource_policy.keypair_policies.get(workload.access_key) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/dependencies.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/dependencies.py index 40780f0f5ca..686d433f55b 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/dependencies.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/dependencies.py @@ -1,5 +1,7 @@ """Validator for session dependencies.""" +from typing import override + from ai.backend.common.types import SessionResult from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot @@ -14,14 +16,17 @@ class DependenciesValidator(ValidatorRule): This corresponds to check_dependencies predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "SessionDependenciesValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "All dependent sessions have completed successfully" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get dependencies for this session dependencies = snapshot.session_dependencies.by_session.get(workload.session_id, []) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/domain_resource_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/domain_resource_limit.py index 3a6f99dc81e..cc20705f6fb 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/domain_resource_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/domain_resource_limit.py @@ -1,5 +1,7 @@ """Validator for domain resource limits.""" +from typing import override + from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot @@ -13,14 +15,17 @@ class DomainResourceLimitValidator(ValidatorRule): This corresponds to check_domain_resource_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "DomainResourceLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Domain has sufficient resource quota for the requested session" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the domain's resource limit domain_limit = snapshot.resource_policy.domain_limits.get(workload.domain_name) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/exceptions.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/exceptions.py index 7366ff75304..39a06016fbe 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/exceptions.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/exceptions.py @@ -5,6 +5,7 @@ from abc import abstractmethod from collections.abc import Sequence from datetime import datetime +from typing import override from aiohttp import web @@ -31,6 +32,7 @@ class SchedulingValidationError(SchedulingError, web.HTTPPreconditionFailed): error_type = "https://api.backend.ai/probs/scheduling-validation-failed" error_title = "Scheduling validation failed." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -46,6 +48,7 @@ def summary(self) -> str: """ raise NotImplementedError + @override def failed_predicates(self) -> list[SchedulingPredicate]: """Return list of failed predicates for this error.""" return [SchedulingPredicate(name=type(self).__name__, msg=str(self))] @@ -69,9 +72,11 @@ def __init__(self, *, max_sessions: int, session_type: str) -> None: self._session_type = session_type super().__init__(self.summary()) + @override def summary(self) -> str: return f"You cannot run more than {self._max_sessions} {self._session_type} sessions" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -92,10 +97,12 @@ def __init__(self, *, pending_dependency_names: Sequence[str]) -> None: self._pending_dependency_names = list(pending_dependency_names) super().__init__(self.summary()) + @override def summary(self) -> str: names = ", ".join(self._pending_dependency_names) return f"Waiting dependency sessions to finish as success. ({names})" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -116,9 +123,11 @@ def __init__(self, *, quota_slots: ResourceSlot) -> None: self._quota_slots = quota_slots super().__init__(self.summary()) + @override def summary(self) -> str: return f"Your keypair resource quota is exceeded. ({_format_slots(self._quota_slots)})" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.KEYPAIR, @@ -139,9 +148,11 @@ def __init__(self, *, quota_slots: ResourceSlot) -> None: self._quota_slots = quota_slots super().__init__(self.summary()) + @override def summary(self) -> str: return f"Your main-keypair resource quota is exceeded. ({_format_slots(self._quota_slots)})" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.USER, @@ -162,9 +173,11 @@ def __init__(self, *, quota_slots: ResourceSlot) -> None: self._quota_slots = quota_slots super().__init__(self.summary()) + @override def summary(self) -> str: return f"Your group resource quota is exceeded. ({_format_slots(self._quota_slots)})" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.GROUP, @@ -185,9 +198,11 @@ def __init__(self, *, quota_slots: ResourceSlot) -> None: self._quota_slots = quota_slots super().__init__(self.summary()) + @override def summary(self) -> str: return f"Your domain resource quota is exceeded. ({_format_slots(self._quota_slots)})" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.DOMAIN, @@ -208,9 +223,11 @@ def __init__(self, *, max_pending_session_count: int) -> None: self._max_pending_session_count = max_pending_session_count super().__init__(self.summary()) + @override def summary(self) -> str: return f"You cannot create more than {self._max_pending_session_count} pending session(s)." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -231,12 +248,14 @@ def __init__(self, *, current_pending_slots: ResourceSlot) -> None: self._current_pending_slots = current_pending_slots super().__init__(self.summary()) + @override def summary(self) -> str: return ( f"Your pending session quota is exceeded." f" ({_format_slots(self._current_pending_slots)})" ) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -257,12 +276,14 @@ def __init__(self, *, scheduled_start: datetime) -> None: self._scheduled_start = scheduled_start super().__init__(self.summary()) + @override def summary(self) -> str: return ( f"Batch session is scheduled to start at {self._scheduled_start.isoformat()};" " current time is before that." ) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.SESSION, @@ -283,10 +304,12 @@ def __init__(self, errors: Sequence[SchedulingValidationError]) -> None: self._errors = list(errors) super().__init__(self.summary()) + @override def summary(self) -> str: lines = [f"- {type(e).__name__}: {e.summary()}" for e in self._errors] return "Multiple validation errors occurred:\n" + "\n".join(lines) + @override def failed_predicates(self) -> list[SchedulingPredicate]: """Return list of all failed predicates from all errors.""" predicates: list[SchedulingPredicate] = [] @@ -294,6 +317,7 @@ def failed_predicates(self) -> list[SchedulingPredicate]: predicates.extend(error.failed_predicates()) return predicates + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/group_resource_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/group_resource_limit.py index 656e1714a06..217b686404d 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/group_resource_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/group_resource_limit.py @@ -1,5 +1,7 @@ """Validator for group resource limits.""" +from typing import override + from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot @@ -13,14 +15,17 @@ class GroupResourceLimitValidator(ValidatorRule): This corresponds to check_group_resource_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "GroupResourceLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Group has sufficient resource quota for the requested session" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the group's resource limit group_limit = snapshot.resource_policy.group_limits.get(workload.group_id) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/keypair_resource_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/keypair_resource_limit.py index d86d82d61d7..22213a80047 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/keypair_resource_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/keypair_resource_limit.py @@ -1,5 +1,7 @@ """Validator for keypair resource limits.""" +from typing import override + from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.sokovan import ( SessionWorkload, @@ -16,14 +18,17 @@ class KeypairResourceLimitValidator(ValidatorRule): This corresponds to check_keypair_resource_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "KeypairResourceLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Keypair has sufficient resource quota for the requested session" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the keypair's resource policy policy = snapshot.resource_policy.keypair_policies.get(workload.access_key) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_count_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_count_limit.py index cf88bc543c5..43212b50291 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_count_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_count_limit.py @@ -1,5 +1,7 @@ """Validator for pending session count limits.""" +from typing import override + from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot from .exceptions import PendingSessionCountLimitExceeded @@ -12,14 +14,17 @@ class PendingSessionCountLimitValidator(ValidatorRule): This corresponds to check_pending_session_count_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "PendingSessionCountLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Pending session count is within the allowed limit for the access key" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the keypair's resource policy policy = snapshot.resource_policy.keypair_policies.get(workload.access_key) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_resource_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_resource_limit.py index 664cd2f7b87..bafe8af2e17 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_resource_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/pending_session_resource_limit.py @@ -1,5 +1,7 @@ """Validator for pending session resource limits.""" +from typing import override + from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot @@ -13,14 +15,17 @@ class PendingSessionResourceLimitValidator(ValidatorRule): This corresponds to check_pending_session_resource_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "PendingSessionResourceLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Pending session resource usage is within the allowed limit for the access key" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the keypair's resource policy policy = snapshot.resource_policy.keypair_policies.get(workload.access_key) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/reserved_batch.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/reserved_batch.py index dd694c61d09..91134b4a1b4 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/reserved_batch.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/reserved_batch.py @@ -1,6 +1,7 @@ """Validator for reserved batch sessions.""" from datetime import datetime +from typing import override from dateutil.tz import tzutc @@ -17,14 +18,17 @@ class ReservedBatchSessionValidator(ValidatorRule): This corresponds to check_reserved_batch_session predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "ReservedBatchSessionValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "Batch session has reached its scheduled start time" + @override def validate(self, _snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Check if this is a batch session with a scheduled start time if workload.session_type == SessionTypes.BATCH and workload.starts_at is not None: diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/user_resource_limit.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/user_resource_limit.py index 03bb8833c4f..f04362c477b 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/user_resource_limit.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/validators/user_resource_limit.py @@ -1,5 +1,7 @@ """Validator for user resource limits.""" +from typing import override + from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.sokovan import SessionWorkload, SystemSnapshot @@ -13,14 +15,17 @@ class UserResourceLimitValidator(ValidatorRule): This corresponds to check_user_resource_limit predicate. """ + @override def name(self) -> str: """Return the validator name for predicates.""" return "UserResourceLimitValidator" + @override def success_message(self) -> str: """Return a message describing successful validation.""" return "User has sufficient resource quota for the requested session" + @override def validate(self, snapshot: SystemSnapshot, workload: SessionWorkload) -> None: # Get the user's resource policy policy = snapshot.resource_policy.user_policies.get(workload.user_uuid) diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_container_user_mapping_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_container_user_mapping_rule.py index 4812afeb647..9fe43b9781a 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_container_user_mapping_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_container_user_mapping_rule.py @@ -17,6 +17,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( SessionSpecDraftRule, @@ -27,9 +29,11 @@ class AssignContainerUserMappingRule(SessionSpecDraftRule): """Fill kernel-draft ``uid`` / ``main_gid`` / ``supplementary_gids`` from the context.""" + @override def name(self) -> str: return "assign_container_user_mapping" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_network_config_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_network_config_rule.py index d4d51faac23..5ed02a10ceb 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_network_config_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_network_config_rule.py @@ -22,6 +22,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.models.network import NetworkType from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( @@ -33,9 +35,11 @@ class AssignNetworkConfigRule(SessionSpecDraftRule): """Resolve ``SessionNetworkDraft.network_type`` using the legacy priority.""" + @override def name(self) -> str: return "assign_network_config" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_user_identity_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_user_identity_rule.py index 7ebb2df91ef..09d4c0a54e0 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_user_identity_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/assign_user_identity_rule.py @@ -13,6 +13,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.contexts.user import current_user from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( @@ -24,9 +26,11 @@ class AssignUserIdentityRule(SessionSpecDraftRule): """Copy the current user's id into ``SessionIdentityDraft.user_uuid``.""" + @override def name(self) -> str: return "assign_user_identity" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/build_internal_data_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/build_internal_data_rule.py index 0ea5a2c881e..19928737088 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/build_internal_data_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/build_internal_data_rule.py @@ -27,7 +27,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, override from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( @@ -39,9 +39,11 @@ class BuildInternalDataRule(SessionSpecDraftRule): """Compose per-kernel ``internal_data`` from draft + context inputs.""" + @override def name(self) -> str: return "build_internal_data" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/compute_kernel_resources_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/compute_kernel_resources_rule.py index 8b93b47eac9..7859df72f26 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/compute_kernel_resources_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/compute_kernel_resources_rule.py @@ -24,7 +24,7 @@ from collections.abc import Mapping from decimal import Decimal -from typing import Any +from typing import Any, override from ai.backend.common.types import BinarySize, ResourceSlotEntry from ai.backend.manager.data.session.creation import ImageInfo @@ -49,9 +49,11 @@ class ComputeKernelResourcesRule(SessionSpecDraftRule): _SHMEM_IMAGE_LABEL = "ai.backend.resource.preferred.shmem" + @override def name(self) -> str: return "compute_kernel_resources" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/expand_kernel_groups_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/expand_kernel_groups_rule.py index 38bd73546e2..d84ee741643 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/expand_kernel_groups_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/expand_kernel_groups_rule.py @@ -20,6 +20,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import ( KernelSpecDraft, SessionSpecDraft, @@ -38,9 +40,11 @@ class ExpandKernelGroupsRule(SessionSpecDraftRule): its attribute-path error message intact. """ + @override def name(self) -> str: return "expand_kernel_groups" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/inject_session_environ_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/inject_session_environ_rule.py index 380185a84a9..f27f60f925f 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/inject_session_environ_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/inject_session_environ_rule.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( SessionSpecDraftRule, @@ -27,9 +29,11 @@ class InjectSessionEnvironRule(SessionSpecDraftRule): """Overlay ``SUDO_SESSION_ENABLED=1`` on every kernel when sudo is enabled.""" + @override def name(self) -> str: return "inject_session_environ" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/merge_resource_group_defaults_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/merge_resource_group_defaults_rule.py index d5fc3dcf170..6da3511cc81 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/merge_resource_group_defaults_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/merge_resource_group_defaults_rule.py @@ -28,6 +28,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import ( KernelExecutionSpecDraft, SessionSpecDraft, @@ -42,9 +44,11 @@ class MergeResourceGroupDefaultsRule(SessionSpecDraftRule): """Overlay RG defaults under caller-supplied draft values.""" + @override def name(self) -> str: return "merge_resource_group_defaults" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/resolve_vfolder_mounts_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/resolve_vfolder_mounts_rule.py index 99498b0aedd..ab317ff8695 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/preparers/resolve_vfolder_mounts_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/preparers/resolve_vfolder_mounts_rule.py @@ -21,6 +21,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.draft import SessionSpecDraft from ai.backend.manager.sokovan.scheduling_controller.preparers.draft_rule import ( SessionSpecDraftRule, @@ -31,9 +33,11 @@ class ResolveVFolderMountsRule(SessionSpecDraftRule): """Copy the per-role resolved vfolder mounts from context onto each kernel.""" + @override def name(self) -> str: return "resolve_vfolder_mounts" + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/concurrent_session_limit_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/concurrent_session_limit_rule.py index cbe68f90616..8f578b910f7 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/concurrent_session_limit_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/concurrent_session_limit_rule.py @@ -10,6 +10,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.kernel import QuotaExceeded from ai.backend.manager.sokovan.scheduling_controller.validators.session_spec_base import ( @@ -21,9 +23,11 @@ class ConcurrentSessionLimitRule(SessionSpecValidatorRule): """Reject enqueue when the keypair is already at its concurrent-session cap.""" + @override def name(self) -> str: return "concurrent_session_limit" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/container_limit_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/container_limit_rule.py index 2b8b2dab29e..0fdd5a02926 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/container_limit_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/container_limit_rule.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.kernel import QuotaExceeded from ai.backend.manager.sokovan.scheduling_controller.validators.session_spec_base import ( @@ -13,9 +15,11 @@ class ContainerLimitRule(SessionSpecValidatorRule): """Session's total kernel count must not exceed the keypair limit.""" + @override def name(self) -> str: return "container_limit" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/dotfile_vfolder_conflict_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/dotfile_vfolder_conflict_rule.py index 4fd1b038c4c..b321e60f069 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/dotfile_vfolder_conflict_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/dotfile_vfolder_conflict_rule.py @@ -10,6 +10,7 @@ from __future__ import annotations from pathlib import PurePosixPath +from typing import override from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.storage import DotfileVFolderPathConflict @@ -22,9 +23,11 @@ class DotfileVFolderConflictRule(SessionSpecValidatorRule): """Dotfile paths must not collide with any kernel's resolved mount path.""" + @override def name(self) -> str: return "dotfile_vfolder_conflict" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/image_slot_type_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/image_slot_type_rule.py index 2a54c9b995e..b24f212e3fc 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/image_slot_type_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/image_slot_type_rule.py @@ -15,6 +15,7 @@ from __future__ import annotations from decimal import Decimal +from typing import override from ai.backend.common.types import SlotName from ai.backend.manager.data.session.spec import SessionSpec @@ -29,9 +30,11 @@ class ImageSlotTypeRule(SessionSpecValidatorRule): """Image-declared slot keys must be served by an agent in the target RG.""" + @override def name(self) -> str: return "image_slot_type" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/inference_model_folder_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/inference_model_folder_rule.py index cc755b5388c..52a3fc99233 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/inference_model_folder_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/inference_model_folder_rule.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.common.types import SessionTypes, VFolderUsageMode from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.api import InvalidAPIParameters @@ -21,9 +23,11 @@ class InferenceModelFolderRule(SessionSpecValidatorRule): needed. """ + @override def name(self) -> str: return "inference_model_folder" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/mount_name_validation_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/mount_name_validation_rule.py index e9de71fecf2..bf6f48cfe48 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/mount_name_validation_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/mount_name_validation_rule.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import override + from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.api import InvalidAPIParameters from ai.backend.manager.models.vfolder import verify_vfolder_name @@ -20,9 +22,11 @@ class MountNameValidationRule(SessionSpecValidatorRule): vfolder at the same path. """ + @override def name(self) -> str: return "mount_name_validation" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/requested_slot_type_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/requested_slot_type_rule.py index 98c6b4ecd23..f7848df191c 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/requested_slot_type_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/requested_slot_type_rule.py @@ -16,6 +16,7 @@ from __future__ import annotations from decimal import Decimal +from typing import override from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.api import InvalidAPIParameters @@ -29,9 +30,11 @@ class RequestedSlotTypeRule(SessionSpecValidatorRule): """Requested slot keys must be served by an agent in the target RG.""" + @override def name(self) -> str: return "requested_slot_type" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/required_resource_slot_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/required_resource_slot_rule.py index 70c77e80249..a774419f14f 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/required_resource_slot_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/required_resource_slot_rule.py @@ -9,6 +9,7 @@ from __future__ import annotations from decimal import Decimal +from typing import override from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.api import InvalidAPIParameters @@ -22,9 +23,11 @@ class RequiredResourceSlotRule(SessionSpecValidatorRule): """Every kernel request must include all globally required slots.""" + @override def name(self) -> str: return "required_resource_slot" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/resource_limit_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/resource_limit_rule.py index b44284d61eb..1611d43d6aa 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/resource_limit_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/resource_limit_rule.py @@ -10,7 +10,7 @@ from collections.abc import Mapping from decimal import Decimal -from typing import Any, cast +from typing import Any, cast, override from ai.backend.common.types import ( BinarySize, @@ -36,9 +36,11 @@ class ResourceLimitRule(SessionSpecValidatorRule): """Per-kernel requested slots must satisfy image min/max + shmem rules.""" + @override def name(self) -> str: return "resource_limit" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/sokovan/scheduling_controller/validators/service_port_rule.py b/src/ai/backend/manager/sokovan/scheduling_controller/validators/service_port_rule.py index 8217b0ecc64..0b07afa5236 100644 --- a/src/ai/backend/manager/sokovan/scheduling_controller/validators/service_port_rule.py +++ b/src/ai/backend/manager/sokovan/scheduling_controller/validators/service_port_rule.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Mapping -from typing import Any +from typing import Any, override from ai.backend.manager.data.session.spec import SessionSpec from ai.backend.manager.errors.api import InvalidAPIParameters @@ -18,9 +18,11 @@ class ServicePortRule(SessionSpecValidatorRule): """Per-kernel preopen_ports must not collide with reserved or service ports.""" + @override def name(self) -> str: return "service_port" + @override def validate( self, spec: SessionSpec, diff --git a/src/ai/backend/manager/tasks/agent_lost_checker.py b/src/ai/backend/manager/tasks/agent_lost_checker.py index 4a56150a878..64ca2e9e5da 100644 --- a/src/ai/backend/manager/tasks/agent_lost_checker.py +++ b/src/ai/backend/manager/tasks/agent_lost_checker.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override from dateutil.tz import tzutc @@ -37,17 +37,21 @@ def __init__( self._event_producer = event_producer @property + @override def name(self) -> str: return "agent_lost_checker" @property + @override def interval(self) -> float: return _CHECK_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: now = datetime.now(tzutc()) timeout = timedelta(seconds=self._config_provider.config.manager.heartbeat_timeout) diff --git a/src/ai/backend/manager/tasks/stats_reporter.py b/src/ai/backend/manager/tasks/stats_reporter.py index 86883a5c02a..011e43ce8f2 100644 --- a/src/ai/backend/manager/tasks/stats_reporter.py +++ b/src/ai/backend/manager/tasks/stats_reporter.py @@ -4,7 +4,7 @@ import asyncio import logging -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, override import sqlalchemy as sa import sqlalchemy.exc @@ -48,17 +48,21 @@ def __init__( self._db = db @property + @override def name(self) -> str: return "stats_task" @property + @override def interval(self) -> float: return _REPORT_INTERVAL @property + @override def initial_delay(self) -> float: return 0.0 + @override async def run(self) -> None: try: await self._stats_monitor.report_metric( diff --git a/src/ai/backend/storage/cli/context.py b/src/ai/backend/storage/cli/context.py index 8bfa409a561..343bf5a8e9b 100644 --- a/src/ai/backend/storage/cli/context.py +++ b/src/ai/backend/storage/cli/context.py @@ -2,7 +2,7 @@ from contextlib import AbstractContextManager from pathlib import Path -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, Self, override import click @@ -19,6 +19,7 @@ def __init__(self, log_level: LogLevel, config_path: Path | None = None) -> None self.config_path = config_path self.log_level = log_level + @override def __enter__(self) -> Self: from ai.backend.logging import LocalLogger @@ -31,6 +32,7 @@ def __enter__(self) -> Self: self._logger.__enter__() return self + @override def __exit__(self, *exc_info: Any) -> None: click_ctx = click.get_current_context() if click_ctx.invoked_subcommand != "start-server": diff --git a/src/ai/backend/storage/data/storage/types.py b/src/ai/backend/storage/data/storage/types.py index a2e7fe1ef55..bf3b44dbcaf 100644 --- a/src/ai/backend/storage/data/storage/types.py +++ b/src/ai/backend/storage/data/storage/types.py @@ -2,7 +2,7 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from ai.backend.common.artifact_storage import AbstractStorage, AbstractStoragePool from ai.backend.common.contexts.request_id import current_request_id @@ -88,6 +88,7 @@ def resolve_storage(self, storage_pool: AbstractStoragePool) -> AbstractStorage: raise InvalidStorageTargetError("StorageTarget has no storage_name or storage") return storage_pool.get_storage(self._storage_name) + @override def __eq__(self, other: object) -> bool: if isinstance(other, StorageTarget): if self._storage_name is not None and other._storage_name is not None: @@ -97,6 +98,7 @@ def __eq__(self, other: object) -> bool: return False return False + @override def __hash__(self) -> int: if self._storage_name is not None: return hash(self._storage_name) diff --git a/src/ai/backend/storage/dependencies/bootstrap/composer.py b/src/ai/backend/storage/dependencies/bootstrap/composer.py index dc01031bb3e..ed7b1a6a914 100644 --- a/src/ai/backend/storage/dependencies/bootstrap/composer.py +++ b/src/ai/backend/storage/dependencies/bootstrap/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.storage.config.unified import StorageProxyUnifiedConfig @@ -29,10 +30,12 @@ class BootstrapComposer(DependencyComposer[BootstrapInput, BootstrapResources]): """Composes bootstrap dependencies.""" @property + @override def stage_name(self) -> str: return "bootstrap" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/storage/dependencies/bootstrap/config.py b/src/ai/backend/storage/dependencies/bootstrap/config.py index 40e7a1e4983..aa889947477 100644 --- a/src/ai/backend/storage/dependencies/bootstrap/config.py +++ b/src/ai/backend/storage/dependencies/bootstrap/config.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.storage.config.loaders import load_local_config @@ -23,10 +24,12 @@ class ConfigProvider( """Provider for storage proxy configuration.""" @property + @override def stage_name(self) -> str: return "config" @asynccontextmanager + @override async def provide( self, setup_input: ConfigProviderInput ) -> AsyncIterator[StorageProxyUnifiedConfig]: diff --git a/src/ai/backend/storage/dependencies/composer.py b/src/ai/backend/storage/dependencies/composer.py index bc5730ab2d6..c6b6b2648c7 100644 --- a/src/ai/backend/storage/dependencies/composer.py +++ b/src/ai/backend/storage/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -36,10 +37,12 @@ class StorageDependencyComposer(DependencyComposer[DependencyInput, DependencyRe """Main composer for all storage proxy dependencies.""" @property + @override def stage_name(self) -> str: return "storage-proxy" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/storage/dependencies/infrastructure/composer.py b/src/ai/backend/storage/dependencies/infrastructure/composer.py index 97869ee0173..9bf9e6762f1 100644 --- a/src/ai/backend/storage/dependencies/infrastructure/composer.py +++ b/src/ai/backend/storage/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.common.etcd import AsyncEtcd @@ -33,10 +34,12 @@ class InfrastructureComposer( """Composer for infrastructure layer dependencies.""" @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/storage/dependencies/infrastructure/etcd.py b/src/ai/backend/storage/dependencies/infrastructure/etcd.py index b76e10fae08..14dab6607bc 100644 --- a/src/ai/backend/storage/dependencies/infrastructure/etcd.py +++ b/src/ai/backend/storage/dependencies/infrastructure/etcd.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import DependencyProvider from ai.backend.common.etcd import AsyncEtcd @@ -15,15 +16,18 @@ class EtcdProvider(DependencyProvider[StorageProxyUnifiedConfig, AsyncEtcd]): """Provider for etcd client connection.""" @property + @override def stage_name(self) -> str: return "etcd" @asynccontextmanager + @override async def provide(self, setup_input: StorageProxyUnifiedConfig) -> AsyncIterator[AsyncEtcd]: """Create and provide an etcd client.""" async with make_etcd(setup_input) as etcd: yield etcd + @override def gen_liveness_checker(self, resource: AsyncEtcd) -> ServiceHealthChecker: """Liveness — stuck etcd connection observed; restart is the recovery path.""" return EtcdHealthChecker(etcd=resource) diff --git a/src/ai/backend/storage/dependencies/infrastructure/redis.py b/src/ai/backend/storage/dependencies/infrastructure/redis.py index 6bde465eac7..d3b36920a84 100644 --- a/src/ai/backend/storage/dependencies/infrastructure/redis.py +++ b/src/ai/backend/storage/dependencies/infrastructure/redis.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_artifact.client import ( ValkeyArtifactDownloadTrackingClient, @@ -32,10 +33,12 @@ class RedisProvider(DependencyProvider[AsyncEtcd, StorageProxyValkeyClients]): """Provider for Redis configuration and Valkey clients.""" @property + @override def stage_name(self) -> str: return "redis" @asynccontextmanager + @override async def provide(self, setup_input: AsyncEtcd) -> AsyncIterator[StorageProxyValkeyClients]: """Load and provide Redis configuration and Valkey clients.""" # Load Redis config from etcd @@ -70,6 +73,7 @@ async def provide(self, setup_input: AsyncEtcd) -> AsyncIterator[StorageProxyVal await bgtask_client.close() await artifact_client.close() + @override def gen_liveness_checker(self, resource: StorageProxyValkeyClients) -> ServiceHealthChecker: """Liveness — Valkey connection-stuck observed; restart recovers.""" return ValkeyHealthChecker( diff --git a/src/ai/backend/storage/dependencies/storage/composer.py b/src/ai/backend/storage/dependencies/storage/composer.py index 2629717815a..2904b83a0f7 100644 --- a/src/ai/backend/storage/dependencies/storage/composer.py +++ b/src/ai/backend/storage/dependencies/storage/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.storage.config.unified import StorageProxyUnifiedConfig @@ -29,10 +30,12 @@ class StorageComposer(DependencyComposer[StorageComposerInput, StorageResources] """Composer for storage layer dependencies.""" @property + @override def stage_name(self) -> str: return "storage" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/storage/dependencies/storage/storage_pool.py b/src/ai/backend/storage/dependencies/storage/storage_pool.py index 10de2ca63ca..d096ff7bc2a 100644 --- a/src/ai/backend/storage/dependencies/storage/storage_pool.py +++ b/src/ai/backend/storage/dependencies/storage/storage_pool.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.dependencies import NonMonitorableDependencyProvider from ai.backend.storage.config.unified import StorageProxyUnifiedConfig @@ -12,10 +13,12 @@ class StoragePoolProvider(NonMonitorableDependencyProvider[StorageProxyUnifiedCo """Provider for storage pool.""" @property + @override def stage_name(self) -> str: return "storage-pool" @asynccontextmanager + @override async def provide(self, setup_input: StorageProxyUnifiedConfig) -> AsyncIterator[StoragePool]: """Create and provide storage pool.""" storage_pool = StoragePool.from_config(setup_input) diff --git a/src/ai/backend/storage/errors/base.py b/src/ai/backend/storage/errors/base.py index c3916bd1423..00f786e4164 100644 --- a/src/ai/backend/storage/errors/base.py +++ b/src/ai/backend/storage/errors/base.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class StorageProxyError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/generic" error_title = "Storage Proxy Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -35,6 +38,7 @@ class ProcessExecutionError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/execution/failed" error_title = "Storage Operation Execution Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -49,6 +53,7 @@ class ExternalStorageServiceError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/external/failed" error_title = "External Operation Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -63,6 +68,7 @@ class NotImplementedAPI(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/api/not-implemented" error_title = "API Not Implemented" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, diff --git a/src/ai/backend/storage/errors/common.py b/src/ai/backend/storage/errors/common.py index 455e2866c54..3628d63845e 100644 --- a/src/ai/backend/storage/errors/common.py +++ b/src/ai/backend/storage/errors/common.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class InvalidAPIParameters(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/invalid-api-params" error_title = "Invalid API parameters" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -35,6 +38,7 @@ class StorageNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/object-not-found" error_title = "Storage Config Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -49,6 +53,7 @@ class StorageTypeInvalidError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/object-type-invalid" error_title = "Storage Config Invalid Type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -63,6 +68,7 @@ class StorageTransferError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/transfer/failed" error_title = "Storage Transfer Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -77,6 +83,7 @@ class StorageStepRequiredStepNotProvided(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/step-mapping-not-provided" error_title = "Storage Step Mapping Not Provided" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -94,6 +101,7 @@ class InvalidPathError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/path/invalid" error_title = "Invalid Path" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -108,6 +116,7 @@ class InvalidSocketPathError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/socket-path/invalid" error_title = "Invalid Socket Path" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -122,6 +131,7 @@ class InvalidConfigurationSourceError(BackendAIError, web.HTTPInternalServerErro error_type = "https://api.backend.ai/probs/storage/config/invalid-source" error_title = "Invalid Configuration Source" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -136,6 +146,7 @@ class InvalidDataLengthError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/data/invalid-length" error_title = "Invalid Data Length" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -150,6 +161,7 @@ class ServiceNotInitializedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/service/not-initialized" error_title = "Service Not Initialized" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -164,6 +176,7 @@ class InvalidStorageTargetError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/invalid-storage-target" error_title = "Invalid Storage Target" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -181,6 +194,7 @@ class UploadOffsetMismatchError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/storage/upload-offset-mismatch" error_title = "Upload Offset Mismatch" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, diff --git a/src/ai/backend/storage/errors/object.py b/src/ai/backend/storage/errors/object.py index db3802b6e3c..34ab41aced3 100644 --- a/src/ai/backend/storage/errors/object.py +++ b/src/ai/backend/storage/errors/object.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -23,6 +25,7 @@ class FileStreamUploadError(ProcessExecutionError): error_type = "https://api.backend.ai/probs/storage/file-stream-upload-failed" error_title = "Failed to upload file stream" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -37,6 +40,7 @@ class FileStreamDownloadError(ProcessExecutionError): error_type = "https://api.backend.ai/probs/storage/file-stream-download-failed" error_title = "Failed to download file stream" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -51,6 +55,7 @@ class PresignedUploadURLGenerationError(ProcessExecutionError): error_type = "https://api.backend.ai/probs/storage/presigned-upload-url-generation-failed" error_title = "Failed to generate presigned upload URL" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -65,6 +70,7 @@ class PresignedDownloadURLGenerationError(ProcessExecutionError): error_type = "https://api.backend.ai/probs/storage/presigned-download-url-generation-failed" error_title = "Failed to generate presigned download URL" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -79,6 +85,7 @@ class ObjectInfoFetchError(ProcessExecutionError): error_type = "https://api.backend.ai/probs/storage/object-info-fetch-failed" error_title = "Failed to fetch object info" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -93,6 +100,7 @@ class ObjectStorageBucketNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/bucket/object-not-found" error_title = "Storage Bucket Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -107,6 +115,7 @@ class StorageBucketFileNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/bucket/file/object-not-found" error_title = "Storage Bucket File Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -121,6 +130,7 @@ class ObjectStorageConfigInvalidError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/object/config/invalid" error_title = "Object Storage Config Invalid" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -135,6 +145,7 @@ class ReservoirStorageConfigInvalidError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/reservoir/config/invalid" error_title = "Reservoir Storage Config Invalid" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFS_STORAGE, @@ -149,6 +160,7 @@ class ObjectStorageObjectDeletionError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/object/deletion/failed" error_title = "Object Deletion Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -166,6 +178,7 @@ class RegistryNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/registries/registry-not-found" error_title = "Registry Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -180,6 +193,7 @@ class HuggingFaceAPIError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/registries/huggingface/api-error" error_title = "HuggingFace API Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -194,6 +208,7 @@ class HuggingFaceModelNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/registries/huggingface/model-not-found" error_title = "HuggingFace Model Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -208,6 +223,7 @@ class ArtifactStorageEmptyError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/artifact/config/invalid" error_title = "Artifact Storage Empty" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -222,6 +238,7 @@ class ArtifactRevisionEmptyError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/artifact/revision/empty" error_title = "Artifact Revision Empty" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -236,6 +253,7 @@ class ArtifactImportError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/artifact/import/failed" error_title = "Artifact Import Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -250,6 +268,7 @@ class ArtifactVerifyStorageTypeInvalid(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/artifact/verify/storage-type/invalid" error_title = "Artifact Verify Storage Type Invalid" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -264,6 +283,7 @@ class ArtifactVerificationFailedError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/artifact/verification/failed" error_title = "Artifact Verification Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -278,6 +298,7 @@ class UnsupportedFileTypeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/unsupported-file-type" error_title = "Unsupported File Type" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFS_STORAGE, diff --git a/src/ai/backend/storage/errors/process.py b/src/ai/backend/storage/errors/process.py index 1590dee8b1a..4349bc0d080 100644 --- a/src/ai/backend/storage/errors/process.py +++ b/src/ai/backend/storage/errors/process.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -23,6 +25,7 @@ class SubprocessStdoutNotAvailableError(BackendAIError, web.HTTPInternalServerEr error_type = "https://api.backend.ai/probs/storage/subprocess/stdout-unavailable" error_title = "Subprocess Stdout Not Available" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -37,6 +40,7 @@ class QuotaCommandFailedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/quota/command-failed" error_title = "Quota Command Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -51,6 +55,7 @@ class CephNotInstalledError(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/storage/ceph/not-installed" error_title = "Ceph Not Installed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -65,6 +70,7 @@ class PureStorageCommandFailedError(BackendAIError, web.HTTPInternalServerError) error_type = "https://api.backend.ai/probs/storage/purestorage/command-failed" error_title = "PureStorage Command Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -79,6 +85,7 @@ class NetAppClientError(ExternalStorageServiceError, web.HTTPServiceUnavailable) error_type = "https://api.backend.ai/probs/storage/netapp/api-error" error_title = "NetApp API Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -93,6 +100,7 @@ class NetAppQTreeNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/netapp/qtree-not-found" error_title = "NetApp QTree Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -107,6 +115,7 @@ class DDNCommandFailedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/ddn/command-failed" error_title = "DDN Command Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -121,6 +130,7 @@ class MetricNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/metric-not-found" error_title = "Metric Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, diff --git a/src/ai/backend/storage/errors/quota.py b/src/ai/backend/storage/errors/quota.py index 2b64d1869e6..239f8921ed3 100644 --- a/src/ai/backend/storage/errors/quota.py +++ b/src/ai/backend/storage/errors/quota.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class QuotaDirectoryNotEmptyError(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/storage/quota-directory-not-empty" error_title = "Quota Directory Not Empty" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -35,6 +38,7 @@ class QuotaScopeNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/quota/scope/not-found" error_title = "Quota Scope Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -49,6 +53,7 @@ class QuotaScopeAlreadyExists(BackendAIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/storage/quota/scope/already-exists" error_title = "Quota Scope Already Exists" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -63,6 +68,7 @@ class InvalidQuotaConfig(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/quota/config/invalid" error_title = "Invalid Quota Config" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -77,6 +83,7 @@ class InvalidQuotaScopeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/quota/scope/invalid" error_title = "Invalid Quota Scope" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -91,6 +98,7 @@ class InvalidQuotaFormatError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/quota/format/invalid" error_title = "Invalid Quota Format" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -105,6 +113,7 @@ class QuotaScopeCreationFailedError(BackendAIError, web.HTTPInternalServerError) error_type = "https://api.backend.ai/probs/storage/quota/scope/creation-failed" error_title = "Quota Scope Creation Failed" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -119,6 +128,7 @@ class QuotaTreeNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/quota/tree/not-found" error_title = "Quota Tree Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, diff --git a/src/ai/backend/storage/errors/vfolder.py b/src/ai/backend/storage/errors/vfolder.py index a3e4439804f..df4376a8bd1 100644 --- a/src/ai/backend/storage/errors/vfolder.py +++ b/src/ai/backend/storage/errors/vfolder.py @@ -5,6 +5,7 @@ from __future__ import annotations from pathlib import PurePosixPath +from typing import override from aiohttp import web @@ -24,6 +25,7 @@ class VFolderNotFoundError(BackendAIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/vfolder/not-found" error_title = "VFolder Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -42,6 +44,7 @@ def __init__(self, vfid: VFolderID, relpath: PurePosixPath) -> None: msg_str = f"Invalid Subpath (vfid={vfid}, relpath={relpath})" super().__init__(extra_msg=msg_str, extra_data=msg_str) + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, diff --git a/src/ai/backend/storage/errors/volume.py b/src/ai/backend/storage/errors/volume.py index a3196ca3291..9642fc1f6a7 100644 --- a/src/ai/backend/storage/errors/volume.py +++ b/src/ai/backend/storage/errors/volume.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class InvalidVolumeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/volume/invalid" error_title = "Invalid Volume" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -35,6 +38,7 @@ class VolumeNotInitializedError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/volume/not-initialized" error_title = "Volume Not Initialized" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -49,6 +53,7 @@ class MetadataTooLargeError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/volume/metadata-too-large" error_title = "Metadata Too Large" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, diff --git a/src/ai/backend/storage/exception.py b/src/ai/backend/storage/exception.py index 36a347ee941..0206029f4bd 100644 --- a/src/ai/backend/storage/exception.py +++ b/src/ai/backend/storage/exception.py @@ -1,4 +1,5 @@ from pathlib import PurePosixPath +from typing import override from aiohttp import web @@ -17,6 +18,7 @@ class StorageProxyError(BackendAIError, web.HTTPInternalServerError): error_title = "Storage Proxy Error" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -30,6 +32,7 @@ class ProcessExecutionError(BackendAIError, web.HTTPInternalServerError): error_title = "Storage Operation Execution Failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -43,6 +46,7 @@ class ExternalStorageServiceError(BackendAIError, web.HTTPInternalServerError): error_title = "External Operation Failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -56,6 +60,7 @@ class NetAppClientError(ExternalStorageServiceError, web.HTTPServiceUnavailable) error_title = "NetApp API Error" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -69,6 +74,7 @@ class VFolderNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "VFolder Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.VFOLDER, @@ -82,6 +88,7 @@ class QuotaDirectoryNotEmptyError(BackendAIError, web.HTTPConflict): error_title = "Quota Directory Not Empty" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -95,6 +102,7 @@ class QuotaScopeNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "Quota Scope Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -108,6 +116,7 @@ class QuotaScopeRequiredError(BackendAIError, web.HTTPPreconditionRequired): error_title = "Quota Scope Required" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.QUOTA_SCOPE, @@ -121,6 +130,7 @@ class QuotaScopeAlreadyExists(BackendAIError, web.HTTPConflict): error_title = "Quota Scope Already Exists" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -134,6 +144,7 @@ class InvalidQuotaConfig(BackendAIError, web.HTTPBadRequest): error_title = "Invalid Quota Config" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -151,6 +162,7 @@ def __init__(self, vfid: VFolderID, relpath: PurePosixPath) -> None: super().__init__(extra_msg=msg_str, extra_data=msg_str) @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -164,6 +176,7 @@ class InvalidQuotaScopeError(BackendAIError, web.HTTPBadRequest): error_title = "Invalid Quota Scope" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -177,6 +190,7 @@ class InvalidVolumeError(BackendAIError, web.HTTPBadRequest): error_title = "Invalid Volume" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -190,6 +204,7 @@ class InvalidAPIParameters(BackendAIError, web.HTTPBadRequest): error_title = "Invalid API parameters" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -203,6 +218,7 @@ class FileStreamUploadError(ProcessExecutionError): error_title = "Failed to upload file stream" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -216,6 +232,7 @@ class FileStreamDownloadError(ProcessExecutionError): error_title = "Failed to download file stream" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -229,6 +246,7 @@ class PresignedUploadURLGenerationError(ProcessExecutionError): error_title = "Failed to generate presigned upload URL" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -242,6 +260,7 @@ class PresignedDownloadURLGenerationError(ProcessExecutionError): error_title = "Failed to generate presigned download URL" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -255,6 +274,7 @@ class ObjectInfoFetchError(ProcessExecutionError): error_title = "Failed to fetch object info" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -268,6 +288,7 @@ class StorageNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "Storage Config Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -281,6 +302,7 @@ class StorageTypeInvalidError(BackendAIError, web.HTTPBadRequest): error_title = "Storage Config Invalid Type" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -294,6 +316,7 @@ class ObjectStorageBucketNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "Storage Bucket Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -307,6 +330,7 @@ class StorageBucketFileNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "Storage Bucket File Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -320,6 +344,7 @@ class RegistryNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "Registry Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -333,6 +358,7 @@ class HuggingFaceAPIError(BackendAIError, web.HTTPInternalServerError): error_title = "HuggingFace API Error" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT_REGISTRY, @@ -346,6 +372,7 @@ class HuggingFaceModelNotFoundError(BackendAIError, web.HTTPNotFound): error_title = "HuggingFace Model Not Found" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -359,6 +386,7 @@ class ObjectStorageConfigInvalidError(BackendAIError, web.HTTPBadRequest): error_title = "Object Storage Config Invalid" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -372,6 +400,7 @@ class ReservoirStorageConfigInvalidError(BackendAIError, web.HTTPBadRequest): error_title = "Reservoir Storage Config Invalid" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -385,6 +414,7 @@ class ArtifactStorageEmptyError(BackendAIError, web.HTTPNotFound): error_title = "Artifact Storage Empty" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -398,6 +428,7 @@ class ArtifactRevisionEmptyError(BackendAIError, web.HTTPBadRequest): error_title = "Artifact Revision Empty" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -411,6 +442,7 @@ class ArtifactImportError(BackendAIError, web.HTTPInternalServerError): error_title = "Artifact Import Failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.ARTIFACT, @@ -424,6 +456,7 @@ class NotImplementedAPI(BackendAIError, web.HTTPBadRequest): error_title = "API Not Implemented" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -437,6 +470,7 @@ class ObjectStorageObjectDeletionError(BackendAIError, web.HTTPBadRequest): error_title = "Object Deletion Failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.OBJECT_STORAGE, @@ -450,6 +484,7 @@ class StorageTransferError(BackendAIError, web.HTTPInternalServerError): error_title = "Storage Transfer Failed" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, @@ -463,6 +498,7 @@ class StorageStepRequiredStepNotProvided(BackendAIError, web.HTTPBadRequest): error_title = "Storage Step Mapping Not Provided" @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE_PROXY, diff --git a/src/ai/backend/storage/plugin.py b/src/ai/backend/storage/plugin.py index b33e7fe9aa5..79a5ec220a7 100644 --- a/src/ai/backend/storage/plugin.py +++ b/src/ai/backend/storage/plugin.py @@ -4,7 +4,7 @@ from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from aiohttp import web @@ -41,6 +41,7 @@ class StoragePluginContext(BasePluginContext[AbstractStoragePlugin]): plugin_group = "backendai_storage_v10" @classmethod + @override def discover_plugins( cls, plugin_group: str, @@ -55,6 +56,7 @@ class StorageArtifactVerifierPluginContext(BasePluginContext[AbstractArtifactVer plugin_group = "backendai_storage_artifact_verifier_v1" @classmethod + @override def discover_plugins( cls, plugin_group: str, diff --git a/src/ai/backend/storage/services/artifacts/common.py b/src/ai/backend/storage/services/artifacts/common.py index ca3025f8707..5df983764a9 100644 --- a/src/ai/backend/storage/services/artifacts/common.py +++ b/src/ai/backend/storage/services/artifacts/common.py @@ -55,11 +55,13 @@ def __init__( self._event_producer = event_producer @property + @override def step_type(self) -> ArtifactStorageImportStep: return ArtifactStorageImportStep.VERIFY @property @abstractmethod + @override def registry_type(self) -> ArtifactRegistryType: """Registry type used for revision resolution""" raise NotImplementedAPI @@ -200,11 +202,13 @@ def __init__(self, transfer_manager: StorageTransferManager) -> None: self._transfer_manager = transfer_manager @property + @override def step_type(self) -> ArtifactStorageImportStep: return ArtifactStorageImportStep.ARCHIVE @property @abstractmethod + @override def registry_type(self) -> ArtifactRegistryType: """Registry type used for revision resolution""" raise NotImplementedAPI diff --git a/src/ai/backend/storage/services/artifacts/huggingface.py b/src/ai/backend/storage/services/artifacts/huggingface.py index 8c54ed35199..f0c7616970a 100644 --- a/src/ai/backend/storage/services/artifacts/huggingface.py +++ b/src/ai/backend/storage/services/artifacts/huggingface.py @@ -755,6 +755,7 @@ def __init__( self._redis_client = redis_client @property + @override def step_type(self) -> ArtifactStorageImportStep: return ArtifactStorageImportStep.DOWNLOAD diff --git a/src/ai/backend/storage/services/artifacts/reservoir.py b/src/ai/backend/storage/services/artifacts/reservoir.py index 32efe783a3b..c064f08032e 100644 --- a/src/ai/backend/storage/services/artifacts/reservoir.py +++ b/src/ai/backend/storage/services/artifacts/reservoir.py @@ -642,6 +642,7 @@ def __init__( self._redis_client = redis_client @property + @override def step_type(self) -> ArtifactStorageImportStep: return ArtifactStorageImportStep.DOWNLOAD diff --git a/src/ai/backend/storage/storages/vfs_storage.py b/src/ai/backend/storage/storages/vfs_storage.py index dcca95bd98b..88db346c1a0 100644 --- a/src/ai/backend/storage/storages/vfs_storage.py +++ b/src/ai/backend/storage/storages/vfs_storage.py @@ -173,6 +173,7 @@ def cleanup_temporary_storage(self) -> None: log.warning("Failed to clean temporary storage {}: {}", self._base_path, e) @property + @override def name(self) -> str: return self._name diff --git a/src/ai/backend/storage/volumes/cephfs/__init__.py b/src/ai/backend/storage/volumes/cephfs/__init__.py index e29d776dcd0..00966c68b2f 100644 --- a/src/ai/backend/storage/volumes/cephfs/__init__.py +++ b/src/ai/backend/storage/volumes/cephfs/__init__.py @@ -3,7 +3,7 @@ import os import shutil from pathlib import Path -from typing import Any +from typing import Any, override import aiofiles.os @@ -25,6 +25,7 @@ class CephDirQuotaModel(BaseQuotaModel): + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -36,6 +37,7 @@ async def create_quota_scope( if options is not None: await self.update_quota_scope(quota_scope_id, options) + @override async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage | None: qspath = self.mangle_qspath(quota_scope_id) if not qspath.exists(): @@ -74,6 +76,7 @@ def read_attrs() -> tuple[int, int]: ) return QuotaUsage(used_bytes=used_bytes, limit_bytes=limit_bytes) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -93,6 +96,7 @@ async def update_quota_scope( ), ) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: qspath = self.mangle_qspath(quota_scope_id) if not qspath.exists(): @@ -108,6 +112,7 @@ async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: class CephFSOpModel(BaseFSOpModel): + @override async def scan_tree_usage(self, path: Path) -> TreeUsage: loop = asyncio.get_running_loop() raw_reports = await loop.run_in_executor( @@ -121,6 +126,7 @@ async def scan_tree_usage(self, path: Path) -> TreeUsage: used_bytes = int(raw_reports[1].strip().decode()) return TreeUsage(file_count=file_count, used_bytes=used_bytes) + @override async def scan_tree_size(self, path: Path) -> BinarySize: loop = asyncio.get_running_loop() raw_report = await loop.run_in_executor( @@ -136,6 +142,7 @@ class CephFSVolume(BaseVolume): registry: dict[str, int] project_id_pool: list[int] + @override async def init(self) -> None: try: await run([b"ceph", b"--version"]) @@ -143,18 +150,22 @@ async def init(self) -> None: raise CephNotInstalledError("Ceph is not installed.") from e await super().init() + @override async def create_quota_model(self) -> AbstractQuotaModel: return CephDirQuotaModel(self.mount_path) + @override async def create_fsop_model(self) -> AbstractFSOpModel: return CephFSOpModel( self.mount_path, self.local_config["storage-proxy"]["scandir-limit"], ) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_QUOTA, CAP_FAST_SIZE]) + @override async def get_fs_usage(self) -> CapacityUsage: (total, used, _) = await asyncio.get_running_loop().run_in_executor( None, diff --git a/src/ai/backend/storage/volumes/ddn/__init__.py b/src/ai/backend/storage/volumes/ddn/__init__.py index 3c8f10cd9c5..a27048e12b2 100644 --- a/src/ai/backend/storage/volumes/ddn/__init__.py +++ b/src/ai/backend/storage/volumes/ddn/__init__.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from pathlib import Path from subprocess import CalledProcessError -from typing import Any, Final +from typing import Any, Final, override import aiofiles import aiofiles.os @@ -140,6 +140,7 @@ async def _get_quota_by_project(self, pid: int, qspath: Path) -> QuotaUsage | No finally: await proc.wait() + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -182,6 +183,7 @@ async def create_quota_scope( if options is not None: await self._set_quota_by_project(project_id, qspath, options) + @override async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage | None: """ $ lfs quota -p @@ -207,6 +209,7 @@ async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage return await self._get_quota_by_project(pid, qspath) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -219,6 +222,7 @@ async def update_quota_scope( raise QuotaScopeNotFoundError await self._set_quota_by_project(pid, qspath, config) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: qspath = self.mangle_qspath(quota_scope_id) pid_path = qspath / PROJECT_ID_FILE_NAME @@ -227,6 +231,7 @@ async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: raise QuotaScopeNotFoundError await self._unset_quota_by_project(pid, qspath) + @override async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None: await self.unset_quota(quota_scope_id) qspath = self.mangle_qspath(quota_scope_id) @@ -236,8 +241,10 @@ async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None: class EXAScalerFSVolume(BaseVolume): name = "exascaler" + @override async def create_quota_model(self) -> AbstractQuotaModel: return EXAScalerQuotaModel(self.mount_path, self.local_config, self.etcd) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_QUOTA]) diff --git a/src/ai/backend/storage/volumes/dellemc/__init__.py b/src/ai/backend/storage/volumes/dellemc/__init__.py index 1a426680fe3..99a4b3a9c94 100644 --- a/src/ai/backend/storage/volumes/dellemc/__init__.py +++ b/src/ai/backend/storage/volumes/dellemc/__init__.py @@ -3,7 +3,7 @@ import logging from collections.abc import Mapping from pathlib import Path -from typing import Any, cast +from typing import Any, cast, override import aiofiles.os @@ -64,6 +64,7 @@ async def _unset_quota_id(self, qspath: Path) -> None: if quota_id_path.exists(): await aiofiles.os.remove(quota_id_path) + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -75,6 +76,7 @@ async def create_quota_scope( if options is not None: await self.update_quota_scope(quota_scope_id, options) + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -99,6 +101,7 @@ async def describe_quota_scope( ) return None + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -124,6 +127,7 @@ async def update_quota_scope( ) await self._set_quota_id(qspath, result["id"]) + @override async def unset_quota( self, quota_scope_id: QuotaScopeID, @@ -139,6 +143,7 @@ async def unset_quota( await self.api_client.delete_quota(quota_id) await self._unset_quota_id(qspath) + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -188,18 +193,22 @@ def __init__( system_name=self.dell_system_name, ) + @override async def shutdown(self) -> None: await self.api_client.aclose() + @override async def create_quota_model(self) -> AbstractQuotaModel: ifs_path = Path(self.config["dell_ifs_path"]) return DellEMCOneFSQuotaModel( self.mount_path, api_client=self.api_client, ifs_path=ifs_path ) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_FAST_FS_SIZE, CAP_VFOLDER, CAP_QUOTA, CAP_METRIC]) + @override async def get_hwinfo(self) -> HardwareMetadata: raw_metadata = await self.api_client.get_metadata() return { @@ -208,6 +217,7 @@ async def get_hwinfo(self) -> HardwareMetadata: "metadata": {**raw_metadata}, } + @override async def get_fs_usage(self) -> CapacityUsage: usage = await self.api_client.get_usage() return CapacityUsage( @@ -215,6 +225,7 @@ async def get_fs_usage(self) -> CapacityUsage: used_bytes=int(usage["used_bytes"]), ) + @override async def get_performance_metric(self) -> FSPerfMetric: protocol_stats = await self.get_protocol_stats() workload = await self.get_workload_stats() diff --git a/src/ai/backend/storage/volumes/dellemc/exceptions.py b/src/ai/backend/storage/volumes/dellemc/exceptions.py index 9c3330b4370..1fdac8e91dc 100644 --- a/src/ai/backend/storage/volumes/dellemc/exceptions.py +++ b/src/ai/backend/storage/volumes/dellemc/exceptions.py @@ -1,3 +1,6 @@ +from typing import override + + class DellError(Exception): pass @@ -14,6 +17,7 @@ def __init__(self, message: str | None = None, *args: object) -> None: self.message = message or "" + @override def __str__(self) -> str: return type(self).__name__ diff --git a/src/ai/backend/storage/volumes/gpfs/__init__.py b/src/ai/backend/storage/volumes/gpfs/__init__.py index 7231feeec5b..eded1ac91fa 100644 --- a/src/ai/backend/storage/volumes/gpfs/__init__.py +++ b/src/ai/backend/storage/volumes/gpfs/__init__.py @@ -1,7 +1,7 @@ import logging from collections.abc import Mapping from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, override from ai.backend.common.etcd import AsyncEtcd from ai.backend.common.events.dispatcher import EventDispatcher, EventProducer @@ -48,6 +48,7 @@ def __init__( def _get_fileset_name(self, quota_scope_id: QuotaScopeID) -> str: return self._fileset_prefix + quota_scope_id.pathname + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -64,6 +65,7 @@ async def create_quota_scope( if options is not None: await self.update_quota_scope(quota_scope_id, options) + @override async def update_quota_scope(self, quota_scope_id: QuotaScopeID, config: QuotaConfig) -> None: await self.api_client.set_quota( self.fs, @@ -71,6 +73,7 @@ async def update_quota_scope(self, quota_scope_id: QuotaScopeID, config: QuotaCo config.limit_bytes, ) + @override async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage | None: if not self.mangle_qspath(quota_scope_id).exists(): return None @@ -100,9 +103,11 @@ async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage limit_bytes=limit_bytes, ) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: await self.api_client.remove_quota(self.fs, self._get_fileset_name(quota_scope_id)) + @override async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None: await self.api_client.remove_fileset(self.fs, self._get_fileset_name(quota_scope_id)) @@ -119,6 +124,7 @@ def __init__( self.api_client = api_client self.fs = fs + @override async def copy_tree( self, src_path: Path, @@ -169,9 +175,11 @@ def __init__( self.fs = self.config["gpfs_fs_name"] self.gpfs_owner = self.config.get("gpfs_owner", "1000:1000") + @override async def init(self) -> None: await super().init() + @override async def create_quota_model(self) -> AbstractQuotaModel: return GPFSQuotaModel( self.mount_path, @@ -181,6 +189,7 @@ async def create_quota_model(self) -> AbstractQuotaModel: fileset_prefix=self.config.get("gpfs_fileset_prefix"), ) + @override async def create_fsop_model(self) -> AbstractFSOpModel: return GPFSOpModel( self.mount_path, @@ -189,9 +198,11 @@ async def create_fsop_model(self) -> AbstractFSOpModel: self.fs, ) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_FAST_FS_SIZE, CAP_VFOLDER, CAP_QUOTA, CAP_METRIC]) + @override async def get_hwinfo(self) -> HardwareMetadata: nodes = await self.api_client.list_nodes() invalid_status = ["FAILED", "DEGRADED", "ERROR"] @@ -224,6 +235,7 @@ async def get_hwinfo(self) -> HardwareMetadata: }, } + @override async def get_fs_usage(self) -> CapacityUsage: storage_pools = await self.api_client.list_fs_pools(self.fs) free, total = 0, 0 @@ -238,6 +250,7 @@ async def get_fs_usage(self) -> CapacityUsage: capacity_bytes=BinarySize(total) * 1024, ) + @override async def get_performance_metric(self) -> FSPerfMetric: # ref: https://www.ibm.com/docs/en/spectrum-scale/5.0.3?topic=2-perfmondata-get query = ( diff --git a/src/ai/backend/storage/volumes/gpfs/exceptions.py b/src/ai/backend/storage/volumes/gpfs/exceptions.py index 8d689ec850c..7597b170f74 100644 --- a/src/ai/backend/storage/volumes/gpfs/exceptions.py +++ b/src/ai/backend/storage/volumes/gpfs/exceptions.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -17,6 +19,7 @@ class GPFSInitError(GPFSError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/gpfs-init-error" error_title = "GPFS initialization failed." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.SETUP, ErrorDetail.NOT_READY) @@ -25,6 +28,7 @@ class GPFSAPIError(GPFSError): error_type = "https://api.backend.ai/probs/gpfs-api-error" error_title = "GPFS API error." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.REQUEST, ErrorDetail.INTERNAL_ERROR) @@ -33,6 +37,7 @@ class GPFSInvalidBodyError(GPFSAPIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/gpfs-bad-request" error_title = "GPFS API bad request." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.REQUEST, ErrorDetail.BAD_REQUEST) @@ -41,6 +46,7 @@ class GPFSUnauthorizedError(GPFSAPIError, web.HTTPUnauthorized): error_type = "https://api.backend.ai/probs/gpfs-authentication-failure" error_title = "GPFS API authentication failed." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.AUTH, ErrorDetail.UNAUTHORIZED) @@ -49,6 +55,7 @@ class GPFSForbiddenError(GPFSAPIError, web.HTTPForbidden): error_type = "https://api.backend.ai/probs/gpfs-forbidden" error_title = "GPFS API access forbidden." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.AUTH, ErrorDetail.FORBIDDEN) @@ -57,6 +64,7 @@ class GPFSNotFoundError(GPFSAPIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/gpfs-not-found" error_title = "GPFS resource not found." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.READ, ErrorDetail.NOT_FOUND) @@ -65,6 +73,7 @@ class GPFSConflictError(GPFSAPIError, web.HTTPConflict): error_type = "https://api.backend.ai/probs/gpfs-conflict" error_title = "GPFS resource conflict." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.CREATE, ErrorDetail.CONFLICT) @@ -73,6 +82,7 @@ class GPFSInternalError(GPFSAPIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/gpfs-internal-error" error_title = "GPFS internal error." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.REQUEST, ErrorDetail.INTERNAL_ERROR) @@ -81,6 +91,7 @@ class GPFSNoMetricError(GPFSError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/gpfs-no-metric" error_title = "GPFS metric not available." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.READ, ErrorDetail.NOT_FOUND) @@ -89,6 +100,7 @@ class GPFSJobFailedError(GPFSError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/gpfs-job-failed" error_title = "GPFS job failed." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.EXECUTE, ErrorDetail.INTERNAL_ERROR) @@ -97,5 +109,6 @@ class GPFSJobCancelledError(GPFSError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/gpfs-job-cancelled" error_title = "GPFS job was cancelled." + @override def error_code(self) -> ErrorCode: return ErrorCode(ErrorDomain.STORAGE, ErrorOperation.EXECUTE, ErrorDetail.CANCELED) diff --git a/src/ai/backend/storage/volumes/hammerspace/exception.py b/src/ai/backend/storage/volumes/hammerspace/exception.py index 1b7c390181a..3d4fc722260 100644 --- a/src/ai/backend/storage/volumes/hammerspace/exception.py +++ b/src/ai/backend/storage/volumes/hammerspace/exception.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -18,6 +20,7 @@ class ConfigurationError(HammerspaceError): error_title = "Hammerspace configuration is invalid." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -31,6 +34,7 @@ class AuthenticationError(HammerspaceError, web.HTTPUnauthorized): error_title = "Hammerspace API authentication failed." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -44,6 +48,7 @@ class ObjectiveNotFound(HammerspaceError, web.HTTPNotFound): error_title = "Hammerspace has no such objective." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -57,6 +62,7 @@ class ShareNotFound(HammerspaceError, web.HTTPNotFound): error_title = "Hammerspace has no such share." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -70,6 +76,7 @@ class VolumeNotFound(HammerspaceError, web.HTTPNotFound): error_title = "Hammerspace has no such volume." @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, diff --git a/src/ai/backend/storage/volumes/netapp/__init__.py b/src/ai/backend/storage/volumes/netapp/__init__.py index 11417df2c2c..f7efe1727eb 100644 --- a/src/ai/backend/storage/volumes/netapp/__init__.py +++ b/src/ai/backend/storage/volumes/netapp/__init__.py @@ -10,7 +10,7 @@ from collections.abc import AsyncIterator from contextlib import aclosing from pathlib import Path -from typing import Any +from typing import Any, override import aiofiles import aiofiles.os @@ -83,6 +83,7 @@ def __init__( self.volume_id = volume_id self.quota_enabled: bool = False + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -127,6 +128,7 @@ async def create_quota_scope( ) # pass if "already on" self.quota_enabled = True + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -136,6 +138,7 @@ async def describe_quota_scope( return None return await self.netapp_client.get_quota_report(self.svm_id, self.volume_id, qspath.name) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -151,11 +154,13 @@ async def update_quota_scope( self.netapp_client.check_job_result(result, []) # FIXME: How do we implement unset_quota() for NetApp? + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: raise InvalidQuotaScopeError( "Unsetting folder limit without removing quota scope is not possible for this backend" ) + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -212,6 +217,7 @@ async def check_license(self) -> bool: xcp_lic_check_path.write_bytes(b"0") return result + @override async def copy_tree( self, src_path: Path, @@ -240,6 +246,7 @@ async def copy_tree( # TODO: line for bgtask pass + @override async def delete_tree(self, path: Path) -> None: relpath = path.relative_to(self.mount_path) nfspath = f"{self.netapp_nfs_host}:{self.nas_path}/{relpath}" @@ -255,6 +262,7 @@ async def delete_tree(self, path: Path) -> None: # TODO: line for bgtask pass + @override def scan_tree( self, path: Path, @@ -369,6 +377,7 @@ async def read_stderr() -> None: return aiter() + @override async def scan_tree_usage( self, path: Path, @@ -417,6 +426,7 @@ async def scan_tree_usage( await proc.wait() return TreeUsage(file_count=total_count, used_bytes=total_size) + @override async def scan_tree_size( self, path: Path, @@ -436,6 +446,7 @@ class NetAppVolume(BaseVolume): volume_id: VolumeID nas_path: Path + @override async def create_quota_model(self) -> AbstractQuotaModel: return QTreeQuotaModel( self.mount_path, @@ -444,6 +455,7 @@ async def create_quota_model(self) -> AbstractQuotaModel: self.volume_id, ) + @override async def create_fsop_model(self) -> AbstractFSOpModel: xcp_fsop_model = XCPFSOpModel( self.mount_path, @@ -464,6 +476,7 @@ async def create_fsop_model(self) -> AbstractFSOpModel: self.local_config["storage-proxy"]["scandir-limit"], ) + @override async def init(self) -> None: self.ontap_endpoint = self.config["netapp_ontap_endpoint"] self.netapp_client = NetAppClient( @@ -505,12 +518,15 @@ async def init(self) -> None: # NOTE: The default qtree in the volume has no explicit path. await super().init() + @override async def shutdown(self) -> None: await self.netapp_client.aclose() + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_FAST_FS_SIZE, CAP_FAST_SIZE, CAP_QUOTA, CAP_METRIC]) + @override async def get_hwinfo(self) -> HardwareMetadata: volume_info = await self.netapp_client.get_volume_by_id(self.volume_id, ["files", "quota"]) metadata = { @@ -519,6 +535,7 @@ async def get_hwinfo(self) -> HardwareMetadata: } return {"status": "healthy", "status_info": None, "metadata": metadata} + @override async def get_fs_usage(self) -> CapacityUsage: volume_info = await self.netapp_client.get_volume_by_id( self.volume_id, ["space.size,space.used"] @@ -530,6 +547,7 @@ async def get_fs_usage(self) -> CapacityUsage: used_bytes=BinarySize(volume_info["space"]["used"]), ) + @override async def get_performance_metric(self) -> FSPerfMetric: volume_info = await self.netapp_client.get_volume_by_id(self.volume_id, ["statistics"]) if "statistics" not in volume_info: diff --git a/src/ai/backend/storage/volumes/noop/__init__.py b/src/ai/backend/storage/volumes/noop/__init__.py index 0fba9e48645..c482d6c6c48 100644 --- a/src/ai/backend/storage/volumes/noop/__init__.py +++ b/src/ai/backend/storage/volumes/noop/__init__.py @@ -37,9 +37,11 @@ class NoopQuotaModel(AbstractQuotaModel): def __init__(self) -> None: pass + @override def mangle_qspath(self, ref: VFolderID | QuotaScopeID | str | None) -> Path: return Path() + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -48,12 +50,14 @@ async def create_quota_scope( ) -> None: pass + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, ) -> QuotaUsage | None: pass + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -61,12 +65,14 @@ async def update_quota_scope( ) -> None: pass + @override async def unset_quota( self, quota_scope_id: QuotaScopeID, ) -> None: pass + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -78,6 +84,7 @@ class NoopFSOpModel(AbstractFSOpModel): def __init__(self) -> None: return + @override async def copy_tree( self, src_path: Path, @@ -85,6 +92,7 @@ async def copy_tree( ) -> None: pass + @override async def move_tree( self, src_path: Path, @@ -92,12 +100,14 @@ async def move_tree( ) -> None: pass + @override async def delete_tree( self, path: Path, ) -> None: pass + @override def scan_tree( self, path: Path, @@ -106,12 +116,14 @@ def scan_tree( ) -> AsyncIterator[DirEntry]: return _return_empty_dir_entry() + @override async def scan_tree_usage( self, path: Path, ) -> TreeUsage: return TreeUsage(0, 0) + @override async def scan_tree_size( self, path: Path, @@ -131,17 +143,21 @@ def info(self) -> VolumeInfo: options=None, ) + @override async def create_quota_model(self) -> AbstractQuotaModel: return NoopQuotaModel() + @override async def create_fsop_model(self) -> AbstractFSOpModel: return NoopFSOpModel() # ------ volume operations ------- + @override async def get_capabilities(self) -> frozenset[str]: return frozenset() + @override async def get_hwinfo(self) -> HardwareMetadata: return { "status": "healthy", @@ -149,6 +165,7 @@ async def get_hwinfo(self) -> HardwareMetadata: "metadata": {}, } + @override async def create_vfolder( self, vfid: VFolderID, @@ -157,9 +174,11 @@ async def create_vfolder( ) -> None: return None + @override async def delete_vfolder(self, vfid: VFolderID) -> None: return None + @override async def clone_vfolder( self, src_vfid: VFolderID, @@ -167,21 +186,27 @@ async def clone_vfolder( ) -> None: return None + @override async def get_vfolder_mount(self, vfid: VFolderID, subpath: str) -> Path: return Path() + @override async def put_metadata(self, vfid: VFolderID, payload: bytes) -> None: pass + @override async def get_metadata(self, vfid: VFolderID) -> bytes: return b"" + @override async def get_performance_metric(self) -> FSPerfMetric: return FSPerfMetric(0, 0, 0, 0, 0.0, 0.0) + @override async def get_fs_usage(self) -> CapacityUsage: return CapacityUsage(0, 0) + @override async def get_usage( self, vfid: VFolderID, @@ -189,11 +214,13 @@ async def get_usage( ) -> TreeUsage: return TreeUsage(0, 0) + @override async def get_used_bytes(self, vfid: VFolderID) -> BinarySize: return BinarySize(0) # ------ vfolder operations ------- + @override def scandir( self, vfid: VFolderID, @@ -203,6 +230,7 @@ def scandir( ) -> AsyncIterator[DirEntry]: return _return_empty_dir_entry() + @override async def mkdir( self, vfid: VFolderID, @@ -213,6 +241,7 @@ async def mkdir( ) -> None: pass + @override async def rmdir( self, vfid: VFolderID, @@ -222,6 +251,7 @@ async def rmdir( ) -> None: pass + @override async def move_file( self, vfid: VFolderID, @@ -230,6 +260,7 @@ async def move_file( ) -> None: pass + @override async def move_tree( self, vfid: VFolderID, @@ -238,6 +269,7 @@ async def move_tree( ) -> None: pass + @override async def copy_file( self, vfid: VFolderID, @@ -246,9 +278,11 @@ async def copy_file( ) -> None: pass + @override async def prepare_upload(self, vfid: VFolderID) -> str: return "" + @override async def add_file( self, vfid: VFolderID, @@ -257,6 +291,7 @@ async def add_file( ) -> None: pass + @override def read_file( self, vfid: VFolderID, @@ -269,6 +304,7 @@ async def _noop() -> AsyncIterator[bytes]: return _noop() + @override async def delete_files( self, vfid: VFolderID, diff --git a/src/ai/backend/storage/volumes/purestorage/__init__.py b/src/ai/backend/storage/volumes/purestorage/__init__.py index ff016920077..faeb55f6eb5 100644 --- a/src/ai/backend/storage/volumes/purestorage/__init__.py +++ b/src/ai/backend/storage/volumes/purestorage/__init__.py @@ -4,7 +4,7 @@ import contextlib import logging import re -from typing import Any +from typing import Any, override from ai.backend.common.types import HardwareMetadata from ai.backend.logging import BraceStyleAdapter @@ -42,6 +42,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._toolkit_version = None + @override async def create_fsop_model(self) -> AbstractFSOpModel: if (await self.get_toolkit_version()) == 2: return RapidFileToolsv2FSOpModel( @@ -87,6 +88,7 @@ async def get_toolkit_version(self) -> int: raise VolumeNotInitializedError("Failed to detect FlashBlade Toolkit version") return self._toolkit_version + @override async def init(self) -> None: toolkit_version = await self.get_toolkit_version() if toolkit_version == -1: @@ -101,9 +103,11 @@ async def init(self) -> None: ) await super().init() + @override async def shutdown(self) -> None: await self.purity_client.aclose() + @override async def get_capabilities(self) -> frozenset[str]: return frozenset( [ @@ -114,6 +118,7 @@ async def get_capabilities(self) -> frozenset[str]: ], ) + @override async def get_hwinfo(self) -> HardwareMetadata: async with self.purity_client as client: metadata = await client.get_metadata() @@ -125,6 +130,7 @@ async def get_hwinfo(self) -> HardwareMetadata: }, } + @override async def get_fs_usage(self) -> CapacityUsage: async with self.purity_client as client: usage = await client.get_usage(self.config["purity_fs_name"]) @@ -133,6 +139,7 @@ async def get_fs_usage(self) -> CapacityUsage: used_bytes=usage["used_bytes"], ) + @override async def get_performance_metric(self) -> FSPerfMetric: async with ( self.purity_client as client, diff --git a/src/ai/backend/storage/volumes/purestorage/exceptions.py b/src/ai/backend/storage/volumes/purestorage/exceptions.py index 30e0429b2d9..515014dd2ff 100644 --- a/src/ai/backend/storage/volumes/purestorage/exceptions.py +++ b/src/ai/backend/storage/volumes/purestorage/exceptions.py @@ -1,3 +1,5 @@ +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -13,6 +15,7 @@ class UnauthorizedPurityClient(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/purity-unauthorized-client" error_title = "Unauthorized Purity Client" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, diff --git a/src/ai/backend/storage/volumes/purestorage/rapidfiles.py b/src/ai/backend/storage/volumes/purestorage/rapidfiles.py index 87124f1a7f7..81ab2fd2335 100644 --- a/src/ai/backend/storage/volumes/purestorage/rapidfiles.py +++ b/src/ai/backend/storage/volumes/purestorage/rapidfiles.py @@ -7,6 +7,7 @@ from pathlib import Path from stat import S_IFDIR, S_IFLNK from subprocess import CalledProcessError +from typing import override from ai.backend.common.json import load_json from ai.backend.common.types import BinarySize @@ -24,6 +25,7 @@ class RapidFileToolsFSOpModel(BaseFSOpModel): + @override async def copy_tree( self, src_path: Path, @@ -44,6 +46,7 @@ async def copy_tree( except CalledProcessError as e: raise PureStorageCommandFailedError(f'"pcp" command failed: {e.stderr}') from e + @override async def delete_tree( self, path: Path, @@ -57,6 +60,7 @@ async def delete_tree( except CalledProcessError as e: raise PureStorageCommandFailedError(f"'prm' command failed: {e.stderr}") from e + @override def scan_tree( self, path: Path, @@ -109,6 +113,7 @@ async def _aiter() -> AsyncIterator[DirEntry]: return _aiter() + @override async def scan_tree_usage( self, path: Path, @@ -146,6 +151,7 @@ async def scan_tree_usage( await proc.wait() return TreeUsage(file_count=total_count, used_bytes=total_size) + @override async def scan_tree_size( self, path: Path, diff --git a/src/ai/backend/storage/volumes/purestorage/rapidfiles_v2.py b/src/ai/backend/storage/volumes/purestorage/rapidfiles_v2.py index 7f02a44bda3..09c7d10e19d 100644 --- a/src/ai/backend/storage/volumes/purestorage/rapidfiles_v2.py +++ b/src/ai/backend/storage/volumes/purestorage/rapidfiles_v2.py @@ -7,6 +7,7 @@ from pathlib import Path from stat import S_IFDIR, S_IFLNK from subprocess import CalledProcessError +from typing import override from ai.backend.common.json import load_json from ai.backend.logging.utils import BraceStyleAdapter @@ -24,6 +25,7 @@ class RapidFileToolsv2FSOpModel(RapidFileToolsFSOpModel): + @override async def copy_tree( self, src_path: Path, @@ -46,6 +48,7 @@ async def copy_tree( except CalledProcessError as e: raise PureStorageCommandFailedError(f'"pcopy" command failed: {e.stderr}') from e + @override def scan_tree( self, path: Path, @@ -96,6 +99,7 @@ async def _aiter() -> AsyncIterator[DirEntry]: return _aiter() + @override async def scan_tree_usage( self, path: Path, diff --git a/src/ai/backend/storage/volumes/types.py b/src/ai/backend/storage/volumes/types.py index 3e3a098b090..781c95cad9b 100644 --- a/src/ai/backend/storage/volumes/types.py +++ b/src/ai/backend/storage/volumes/types.py @@ -1,7 +1,7 @@ from abc import ABCMeta, abstractmethod from dataclasses import dataclass from pathlib import Path, PurePath -from typing import Self +from typing import Self, override from ai.backend.common.dto.storage.field import VFolderMetaField, VolumeMetaField from ai.backend.common.dto.storage.path import QuotaScopeKeyPath, VFolderKeyPath @@ -27,6 +27,7 @@ def from_quota_scope_path(cls, path: QuotaScopeKeyPath) -> Self: quota_scope_id=QuotaScopeID(path.quota_scope_type, path.quota_scope_uuid), ) + @override def to_logging_str(self) -> str: return f"QuotaScopeKey(volume_id={self.volume_id}, quota_scope_id={self.quota_scope_id})" @@ -44,6 +45,7 @@ def from_vfolder_path(cls, path: VFolderKeyPath) -> Self: vfolder_id=VFolderID(quota_scope_id, path.folder_uuid), ) + @override def to_logging_str(self) -> str: return f"VFolderKey(volume_id={self.volume_id}, vfolder_id={self.vfolder_id})" @@ -65,6 +67,7 @@ def to_field(self) -> VolumeMetaField: capabilities=self.capabilities, ) + @override def to_logging_str(self) -> str: return ( f"VolumeMeta(volume_id={self.volume_id}, backend={self.backend}, path={self.path}, " @@ -89,6 +92,7 @@ def to_field(self) -> VFolderMetaField: fs_used_bytes=self.fs_used_bytes, ) + @override def to_logging_str(self) -> str: return ( f"VFolderMeta(mount_path={self.mount_path}, file_count={self.file_count}, " @@ -108,5 +112,6 @@ def to_response(self) -> QuotaScopeResponse: limit_bytes=self.limit_bytes, ) + @override def to_logging_str(self) -> str: return f"QuotaScopeMeta(used_bytes={self.used_bytes}, limit_bytes={self.limit_bytes})" diff --git a/src/ai/backend/storage/volumes/vast/__init__.py b/src/ai/backend/storage/volumes/vast/__init__.py index 15381b1c4f7..c6887974b6a 100644 --- a/src/ai/backend/storage/volumes/vast/__init__.py +++ b/src/ai/backend/storage/volumes/vast/__init__.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from dataclasses import asdict from pathlib import Path -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, cast, override import aiofiles import aiofiles.os @@ -99,6 +99,7 @@ async def _modify_quota_scope( f"VAST API unknown error during quota modification: {e}" ) from e + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -166,6 +167,7 @@ async def _set_quota(_options: QuotaConfig) -> VASTQuota: quota = await _set_quota(options) await self._set_vast_quota_id(quota_scope_id, quota.id) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -176,6 +178,7 @@ async def update_quota_scope( raise QuotaScopeNotFoundError await self._modify_quota_scope(vast_quota_id, config) + @override async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage | None: if (vast_quota_id := await self._get_vast_quota_id(quota_scope_id)) is None: return None @@ -194,6 +197,7 @@ async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage limit_bytes=quota.hard_limit, ) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: vast_quota_id = await self._get_vast_quota_id(quota_scope_id) if vast_quota_id is None: @@ -204,6 +208,7 @@ async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: raise QuotaScopeNotFoundError from e await self._rm_vast_quota_id(quota_scope_id) + @override async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None: await self.unset_quota(quota_scope_id) qspath = self.mangle_qspath(quota_scope_id) @@ -248,15 +253,19 @@ def __init__( cluster_info_cache_ttl=self.config["vast_cluster_info_cache_ttl"], ) + @override async def shutdown(self) -> None: pass + @override async def create_quota_model(self) -> VASTQuotaModel: return VASTQuotaModel(self.mount_path, self.api_client) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_METRIC, CAP_QUOTA, CAP_FAST_FS_SIZE, CAP_FAST_SIZE]) + @override async def get_hwinfo(self) -> HardwareMetadata: cluster_id: int = self.config["vast_cluster_id"] try: @@ -287,6 +296,7 @@ async def get_hwinfo(self) -> HardwareMetadata: }, } + @override async def get_performance_metric(self) -> FSPerfMetric: cluster_id: int = self.config["vast_cluster_id"] try: @@ -311,6 +321,7 @@ async def get_performance_metric(self) -> FSPerfMetric: io_bytes_write=clsuter_info.wr_bw, ) + @override async def get_fs_usage(self) -> CapacityUsage: cluster_id: int = self.config["vast_cluster_id"] try: diff --git a/src/ai/backend/storage/volumes/vast/exceptions.py b/src/ai/backend/storage/volumes/vast/exceptions.py index d022187a810..6e7b0a75ce8 100644 --- a/src/ai/backend/storage/volumes/vast/exceptions.py +++ b/src/ai/backend/storage/volumes/vast/exceptions.py @@ -1,3 +1,6 @@ +from typing import override + + class VASTError(Exception): pass @@ -10,6 +13,7 @@ def __init__(self, message: str | None = None, *args: object) -> None: self.message = message or "" + @override def __str__(self) -> str: return type(self).__name__ diff --git a/src/ai/backend/storage/volumes/vfs/__init__.py b/src/ai/backend/storage/volumes/vfs/__init__.py index 9113d817698..966b9973af9 100644 --- a/src/ai/backend/storage/volumes/vfs/__init__.py +++ b/src/ai/backend/storage/volumes/vfs/__init__.py @@ -68,6 +68,7 @@ class BaseQuotaModel(AbstractQuotaModel): def __init__(self, mount_path: Path) -> None: self.mount_path = mount_path + @override def mangle_qspath(self, ref: VFolderID | QuotaScopeID | str | None) -> Path: try: match ref: @@ -89,6 +90,7 @@ def mangle_qspath(self, ref: VFolderID | QuotaScopeID | str | None) -> Path: except t.DataError as e: raise InvalidQuotaScopeError(f"Invalid value format for quota scope ID: {ref!r}") from e + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -105,6 +107,7 @@ async def create_quota_scope( except FileExistsError as e: raise QuotaScopeAlreadyExists(f"Quota scope '{quota_scope_id}' already exists") from e + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -114,6 +117,7 @@ async def describe_quota_scope( return QuotaUsage(-1, -1) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -122,6 +126,7 @@ async def update_quota_scope( # This is a no-op. pass + @override async def unset_quota( self, quota_scope_id: QuotaScopeID, @@ -129,6 +134,7 @@ async def unset_quota( # This is a no-op. pass + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -152,6 +158,7 @@ class SetGIDQuotaModel(BaseQuotaModel): or each project. """ + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -169,6 +176,7 @@ async def create_quota_scope( raise QuotaScopeAlreadyExists(f"Quota scope '{quota_scope_id}' already exists") from e # TODO: setgid impl. + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -178,6 +186,7 @@ async def describe_quota_scope( # TODO: setgid impl. return QuotaUsage(-1, -1) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -186,6 +195,7 @@ async def update_quota_scope( # TODO: setgid impl. raise NotImplementedError + @override async def unset_quota( self, quota_scope_id: QuotaScopeID, @@ -193,6 +203,7 @@ async def unset_quota( # TODO: setgid impl. raise NotImplementedError + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -218,6 +229,7 @@ def __init__( self.scandir_limit = scandir_limit self.watcher = watcher + @override async def copy_tree( self, src_path: Path, @@ -234,6 +246,7 @@ async def copy_tree( ), ) + @override async def move_tree( self, src_path: Path, @@ -245,6 +258,7 @@ async def move_tree( lambda: shutil.move(str(src_path), str(dst_path)), ) + @override async def delete_tree( self, path: Path, @@ -255,6 +269,7 @@ async def delete_tree( except FileNotFoundError: pass + @override def scan_tree( self, path: Path, @@ -337,6 +352,7 @@ async def _aiter() -> AsyncIterator[DirEntry]: return _aiter() + @override async def scan_tree_usage( self, path: Path, @@ -377,6 +393,7 @@ def _calc_usage(path: Path) -> None: total_count = -1 return TreeUsage(file_count=total_count, used_bytes=total_size) + @override async def scan_tree_size( self, path: Path, @@ -398,9 +415,11 @@ def info(self) -> VolumeInfo: options=self.config, ) + @override async def create_quota_model(self) -> AbstractQuotaModel: return BaseQuotaModel(self.mount_path) + @override async def create_fsop_model(self) -> AbstractFSOpModel: return BaseFSOpModel( self.mount_path, @@ -408,9 +427,11 @@ async def create_fsop_model(self) -> AbstractFSOpModel: self.watcher, ) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER]) + @override async def get_hwinfo(self) -> HardwareMetadata: return { "status": "healthy", @@ -419,6 +440,7 @@ async def get_hwinfo(self) -> HardwareMetadata: } @final + @override async def create_vfolder( self, vfid: VFolderID, @@ -436,6 +458,7 @@ async def create_vfolder( vfpath.chmod(mode) @final + @override async def delete_vfolder(self, vfid: VFolderID) -> None: vfpath = self.mangle_vfpath(vfid) await self.fsop_model.delete_tree(vfpath) @@ -452,6 +475,7 @@ async def delete_vfolder(self, vfid: VFolderID) -> None: raise @final + @override async def clone_vfolder( self, src_vfid: VFolderID, @@ -489,16 +513,19 @@ async def clone_vfolder( raise ProcessExecutionError("Copying files from source directories failed.") from e @final + @override async def get_vfolder_mount(self, vfid: VFolderID, subpath: str) -> Path: self.sanitize_vfpath(vfid, PurePosixPath(subpath)) return self.mangle_vfpath(vfid).resolve() + @override async def put_metadata(self, vfid: VFolderID, payload: bytes) -> None: vfpath = self.mangle_vfpath(vfid) metadata_path = vfpath / "metadata.json" loop = asyncio.get_running_loop() await loop.run_in_executor(None, metadata_path.write_bytes, payload) + @override async def get_metadata(self, vfid: VFolderID) -> bytes: vfpath = self.mangle_vfpath(vfid) metadata_path = vfpath / "metadata.json" @@ -512,9 +539,11 @@ async def get_metadata(self, vfid: VFolderID) -> bytes: return b"" # Other IO errors should be bubbled up. + @override async def get_performance_metric(self) -> FSPerfMetric: raise NotImplementedError + @override async def get_fs_usage(self) -> CapacityUsage: loop = asyncio.get_running_loop() stat = await loop.run_in_executor(None, os.statvfs, self.mount_path) @@ -524,6 +553,7 @@ async def get_fs_usage(self) -> CapacityUsage: ) @final + @override async def get_usage( self, vfid: VFolderID, @@ -533,6 +563,7 @@ async def get_usage( return await self.fsop_model.scan_tree_usage(target_path) @final + @override async def get_used_bytes(self, vfid: VFolderID) -> BinarySize: vfpath = self.mangle_vfpath(vfid) return await self.fsop_model.scan_tree_size(vfpath) @@ -540,6 +571,7 @@ async def get_used_bytes(self, vfid: VFolderID) -> BinarySize: # ------ vfolder internal operations ------- @final + @override def scandir( self, vfid: VFolderID, @@ -550,6 +582,7 @@ def scandir( target_path = self.sanitize_vfpath(vfid, relpath) return self.fsop_model.scan_tree(target_path, recursive=recursive) + @override async def mkdir( self, vfid: VFolderID, @@ -565,6 +598,7 @@ async def mkdir( lambda: target_path.mkdir(0o755, parents=parents, exist_ok=exist_ok), ) + @override async def rmdir( self, vfid: VFolderID, @@ -576,6 +610,7 @@ async def rmdir( loop = asyncio.get_running_loop() await loop.run_in_executor(None, target_path.rmdir) + @override async def move_file( self, vfid: VFolderID, @@ -586,6 +621,7 @@ async def move_file( dst_path = self.sanitize_vfpath(vfid, dst) await self.fsop_model.move_tree(src_path, dst_path) + @override async def move_tree( self, vfid: VFolderID, @@ -600,6 +636,7 @@ async def move_tree( dst_path = self.sanitize_vfpath(vfid, dst) await self.fsop_model.move_tree(src_path, dst_path) + @override async def copy_file( self, vfid: VFolderID, @@ -617,6 +654,7 @@ async def copy_file( ) await self.fsop_model.copy_tree(src_path, dst_path) + @override async def prepare_upload(self, vfid: VFolderID) -> str: vfpath = self.mangle_vfpath(vfid) session_id = secrets.token_hex(16) @@ -631,6 +669,7 @@ def _create_target() -> None: await loop.run_in_executor(None, _create_target) return session_id + @override async def add_file( self, vfid: VFolderID, @@ -661,6 +700,7 @@ def _write(q: janus._SyncQueueProxy[bytes]) -> None: finally: await write_fut + @override def read_file( self, vfid: VFolderID, @@ -714,6 +754,7 @@ async def _aiter() -> AsyncIterator[bytes]: return _aiter() + @override async def delete_files( self, vfid: VFolderID, diff --git a/src/ai/backend/storage/volumes/weka/__init__.py b/src/ai/backend/storage/volumes/weka/__init__.py index 16b318bb991..c1196e13bab 100644 --- a/src/ai/backend/storage/volumes/weka/__init__.py +++ b/src/ai/backend/storage/volumes/weka/__init__.py @@ -5,7 +5,7 @@ from collections.abc import Mapping from datetime import UTC, datetime, timedelta from pathlib import Path -from typing import Any +from typing import Any, override import aiofiles.os @@ -50,6 +50,7 @@ async def _get_inode_id(self, path: Path) -> int: lambda: path.stat().st_ino, ) + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -63,6 +64,7 @@ async def create_quota_scope( if options is not None: await self.update_quota_scope(quota_scope_id, options) + @override async def update_quota_scope(self, quota_scope_id: QuotaScopeID, config: QuotaConfig) -> None: qspath = self.mangle_qspath(quota_scope_id) inode_id = await self._get_inode_id(qspath) @@ -73,6 +75,7 @@ async def update_quota_scope(self, quota_scope_id: QuotaScopeID, config: QuotaCo qs_relpath, inode_id, soft_limit=config.limit_bytes, hard_limit=config.limit_bytes ) + @override async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage | None: qspath = self.mangle_qspath(quota_scope_id) if not qspath.exists(): @@ -95,6 +98,7 @@ async def describe_quota_scope(self, quota_scope_id: QuotaScopeID) -> QuotaUsage limit_bytes=limit_bytes, ) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: qspath = self.mangle_qspath(quota_scope_id) inode_id = await self._get_inode_id(qspath) @@ -103,6 +107,7 @@ async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: except WekaNotFoundError: pass + @override async def delete_quota_scope(self, quota_scope_id: QuotaScopeID) -> None: qspath = self.mangle_qspath(quota_scope_id) inode_id = await self._get_inode_id(qspath) @@ -149,6 +154,7 @@ def __init__( ssl=ssl_verify, ) + @override async def init(self) -> None: for fs in await self.api_client.list_fs(): if fs.name == self.config["weka_fs_name"]: @@ -158,12 +164,15 @@ async def init(self) -> None: raise WekaInitError(f"FileSystem {self.config['weka_fs_name']} not found") await super().init() + @override async def create_quota_model(self) -> AbstractQuotaModel: return WekaQuotaModel(self.mount_path, self._fs_uid, self.api_client) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_QUOTA, CAP_METRIC, CAP_FAST_FS_SIZE]) + @override async def get_hwinfo(self) -> HardwareMetadata: if self._fs_uid is None: raise VolumeNotInitializedError("Weka filesystem UID is not initialized") @@ -188,6 +197,7 @@ async def get_hwinfo(self) -> HardwareMetadata: "metadata": {}, } + @override async def get_fs_usage(self) -> CapacityUsage: if self._fs_uid is None: raise VolumeNotInitializedError("Weka filesystem UID is not initialized") @@ -197,6 +207,7 @@ async def get_fs_usage(self) -> CapacityUsage: used_bytes=fs.used_total, ) + @override async def get_performance_metric(self) -> FSPerfMetric: start_time = datetime.now(UTC).replace(second=0, microsecond=0) - timedelta( minutes=1, diff --git a/src/ai/backend/storage/volumes/weka/exceptions.py b/src/ai/backend/storage/volumes/weka/exceptions.py index bda229a55aa..16e41fd1eda 100644 --- a/src/ai/backend/storage/volumes/weka/exceptions.py +++ b/src/ai/backend/storage/volumes/weka/exceptions.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class WekaError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/storage/weka/generic" error_title = "Weka Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -35,6 +38,7 @@ class WekaInitError(WekaError): error_type = "https://api.backend.ai/probs/storage/weka/init-error" error_title = "Weka Initialization Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -49,6 +53,7 @@ class WekaAPIError(WekaError): error_type = "https://api.backend.ai/probs/storage/weka/api-error" error_title = "Weka API Error" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -63,6 +68,7 @@ class WekaInvalidBodyError(WekaAPIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/storage/weka/invalid-body" error_title = "Weka Invalid Body" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -77,6 +83,7 @@ class WekaUnauthorizedError(WekaAPIError): error_type = "https://api.backend.ai/probs/storage/weka/unauthorized" error_title = "Weka Unauthorized" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -91,6 +98,7 @@ class WekaNotFoundError(WekaAPIError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/weka/not-found" error_title = "Weka Not Found" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, @@ -112,6 +120,7 @@ class WekaNoMetricError(WekaError, web.HTTPNotFound): error_type = "https://api.backend.ai/probs/storage/weka/no-metric" error_title = "Weka No Metric" + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.STORAGE, diff --git a/src/ai/backend/storage/volumes/xfs/__init__.py b/src/ai/backend/storage/volumes/xfs/__init__.py index 7417d8221e6..b9fc7c773bb 100644 --- a/src/ai/backend/storage/volumes/xfs/__init__.py +++ b/src/ai/backend/storage/volumes/xfs/__init__.py @@ -4,7 +4,7 @@ from collections.abc import Mapping from pathlib import Path from tempfile import NamedTemporaryFile -from typing import Any +from typing import Any, override import aiofiles import aiofiles.os @@ -158,6 +158,7 @@ def __init__( self.block_size = stat_vfs.f_bsize self._lock_path = lock_path + @override async def create_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -189,6 +190,7 @@ async def create_quota_scope( if options is not None: await self.update_quota_scope(quota_scope_id, options) + @override async def describe_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -224,6 +226,7 @@ async def describe_quota_scope( ) return QuotaUsage(used_bytes, hard_limit_bytes) + @override async def update_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -257,11 +260,13 @@ async def update_quota_scope( ], ) + @override async def unset_quota(self, quota_scope_id: QuotaScopeID) -> None: raise InvalidQuotaScopeError( "Unsetting folder limit without removing quota scope is not possible for this backend" ) + @override async def delete_quota_scope( self, quota_scope_id: QuotaScopeID, @@ -325,11 +330,13 @@ def __init__( f"Cannot create XFS backend lock file at {self._lock_path}" ) from e + @override async def init(self) -> None: self.project_registry = XfsProjectRegistry() await self.project_registry.init(self) await super().init() + @override async def create_quota_model(self) -> AbstractQuotaModel: return XFSProjectQuotaModel( self.mount_path, @@ -337,5 +344,6 @@ async def create_quota_model(self) -> AbstractQuotaModel: self._lock_path, ) + @override async def get_capabilities(self) -> frozenset[str]: return frozenset([CAP_VFOLDER, CAP_QUOTA]) diff --git a/src/ai/backend/storage/watcher.py b/src/ai/backend/storage/watcher.py index c552819f6f0..792396b6a84 100644 --- a/src/ai/backend/storage/watcher.py +++ b/src/ai/backend/storage/watcher.py @@ -8,7 +8,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Sequence from pathlib import Path -from typing import Any, ClassVar, Self +from typing import Any, ClassVar, Self, override import aiofiles.os import attrs @@ -109,9 +109,11 @@ class ChownTask(AbstractTask): uid: int = attrs.field() gid: int = attrs.field() + @override async def run(self) -> Any: os.chown(self.directory, self.uid, self.gid) + @override def serialize(self) -> bytes: return msgpack.packb(( self.directory, @@ -120,6 +122,7 @@ def serialize(self) -> bytes: )) @classmethod + @override def deserialize(cls, values: tuple[Any, ...]) -> Self: return cls( values[0], @@ -146,6 +149,7 @@ class MountTask(AbstractTask): fstab_path: str = attrs.field(default="/etc/fstab") mount_prefix: str | None = attrs.field(default=None) + @override async def run(self) -> Any: return await _mount( self.mount_path, @@ -173,6 +177,7 @@ def from_event( mount_prefix, ) + @override def serialize(self) -> bytes: return msgpack.packb(( self.mount_path, @@ -187,6 +192,7 @@ def serialize(self) -> bytes: )) @classmethod + @override def deserialize(cls, values: tuple[Any, ...]) -> Self: return cls( values[0], @@ -216,6 +222,7 @@ class UmountTask(AbstractTask): mount_prefix: str | None = attrs.field(default=None) timeout: float | None = attrs.field(default=None) + @override async def run(self) -> Any: did_umount = await _umount( self.mount_path, @@ -246,6 +253,7 @@ def from_event( timeout, ) + @override def serialize(self) -> bytes: return msgpack.packb(( self.mount_path, @@ -258,6 +266,7 @@ def serialize(self) -> bytes: )) @classmethod + @override def deserialize(cls, values: tuple[Any, ...]) -> Self: return cls( values[0], @@ -275,6 +284,7 @@ class DeletePathTask(AbstractTask): name = "delete-path" path: Path + @override async def run(self) -> Any: if self.path.is_dir(): loop = asyncio.get_running_loop() @@ -285,10 +295,12 @@ async def run(self) -> Any: else: await aiofiles.os.remove(self.path) + @override def serialize(self) -> bytes: return msgpack.packb((self.path,)) @classmethod + @override def deserialize(cls, values: tuple[Any, ...]) -> Self: return cls( values[0], diff --git a/src/ai/backend/test/cli/utils.py b/src/ai/backend/test/cli/utils.py index 5f48d846b31..3a7df96c7d7 100644 --- a/src/ai/backend/test/cli/utils.py +++ b/src/ai/backend/test/cli/utils.py @@ -1,11 +1,12 @@ from __future__ import annotations -from typing import Any +from typing import Any, override import click class CommaSeparatedChoice(click.Choice): + @override def convert( self, value: str, @@ -21,6 +22,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._usage_args = kwargs.pop("usage_args") super().__init__(*args, **kwargs) + @override def format_usage(self, ctx: click.Context, formatter: click.HelpFormatter) -> None: if self._usage_args: formatter.write_usage(ctx.command_path, self._usage_args) diff --git a/src/ai/backend/test/contexts/auth.py b/src/ai/backend/test/contexts/auth.py index 075c131d97a..2ff972df67a 100644 --- a/src/ai/backend/test/contexts/auth.py +++ b/src/ai/backend/test/contexts/auth.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.tester.dependency import ( EndpointDep, @@ -8,17 +10,20 @@ class KeypairContext(BaseTestContext[KeyPairDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.KEYPAIR class LoginCredentialContext(BaseTestContext[LoginCredentialDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.LOGIN_CREDENTIAL class EndpointContext(BaseTestContext[EndpointDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.ENDPOINT diff --git a/src/ai/backend/test/contexts/client_session.py b/src/ai/backend/test/contexts/client_session.py index 0aab9885ce9..55ec03b62a4 100644 --- a/src/ai/backend/test/contexts/client_session.py +++ b/src/ai/backend/test/contexts/client_session.py @@ -1,14 +1,18 @@ +from typing import override + from ai.backend.client.session import AsyncSession from ai.backend.test.contexts.context import BaseTestContext, ContextName class ClientSessionContext(BaseTestContext[AsyncSession]): @classmethod + @override def name(cls) -> ContextName: return ContextName.CLIENT_SESSION class CreatedUserClientSessionContext(BaseTestContext[AsyncSession]): @classmethod + @override def name(cls) -> ContextName: return ContextName.CREATED_USER_CLIENT_SESSION diff --git a/src/ai/backend/test/contexts/context.py b/src/ai/backend/test/contexts/context.py index bdaed1113cf..c5b3fc93f1d 100644 --- a/src/ai/backend/test/contexts/context.py +++ b/src/ai/backend/test/contexts/context.py @@ -3,7 +3,7 @@ from collections.abc import Iterator, Mapping, MutableMapping from contextlib import contextmanager from contextvars import ContextVar -from typing import TypeVar, final +from typing import TypeVar, final, override T = TypeVar("T") @@ -64,6 +64,7 @@ class BaseTestContext[T]: _ctxvar: ContextVar[T | None] | None = None _used: MutableMapping[ContextName, "BaseTestContext[T]"] = {} + @override def __init_subclass__(cls) -> None: if cls._ctxvar is not None: raise RuntimeError(f"{cls.__name__} is already initialized") diff --git a/src/ai/backend/test/contexts/domain.py b/src/ai/backend/test/contexts/domain.py index 959c5ce4ec6..401f3f2a540 100644 --- a/src/ai/backend/test/contexts/domain.py +++ b/src/ai/backend/test/contexts/domain.py @@ -1,8 +1,11 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.tester.dependency import DomainDep class DomainContext(BaseTestContext[DomainDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.DOMAIN diff --git a/src/ai/backend/test/contexts/group.py b/src/ai/backend/test/contexts/group.py index 2519ec86a44..c2ec9f20653 100644 --- a/src/ai/backend/test/contexts/group.py +++ b/src/ai/backend/test/contexts/group.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.data.group import CreatedGroupMeta from ai.backend.test.tester.dependency import GroupDep @@ -5,11 +7,13 @@ class GroupContext(BaseTestContext[GroupDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.GROUP class CreatedGroupContext(BaseTestContext[CreatedGroupMeta]): @classmethod + @override def name(cls) -> ContextName: return ContextName.CREATED_GROUP diff --git a/src/ai/backend/test/contexts/model_service.py b/src/ai/backend/test/contexts/model_service.py index ce2cc085fd2..b3ad829a8ac 100644 --- a/src/ai/backend/test/contexts/model_service.py +++ b/src/ai/backend/test/contexts/model_service.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.data.model_service import ModelServiceEndpointMeta from ai.backend.test.tester.dependency import ModelServiceDep @@ -5,17 +7,20 @@ class ModelServiceContext(BaseTestContext[ModelServiceDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.MODEL_SERVICE class CreatedModelServiceEndpointMetaContext(BaseTestContext[ModelServiceEndpointMeta]): @classmethod + @override def name(cls) -> ContextName: return ContextName.CREATED_MODEL_SERVICE_ENDPOINT class CreatedModelServiceTokenContext(BaseTestContext[str]): @classmethod + @override def name(cls) -> ContextName: return ContextName.CREATED_MODEL_SERVICE_TOKEN diff --git a/src/ai/backend/test/contexts/resource_policy.py b/src/ai/backend/test/contexts/resource_policy.py index 146db4c4c52..d1af2a2e8fd 100644 --- a/src/ai/backend/test/contexts/resource_policy.py +++ b/src/ai/backend/test/contexts/resource_policy.py @@ -1,8 +1,11 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.tester.dependency import UserResourcePolicyDep class UserResourcePolicyContext(BaseTestContext[UserResourcePolicyDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.USER_RESOURCE_POLICY diff --git a/src/ai/backend/test/contexts/scaling_group.py b/src/ai/backend/test/contexts/scaling_group.py index 292ae2fbe01..679afc5c9db 100644 --- a/src/ai/backend/test/contexts/scaling_group.py +++ b/src/ai/backend/test/contexts/scaling_group.py @@ -1,8 +1,11 @@ +from typing import override + from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.tester.dependency import ScalingGroupDep class ScalingGroupContext(BaseTestContext[ScalingGroupDep]): @classmethod + @override def name(cls) -> ContextName: return ContextName.SCALING_GROUP diff --git a/src/ai/backend/test/contexts/tester.py b/src/ai/backend/test/contexts/tester.py index b1e82e35e7b..34e632ca586 100644 --- a/src/ai/backend/test/contexts/tester.py +++ b/src/ai/backend/test/contexts/tester.py @@ -1,4 +1,5 @@ from pathlib import Path +from typing import override from ai.backend.test.contexts.context import BaseTestContext, ContextName from ai.backend.test.data.tester import TestSpecMeta @@ -6,11 +7,13 @@ class TestSpecMetaContext(BaseTestContext[TestSpecMeta]): @classmethod + @override def name(cls) -> ContextName: return ContextName.TEST_SPEC_META class TestErrorOutputDirectoryContext(BaseTestContext[Path]): @classmethod + @override def name(cls) -> ContextName: return ContextName.TEST_ERROR_OUTPUT_DIRECTORY diff --git a/src/ai/backend/test/templates/auth/keypair.py b/src/ai/backend/test/templates/auth/keypair.py index ba1179f77a3..3a86b7cd4f8 100644 --- a/src/ai/backend/test/templates/auth/keypair.py +++ b/src/ai/backend/test/templates/auth/keypair.py @@ -17,6 +17,7 @@ class KeypairAuthTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "keypair_auth" @@ -39,6 +40,7 @@ async def _context(self) -> AsyncIterator[None]: class KeypairAuthAsCreatedUserTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "keypair_auth_as_created_user" diff --git a/src/ai/backend/test/templates/auth/login.py b/src/ai/backend/test/templates/auth/login.py index bd94c320267..785f9e8ac9b 100644 --- a/src/ai/backend/test/templates/auth/login.py +++ b/src/ai/backend/test/templates/auth/login.py @@ -17,6 +17,7 @@ class LoginTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "login" diff --git a/src/ai/backend/test/templates/group/add_group_to_user.py b/src/ai/backend/test/templates/group/add_group_to_user.py index ee9a3e05c48..b2673c212de 100644 --- a/src/ai/backend/test/templates/group/add_group_to_user.py +++ b/src/ai/backend/test/templates/group/add_group_to_user.py @@ -18,6 +18,7 @@ class _UserInfo: class AddGroupToKeypairTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "add_group" diff --git a/src/ai/backend/test/templates/group/group.py b/src/ai/backend/test/templates/group/group.py index a99806d8332..4107ba31070 100644 --- a/src/ai/backend/test/templates/group/group.py +++ b/src/ai/backend/test/templates/group/group.py @@ -14,6 +14,7 @@ class GroupTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "create_group" diff --git a/src/ai/backend/test/templates/model_service/endpoint.py b/src/ai/backend/test/templates/model_service/endpoint.py index e120cc347f9..f70ed01525b 100644 --- a/src/ai/backend/test/templates/model_service/endpoint.py +++ b/src/ai/backend/test/templates/model_service/endpoint.py @@ -131,6 +131,7 @@ async def _context(self) -> AsyncIterator[None]: class EndpointTemplate(_BaseEndpointTemplate): @property + @override def name(self) -> str: return "endpoint_template" @@ -141,6 +142,7 @@ def _extra_service_params(self) -> dict[str, Any]: class PublicEndpointTemplate(_BaseEndpointTemplate): @property + @override def name(self) -> str: return "public_endpoint_template" @@ -151,6 +153,7 @@ def _extra_service_params(self) -> dict[str, Any]: class CustomGroupEndpointTemplate(_BaseEndpointTemplate): @property + @override def name(self) -> str: return "custom_group_endpoint_template" diff --git a/src/ai/backend/test/templates/model_service/jwt_token.py b/src/ai/backend/test/templates/model_service/jwt_token.py index 63f84b2399b..f67fc908130 100644 --- a/src/ai/backend/test/templates/model_service/jwt_token.py +++ b/src/ai/backend/test/templates/model_service/jwt_token.py @@ -16,6 +16,7 @@ class ModelServiceTokenTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "generate_model_service_token" diff --git a/src/ai/backend/test/templates/session/batch_session.py b/src/ai/backend/test/templates/session/batch_session.py index fb11aaed048..1d2001b3641 100644 --- a/src/ai/backend/test/templates/session/batch_session.py +++ b/src/ai/backend/test/templates/session/batch_session.py @@ -26,6 +26,7 @@ class BatchSessionTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "batch_session" diff --git a/src/ai/backend/test/templates/session/dependent_session.py b/src/ai/backend/test/templates/session/dependent_session.py index 9b629669677..d8559717a51 100644 --- a/src/ai/backend/test/templates/session/dependent_session.py +++ b/src/ai/backend/test/templates/session/dependent_session.py @@ -22,6 +22,7 @@ class DependentSessionTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "dependent_session" diff --git a/src/ai/backend/test/templates/session/interactive_session.py b/src/ai/backend/test/templates/session/interactive_session.py index 31f02ad1190..3777a1d554a 100644 --- a/src/ai/backend/test/templates/session/interactive_session.py +++ b/src/ai/backend/test/templates/session/interactive_session.py @@ -126,6 +126,7 @@ async def _context(self) -> AsyncIterator[None]: class InteractiveSessionTemplate(_BaseInteractiveSessionTemplate): @property + @override def name(self) -> str: return "interactive_session" @@ -136,6 +137,7 @@ def _extra_session_params(self) -> dict[str, Any]: class InteractiveSessionWithBootstrapScriptTemplate(_BaseInteractiveSessionTemplate): @property + @override def name(self) -> str: return "interactive_session_with_bootstrap_script" @@ -147,6 +149,7 @@ def _extra_session_params(self) -> dict[str, Any]: class InteractiveSessionWithCustomGroupTemplate(_BaseInteractiveSessionTemplate): @property + @override def name(self) -> str: return "interactive_session_with_custom_group" diff --git a/src/ai/backend/test/templates/session/session_template.py b/src/ai/backend/test/templates/session/session_template.py index a8f489e6076..2ab17fc4a42 100644 --- a/src/ai/backend/test/templates/session/session_template.py +++ b/src/ai/backend/test/templates/session/session_template.py @@ -31,6 +31,7 @@ class SessionTemplateTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "session_template" @@ -99,6 +100,7 @@ async def _context(self) -> AsyncIterator[None]: class BatchSessionFromTemplateTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "batch_session_from_template" @@ -159,6 +161,7 @@ async def _context(self) -> AsyncIterator[None]: class InteractiveSessionFromTemplateTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "interactive_session_from_template" diff --git a/src/ai/backend/test/templates/session/vfolder_mounted_interactive_session.py b/src/ai/backend/test/templates/session/vfolder_mounted_interactive_session.py index dfafb65d1cd..11a93c445e3 100644 --- a/src/ai/backend/test/templates/session/vfolder_mounted_interactive_session.py +++ b/src/ai/backend/test/templates/session/vfolder_mounted_interactive_session.py @@ -22,6 +22,7 @@ class VFolderMountedInteractiveSessionTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "vfolder_mounted_interactive_session" diff --git a/src/ai/backend/test/templates/template.py b/src/ai/backend/test/templates/template.py index a4536c7d123..85865111b5f 100644 --- a/src/ai/backend/test/templates/template.py +++ b/src/ai/backend/test/templates/template.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterator from contextlib import asynccontextmanager as actxmgr -from typing import Protocol, final +from typing import Protocol, final, override from ai.backend.test.tester.exporter import TestExporter @@ -27,6 +27,7 @@ async def test(self) -> None: class NopTestCode(TestCode): + @override async def test(self) -> None: """ A no-operation test code that does nothing. @@ -79,9 +80,11 @@ def __init__(self, testcode: TestCode) -> None: self._testcode = testcode @property + @override def name(self) -> str: return "basic" + @override async def run_test(self, exporter: TestExporter) -> None: await self._testcode.test() @@ -118,6 +121,7 @@ async def _context(self) -> AsyncIterator[None]: yield # type: ignore[unreachable] # Required for async generator typing @final + @override async def run_test(self, exporter: TestExporter) -> None: try: async with self._context(): @@ -141,6 +145,7 @@ def __init__( self._templates = templates @final + @override async def run_test(self, exporter: TestExporter) -> None: """ Run the test case by executing each template in sequence. diff --git a/src/ai/backend/test/templates/user/user.py b/src/ai/backend/test/templates/user/user.py index 31d85109b4b..761298d8684 100644 --- a/src/ai/backend/test/templates/user/user.py +++ b/src/ai/backend/test/templates/user/user.py @@ -17,6 +17,7 @@ class UserTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "create_user" diff --git a/src/ai/backend/test/templates/vfolder/file_uploader.py b/src/ai/backend/test/templates/vfolder/file_uploader.py index bfe84816825..fbf5006c2b0 100644 --- a/src/ai/backend/test/templates/vfolder/file_uploader.py +++ b/src/ai/backend/test/templates/vfolder/file_uploader.py @@ -17,6 +17,7 @@ class PlainTextFilesUploader(WrapperTestTemplate): @property + @override def name(self) -> str: return "upload_plain_text_files" diff --git a/src/ai/backend/test/templates/vfolder/general_vfolder.py b/src/ai/backend/test/templates/vfolder/general_vfolder.py index 0d71eaa6436..95a2a294a6b 100644 --- a/src/ai/backend/test/templates/vfolder/general_vfolder.py +++ b/src/ai/backend/test/templates/vfolder/general_vfolder.py @@ -11,6 +11,7 @@ class UserVFolderTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "general_user_vfolder" @@ -44,6 +45,7 @@ async def _context(self) -> AsyncIterator[None]: class ProjectVFolderTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "general_project_vfolder" diff --git a/src/ai/backend/test/templates/vfolder/invite.py b/src/ai/backend/test/templates/vfolder/invite.py index 2ba514bc985..80e0390621d 100644 --- a/src/ai/backend/test/templates/vfolder/invite.py +++ b/src/ai/backend/test/templates/vfolder/invite.py @@ -18,6 +18,7 @@ class VFolderInviteTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "vfolder_invite" @@ -47,6 +48,7 @@ async def _context(self) -> AsyncIterator[None]: class AcceptInvitationTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "accept_invitation" @@ -76,6 +78,7 @@ async def _context(self) -> AsyncIterator[None]: class RejectInvitationTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "reject_invitation" diff --git a/src/ai/backend/test/templates/vfolder/model_vfolder.py b/src/ai/backend/test/templates/vfolder/model_vfolder.py index e1724498a37..6a7d02349cd 100644 --- a/src/ai/backend/test/templates/vfolder/model_vfolder.py +++ b/src/ai/backend/test/templates/vfolder/model_vfolder.py @@ -11,6 +11,7 @@ class ModelVFolderTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "model_vfolder" diff --git a/src/ai/backend/test/templates/vfolder/share.py b/src/ai/backend/test/templates/vfolder/share.py index 4badd3841d4..3db9a3bf063 100644 --- a/src/ai/backend/test/templates/vfolder/share.py +++ b/src/ai/backend/test/templates/vfolder/share.py @@ -16,6 +16,7 @@ class ShareVFolderTemplate(WrapperTestTemplate): @property + @override def name(self) -> str: return "share_vfolder" diff --git a/src/ai/backend/test/testcases/session/commit.py b/src/ai/backend/test/testcases/session/commit.py index 39832b5ad83..f6063060db1 100644 --- a/src/ai/backend/test/testcases/session/commit.py +++ b/src/ai/backend/test/testcases/session/commit.py @@ -1,4 +1,5 @@ import asyncio +from typing import override from ai.backend.common.bgtask.types import BgtaskStatus from ai.backend.test.contexts.client_session import ClientSessionContext @@ -9,6 +10,7 @@ class InteractiveSessionImagifySuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() @@ -36,6 +38,7 @@ async def test(self) -> None: class InteractiveSessionCommitSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() diff --git a/src/ai/backend/test/testcases/session/creation_failure_command_timeout.py b/src/ai/backend/test/testcases/session/creation_failure_command_timeout.py index 5f5cecae7a5..155afcd115c 100644 --- a/src/ai/backend/test/testcases/session/creation_failure_command_timeout.py +++ b/src/ai/backend/test/testcases/session/creation_failure_command_timeout.py @@ -1,4 +1,5 @@ import asyncio +from typing import override from ai.backend.common.events.event_types.kernel.types import KernelLifecycleEventReason from ai.backend.common.types import ClusterMode @@ -13,6 +14,7 @@ class BatchSessionCreationFailureTimeout(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() spec_meta = TestSpecMetaContext.current() diff --git a/src/ai/backend/test/testcases/session/creation_failure_low_resources.py b/src/ai/backend/test/testcases/session/creation_failure_low_resources.py index 0af621d89e1..790bd101823 100644 --- a/src/ai/backend/test/testcases/session/creation_failure_low_resources.py +++ b/src/ai/backend/test/testcases/session/creation_failure_low_resources.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.exceptions import BackendAPIError from ai.backend.client.output.fields import image_fields from ai.backend.common.types import ClusterMode @@ -9,6 +11,7 @@ class SessionCreationFailureLowResources(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() image_dep = ImageContext.current() diff --git a/src/ai/backend/test/testcases/session/creation_failure_schedule_timeout.py b/src/ai/backend/test/testcases/session/creation_failure_schedule_timeout.py index 96640c610c1..e0ce643add0 100644 --- a/src/ai/backend/test/testcases/session/creation_failure_schedule_timeout.py +++ b/src/ai/backend/test/testcases/session/creation_failure_schedule_timeout.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.func.session import ComputeSession from ai.backend.common.types import ClusterMode from ai.backend.test.contexts.client_session import ClientSessionContext @@ -7,6 +9,7 @@ class InteractiveSessionCreationFailureScheduleTimeout(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() image_dep = ImageContext.current() diff --git a/src/ai/backend/test/testcases/session/creation_failure_too_many_container.py b/src/ai/backend/test/testcases/session/creation_failure_too_many_container.py index efc099adf83..585889935bd 100644 --- a/src/ai/backend/test/testcases/session/creation_failure_too_many_container.py +++ b/src/ai/backend/test/testcases/session/creation_failure_too_many_container.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.exceptions import BackendAPIError from ai.backend.client.output.fields import keypair_fields, keypair_resource_policy_fields from ai.backend.common.types import ClusterMode @@ -11,6 +13,7 @@ class SessionCreationFailureTooManyContainer(TestCode): + @override async def test(self) -> None: keypair = KeypairContext.current() client_session = ClientSessionContext.current() diff --git a/src/ai/backend/test/testcases/session/creation_failure_wrong_command.py b/src/ai/backend/test/testcases/session/creation_failure_wrong_command.py index a308edad2c8..245aaa4035f 100644 --- a/src/ai/backend/test/testcases/session/creation_failure_wrong_command.py +++ b/src/ai/backend/test/testcases/session/creation_failure_wrong_command.py @@ -1,4 +1,5 @@ import asyncio +from typing import override from ai.backend.common.types import ClusterMode from ai.backend.test.contexts.client_session import ClientSessionContext @@ -13,6 +14,7 @@ class BatchSessionCreationFailureWrongCommand(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() image_dep = ImageContext.current() diff --git a/src/ai/backend/test/testcases/session/execution.py b/src/ai/backend/test/testcases/session/execution.py index 037a39028dc..fc6d859c15f 100644 --- a/src/ai/backend/test/testcases/session/execution.py +++ b/src/ai/backend/test/testcases/session/execution.py @@ -1,9 +1,12 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.session import CodeExecutionContext, CreatedSessionMetaContext from ai.backend.test.templates.template import TestCode class InteractiveSessionExecuteCodeSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() @@ -40,6 +43,7 @@ async def test(self) -> None: class InteractiveSessionExecuteCodeFailureWrongCommand(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() @@ -76,6 +80,7 @@ async def test(self) -> None: class InteractiveSessionExecuteCodeFailureWithCustomExitCode(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() diff --git a/src/ai/backend/test/testcases/session/filecheck.py b/src/ai/backend/test/testcases/session/filecheck.py index 14ea5f52ed3..9833ca0ede3 100644 --- a/src/ai/backend/test/testcases/session/filecheck.py +++ b/src/ai/backend/test/testcases/session/filecheck.py @@ -1,4 +1,5 @@ from pathlib import Path +from typing import override from ai.backend.common.json import load_json from ai.backend.test.contexts.client_session import ClientSessionContext @@ -15,6 +16,7 @@ def __init__(self, path: Path, checklist: list[str]) -> None: self._path = path self._checklist = checklist + @override async def test(self) -> None: client_session = ClientSessionContext.current() session_meta = CreatedSessionMetaContext.current() diff --git a/src/ai/backend/test/testcases/spec_manager.py b/src/ai/backend/test/testcases/spec_manager.py index c13c9d516e9..943e3c15ac3 100644 --- a/src/ai/backend/test/testcases/spec_manager.py +++ b/src/ai/backend/test/testcases/spec_manager.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from itertools import product -from typing import Any +from typing import Any, override from ai.backend.test.contexts.context import ContextName from ai.backend.test.templates.template import ( @@ -47,9 +47,11 @@ class TestSpec: template: TestTemplate parametrizes: Mapping[ContextName, Sequence[Any]] | None = None + @override def __hash__(self) -> int: return hash(self.name) + @override def __eq__(self, other: object) -> bool: if not isinstance(other, TestSpec): return NotImplemented diff --git a/src/ai/backend/test/testcases/vfolder/access.py b/src/ai/backend/test/testcases/vfolder/access.py index 089e15fd617..7f61a805b76 100644 --- a/src/ai/backend/test/testcases/vfolder/access.py +++ b/src/ai/backend/test/testcases/vfolder/access.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.exceptions import BackendAPIError from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import ( @@ -13,6 +15,7 @@ class VFolderAccessSuccess(TestCode): Test whether logged-in user can access the vfolder properly. """ + @override async def test(self) -> None: vfolder_meta = CreatedVFolderMetaContext.current() client_session = ClientSessionContext.current() @@ -35,6 +38,7 @@ async def test(self) -> None: class VFolderAccessFailure(TestCode): + @override async def test(self) -> None: vfolder_meta = CreatedVFolderMetaContext.current() client_session = ClientSessionContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/clone.py b/src/ai/backend/test/testcases/vfolder/clone.py index 451f675516e..4b7c4cfcc6a 100644 --- a/src/ai/backend/test/testcases/vfolder/clone.py +++ b/src/ai/backend/test/testcases/vfolder/clone.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.common.bgtask.types import BgtaskStatus from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext @@ -7,6 +9,7 @@ class VFolderCloneSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/delete_files.py b/src/ai/backend/test/testcases/vfolder/delete_files.py index be70419e4e5..f4725f10790 100644 --- a/src/ai/backend/test/testcases/vfolder/delete_files.py +++ b/src/ai/backend/test/testcases/vfolder/delete_files.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext, UploadedFilesContext from ai.backend.test.templates.template import TestCode @@ -5,6 +7,7 @@ class VFolderFilesDeletionSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() @@ -29,6 +32,7 @@ class VFolderFilesRecursiveDeletionSuccess(TestCode): def __init__(self, dirname_to_delete: str) -> None: self._dirname_to_delete = dirname_to_delete + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/download.py b/src/ai/backend/test/testcases/vfolder/download.py index c7ee0ebde13..c6a17933014 100644 --- a/src/ai/backend/test/testcases/vfolder/download.py +++ b/src/ai/backend/test/testcases/vfolder/download.py @@ -1,5 +1,6 @@ import tempfile from pathlib import Path +from typing import override from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext, UploadedFilesContext @@ -7,6 +8,7 @@ class VFolderDownloadSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/invite_failures.py b/src/ai/backend/test/testcases/vfolder/invite_failures.py index fcd573e297b..dd48250c1cd 100644 --- a/src/ai/backend/test/testcases/vfolder/invite_failures.py +++ b/src/ai/backend/test/testcases/vfolder/invite_failures.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.exceptions import BackendAPIError from ai.backend.test.contexts.client_session import ( ClientSessionContext, @@ -13,6 +15,7 @@ class VFolderInviteFailure(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() @@ -28,6 +31,7 @@ async def test(self) -> None: class VFolderAcceptDuplicatedInvitation(TestCode): + @override async def test(self) -> None: client_session = CreatedUserClientSessionContext.current() created_user_meta = CreatedUserContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/list_files.py b/src/ai/backend/test/testcases/vfolder/list_files.py index c638989ea2b..51440a5ec2b 100644 --- a/src/ai/backend/test/testcases/vfolder/list_files.py +++ b/src/ai/backend/test/testcases/vfolder/list_files.py @@ -1,9 +1,12 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext, UploadedFilesContext from ai.backend.test.templates.template import TestCode class VFolderListFilesSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/move.py b/src/ai/backend/test/testcases/vfolder/move.py index 6ec6e438c72..f6dad384363 100644 --- a/src/ai/backend/test/testcases/vfolder/move.py +++ b/src/ai/backend/test/testcases/vfolder/move.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext, UploadedFilesContext from ai.backend.test.templates.template import TestCode @@ -5,6 +7,7 @@ class VFolderFileMoveSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/purge.py b/src/ai/backend/test/testcases/vfolder/purge.py index 33dba5d7125..bd67746e236 100644 --- a/src/ai/backend/test/testcases/vfolder/purge.py +++ b/src/ai/backend/test/testcases/vfolder/purge.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.tester import TestSpecMetaContext from ai.backend.test.contexts.vfolder import VFolderContext @@ -6,6 +8,7 @@ class VFolderPurgeSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() spec_meta = TestSpecMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/rename_file.py b/src/ai/backend/test/testcases/vfolder/rename_file.py index 414a14a802d..887d0a0107a 100644 --- a/src/ai/backend/test/testcases/vfolder/rename_file.py +++ b/src/ai/backend/test/testcases/vfolder/rename_file.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext, UploadedFilesContext from ai.backend.test.templates.template import TestCode @@ -5,6 +7,7 @@ class VFolderFileRenameSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/restore.py b/src/ai/backend/test/testcases/vfolder/restore.py index acd19bea883..86291b5cb7d 100644 --- a/src/ai/backend/test/testcases/vfolder/restore.py +++ b/src/ai/backend/test/testcases/vfolder/restore.py @@ -1,9 +1,12 @@ +from typing import override + from ai.backend.test.contexts.client_session import ClientSessionContext from ai.backend.test.contexts.vfolder import CreatedVFolderMetaContext from ai.backend.test.templates.template import TestCode class VFolderDeleteAndRestoreSuccess(TestCode): + @override async def test(self) -> None: client_session = ClientSessionContext.current() vfolder_meta = CreatedVFolderMetaContext.current() diff --git a/src/ai/backend/test/testcases/vfolder/share.py b/src/ai/backend/test/testcases/vfolder/share.py index 860f3360320..b752e4ec178 100644 --- a/src/ai/backend/test/testcases/vfolder/share.py +++ b/src/ai/backend/test/testcases/vfolder/share.py @@ -1,3 +1,5 @@ +from typing import override + from ai.backend.client.session import set_api_context from ai.backend.test.contexts.client_session import ( ClientSessionContext, @@ -13,6 +15,7 @@ class VFolderSharePermissionOverrideSuccess(TestCode): + @override async def test(self) -> None: vfolder_meta = CreatedVFolderMetaContext.current() # The session of the user who shared the vfolder diff --git a/src/ai/backend/test/tester/exporter.py b/src/ai/backend/test/tester/exporter.py index 441c0750738..3b61dea3068 100644 --- a/src/ai/backend/test/tester/exporter.py +++ b/src/ai/backend/test/tester/exporter.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from pathlib import Path +from typing import override from ai.backend.common.json import pretty_json_str from ai.backend.test.contexts.tester import TestSpecMetaContext @@ -85,20 +86,24 @@ def __init__(self, sub_name: str | None, started: float, last_stage_done: float) self._output_directory.mkdir(parents=True, exist_ok=True) @classmethod + @override async def create(cls, sub_name: str | None = None) -> TestExporter: return cls(sub_name, started=0.0, last_stage_done=0.0) + @override async def export_start(self) -> None: current = time.perf_counter() self._started = current self._last_stage_done = current print(f"Starting '{self._spec_name()}'...") + @override async def export_done(self) -> None: current = time.perf_counter() elapsed = current - self._started print(f"\033[92m✓ '{self._spec_name()}' completed in {elapsed:.2f} seconds.\033[0m") + @override async def export_stage_done(self, stage: str) -> None: current = time.perf_counter() elapsed = current - self._last_stage_done @@ -106,6 +111,7 @@ async def export_stage_done(self, stage: str) -> None: self._last_stage = stage print(f"\033[92m✓ '{self._stage_name(stage)}' completed in {elapsed:.2f} seconds.\033[0m") + @override async def export_stage_exception(self, stage: str, exception: BaseException) -> None: current = time.perf_counter() elapsed = current - self._last_stage_done @@ -115,6 +121,7 @@ async def export_stage_exception(self, stage: str, exception: BaseException) -> f"{exception.__class__.__name__}\033[0m" ) + @override async def export_exception(self, exception: BaseException) -> None: current = time.perf_counter() elapsed = current - self._started diff --git a/src/ai/backend/testutils/db.py b/src/ai/backend/testutils/db.py index e78fa7c8a1d..3e1fd991381 100644 --- a/src/ai/backend/testutils/db.py +++ b/src/ai/backend/testutils/db.py @@ -2,7 +2,7 @@ from collections.abc import AsyncGenerator, Callable, Sequence from contextlib import asynccontextmanager -from typing import Any, Protocol, cast +from typing import Any, Protocol, cast, override from sqlalchemy import Table, text from sqlalchemy.ext.asyncio import AsyncEngine @@ -37,6 +37,7 @@ def _make_subset_schema_generator(table_names: frozenset[str]) -> type[SchemaGen """ class _SubsetSchemaGenerator(SchemaGenerator): + @override def visit_foreign_key_constraint(self, constraint: ForeignKeyConstraint) -> None: elements = constraint.elements if elements: diff --git a/src/ai/backend/web/cli/context.py b/src/ai/backend/web/cli/context.py index 8bfa409a561..343bf5a8e9b 100644 --- a/src/ai/backend/web/cli/context.py +++ b/src/ai/backend/web/cli/context.py @@ -2,7 +2,7 @@ from contextlib import AbstractContextManager from pathlib import Path -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, Self, override import click @@ -19,6 +19,7 @@ def __init__(self, log_level: LogLevel, config_path: Path | None = None) -> None self.config_path = config_path self.log_level = log_level + @override def __enter__(self) -> Self: from ai.backend.logging import LocalLogger @@ -31,6 +32,7 @@ def __enter__(self) -> Self: self._logger.__enter__() return self + @override def __exit__(self, *exc_info: Any) -> None: click_ctx = click.get_current_context() if click_ctx.invoked_subcommand != "start-server": diff --git a/src/ai/backend/web/clients/apollo_router_pool/health_checker.py b/src/ai/backend/web/clients/apollo_router_pool/health_checker.py index 6dc954ed971..ef32c93f1f6 100644 --- a/src/ai/backend/web/clients/apollo_router_pool/health_checker.py +++ b/src/ai/backend/web/clients/apollo_router_pool/health_checker.py @@ -17,6 +17,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.common.health_checker.abc import ServiceHealthChecker from ai.backend.common.health_checker.types import ( @@ -45,13 +46,16 @@ def __init__(self, pool: HealthyEndpointPool, *, timeout: float = _DEFAULT_TIMEO self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: return APOLLO_ROUTER @property + @override def timeout(self) -> float: return self._timeout + @override async def check_service(self) -> ServiceHealth: has_healthy = self._pool.has_any_healthy() return ServiceHealth( @@ -76,13 +80,16 @@ def __init__(self, pool: HealthyEndpointPool, *, timeout: float = _DEFAULT_TIMEO self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: return APOLLO_ROUTER_ENDPOINTS @property + @override def timeout(self) -> float: return self._timeout + @override async def check_service(self) -> ServiceHealth: now = datetime.now(UTC) results: dict[ComponentId, ComponentHealthStatus] = {} diff --git a/src/ai/backend/web/clients/manager_pool/health_checker.py b/src/ai/backend/web/clients/manager_pool/health_checker.py index 1eb76aacda0..c143c557888 100644 --- a/src/ai/backend/web/clients/manager_pool/health_checker.py +++ b/src/ai/backend/web/clients/manager_pool/health_checker.py @@ -18,6 +18,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import override from ai.backend.common.health_checker.abc import ServiceHealthChecker from ai.backend.common.health_checker.types import ( @@ -48,13 +49,16 @@ def __init__(self, pool: HealthyEndpointPool, *, timeout: float = _DEFAULT_TIMEO self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: return MANAGER @property + @override def timeout(self) -> float: return self._timeout + @override async def check_service(self) -> ServiceHealth: has_healthy = self._pool.has_any_healthy() return ServiceHealth( @@ -79,13 +83,16 @@ def __init__(self, pool: HealthyEndpointPool, *, timeout: float = _DEFAULT_TIMEO self._timeout = timeout @property + @override def target_service_group(self) -> ServiceGroup: return MANAGER_ENDPOINTS @property + @override def timeout(self) -> float: return self._timeout + @override async def check_service(self) -> ServiceHealth: now = datetime.now(UTC) results: dict[ComponentId, ComponentHealthStatus] = {} diff --git a/src/ai/backend/web/dependencies/components/composer.py b/src/ai/backend/web/dependencies/components/composer.py index a927bf5617f..839bff50718 100644 --- a/src/ai/backend/web/dependencies/components/composer.py +++ b/src/ai/backend/web/dependencies/components/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.dependencies import DependencyComposer, DependencyStack from ai.backend.web.config.unified import WebServerUnifiedConfig @@ -27,10 +28,12 @@ class ComponentComposer(DependencyComposer[WebServerUnifiedConfig, ComponentReso """ @property + @override def stage_name(self) -> str: return "components" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/web/dependencies/components/hive_router_client.py b/src/ai/backend/web/dependencies/components/hive_router_client.py index 4b0ab30864c..6204c56c0bf 100644 --- a/src/ai/backend/web/dependencies/components/hive_router_client.py +++ b/src/ai/backend/web/dependencies/components/hive_router_client.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override import aiohttp @@ -27,10 +28,12 @@ class HiveRouterClientProvider( """ @property + @override def stage_name(self) -> str: return "hive-router-client" @asynccontextmanager + @override async def provide( self, setup_input: WebServerUnifiedConfig ) -> AsyncIterator[HiveRouterClientInfo]: diff --git a/src/ai/backend/web/dependencies/components/manager_client.py b/src/ai/backend/web/dependencies/components/manager_client.py index f14f67dc4f2..4298a1b340e 100644 --- a/src/ai/backend/web/dependencies/components/manager_client.py +++ b/src/ai/backend/web/dependencies/components/manager_client.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override import aiohttp @@ -27,10 +28,12 @@ class ManagerClientProvider( """ @property + @override def stage_name(self) -> str: return "manager-client" @asynccontextmanager + @override async def provide( self, setup_input: WebServerUnifiedConfig ) -> AsyncIterator[ManagerClientInfo]: diff --git a/src/ai/backend/web/dependencies/composer.py b/src/ai/backend/web/dependencies/composer.py index da371b7a1f1..1d1a0f1da41 100644 --- a/src/ai/backend/web/dependencies/composer.py +++ b/src/ai/backend/web/dependencies/composer.py @@ -4,6 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path +from typing import override import tomli @@ -51,10 +52,12 @@ class WebDependencyComposer(DependencyComposer[DependencyInput, DependencyResour """ @property + @override def stage_name(self) -> str: return "web-dependencies" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/web/dependencies/infrastructure/composer.py b/src/ai/backend/web/dependencies/infrastructure/composer.py index 3e2be59c189..ecd73bdee10 100644 --- a/src/ai/backend/web/dependencies/infrastructure/composer.py +++ b/src/ai/backend/web/dependencies/infrastructure/composer.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass +from typing import override from ai.backend.common.clients.valkey_client.valkey_session.client import ValkeySessionClient from ai.backend.common.dependencies import DependencyComposer, DependencyStack @@ -26,10 +27,12 @@ class InfrastructureComposer(DependencyComposer[WebServerUnifiedConfig, Infrastr """ @property + @override def stage_name(self) -> str: return "infrastructure" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, diff --git a/src/ai/backend/web/dependencies/infrastructure/redis.py b/src/ai/backend/web/dependencies/infrastructure/redis.py index 15f05f5444d..ba1da56b1cc 100644 --- a/src/ai/backend/web/dependencies/infrastructure/redis.py +++ b/src/ai/backend/web/dependencies/infrastructure/redis.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from ai.backend.common.clients.valkey_client.valkey_session.client import ValkeySessionClient from ai.backend.common.defs import REDIS_STATISTICS_DB, RedisRole @@ -17,10 +18,12 @@ class RedisProvider(DependencyProvider[WebServerUnifiedConfig, ValkeySessionClie """ @property + @override def stage_name(self) -> str: return "redis" @asynccontextmanager + @override async def provide( self, setup_input: WebServerUnifiedConfig ) -> AsyncIterator[ValkeySessionClient]: @@ -42,6 +45,7 @@ async def provide( finally: await valkey_session_client.close() + @override def gen_liveness_checker(self, resource: ValkeySessionClient) -> ServiceHealthChecker: """Liveness — Valkey connection-stuck observed; restart recovers.""" return ValkeyHealthChecker(clients={CID_REDIS_SESSION: resource}) diff --git a/src/ai/backend/web/errors.py b/src/ai/backend/web/errors.py index aa5f1889214..bbc828b2e46 100644 --- a/src/ai/backend/web/errors.py +++ b/src/ai/backend/web/errors.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import override + from aiohttp import web from ai.backend.common.exception import ( @@ -21,6 +23,7 @@ class InvalidAPIConfigurationError(BackendAIError, web.HTTPInternalServerError): error_type = "https://api.backend.ai/probs/webserver/invalid-api-configuration" error_title = "Invalid API configuration state." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -35,6 +38,7 @@ class InvalidTemplateValueError(BackendAIError, web.HTTPBadRequest): error_type = "https://api.backend.ai/probs/webserver/invalid-template-value" error_title = "Invalid value in template." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, @@ -49,6 +53,7 @@ class ManagerConnectionUnavailable(BackendAIError, web.HTTPServiceUnavailable): error_type = "https://api.backend.ai/probs/webserver/manager-connection-unavailable" error_title = "No healthy Manager endpoint is available." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.API, diff --git a/src/ai/backend/web/logging.py b/src/ai/backend/web/logging.py index e4b8848f597..232b96139e7 100644 --- a/src/ai/backend/web/logging.py +++ b/src/ai/backend/web/logging.py @@ -3,7 +3,7 @@ import logging from collections.abc import Mapping from types import TracebackType -from typing import Any +from typing import Any, override class BraceMessage: @@ -13,6 +13,7 @@ def __init__(self, fmt: str, args: tuple[Any, ...]) -> None: self.fmt = fmt self.args = args + @override def __str__(self) -> str: return self.fmt.format(*self.args) @@ -21,6 +22,7 @@ class BraceStyleAdapter(logging.LoggerAdapter[logging.Logger]): def __init__(self, logger: logging.Logger, extra: Mapping[str, Any] | None = None) -> None: super().__init__(logger, extra) + @override def log( self, level: int, diff --git a/src/ai/backend/web/template.py b/src/ai/backend/web/template.py index 04cf1d58f9e..69743c7a1b3 100644 --- a/src/ai/backend/web/template.py +++ b/src/ai/backend/web/template.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Any +from typing import Any, override from jinja2 import nodes from jinja2.ext import Extension @@ -13,6 +13,7 @@ class TOMLField(Extension): tags = {"toml_field"} + @override def parse(self, parser: Parser) -> nodes.Node | list[nodes.Node]: tag_name = list(self.tags)[0] lineno = parser.stream.expect(f"name:{tag_name}").lineno @@ -55,6 +56,7 @@ def _transform(self, field_value: nodes.Expr, lineno: int | None = None) -> node class TOMLStringListField(TOMLField): tags = {"toml_strlist_field"} + @override def _transform(self, field_value: nodes.Expr, lineno: int | None = None) -> nodes.Expr: return nodes.Filter(field_value, "toml_scalar", [], [], None, None, lineno=lineno) diff --git a/stubs/trafaret/base.pyi b/stubs/trafaret/base.pyi index bb2d981d31c..b4d4e0b2d85 100644 --- a/stubs/trafaret/base.pyi +++ b/stubs/trafaret/base.pyi @@ -1,10 +1,11 @@ from collections.abc import Hashable from typing import Any as _Any -from typing import NoReturn +from typing import NoReturn, override from trafaret.dataerror import DataError class TrafaretMeta(type): + @override def __or__(cls, other: _Any) -> _Any: ... def __and__(cls, other: _Any) -> _Any: ... def __rshift__(cls, other: str) -> _Any: ... diff --git a/stubs/trafaret/constructor.pyi b/stubs/trafaret/constructor.pyi index b0084e90d2d..86405a09b2b 100644 --- a/stubs/trafaret/constructor.pyi +++ b/stubs/trafaret/constructor.pyi @@ -1,8 +1,10 @@ from typing import Any as _Any +from typing import override from .base import Key, Trafaret class ConstructMeta(type): + @override def __or__(self, other: _Any) -> _Any: ... def __and__(self, other: _Any) -> _Any: ... diff --git a/tests/unit/agent/test_affinity_map.py b/tests/unit/agent/test_affinity_map.py index 76de0526ac5..b5596786acb 100644 --- a/tests/unit/agent/test_affinity_map.py +++ b/tests/unit/agent/test_affinity_map.py @@ -4,7 +4,7 @@ from collections.abc import MutableMapping, Sequence from decimal import Decimal from pprint import pprint -from typing import Any +from typing import Any, override import pytest @@ -27,9 +27,11 @@ def __init__(self, *args: Any, extra_prop1: str = "zzz", **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.extra_prop1 = extra_prop1 + @override def __str__(self) -> str: return self.device_id + @override def __repr__(self) -> str: # for simpler output of debug prints return self.device_id @@ -42,9 +44,11 @@ def __init__(self, *args: Any, extra_prop1: str = "yyy", **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.extra_prop1 = extra_prop1 + @override def __str__(self) -> str: return self.device_id + @override def __repr__(self) -> str: # for simpler output of debug prints return self.device_id diff --git a/tests/unit/client/v2/test_anonymous_client.py b/tests/unit/client/v2/test_anonymous_client.py index 56b89b710be..ed60af68d0d 100644 --- a/tests/unit/client/v2/test_anonymous_client.py +++ b/tests/unit/client/v2/test_anonymous_client.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from datetime import datetime +from typing import override from unittest.mock import MagicMock import aiohttp @@ -22,6 +23,7 @@ def config() -> ClientConfig: class FakeAuth(AuthStrategy): + @override def sign( self, method: str, diff --git a/tests/unit/client/v2/test_infra_client.py b/tests/unit/client/v2/test_infra_client.py index d9cb65e828f..9a640c1731a 100644 --- a/tests/unit/client/v2/test_infra_client.py +++ b/tests/unit/client/v2/test_infra_client.py @@ -10,7 +10,7 @@ from collections.abc import Mapping from datetime import datetime -from typing import Any +from typing import Any, override from unittest.mock import MagicMock from uuid import UUID @@ -26,6 +26,7 @@ class FakeAuth(AuthStrategy): + @override def sign( self, method: str, diff --git a/tests/unit/client/v2/test_session_client.py b/tests/unit/client/v2/test_session_client.py index eccb52ec268..0d263fd42a6 100644 --- a/tests/unit/client/v2/test_session_client.py +++ b/tests/unit/client/v2/test_session_client.py @@ -11,7 +11,7 @@ from collections.abc import Mapping from datetime import datetime -from typing import Any +from typing import Any, override from unittest.mock import MagicMock from uuid import UUID @@ -32,6 +32,7 @@ class FakeAuth(AuthStrategy): + @override def sign( self, method: str, diff --git a/tests/unit/client_v2/conftest.py b/tests/unit/client_v2/conftest.py index b09d9139588..4716ebbab73 100644 --- a/tests/unit/client_v2/conftest.py +++ b/tests/unit/client_v2/conftest.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from datetime import datetime +from typing import override from uuid import UUID, uuid4 import pytest @@ -12,6 +13,7 @@ class MockAuth(AuthStrategy): """A mock auth strategy for testing that returns predictable headers.""" + @override def sign( self, method: str, diff --git a/tests/unit/client_v2/test_base_client.py b/tests/unit/client_v2/test_base_client.py index e0012f6d886..f71538fb397 100644 --- a/tests/unit/client_v2/test_base_client.py +++ b/tests/unit/client_v2/test_base_client.py @@ -1,5 +1,6 @@ from collections.abc import AsyncIterator, Mapping from datetime import datetime +from typing import override from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import aiohttp @@ -62,6 +63,7 @@ class _RecordingAuth(MockAuth): def __init__(self) -> None: self.signed_rel_urls: list[str] = [] + @override def sign( self, method: str, diff --git a/tests/unit/common/configs/loader/test_loader_chain.py b/tests/unit/common/configs/loader/test_loader_chain.py index 3bf1e306903..56ab1278d6b 100644 --- a/tests/unit/common/configs/loader/test_loader_chain.py +++ b/tests/unit/common/configs/loader/test_loader_chain.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, override from ai.backend.common.configs.loader import LoaderChain from ai.backend.common.configs.loader.types import AbstractConfigLoader @@ -8,6 +8,7 @@ class DummyLoader(AbstractConfigLoader): def __init__(self, payload: dict[str, Any]) -> None: self.payload = payload + @override async def load(self) -> dict[str, Any]: return self.payload diff --git a/tests/unit/common/configs/loader/test_source_name.py b/tests/unit/common/configs/loader/test_source_name.py index 7acb27b40b5..86f759ccd95 100644 --- a/tests/unit/common/configs/loader/test_source_name.py +++ b/tests/unit/common/configs/loader/test_source_name.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Any +from typing import Any, override from unittest.mock import AsyncMock from ai.backend.common.configs.loader import ( @@ -14,6 +14,7 @@ class CustomLoader(AbstractConfigLoader): + @override async def load(self) -> dict[str, Any]: return {} diff --git a/tests/unit/common/cron/test_local_cron.py b/tests/unit/common/cron/test_local_cron.py index c525e377054..da8223b6169 100644 --- a/tests/unit/common/cron/test_local_cron.py +++ b/tests/unit/common/cron/test_local_cron.py @@ -4,7 +4,7 @@ import asyncio from collections.abc import Awaitable, Callable -from typing import cast +from typing import cast, override from unittest.mock import AsyncMock import pytest @@ -30,24 +30,29 @@ def __init__( self.run_func = run_func or AsyncMock() self.run_count = 0 + @override async def run(self) -> None: """Execute the mock task.""" self.run_count += 1 await self.run_func() @property + @override def name(self) -> str: return self._name @property + @override def interval(self) -> float: return self._interval @property + @override def initial_delay(self) -> float: return self._initial_delay @property + @override def run_timeout(self) -> float | None: return self._run_timeout diff --git a/tests/unit/common/dependencies/stacks/test_visualizing.py b/tests/unit/common/dependencies/stacks/test_visualizing.py index 2bbe66175e2..bcc671880d2 100644 --- a/tests/unit/common/dependencies/stacks/test_visualizing.py +++ b/tests/unit/common/dependencies/stacks/test_visualizing.py @@ -4,7 +4,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from io import StringIO -from typing import Any +from typing import Any, override import pytest @@ -32,6 +32,7 @@ def __init__( self._raise_on_enter = raise_on_enter @property + @override def stage_name(self) -> str: """ Get the stage name. @@ -39,6 +40,7 @@ def stage_name(self) -> str: return self._stage_name @asynccontextmanager + @override async def provide(self, setup_input: Any) -> AsyncIterator[str]: """ Provide a resource. @@ -95,10 +97,12 @@ class ComposerResources: class TestComposer(DependencyComposer[str, ComposerResources]): @property + @override def stage_name(self) -> str: return "test-composer" @asynccontextmanager + @override async def compose( self, stack: DependencyStack, setup_input: str ) -> AsyncIterator[ComposerResources]: @@ -141,10 +145,12 @@ async def test_nested_composer_indentation(self) -> None: class InnerComposer(DependencyComposer[str, str]): @property + @override def stage_name(self) -> str: return "inner-composer" @asynccontextmanager + @override async def compose(self, stack: DependencyStack, setup_input: str) -> AsyncIterator[str]: provider = SimpleDependencyProvider("inner-dep", cleanup_tracker) resource = await stack.enter_dependency(provider, setup_input) @@ -152,10 +158,12 @@ async def compose(self, stack: DependencyStack, setup_input: str) -> AsyncIterat class OuterComposer(DependencyComposer[str, str]): @property + @override def stage_name(self) -> str: return "outer-composer" @asynccontextmanager + @override async def compose(self, stack: DependencyStack, setup_input: str) -> AsyncIterator[str]: inner_composer = InnerComposer() resource = await stack.enter_composer(inner_composer, setup_input) @@ -258,10 +266,12 @@ async def test_composer_failure_visualization(self) -> None: class FailingComposer(DependencyComposer[str, str]): @property + @override def stage_name(self) -> str: return "failing-composer" @asynccontextmanager + @override async def compose(self, stack: DependencyStack, setup_input: str) -> AsyncIterator[str]: provider1 = SimpleDependencyProvider("dep1", cleanup_tracker) await stack.enter_dependency(provider1, setup_input) diff --git a/tests/unit/common/dependencies/test_base.py b/tests/unit/common/dependencies/test_base.py index f8bab00f37f..5d32937b558 100644 --- a/tests/unit/common/dependencies/test_base.py +++ b/tests/unit/common/dependencies/test_base.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override import pytest @@ -65,6 +66,7 @@ class SimpleDependencyProvider(NonMonitorableDependencyProvider[str, str]): """ @property + @override def stage_name(self) -> str: """ @@ -73,6 +75,7 @@ def stage_name(self) -> str: return "simple" @asynccontextmanager + @override async def provide(self, setup_input: str) -> AsyncIterator[str]: """ @@ -111,10 +114,12 @@ class CleanupTracker: class ProviderWithCleanup(NonMonitorableDependencyProvider[str, str]): @property + @override def stage_name(self) -> str: return "cleanup-test" @asynccontextmanager + @override async def provide(self, setup_input: str) -> AsyncIterator[str]: try: yield "resource" diff --git a/tests/unit/common/events/test_dispatcher.py b/tests/unit/common/events/test_dispatcher.py index 97a11fbb108..13aee72c288 100644 --- a/tests/unit/common/events/test_dispatcher.py +++ b/tests/unit/common/events/test_dispatcher.py @@ -4,7 +4,7 @@ import base64 from collections.abc import AsyncGenerator from dataclasses import dataclass -from typing import Any +from typing import Any, override import msgpack import pytest @@ -27,24 +27,30 @@ class DummyAnycastEvent(AbstractAnycastEvent): value: int + @override def serialize(self) -> tuple[Any, ...]: return (self.value,) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> DummyAnycastEvent: return cls(value[0]) @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.AGENT + @override def domain_id(self) -> str | None: return None + @override def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "test_anycast" @@ -53,24 +59,30 @@ def event_name(cls) -> str: class DummyBroadcastEvent(AbstractBroadcastEvent): value: int + @override def serialize(self) -> tuple[Any, ...]: return (self.value,) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> DummyBroadcastEvent: return cls(value[0]) @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.AGENT + @override def domain_id(self) -> str | None: return None + @override def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "test_broadcast" diff --git a/tests/unit/common/events/test_hub.py b/tests/unit/common/events/test_hub.py index 0f3969a4408..b12f13f988b 100644 --- a/tests/unit/common/events/test_hub.py +++ b/tests/unit/common/events/test_hub.py @@ -1,7 +1,7 @@ from __future__ import annotations import uuid -from typing import Self +from typing import Self, override from unittest.mock import AsyncMock from ai.backend.common.events.hub import WILDCARD, EventHub, EventPropagator @@ -13,36 +13,44 @@ class DummyBaseEvent(AbstractEvent): def __init__(self, domain_id: str) -> None: self._domain_id = domain_id + @override def domain_id(self) -> str | None: return self._domain_id @classmethod + @override def delivery_pattern(cls) -> DeliveryPattern: return DeliveryPattern.BROADCAST + @override def serialize(self) -> tuple[bytes, ...]: raise NotImplementedError @classmethod + @override def deserialize(cls, value: tuple[bytes, ...]) -> Self: raise NotImplementedError @classmethod + @override def event_name(cls) -> str: return "dummy" + @override def user_event(self) -> UserEvent | None: return None class DummySessionEvent(DummyBaseEvent): @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.SESSION class DummyKernelEvent(DummyBaseEvent): @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.KERNEL @@ -54,12 +62,15 @@ def __init__(self) -> None: self._id = uuid.uuid4() self.records = [] + @override def id(self) -> uuid.UUID: return self._id + @override async def propagate_event(self, event: AbstractEvent) -> None: self.records.append(event.domain_id() or "") + @override async def close(self) -> None: pass diff --git a/tests/unit/common/resilience/test_policy_composition.py b/tests/unit/common/resilience/test_policy_composition.py index 9449c5eb9e5..9ffc646ac0a 100644 --- a/tests/unit/common/resilience/test_policy_composition.py +++ b/tests/unit/common/resilience/test_policy_composition.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Awaitable, Callable -from typing import ParamSpec, TypeVar +from typing import ParamSpec, TypeVar, override import pytest @@ -18,6 +18,7 @@ async def test_single_policy_execution(self) -> None: execute_called = False class MockPolicy(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -44,6 +45,7 @@ async def test_multiple_policies_execution_order(self) -> None: execution_order: list[str] = [] class MockPolicy1(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -56,6 +58,7 @@ async def execute( return result class MockPolicy2(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -68,6 +71,7 @@ async def execute( return result class MockPolicy3(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -103,6 +107,7 @@ async def test_exception_propagation_through_policies(self) -> None: execution_order: list[str] = [] class MockPolicy1(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -116,6 +121,7 @@ async def execute( execution_order.append("policy1_after") class MockPolicy2(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], @@ -159,6 +165,7 @@ async def test_policy_exception_handling(self) -> None: """Test that exceptions from policies themselves are propagated.""" class FailingMockPolicy(Policy): + @override async def execute( self, next_call: Callable[P, Awaitable[R]], diff --git a/tests/unit/common/test_api_handlers.py b/tests/unit/common/test_api_handlers.py index b638740094b..96db9437124 100644 --- a/tests/unit/common/test_api_handlers.py +++ b/tests/unit/common/test_api_handlers.py @@ -3,7 +3,7 @@ import uuid from collections.abc import Awaitable, Callable from http import HTTPStatus -from typing import Any, Self +from typing import Any, Self, override from aiohttp import web from pydantic import AliasChoices, Field @@ -281,6 +281,7 @@ class TestAuthInfo(MiddlewareParam): is_authorized: bool = Field(default=False) @classmethod + @override async def from_request(cls, request: web.Request) -> Self: return cls(is_authorized=request.get("is_authorized", False)) @@ -356,6 +357,7 @@ class TestMiddlewareModel(MiddlewareParam): is_authorized: bool @classmethod + @override async def from_request(cls, request: web.Request) -> Self: return cls(is_authorized=request.get("is_authorized", False)) diff --git a/tests/unit/common/test_distributed.py b/tests/unit/common/test_distributed.py index 4f20a45e9e6..e9dad2bd182 100644 --- a/tests/unit/common/test_distributed.py +++ b/tests/unit/common/test_distributed.py @@ -10,7 +10,7 @@ from decimal import Decimal from functools import partial from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, override import aiotools import pytest @@ -75,24 +75,30 @@ def dslice(start: Decimal, stop: Decimal, num: int) -> Iterable[Decimal]: class NoopAnycastEvent(AbstractAnycastEvent): test_case_ns: str + @override def serialize(self) -> tuple[Any, ...]: return (self.test_case_ns,) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> NoopAnycastEvent: return cls(value[0]) @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.AGENT + @override def domain_id(self) -> str | None: return None + @override def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "noop" @@ -317,6 +323,7 @@ async def _tick(context: Any, source: AgentId, event: NoopAnycastEvent) -> None: await event_dispatcher.close() await event_producer.close() + @override def run(self) -> None: asyncio.run(self.timer_node_async()) diff --git a/tests/unit/common/test_events.py b/tests/unit/common/test_events.py index e03e72e5acd..3ae28767199 100644 --- a/tests/unit/common/test_events.py +++ b/tests/unit/common/test_events.py @@ -3,7 +3,7 @@ import asyncio from dataclasses import dataclass from types import TracebackType -from typing import Any +from typing import Any, override import aiotools @@ -26,24 +26,30 @@ class DummyBroadcastEvent(AbstractBroadcastEvent): value: int + @override def serialize(self) -> tuple[Any, ...]: return (self.value + 1,) @classmethod + @override def deserialize(cls, value: tuple[Any, ...]) -> DummyBroadcastEvent: return cls(value[0] + 1) @classmethod + @override def event_domain(cls) -> EventDomain: return EventDomain.AGENT + @override def domain_id(self) -> str | None: return None + @override def user_event(self) -> UserEvent | None: return None @classmethod + @override def event_name(cls) -> str: return "testing" diff --git a/tests/unit/common/test_logging.py b/tests/unit/common/test_logging.py index 8bbe7e89f3d..f75d1d99ea8 100644 --- a/tests/unit/common/test_logging.py +++ b/tests/unit/common/test_logging.py @@ -5,6 +5,7 @@ import threading import time from pathlib import Path +from typing import override import pytest import trafaret as t @@ -43,6 +44,7 @@ class NotPicklableClass: """A class that cannot be pickled.""" + @override def __reduce__(self) -> tuple[type, tuple[()]]: raise TypeError("this is not picklable") @@ -54,6 +56,7 @@ def __init__(self, x: int) -> None: if x == 1: raise TypeError("this is not unpicklable") + @override def __reduce__(self) -> tuple[type[NotUnpicklableClass], tuple[int, ...]]: return type(self), (1,) diff --git a/tests/unit/common/test_plugin.py b/tests/unit/common/test_plugin.py index c15fb4a4095..6cba5e3dd9b 100644 --- a/tests/unit/common/test_plugin.py +++ b/tests/unit/common/test_plugin.py @@ -2,7 +2,7 @@ import functools from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass -from typing import Any +from typing import Any, override from unittest.mock import AsyncMock import pytest @@ -27,12 +27,15 @@ class DummyPlugin(AbstractPlugin): def __init__(self, plugin_config: Mapping[str, Any], local_config: Mapping[str, Any]) -> None: super().__init__(plugin_config, local_config) + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: pass @@ -187,12 +190,15 @@ class DummyExplicitAllowPlugin(AbstractPlugin): require_explicit_allow = True _entrypoint_name = "explicit-only" + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, plugin_config: Mapping[str, Any]) -> None: pass @@ -275,18 +281,22 @@ class DummyHookPassingPlugin(HookPlugin): _entrypoint_name = "hook-p" + @override def get_handlers(self) -> list[tuple[str, Callable[..., Any]]]: return [ ("HOOK1", self.hook1_handler), ("HOOK2", self.hook2_handler), ] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_config: Mapping[str, Any]) -> None: pass @@ -306,18 +316,22 @@ class DummyHookRejectingPlugin(HookPlugin): _entrypoint_name = "hook-r" + @override def get_handlers(self) -> list[tuple[str, Callable[..., Any]]]: return [ ("HOOK1", self.hook1_handler), ("HOOK2", self.hook2_handler), ] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_config: Mapping[str, Any]) -> None: pass @@ -337,17 +351,21 @@ class DummyHookErrorPlugin(HookPlugin): _entrypoint_name = "hook-e" + @override def get_handlers(self) -> list[tuple[str, Callable[..., Any]]]: return [ ("HOOK3", self.hook3_handler), ] + @override async def init(self, context: Any | None = None) -> None: pass + @override async def cleanup(self) -> None: pass + @override async def update_plugin_config(self, new_config: Mapping[str, Any]) -> None: pass diff --git a/tests/unit/common/test_types.py b/tests/unit/common/test_types.py index 34d89ce22e4..fda06d2637f 100644 --- a/tests/unit/common/test_types.py +++ b/tests/unit/common/test_types.py @@ -2,7 +2,7 @@ import uuid from dataclasses import dataclass from decimal import Decimal -from typing import Any +from typing import Any, override import pytest from typeguard import TypeCheckError @@ -39,6 +39,7 @@ def __init__(self, x: int) -> None: init_count += 1 self.x = x + @override async def __ainit__(self) -> None: await asyncio.sleep(0.01) nonlocal ainit_count @@ -51,6 +52,7 @@ def __init__(self, x: int, y: int) -> None: init_count += 1 self.y = y + @override async def __ainit__(self) -> None: await super().__ainit__() await asyncio.sleep(0.01) diff --git a/tests/unit/common/valkey_leader/test_leader_cron.py b/tests/unit/common/valkey_leader/test_leader_cron.py index 4d3a76af43d..eab9f465b05 100644 --- a/tests/unit/common/valkey_leader/test_leader_cron.py +++ b/tests/unit/common/valkey_leader/test_leader_cron.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from typing import cast +from typing import cast, override from unittest.mock import AsyncMock, MagicMock, create_autospec import pytest @@ -35,22 +35,26 @@ def __init__( self.run_func = run_func or AsyncMock() self.run_count = 0 + @override async def run(self) -> None: """Execute the mock task.""" self.run_count += 1 await self.run_func() @property + @override def name(self) -> str: """Task name.""" return self._name @property + @override def interval(self) -> float: """Interval.""" return self._interval @property + @override def initial_delay(self) -> float: """Initial delay.""" return self._initial_delay diff --git a/tests/unit/common/valkey_leader/test_leader_task_integration.py b/tests/unit/common/valkey_leader/test_leader_task_integration.py index 0a405eb0009..409baddb7b1 100644 --- a/tests/unit/common/valkey_leader/test_leader_task_integration.py +++ b/tests/unit/common/valkey_leader/test_leader_task_integration.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from typing import Any +from typing import Any, override from unittest.mock import AsyncMock, MagicMock import pytest @@ -31,6 +31,7 @@ def __init__(self) -> None: self._task: asyncio.Task[Any] | None = None self._stopped = False + @override async def start(self, leadership_checker: LeadershipChecker) -> None: self.start_called = True self.leadership_checker = leadership_checker @@ -44,6 +45,7 @@ async def _run(self) -> None: self.execution_count += 1 await asyncio.sleep(0.1) + @override async def stop(self) -> None: self.stop_called = True self._stopped = True diff --git a/tests/unit/common/valkey_leader/test_valkey_leader_client.py b/tests/unit/common/valkey_leader/test_valkey_leader_client.py index be3d8ee555e..01e3377f8a4 100644 --- a/tests/unit/common/valkey_leader/test_valkey_leader_client.py +++ b/tests/unit/common/valkey_leader/test_valkey_leader_client.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from typing import override from unittest.mock import AsyncMock import pytest @@ -112,6 +113,7 @@ async def test_handle_exception_during_acquire( # Mock the Lua script to raise an exception class MockBackendAIError(BackendAIError): @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode.default() diff --git a/tests/unit/manager/actions/test_processor.py b/tests/unit/manager/actions/test_processor.py index d08c12ae68e..22a14b165d0 100644 --- a/tests/unit/manager/actions/test_processor.py +++ b/tests/unit/manager/actions/test_processor.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from typing import Final +from typing import Final, override from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 @@ -30,14 +30,17 @@ class MockAction(BaseAction): type: str operation: str + @override def entity_id(self) -> str | None: return self.id @classmethod + @override def entity_type(cls) -> EntityType: return EntityType.SESSION @classmethod + @override def operation_type(cls) -> ActionOperationType: return ActionOperationType.CREATE @@ -46,6 +49,7 @@ def operation_type(cls) -> ActionOperationType: class MockActionResult(BaseActionResult): id: str + @override def entity_id(self) -> str | None: return self.id @@ -73,10 +77,12 @@ def __init__( self.expected_done_action = expected_done_action self.expected_done_result = expected_done_result + @override async def prepare(self, action: BaseAction, meta: BaseActionTriggerMeta) -> None: assert isinstance(action, MockAction) assert action == self.expected_prepare_action + @override async def done(self, action: BaseAction, result: ProcessResult) -> None: assert isinstance(action, MockAction) assert action == self.expected_done_action diff --git a/tests/unit/manager/api/test_exceptions.py b/tests/unit/manager/api/test_exceptions.py index dfc877728b3..cba9081e537 100644 --- a/tests/unit/manager/api/test_exceptions.py +++ b/tests/unit/manager/api/test_exceptions.py @@ -1,5 +1,6 @@ import json import pickle +from typing import override from ai.backend.common.exception import BackendAIError, ErrorCode from ai.backend.common.json import dump_json_str @@ -9,6 +10,7 @@ class TestError(BackendAIError): @classmethod + @override def error_code(cls) -> ErrorCode: return ErrorCode.default() diff --git a/tests/unit/manager/api/test_routing.py b/tests/unit/manager/api/test_routing.py index 7314d8e3b40..606c8e66057 100644 --- a/tests/unit/manager/api/test_routing.py +++ b/tests/unit/manager/api/test_routing.py @@ -6,7 +6,7 @@ import functools import json from collections.abc import Awaitable, Callable -from typing import Self +from typing import Self, override from unittest.mock import AsyncMock, MagicMock import aiohttp_cors @@ -283,6 +283,7 @@ class MockMiddlewareParam(MiddlewareParam): user_id: str @classmethod + @override async def from_request(cls, request: web.Request) -> Self: return cls(user_id="test-user-123") diff --git a/tests/unit/manager/models/test_base.py b/tests/unit/manager/models/test_base.py index eb247c92a0d..bd8e31ddcc2 100644 --- a/tests/unit/manager/models/test_base.py +++ b/tests/unit/manager/models/test_base.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from decimal import Decimal from types import SimpleNamespace -from typing import Any +from typing import Any, override from unittest.mock import AsyncMock, MagicMock import sqlalchemy as sa @@ -94,6 +94,7 @@ class _Shape(ABCColumnPayload): """Test-only polymorphic payload base; `load` dispatches by the `kind` tag.""" @classmethod + @override def load(cls, raw: dict[str, Any]) -> _Shape: kind = raw["kind"] if kind == "circle": @@ -107,6 +108,7 @@ def load(cls, raw: dict[str, Any]) -> _Shape: class _Circle(_Shape): radius: float + @override def serialize(self) -> dict[str, Any]: return {"kind": "circle", "radius": self.radius} @@ -115,6 +117,7 @@ def serialize(self) -> dict[str, Any]: class _Square(_Shape): side: float + @override def serialize(self) -> dict[str, Any]: return {"kind": "square", "side": self.side} diff --git a/tests/unit/manager/repositories/base/rbac/test_entity_creator.py b/tests/unit/manager/repositories/base/rbac/test_entity_creator.py index 76c63734ee7..7a33a145c87 100644 --- a/tests/unit/manager/repositories/base/rbac/test_entity_creator.py +++ b/tests/unit/manager/repositories/base/rbac/test_entity_creator.py @@ -4,7 +4,7 @@ import uuid from collections.abc import AsyncGenerator, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from uuid import UUID import aiohttp.web @@ -101,6 +101,7 @@ def __init__( self._scope_id = scope_id self._entity_id = entity_id + @override def build_row(self) -> RBACEntityCreatorTestRow: row_kwargs: dict[str, Any] = { "name": self._name, @@ -120,6 +121,7 @@ def __init__(self, tenant_id: int, item_id: int, name: str) -> None: self._item_id = item_id self._name = name + @override def build_row(self) -> CompositePKTestRow: return CompositePKTestRow( tenant_id=self._tenant_id, @@ -871,6 +873,7 @@ class _TestRBACDuplicateNameError(BackendAIError, aiohttp.web.HTTPConflict): error_type = "https://api.backend.ai/probs/test-rbac-duplicate-name" error_title = "Duplicate name." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -886,6 +889,7 @@ def __init__(self, name: str) -> None: self._name = name @property + @override def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: return ( IntegrityErrorCheck( @@ -895,6 +899,7 @@ def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: ), ) + @override def build_row(self) -> RBACCreatorUniqueTestRow: return RBACCreatorUniqueTestRow(name=self._name) diff --git a/tests/unit/manager/repositories/base/rbac/test_entity_purger.py b/tests/unit/manager/repositories/base/rbac/test_entity_purger.py index cc59a357cd6..37b2e186a29 100644 --- a/tests/unit/manager/repositories/base/rbac/test_entity_purger.py +++ b/tests/unit/manager/repositories/base/rbac/test_entity_purger.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from uuid import UUID import pytest @@ -108,9 +108,11 @@ class SimpleRBACEntityPurgerSpec(RBACEntityPurgerSpec): def __init__(self, entity_uuid: UUID) -> None: self._entity_uuid = entity_uuid + @override def element_type(self) -> RBACElementType: return RBACElementType.VFOLDER + @override def entity_ref(self) -> RBACElementRef: return RBACElementRef( element_type=self.element_type(), @@ -675,9 +677,11 @@ async def test_purger_preserves_unrelated_scope_associations( class TestEntityBatchPurgerSpec(RBACEntityBatchPurgerSpec[RBACEntityPurgerTestRow]): """Test spec for batch purging entities.""" + @override def build_subquery(self) -> sa.sql.Select[tuple[RBACEntityPurgerTestRow]]: return sa.select(RBACEntityPurgerTestRow) + @override def element_type(self) -> RBACElementType: return RBACElementType.VFOLDER @@ -919,9 +923,11 @@ class CompositePKPurgerSpec(RBACEntityPurgerSpec): def __init__(self, entity_uuid: str) -> None: self._entity_uuid = entity_uuid + @override def element_type(self) -> RBACElementType: return RBACElementType.VFOLDER + @override def entity_ref(self) -> RBACElementRef: return RBACElementRef( element_type=self.element_type(), @@ -932,9 +938,11 @@ def entity_ref(self) -> RBACElementRef: class CompositePKBatchPurgerSpec(RBACEntityBatchPurgerSpec[CompositePKPurgerTestRow]): """Batch purger spec for composite PK testing.""" + @override def build_subquery(self) -> sa.Select[Any]: return sa.select(CompositePKPurgerTestRow) + @override def element_type(self) -> RBACElementType: return RBACElementType.VFOLDER diff --git a/tests/unit/manager/repositories/base/rbac/test_scope_binder.py b/tests/unit/manager/repositories/base/rbac/test_scope_binder.py index f6ba4c80a19..3ff8b4f16c8 100644 --- a/tests/unit/manager/repositories/base/rbac/test_scope_binder.py +++ b/tests/unit/manager/repositories/base/rbac/test_scope_binder.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from uuid import UUID import pytest @@ -66,6 +66,7 @@ def __init__(self, entity_id: str, scope_id: str) -> None: self._entity_id = entity_id self._scope_id = scope_id + @override def build_row(self) -> ScopeBinderMappingRow: return ScopeBinderMappingRow( id=uuid.uuid4(), diff --git a/tests/unit/manager/repositories/base/rbac/test_scope_unbinder.py b/tests/unit/manager/repositories/base/rbac/test_scope_unbinder.py index bca70691da3..edf1eef763f 100644 --- a/tests/unit/manager/repositories/base/rbac/test_scope_unbinder.py +++ b/tests/unit/manager/repositories/base/rbac/test_scope_unbinder.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import pytest import sqlalchemy as sa @@ -71,11 +71,13 @@ def __init__( self._scope_id = scope_id self._scope_element_type = scope_element_type + @override def build_purger_spec(self) -> BatchPurgerSpec[ScopeUnbinderMappingRow]: entity_ids = self._entity_ids_value scope_id = self._scope_id class _Spec(BatchPurgerSpec[ScopeUnbinderMappingRow]): + @override def build_subquery(self) -> sa.sql.Select[tuple[ScopeUnbinderMappingRow]]: stmt = sa.select(ScopeUnbinderMappingRow).where( ScopeUnbinderMappingRow.scope_id == scope_id, @@ -89,14 +91,17 @@ def build_subquery(self) -> sa.sql.Select[tuple[ScopeUnbinderMappingRow]]: return _Spec() @property + @override def entity_type(self) -> RBACElementType: return self._entity_element_type @property + @override def scope_ref(self) -> RBACElementRef: return RBACElementRef(self._scope_element_type, self._scope_id) @property + @override def entity_ids(self) -> Sequence[str] | None: return self._entity_ids_value diff --git a/tests/unit/manager/repositories/base/test_creator.py b/tests/unit/manager/repositories/base/test_creator.py index 452d376366c..e35ac5be2c8 100644 --- a/tests/unit/manager/repositories/base/test_creator.py +++ b/tests/unit/manager/repositories/base/test_creator.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator, Sequence from datetime import datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override import aiohttp.web import pytest @@ -59,6 +59,7 @@ def __init__(self, name: str, value: str | None = None) -> None: self._name = name self._value = value + @override def build_row(self) -> CreatorTestRow: return CreatorTestRow(name=self._name, value=self._value) @@ -70,6 +71,7 @@ def __init__(self, should_fail: bool = True, name: str = "test") -> None: self._should_fail = should_fail self._name = name + @override def build_row(self) -> CreatorTestRow: if self._should_fail: raise ValueError(f"Simulated failure for {self._name}") @@ -341,6 +343,7 @@ class DefaultsCreatorSpec(CreatorSpec[CreatorTestRowWithDefaults]): def __init__(self, name: str) -> None: self._name = name + @override def build_row(self) -> CreatorTestRowWithDefaults: return CreatorTestRowWithDefaults(name=self._name) @@ -479,6 +482,7 @@ class PythonDefaultsCreatorSpec(CreatorSpec[CreatorTestRowWithPythonDefaults]): def __init__(self, name: str) -> None: self._name = name + @override def build_row(self) -> CreatorTestRowWithPythonDefaults: return CreatorTestRowWithPythonDefaults(name=self._name) @@ -731,6 +735,7 @@ class _TestDuplicateNameError(BackendAIError, aiohttp.web.HTTPConflict): error_type = "https://api.backend.ai/probs/test-duplicate-name" error_title = "Duplicate name." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -746,6 +751,7 @@ def __init__(self, name: str) -> None: self._name = name @property + @override def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: return ( IntegrityErrorCheck( @@ -755,6 +761,7 @@ def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: ), ) + @override def build_row(self) -> CreatorTestRowWithUnique: return CreatorTestRowWithUnique(name=self._name) @@ -765,6 +772,7 @@ class PlainUniqueCreatorSpec(CreatorSpec[CreatorTestRowWithUnique]): def __init__(self, name: str) -> None: self._name = name + @override def build_row(self) -> CreatorTestRowWithUnique: return CreatorTestRowWithUnique(name=self._name) @@ -776,6 +784,7 @@ def __init__(self, name: str) -> None: self._name = name @property + @override def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: return ( IntegrityErrorCheck( @@ -785,6 +794,7 @@ def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: ), ) + @override def build_row(self) -> CreatorTestRowWithUnique: return CreatorTestRowWithUnique(name=self._name) diff --git a/tests/unit/manager/repositories/base/test_purger.py b/tests/unit/manager/repositories/base/test_purger.py index 96fd4faaea8..2397aab5b97 100644 --- a/tests/unit/manager/repositories/base/test_purger.py +++ b/tests/unit/manager/repositories/base/test_purger.py @@ -4,7 +4,7 @@ import uuid from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, override from uuid import UUID import pytest @@ -237,6 +237,7 @@ def __init__( self._row_class = row_class self._conditions = conditions or [] + @override def build_subquery(self) -> sa.sql.Select[tuple[Base]]: query = sa.select(self._row_class) for cond in self._conditions: @@ -536,6 +537,7 @@ def __init__( self._row_class = row_class self._conditions = conditions or [] + @override def build_subquery(self) -> sa.sql.Select[tuple[Base]]: query = sa.select(self._row_class) for cond in self._conditions: @@ -648,6 +650,7 @@ def __init__( ) -> None: self._conditions = conditions or [] + @override def build_subquery(self) -> sa.sql.Select[tuple[BatchPurgerCompositePKRow]]: query = sa.select(BatchPurgerCompositePKRow) for cond in self._conditions: diff --git a/tests/unit/manager/repositories/base/test_querier.py b/tests/unit/manager/repositories/base/test_querier.py index 05820a9471e..86d5f99be06 100644 --- a/tests/unit/manager/repositories/base/test_querier.py +++ b/tests/unit/manager/repositories/base/test_querier.py @@ -4,7 +4,7 @@ from collections.abc import AsyncGenerator, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from uuid import UUID, uuid4 import pytest @@ -352,6 +352,7 @@ class TestScopeValidationError(BackendAIError): error_type = "https://api.backend.ai/probs/test-scope-error" error_title = "Test scope validation error." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -366,6 +367,7 @@ class TestScopeValidationError2(BackendAIError): error_type = "https://api.backend.ai/probs/test-scope-error-2" error_title = "Test scope validation error 2." + @override def error_code(self) -> ErrorCode: return ErrorCode( domain=ErrorDomain.BACKENDAI, @@ -392,6 +394,7 @@ class MockSearchScope(SearchScope): checks: Sequence[ExistenceCheck[Any]] filter_category: str | None = None + @override def to_condition(self) -> QueryCondition: """Convert scope to a query condition.""" category = self.filter_category @@ -404,6 +407,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: """Return existence checks for scope validation.""" return self.checks diff --git a/tests/unit/manager/repositories/base/test_updater.py b/tests/unit/manager/repositories/base/test_updater.py index ca96e888a1d..b6b5e8651ae 100644 --- a/tests/unit/manager/repositories/base/test_updater.py +++ b/tests/unit/manager/repositories/base/test_updater.py @@ -4,7 +4,7 @@ import uuid from collections.abc import AsyncGenerator, Callable, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from uuid import UUID import pytest @@ -83,9 +83,11 @@ def __init__(self, new_status: str, new_value: int | None = None) -> None: self._new_value = new_value @property + @override def row_class(self) -> type[UpdaterTestRowInt]: return UpdaterTestRowInt + @override def build_values(self) -> dict[str, Any]: values: dict[str, Any] = {"status": self._new_status} if self._new_value is not None: @@ -97,9 +99,11 @@ class IntPKNoValuesUpdaterSpec(UpdaterSpec[UpdaterTestRowInt]): """Updater spec that produces no column changes (empty build_values).""" @property + @override def row_class(self) -> type[UpdaterTestRowInt]: return UpdaterTestRowInt + @override def build_values(self) -> dict[str, Any]: return {} @@ -111,9 +115,11 @@ def __init__(self, new_status: str) -> None: self._new_status = new_status @property + @override def row_class(self) -> type[UpdaterTestRowUUID]: return UpdaterTestRowUUID + @override def build_values(self) -> dict[str, Any]: return {"status": self._new_status} @@ -125,9 +131,11 @@ def __init__(self, new_status: str) -> None: self._new_status = new_status @property + @override def row_class(self) -> type[UpdaterTestRowStr]: return UpdaterTestRowStr + @override def build_values(self) -> dict[str, Any]: return {"status": self._new_status} @@ -142,9 +150,11 @@ def __init__(self, new_status: str) -> None: self._new_status = new_status @property + @override def row_class(self) -> type[UpdaterTestRowInt]: return UpdaterTestRowInt + @override def build_values(self) -> dict[str, Any]: return {"status": self._new_status} @@ -714,13 +724,16 @@ def __init__( self._checks = checks @property + @override def row_class(self) -> type[UpdaterTestRowWithUnique]: return UpdaterTestRowWithUnique + @override def build_values(self) -> dict[str, Any]: return {"name": self._new_name} @property + @override def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: return self._checks @@ -737,13 +750,16 @@ def __init__( self._checks = checks @property + @override def row_class(self) -> type[UpdaterTestRowWithUnique]: return UpdaterTestRowWithUnique + @override def build_values(self) -> dict[str, Any]: return {"name": self._new_name} @property + @override def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]: return self._checks @@ -941,9 +957,11 @@ def __init__(self, **values: Any) -> None: self._values = values @property + @override def row_class(self) -> type[BulkUpdaterPartialTestRow]: return BulkUpdaterPartialTestRow + @override def build_values(self) -> dict[str, Any]: return dict(self._values) diff --git a/tests/unit/manager/repositories/base/test_upserter.py b/tests/unit/manager/repositories/base/test_upserter.py index 1a7e3387bef..c844d14e3c2 100644 --- a/tests/unit/manager/repositories/base/test_upserter.py +++ b/tests/unit/manager/repositories/base/test_upserter.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override import pytest import sqlalchemy as sa @@ -41,12 +41,15 @@ def __init__(self, name: str, value: str | None = None) -> None: self._value = value @property + @override def row_class(self) -> type[UpserterTestRow]: return UpserterTestRow + @override def build_insert_values(self) -> dict[str, Any]: return {"name": self._name, "value": self._value} + @override def build_update_values(self) -> dict[str, Any]: # On conflict (PK), only update value return {"value": self._value} diff --git a/tests/unit/manager/repositories/ops/test_provider.py b/tests/unit/manager/repositories/ops/test_provider.py index 181abf91840..307e698bee6 100644 --- a/tests/unit/manager/repositories/ops/test_provider.py +++ b/tests/unit/manager/repositories/ops/test_provider.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator, Sequence from dataclasses import dataclass -from typing import Any +from typing import Any, override from unittest.mock import AsyncMock import pytest @@ -62,6 +62,7 @@ class ParentCreatorSpec(CreatorSpec[OpsTestParentRow]): name: str domain_name: str + @override def build_row(self) -> OpsTestParentRow: return OpsTestParentRow(name=self.name, domain_name=self.domain_name) @@ -71,9 +72,11 @@ class ParentUpdaterSpec(UpdaterSpec[OpsTestParentRow]): new_name: str @property + @override def row_class(self) -> type[OpsTestParentRow]: return OpsTestParentRow + @override def build_values(self) -> dict[str, Any]: return {"name": self.new_name} @@ -87,6 +90,7 @@ class ChildDependency: class ChildDependentCreatorSpec(DependentCreatorSpec[ChildDependency, OpsTestChildRow]): label: str + @override def build_row(self, dependency: ChildDependency) -> OpsTestChildRow: return OpsTestChildRow(parent_id=dependency.parent_id, label=self.label) @@ -95,10 +99,12 @@ def build_row(self, dependency: ChildDependency) -> OpsTestChildRow: class ParentDomainScope(SearchScope): domain_name: str + @override def to_condition(self) -> QueryCondition: return lambda: OpsTestParentRow.domain_name == self.domain_name @property + @override def existence_checks(self) -> Sequence[ExistenceCheck[Any]]: return () diff --git a/tests/unit/manager/services/deployment/test_deployment_service.py b/tests/unit/manager/services/deployment/test_deployment_service.py index 3ff5e8d5b1c..c90b280d457 100644 --- a/tests/unit/manager/services/deployment/test_deployment_service.py +++ b/tests/unit/manager/services/deployment/test_deployment_service.py @@ -9,7 +9,7 @@ import uuid from collections.abc import Callable, Iterator from datetime import UTC, datetime -from typing import cast +from typing import cast, override from unittest.mock import AsyncMock, MagicMock import pytest @@ -635,6 +635,7 @@ def configure_repository( return mock_deployment_repository @pytest.fixture + @override def mock_appproxy_client_pool(self, sample_coordinator_jwt: str) -> MagicMock: client = MagicMock(spec=AppProxyClient) client.mint_endpoint_token = AsyncMock( @@ -645,6 +646,7 @@ def mock_appproxy_client_pool(self, sample_coordinator_jwt: str) -> MagicMock: return pool @pytest.fixture + @override def deployment_service( self, mock_deployment_controller: MagicMock, diff --git a/tests/unit/manager/sokovan/reconciler/test_coordinator.py b/tests/unit/manager/sokovan/reconciler/test_coordinator.py index c3739e7f1dc..02bab1c35d5 100644 --- a/tests/unit/manager/sokovan/reconciler/test_coordinator.py +++ b/tests/unit/manager/sokovan/reconciler/test_coordinator.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from datetime import UTC, datetime from enum import StrEnum -from typing import cast +from typing import cast, override from unittest.mock import MagicMock from uuid import UUID, uuid4 @@ -52,9 +52,11 @@ class FakeTargetStatuses(BaseReconcilerTargetStatuses): class FakeInfo(BaseReconcilerInfo): ids: list[UUID] = field(default_factory=list) + @override def entity_ids(self) -> Sequence[UUID]: return self.ids + @override def now(self) -> datetime: return datetime(2026, 1, 1, tzinfo=UTC) @@ -64,12 +66,15 @@ class FakeResult(BaseReconcilerResult): processed: int = 0 failed: int = 0 + @override def processed_count(self) -> int: return self.processed + @override def failed_count(self) -> int: return self.failed + @override def decisions(self) -> Sequence[ReconcilerDecision]: return [] @@ -84,6 +89,7 @@ def __init__(self, call_log: CallLog, info: FakeInfo) -> None: self._call_log = call_log self._info = info + @override async def fetch_reconcile_info( self, category: FakeCategory, target_statuses: FakeTargetStatuses ) -> FakeInfo: @@ -96,10 +102,12 @@ def __init__(self, call_log: CallLog, result: FakeResult) -> None: self._call_log = call_log self._result = result + @override async def execute(self, reconcile_info: FakeInfo) -> FakeResult: self._call_log.events.append("execute") return self._result + @override async def post_process(self, result: FakeResult) -> None: self._call_log.events.append("post_process") @@ -110,6 +118,7 @@ class FakeApplier( def __init__(self, call_log: CallLog) -> None: self._call_log = call_log + @override async def apply( self, apply_input: ReconcilerApplyInput[ @@ -123,6 +132,7 @@ class FakeFlag(ReconcilerFlag): def __init__(self, needed: bool) -> None: self._needed = needed + @override async def check_mark_needed(self, reconcile_type: str) -> bool: return self._needed diff --git a/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py b/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py index 7707f43d273..e5d19e8bfab 100644 --- a/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py +++ b/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py @@ -18,6 +18,7 @@ import uuid from collections.abc import Callable +from typing import override import pytest @@ -69,9 +70,11 @@ def __init__( self._log = log self._transform = transform + @override def name(self) -> str: return self._tag + @override async def prepare( self, draft: SessionSpecDraft, diff --git a/tests/unit/storage-proxy/storages/test_s3_client.py b/tests/unit/storage-proxy/storages/test_s3_client.py index a3fcee6cda6..bc67a0d93cf 100644 --- a/tests/unit/storage-proxy/storages/test_s3_client.py +++ b/tests/unit/storage-proxy/storages/test_s3_client.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator +from typing import override import pytest from botocore.exceptions import ClientError @@ -11,10 +12,12 @@ class TestStreamReader(StreamReader): def __init__(self, data_chunks: list[bytes]): self._data_chunks = data_chunks + @override async def read(self) -> AsyncIterator[bytes]: for chunk in self._data_chunks: yield chunk + @override def content_type(self) -> str | None: return None diff --git a/tests/unit/storage-proxy/storages/test_service.py b/tests/unit/storage-proxy/storages/test_service.py index 758a50cdaeb..a3e74103e86 100644 --- a/tests/unit/storage-proxy/storages/test_service.py +++ b/tests/unit/storage-proxy/storages/test_service.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator -from typing import Any +from typing import Any, override import pytest @@ -23,10 +23,12 @@ class TestStreamReader(StreamReader): def __init__(self, data_chunks: list[bytes]): self._data_chunks = data_chunks + @override async def read(self) -> AsyncIterator[bytes]: for chunk in self._data_chunks: yield chunk + @override def content_type(self) -> str | None: return None From 7c5241b74f2a43057ae5f0462c5d20e97588c966 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Thu, 2 Jul 2026 14:23:29 +0900 Subject: [PATCH 2/2] changes: add news fragment for PR #12451 Co-Authored-By: Claude Opus 4.8 --- changes/12451.misc.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/12451.misc.md diff --git a/changes/12451.misc.md b/changes/12451.misc.md new file mode 100644 index 00000000000..22bc2f5252c --- /dev/null +++ b/changes/12451.misc.md @@ -0,0 +1 @@ +Enable mypy's `explicit-override` error code and add the missing `@typing.override` decorators across the codebase, so every method overriding a base-class method is explicitly marked. `ai.backend.client.cli.pretty` is exempted per-module because mypy cannot recognise `@override` on a property overriding a base instance attribute.