Skip to content

Commit b58ad40

Browse files
r4victorr4victor
authored
Refactor backends module to avoid importing deps on models import (#2923)
1 parent 7a64276 commit b58ad40

8 files changed

Lines changed: 78 additions & 79 deletions

File tree

‎src/dstack/_internal/core/backends/__init__.py‎

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +0,0 @@
1-
from dstack._internal.core.backends.base.compute import (
2-
ComputeWithCreateInstanceSupport,
3-
ComputeWithGatewaySupport,
4-
ComputeWithMultinodeSupport,
5-
ComputeWithPlacementGroupSupport,
6-
ComputeWithPrivateGatewaySupport,
7-
ComputeWithReservationSupport,
8-
ComputeWithVolumeSupport,
9-
)
10-
from dstack._internal.core.backends.base.configurator import Configurator
11-
from dstack._internal.core.backends.configurators import list_available_configurator_classes
12-
from dstack._internal.core.backends.local.compute import LocalCompute
13-
from dstack._internal.core.models.backends.base import BackendType
14-
from dstack._internal.settings import LOCAL_BACKEND_ENABLED
15-
16-
17-
def _get_backends_with_compute_feature(
18-
configurator_classes: list[type[Configurator]],
19-
compute_feature_class: type,
20-
) -> list[BackendType]:
21-
backend_types_and_computes = [
22-
(configurator_class.TYPE, configurator_class.BACKEND_CLASS.COMPUTE_CLASS)
23-
for configurator_class in configurator_classes
24-
]
25-
if LOCAL_BACKEND_ENABLED:
26-
backend_types_and_computes.append((BackendType.LOCAL, LocalCompute))
27-
backend_types = []
28-
for backend_type, compute_class in backend_types_and_computes:
29-
if issubclass(compute_class, compute_feature_class):
30-
backend_types.append(backend_type)
31-
return backend_types
32-
33-
34-
_configurator_classes = list_available_configurator_classes()
35-
36-
37-
# The following backend lists do not include unavailable backends (i.e. backends missing deps).
38-
BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
39-
configurator_classes=_configurator_classes,
40-
compute_feature_class=ComputeWithCreateInstanceSupport,
41-
)
42-
BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
43-
configurator_classes=_configurator_classes,
44-
compute_feature_class=ComputeWithMultinodeSupport,
45-
)
46-
BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature(
47-
configurator_classes=_configurator_classes,
48-
compute_feature_class=ComputeWithPlacementGroupSupport,
49-
)
50-
BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature(
51-
configurator_classes=_configurator_classes,
52-
compute_feature_class=ComputeWithReservationSupport,
53-
)
54-
BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
55-
configurator_classes=_configurator_classes,
56-
compute_feature_class=ComputeWithGatewaySupport,
57-
)
58-
BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
59-
configurator_classes=_configurator_classes,
60-
compute_feature_class=ComputeWithPrivateGatewaySupport,
61-
)
62-
BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature(
63-
configurator_classes=_configurator_classes,
64-
compute_feature_class=ComputeWithVolumeSupport,
65-
)

‎src/dstack/_internal/core/backends/features.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
from dstack._internal.core.backends.base.compute import (
2+
ComputeWithCreateInstanceSupport,
3+
ComputeWithGatewaySupport,
4+
ComputeWithMultinodeSupport,
5+
ComputeWithPlacementGroupSupport,
6+
ComputeWithPrivateGatewaySupport,
7+
ComputeWithReservationSupport,
8+
ComputeWithVolumeSupport,
9+
)
10+
from dstack._internal.core.backends.base.configurator import Configurator
11+
from dstack._internal.core.backends.configurators import list_available_configurator_classes
12+
from dstack._internal.core.backends.local.compute import LocalCompute
13+
from dstack._internal.core.models.backends.base import BackendType
14+
from dstack._internal.settings import LOCAL_BACKEND_ENABLED
15+
16+
_configurator_classes = list_available_configurator_classes()
17+
18+
19+
def _get_backends_with_compute_feature(
20+
configurator_classes: list[type[Configurator]],
21+
compute_feature_class: type,
22+
) -> list[BackendType]:
23+
backend_types_and_computes = [
24+
(configurator_class.TYPE, configurator_class.BACKEND_CLASS.COMPUTE_CLASS)
25+
for configurator_class in configurator_classes
26+
]
27+
if LOCAL_BACKEND_ENABLED:
28+
backend_types_and_computes.append((BackendType.LOCAL, LocalCompute))
29+
backend_types = []
30+
for backend_type, compute_class in backend_types_and_computes:
31+
if issubclass(compute_class, compute_feature_class):
32+
backend_types.append(backend_type)
33+
return backend_types
34+
35+
36+
# The following backend lists do not include unavailable backends (i.e. backends missing deps).
37+
BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
38+
configurator_classes=_configurator_classes,
39+
compute_feature_class=ComputeWithCreateInstanceSupport,
40+
)
41+
BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
42+
configurator_classes=_configurator_classes,
43+
compute_feature_class=ComputeWithMultinodeSupport,
44+
)
45+
BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = _get_backends_with_compute_feature(
46+
configurator_classes=_configurator_classes,
47+
compute_feature_class=ComputeWithPlacementGroupSupport,
48+
)
49+
BACKENDS_WITH_RESERVATION_SUPPORT = _get_backends_with_compute_feature(
50+
configurator_classes=_configurator_classes,
51+
compute_feature_class=ComputeWithReservationSupport,
52+
)
53+
BACKENDS_WITH_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
54+
configurator_classes=_configurator_classes,
55+
compute_feature_class=ComputeWithGatewaySupport,
56+
)
57+
BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT = _get_backends_with_compute_feature(
58+
configurator_classes=_configurator_classes,
59+
compute_feature_class=ComputeWithPrivateGatewaySupport,
60+
)
61+
BACKENDS_WITH_VOLUMES_SUPPORT = _get_backends_with_compute_feature(
62+
configurator_classes=_configurator_classes,
63+
compute_feature_class=ComputeWithVolumeSupport,
64+
)

‎src/dstack/_internal/server/background/tasks/process_instances.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,6 @@
1212
from sqlalchemy.orm import joinedload, lazyload
1313

1414
from dstack._internal import settings
15-
from dstack._internal.core.backends import (
16-
BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
17-
BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT,
18-
)
1915
from dstack._internal.core.backends.base.compute import (
2016
ComputeWithCreateInstanceSupport,
2117
ComputeWithPlacementGroupSupport,
@@ -27,6 +23,10 @@
2723
get_shim_env,
2824
get_shim_pre_start_commands,
2925
)
26+
from dstack._internal.core.backends.features import (
27+
BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
28+
BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT,
29+
)
3030
from dstack._internal.core.backends.remote.provisioning import (
3131
detect_cpu_arch,
3232
get_host_info,

‎src/dstack/_internal/server/services/fleets.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88
from sqlalchemy.ext.asyncio import AsyncSession
99
from sqlalchemy.orm import joinedload, selectinload
1010

11-
from dstack._internal.core.backends import BACKENDS_WITH_CREATE_INSTANCE_SUPPORT
1211
from dstack._internal.core.backends.base.backend import Backend
12+
from dstack._internal.core.backends.features import BACKENDS_WITH_CREATE_INSTANCE_SUPPORT
1313
from dstack._internal.core.errors import (
1414
ForbiddenError,
1515
ResourceExistsError,

‎src/dstack/_internal/server/services/gateways/__init__.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,16 @@
1111
from sqlalchemy.orm import selectinload
1212

1313
import dstack._internal.utils.random_names as random_names
14-
from dstack._internal.core.backends import (
15-
BACKENDS_WITH_GATEWAY_SUPPORT,
16-
BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT,
17-
)
1814
from dstack._internal.core.backends.base.compute import (
1915
Compute,
2016
ComputeWithGatewaySupport,
2117
get_dstack_gateway_wheel,
2218
get_dstack_runner_version,
2319
)
20+
from dstack._internal.core.backends.features import (
21+
BACKENDS_WITH_GATEWAY_SUPPORT,
22+
BACKENDS_WITH_PRIVATE_GATEWAY_SUPPORT,
23+
)
2424
from dstack._internal.core.errors import (
2525
GatewayError,
2626
ResourceNotExistsError,

‎src/dstack/_internal/server/services/instances.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,11 @@
88
from sqlalchemy.ext.asyncio import AsyncSession
99
from sqlalchemy.orm import joinedload
1010

11-
from dstack._internal.core.backends import BACKENDS_WITH_MULTINODE_SUPPORT
1211
from dstack._internal.core.backends.base.offers import (
1312
offer_to_catalog_item,
1413
requirements_to_query_filter,
1514
)
15+
from dstack._internal.core.backends.features import BACKENDS_WITH_MULTINODE_SUPPORT
1616
from dstack._internal.core.models.backends.base import BackendType
1717
from dstack._internal.core.models.envs import Env
1818
from dstack._internal.core.models.instances import (

‎src/dstack/_internal/server/services/offers.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,13 @@
22

33
import gpuhunt
44

5-
from dstack._internal.core.backends import (
5+
from dstack._internal.core.backends.base.backend import Backend
6+
from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
7+
from dstack._internal.core.backends.features import (
68
BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
79
BACKENDS_WITH_MULTINODE_SUPPORT,
810
BACKENDS_WITH_RESERVATION_SUPPORT,
911
)
10-
from dstack._internal.core.backends.base.backend import Backend
11-
from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
1212
from dstack._internal.core.models.backends.base import BackendType
1313
from dstack._internal.core.models.instances import (
1414
InstanceOfferWithAvailability,

‎src/dstack/_internal/server/services/volumes.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
from sqlalchemy.ext.asyncio import AsyncSession
77
from sqlalchemy.orm import joinedload, selectinload
88

9-
from dstack._internal.core.backends import BACKENDS_WITH_VOLUMES_SUPPORT
109
from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
10+
from dstack._internal.core.backends.features import BACKENDS_WITH_VOLUMES_SUPPORT
1111
from dstack._internal.core.errors import (
1212
BackendNotAvailable,
1313
ResourceExistsError,

0 commit comments

Comments
 (0)