From ef8ce3bacd0e0d4b12ef4010f4ac1db49385ce07 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 22 Jul 2025 14:26:09 +0500 Subject: [PATCH] Refactor backends module to avoid importing deps on models import --- .../_internal/core/backends/__init__.py | 65 ------------------- .../_internal/core/backends/features.py | 64 ++++++++++++++++++ .../background/tasks/process_instances.py | 8 +-- .../_internal/server/services/fleets.py | 2 +- .../server/services/gateways/__init__.py | 8 +-- .../_internal/server/services/instances.py | 2 +- .../_internal/server/services/offers.py | 6 +- .../_internal/server/services/volumes.py | 2 +- 8 files changed, 78 insertions(+), 79 deletions(-) create mode 100644 src/dstack/_internal/core/backends/features.py diff --git a/src/dstack/_internal/core/backends/__init__.py b/src/dstack/_internal/core/backends/__init__.py index c384d66311..e69de29bb2 100644 --- a/src/dstack/_internal/core/backends/__init__.py +++ b/src/dstack/_internal/core/backends/__init__.py @@ -1,65 +0,0 @@ -from dstack._internal.core.backends.base.compute import ( - ComputeWithCreateInstanceSupport, - ComputeWithGatewaySupport, - ComputeWithMultinodeSupport, - ComputeWithPlacementGroupSupport, - ComputeWithPrivateGatewaySupport, - ComputeWithReservationSupport, - ComputeWithVolumeSupport, -) -from dstack._internal.core.backends.base.configurator import Configurator -from dstack._internal.core.backends.configurators import list_available_configurator_classes -from dstack._internal.core.backends.local.compute import LocalCompute -from dstack._internal.core.models.backends.base import BackendType -from dstack._internal.settings import LOCAL_BACKEND_ENABLED - - -def _get_backends_with_compute_feature( - configurator_classes: list[type[Configurator]], - compute_feature_class: type, -) -> list[BackendType]: - backend_types_and_computes = [ - (configurator_class.TYPE, configurator_class.BACKEND_CLASS.COMPUTE_CLASS) - for configurator_class in configurator_classes - ] - if LOCAL_BACKEND_ENABLED: - backend_types_and_computes.append((BackendType.LOCAL, LocalCompute)) - backend_types = [] - for backend_type, compute_class in backend_types_and_computes: - if issubclass(compute_class, compute_feature_class): - backend_types.append(backend_type) - return backend_types - - -_configurator_classes = list_available_configurator_classes() - - -# The following backend lists do not include unavailable backends (i.e. backends missing deps). -BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithCreateInstanceSupport, -) -BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithMultinodeSupport, -) -BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithPlacementGroupSupport, -) -BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithReservationSupport, -) -BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithGatewaySupport, -) -BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithPrivateGatewaySupport, -) -BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature( - configurator_classes=_configurator_classes, - compute_feature_class=ComputeWithVolumeSupport, -) diff --git a/src/dstack/_internal/core/backends/features.py b/src/dstack/_internal/core/backends/features.py new file mode 100644 index 0000000000..780d6203c3 --- /dev/null +++ b/src/dstack/_internal/core/backends/features.py @@ -0,0 +1,64 @@ +from dstack._internal.core.backends.base.compute import ( + ComputeWithCreateInstanceSupport, + ComputeWithGatewaySupport, + ComputeWithMultinodeSupport, + ComputeWithPlacementGroupSupport, + ComputeWithPrivateGatewaySupport, + ComputeWithReservationSupport, + ComputeWithVolumeSupport, +) +from dstack._internal.core.backends.base.configurator import Configurator +from dstack._internal.core.backends.configurators import list_available_configurator_classes +from dstack._internal.core.backends.local.compute import LocalCompute +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.settings import LOCAL_BACKEND_ENABLED + +_configurator_classes = list_available_configurator_classes() + + +def _get_backends_with_compute_feature( + configurator_classes: list[type[Configurator]], + compute_feature_class: type, +) -> list[BackendType]: + backend_types_and_computes = [ + (configurator_class.TYPE, configurator_class.BACKEND_CLASS.COMPUTE_CLASS) + for configurator_class in configurator_classes + ] + if LOCAL_BACKEND_ENABLED: + backend_types_and_computes.append((BackendType.LOCAL, LocalCompute)) + backend_types = [] + for backend_type, compute_class in backend_types_and_computes: + if issubclass(compute_class, compute_feature_class): + backend_types.append(backend_type) + return backend_types + + +# The following backend lists do not include unavailable backends (i.e. backends missing deps). +BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithCreateInstanceSupport, +) +BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithMultinodeSupport, +) +BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithPlacementGroupSupport, +) +BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithReservationSupport, +) +BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithGatewaySupport, +) +BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithPrivateGatewaySupport, +) +BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature( + configurator_classes=_configurator_classes, + compute_feature_class=ComputeWithVolumeSupport, +) diff --git a/src/dstack/_internal/server/background/tasks/process_instances.py b/src/dstack/_internal/server/background/tasks/process_instances.py index 75bbba61a9..3a9879aca1 100644 --- a/src/dstack/_internal/server/background/tasks/process_instances.py +++ b/src/dstack/_internal/server/background/tasks/process_instances.py @@ -12,10 +12,6 @@ from sqlalchemy.orm import joinedload, lazyload from dstack._internal import settings -from dstack._internal.core.backends import ( - BACKENDS_WITH_CREATE_INSTANCE_SUPPORT, - BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT, -) from dstack._internal.core.backends.base.compute import ( ComputeWithCreateInstanceSupport, ComputeWithPlacementGroupSupport, @@ -27,6 +23,10 @@ get_shim_env, get_shim_pre_start_commands, ) +from dstack._internal.core.backends.features import ( + BACKENDS_WITH_CREATE_INSTANCE_SUPPORT, + BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT, +) from dstack._internal.core.backends.remote.provisioning import ( detect_cpu_arch, get_host_info, diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index b9c7aabd3a..5c0e2ae5c6 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -8,8 +8,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, selectinload -from dstack._internal.core.backends import BACKENDS_WITH_CREATE_INSTANCE_SUPPORT from dstack._internal.core.backends.base.backend import Backend +from dstack._internal.core.backends.features import BACKENDS_WITH_CREATE_INSTANCE_SUPPORT from dstack._internal.core.errors import ( ForbiddenError, ResourceExistsError, diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 224184d9dc..8493f0284f 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -11,16 +11,16 @@ from sqlalchemy.orm import selectinload import dstack._internal.utils.random_names as random_names -from dstack._internal.core.backends import ( - BACKENDS_WITH_GATEWAY_SUPPORT, - BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT, -) from dstack._internal.core.backends.base.compute import ( Compute, ComputeWithGatewaySupport, get_dstack_gateway_wheel, get_dstack_runner_version, ) +from dstack._internal.core.backends.features import ( + BACKENDS_WITH_GATEWAY_SUPPORT, + BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT, +) from dstack._internal.core.errors import ( GatewayError, ResourceNotExistsError, diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index 40dd3c3f15..4bbb3de269 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -8,11 +8,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload -from dstack._internal.core.backends import BACKENDS_WITH_MULTINODE_SUPPORT from dstack._internal.core.backends.base.offers import ( offer_to_catalog_item, requirements_to_query_filter, ) +from dstack._internal.core.backends.features import BACKENDS_WITH_MULTINODE_SUPPORT from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.envs import Env from dstack._internal.core.models.instances import ( diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 7e5bd443f0..1e1b6ff586 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -2,13 +2,13 @@ import gpuhunt -from dstack._internal.core.backends import ( +from dstack._internal.core.backends.base.backend import Backend +from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport +from dstack._internal.core.backends.features import ( BACKENDS_WITH_CREATE_INSTANCE_SUPPORT, BACKENDS_WITH_MULTINODE_SUPPORT, BACKENDS_WITH_RESERVATION_SUPPORT, ) -from dstack._internal.core.backends.base.backend import Backend -from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( InstanceOfferWithAvailability, diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py index c58ca38dfe..c70b33c87c 100644 --- a/src/dstack/_internal/server/services/volumes.py +++ b/src/dstack/_internal/server/services/volumes.py @@ -6,8 +6,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, selectinload -from dstack._internal.core.backends import BACKENDS_WITH_VOLUMES_SUPPORT from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport +from dstack._internal.core.backends.features import BACKENDS_WITH_VOLUMES_SUPPORT from dstack._internal.core.errors import ( BackendNotAvailable, ResourceExistsError,