dstack 0.19.25__py3-none-any.whl → 0.19.26__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of dstack might be problematic. Click here for more details.

Files changed (128) hide show
  1. dstack/_internal/cli/commands/__init__.py +2 -2
  2. dstack/_internal/cli/commands/apply.py +3 -61
  3. dstack/_internal/cli/commands/attach.py +1 -1
  4. dstack/_internal/cli/commands/completion.py +1 -1
  5. dstack/_internal/cli/commands/delete.py +2 -2
  6. dstack/_internal/cli/commands/fleet.py +1 -1
  7. dstack/_internal/cli/commands/gateway.py +2 -2
  8. dstack/_internal/cli/commands/init.py +56 -24
  9. dstack/_internal/cli/commands/logs.py +1 -1
  10. dstack/_internal/cli/commands/metrics.py +1 -1
  11. dstack/_internal/cli/commands/offer.py +45 -7
  12. dstack/_internal/cli/commands/project.py +2 -2
  13. dstack/_internal/cli/commands/secrets.py +2 -2
  14. dstack/_internal/cli/commands/server.py +1 -1
  15. dstack/_internal/cli/commands/stop.py +1 -1
  16. dstack/_internal/cli/commands/volume.py +1 -1
  17. dstack/_internal/cli/main.py +2 -2
  18. dstack/_internal/cli/services/completion.py +2 -2
  19. dstack/_internal/cli/services/configurators/__init__.py +6 -2
  20. dstack/_internal/cli/services/configurators/base.py +6 -7
  21. dstack/_internal/cli/services/configurators/fleet.py +1 -3
  22. dstack/_internal/cli/services/configurators/gateway.py +2 -4
  23. dstack/_internal/cli/services/configurators/run.py +195 -55
  24. dstack/_internal/cli/services/configurators/volume.py +2 -4
  25. dstack/_internal/cli/services/profile.py +1 -1
  26. dstack/_internal/cli/services/repos.py +51 -47
  27. dstack/_internal/core/backends/aws/configurator.py +11 -7
  28. dstack/_internal/core/backends/azure/configurator.py +11 -7
  29. dstack/_internal/core/backends/base/configurator.py +25 -13
  30. dstack/_internal/core/backends/cloudrift/configurator.py +13 -7
  31. dstack/_internal/core/backends/cudo/configurator.py +11 -7
  32. dstack/_internal/core/backends/datacrunch/compute.py +5 -1
  33. dstack/_internal/core/backends/datacrunch/configurator.py +13 -7
  34. dstack/_internal/core/backends/gcp/configurator.py +11 -7
  35. dstack/_internal/core/backends/hotaisle/configurator.py +13 -7
  36. dstack/_internal/core/backends/kubernetes/configurator.py +13 -7
  37. dstack/_internal/core/backends/lambdalabs/configurator.py +11 -7
  38. dstack/_internal/core/backends/nebius/compute.py +1 -1
  39. dstack/_internal/core/backends/nebius/configurator.py +11 -7
  40. dstack/_internal/core/backends/nebius/resources.py +21 -11
  41. dstack/_internal/core/backends/oci/configurator.py +11 -7
  42. dstack/_internal/core/backends/runpod/configurator.py +11 -7
  43. dstack/_internal/core/backends/template/configurator.py.jinja +11 -7
  44. dstack/_internal/core/backends/tensordock/configurator.py +13 -7
  45. dstack/_internal/core/backends/vastai/configurator.py +11 -7
  46. dstack/_internal/core/backends/vultr/configurator.py +11 -4
  47. dstack/_internal/core/compatibility/gpus.py +13 -0
  48. dstack/_internal/core/compatibility/runs.py +1 -0
  49. dstack/_internal/core/models/common.py +3 -3
  50. dstack/_internal/core/models/configurations.py +172 -27
  51. dstack/_internal/core/models/files.py +1 -1
  52. dstack/_internal/core/models/fleets.py +5 -1
  53. dstack/_internal/core/models/profiles.py +41 -11
  54. dstack/_internal/core/models/resources.py +46 -42
  55. dstack/_internal/core/models/runs.py +4 -0
  56. dstack/_internal/core/services/configs/__init__.py +2 -2
  57. dstack/_internal/core/services/profiles.py +2 -2
  58. dstack/_internal/core/services/repos.py +5 -3
  59. dstack/_internal/core/services/ssh/ports.py +1 -1
  60. dstack/_internal/proxy/lib/deps.py +6 -2
  61. dstack/_internal/server/app.py +22 -17
  62. dstack/_internal/server/background/tasks/process_gateways.py +4 -1
  63. dstack/_internal/server/background/tasks/process_instances.py +10 -2
  64. dstack/_internal/server/background/tasks/process_probes.py +1 -1
  65. dstack/_internal/server/background/tasks/process_running_jobs.py +10 -4
  66. dstack/_internal/server/background/tasks/process_runs.py +1 -1
  67. dstack/_internal/server/background/tasks/process_submitted_jobs.py +54 -43
  68. dstack/_internal/server/background/tasks/process_terminating_jobs.py +2 -2
  69. dstack/_internal/server/background/tasks/process_volumes.py +1 -1
  70. dstack/_internal/server/db.py +8 -4
  71. dstack/_internal/server/models.py +1 -0
  72. dstack/_internal/server/routers/gpus.py +1 -6
  73. dstack/_internal/server/schemas/runner.py +10 -0
  74. dstack/_internal/server/services/backends/__init__.py +14 -8
  75. dstack/_internal/server/services/backends/handlers.py +6 -1
  76. dstack/_internal/server/services/docker.py +5 -5
  77. dstack/_internal/server/services/fleets.py +14 -13
  78. dstack/_internal/server/services/gateways/__init__.py +2 -0
  79. dstack/_internal/server/services/gateways/client.py +5 -2
  80. dstack/_internal/server/services/gateways/connection.py +1 -1
  81. dstack/_internal/server/services/gpus.py +50 -49
  82. dstack/_internal/server/services/instances.py +41 -1
  83. dstack/_internal/server/services/jobs/__init__.py +15 -4
  84. dstack/_internal/server/services/jobs/configurators/base.py +7 -11
  85. dstack/_internal/server/services/jobs/configurators/dev.py +5 -0
  86. dstack/_internal/server/services/jobs/configurators/extensions/cursor.py +3 -3
  87. dstack/_internal/server/services/jobs/configurators/extensions/vscode.py +3 -3
  88. dstack/_internal/server/services/jobs/configurators/service.py +1 -0
  89. dstack/_internal/server/services/jobs/configurators/task.py +3 -0
  90. dstack/_internal/server/services/locking.py +5 -5
  91. dstack/_internal/server/services/logging.py +10 -2
  92. dstack/_internal/server/services/logs/__init__.py +8 -6
  93. dstack/_internal/server/services/logs/aws.py +330 -327
  94. dstack/_internal/server/services/logs/filelog.py +7 -6
  95. dstack/_internal/server/services/logs/gcp.py +141 -139
  96. dstack/_internal/server/services/plugins.py +1 -1
  97. dstack/_internal/server/services/projects.py +2 -5
  98. dstack/_internal/server/services/proxy/repo.py +5 -1
  99. dstack/_internal/server/services/requirements/__init__.py +0 -0
  100. dstack/_internal/server/services/requirements/combine.py +259 -0
  101. dstack/_internal/server/services/runner/client.py +7 -0
  102. dstack/_internal/server/services/runs.py +1 -1
  103. dstack/_internal/server/services/services/__init__.py +8 -2
  104. dstack/_internal/server/services/services/autoscalers.py +2 -0
  105. dstack/_internal/server/services/ssh.py +2 -1
  106. dstack/_internal/server/services/storage/__init__.py +5 -6
  107. dstack/_internal/server/services/storage/gcs.py +49 -49
  108. dstack/_internal/server/services/storage/s3.py +52 -52
  109. dstack/_internal/server/statics/index.html +1 -1
  110. dstack/_internal/server/testing/common.py +1 -1
  111. dstack/_internal/server/utils/logging.py +3 -3
  112. dstack/_internal/server/utils/provisioning.py +3 -3
  113. dstack/_internal/utils/json_schema.py +3 -1
  114. dstack/_internal/utils/typing.py +14 -0
  115. dstack/api/_public/repos.py +21 -2
  116. dstack/api/_public/runs.py +5 -7
  117. dstack/api/server/__init__.py +17 -19
  118. dstack/api/server/_gpus.py +2 -1
  119. dstack/api/server/_group.py +4 -3
  120. dstack/api/server/_repos.py +20 -3
  121. dstack/plugins/builtin/rest_plugin/_plugin.py +1 -0
  122. dstack/version.py +1 -1
  123. {dstack-0.19.25.dist-info → dstack-0.19.26.dist-info}/METADATA +1 -1
  124. {dstack-0.19.25.dist-info → dstack-0.19.26.dist-info}/RECORD +127 -124
  125. dstack/api/huggingface/__init__.py +0 -73
  126. {dstack-0.19.25.dist-info → dstack-0.19.26.dist-info}/WHEEL +0 -0
  127. {dstack-0.19.25.dist-info → dstack-0.19.26.dist-info}/entry_points.txt +0 -0
  128. {dstack-0.19.25.dist-info → dstack-0.19.26.dist-info}/licenses/LICENSE.md +0 -0
@@ -1,11 +1,11 @@
1
1
  from abc import ABC, abstractmethod
2
- from typing import Any, ClassVar, List, Optional
2
+ from typing import Any, ClassVar, Generic, List, Optional, TypeVar
3
3
  from uuid import UUID
4
4
 
5
5
  from dstack._internal.core.backends.base.backend import Backend
6
6
  from dstack._internal.core.backends.models import (
7
- AnyBackendConfig,
8
7
  AnyBackendConfigWithCreds,
8
+ AnyBackendConfigWithoutCreds,
9
9
  )
10
10
  from dstack._internal.core.errors import BackendInvalidCredentialsError
11
11
  from dstack._internal.core.models.backends.base import BackendType
@@ -15,6 +15,11 @@ from dstack._internal.core.models.common import CoreModel
15
15
  # We'll introduce our own base limit that can be customized per backend if required.
16
16
  TAGS_MAX_NUM = 25
17
17
 
18
+ BackendConfigWithoutCredsT = TypeVar(
19
+ "BackendConfigWithoutCredsT", bound=AnyBackendConfigWithoutCreds
20
+ )
21
+ BackendConfigWithCredsT = TypeVar("BackendConfigWithCredsT", bound=AnyBackendConfigWithCreds)
22
+
18
23
 
19
24
  class BackendRecord(CoreModel):
20
25
  """
@@ -39,7 +44,7 @@ class StoredBackendRecord(BackendRecord):
39
44
  backend_id: UUID
40
45
 
41
46
 
42
- class Configurator(ABC):
47
+ class Configurator(ABC, Generic[BackendConfigWithoutCredsT, BackendConfigWithCredsT]):
43
48
  """
44
49
  `Configurator` is responsible for configuring backends
45
50
  and initializing `Backend` instances from backend configs.
@@ -52,7 +57,7 @@ class Configurator(ABC):
52
57
  BACKEND_CLASS: ClassVar[type[Backend]]
53
58
 
54
59
  @abstractmethod
55
- def validate_config(self, config: AnyBackendConfigWithCreds, default_creds_enabled: bool):
60
+ def validate_config(self, config: BackendConfigWithCredsT, default_creds_enabled: bool):
56
61
  """
57
62
  Validates backend config including backend creds and other parameters.
58
63
  Raises `ServerClientError` or its subclass if config is invalid.
@@ -61,9 +66,7 @@ class Configurator(ABC):
61
66
  pass
62
67
 
63
68
  @abstractmethod
64
- def create_backend(
65
- self, project_name: str, config: AnyBackendConfigWithCreds
66
- ) -> BackendRecord:
69
+ def create_backend(self, project_name: str, config: BackendConfigWithCredsT) -> BackendRecord:
67
70
  """
68
71
  Sets up backend given backend config and returns
69
72
  text-encoded config and creds to be stored in the DB.
@@ -78,13 +81,22 @@ class Configurator(ABC):
78
81
  pass
79
82
 
80
83
  @abstractmethod
81
- def get_backend_config(
82
- self, record: StoredBackendRecord, include_creds: bool
83
- ) -> AnyBackendConfig:
84
+ def get_backend_config_with_creds(
85
+ self, record: StoredBackendRecord
86
+ ) -> BackendConfigWithCredsT:
87
+ """
88
+ Constructs `BackendConfig` with credentials included.
89
+ Used internally and when project admins need to see backend's creds.
90
+ """
91
+ pass
92
+
93
+ @abstractmethod
94
+ def get_backend_config_without_creds(
95
+ self, record: StoredBackendRecord
96
+ ) -> BackendConfigWithoutCredsT:
84
97
  """
85
- Constructs `BackendConfig` to be returned in API responses.
86
- Project admins may need to see backend's creds. In this case `include_creds` will be `True`.
87
- Otherwise, no sensitive information should be included.
98
+ Constructs `BackendConfig` without sensitive information.
99
+ Used for API responses where creds should not be exposed.
88
100
  """
89
101
  pass
90
102
 
@@ -8,7 +8,6 @@ from dstack._internal.core.backends.base.configurator import (
8
8
  from dstack._internal.core.backends.cloudrift.api_client import RiftClient
9
9
  from dstack._internal.core.backends.cloudrift.backend import CloudRiftBackend
10
10
  from dstack._internal.core.backends.cloudrift.models import (
11
- AnyCloudRiftBackendConfig,
12
11
  AnyCloudRiftCreds,
13
12
  CloudRiftBackendConfig,
14
13
  CloudRiftBackendConfigWithCreds,
@@ -21,7 +20,12 @@ from dstack._internal.core.models.backends.base import (
21
20
  )
22
21
 
23
22
 
24
- class CloudRiftConfigurator(Configurator):
23
+ class CloudRiftConfigurator(
24
+ Configurator[
25
+ CloudRiftBackendConfig,
26
+ CloudRiftBackendConfigWithCreds,
27
+ ]
28
+ ):
25
29
  TYPE = BackendType.CLOUDRIFT
26
30
  BACKEND_CLASS = CloudRiftBackend
27
31
 
@@ -40,12 +44,14 @@ class CloudRiftConfigurator(Configurator):
40
44
  auth=CloudRiftCreds.parse_obj(config.creds).json(),
41
45
  )
42
46
 
43
- def get_backend_config(
44
- self, record: BackendRecord, include_creds: bool
45
- ) -> AnyCloudRiftBackendConfig:
47
+ def get_backend_config_with_creds(
48
+ self, record: BackendRecord
49
+ ) -> CloudRiftBackendConfigWithCreds:
50
+ config = self._get_config(record)
51
+ return CloudRiftBackendConfigWithCreds.__response__.parse_obj(config)
52
+
53
+ def get_backend_config_without_creds(self, record: BackendRecord) -> CloudRiftBackendConfig:
46
54
  config = self._get_config(record)
47
- if include_creds:
48
- return CloudRiftBackendConfigWithCreds.__response__.parse_obj(config)
49
55
  return CloudRiftBackendConfig.__response__.parse_obj(config)
50
56
 
51
57
  def get_backend(self, record: BackendRecord) -> CloudRiftBackend:
@@ -8,7 +8,6 @@ from dstack._internal.core.backends.base.configurator import (
8
8
  from dstack._internal.core.backends.cudo import api_client
9
9
  from dstack._internal.core.backends.cudo.backend import CudoBackend
10
10
  from dstack._internal.core.backends.cudo.models import (
11
- AnyCudoBackendConfig,
12
11
  CudoBackendConfig,
13
12
  CudoBackendConfigWithCreds,
14
13
  CudoConfig,
@@ -18,7 +17,12 @@ from dstack._internal.core.backends.cudo.models import (
18
17
  from dstack._internal.core.models.backends.base import BackendType
19
18
 
20
19
 
21
- class CudoConfigurator(Configurator):
20
+ class CudoConfigurator(
21
+ Configurator[
22
+ CudoBackendConfig,
23
+ CudoBackendConfigWithCreds,
24
+ ]
25
+ ):
22
26
  TYPE = BackendType.CUDO
23
27
  BACKEND_CLASS = CudoBackend
24
28
 
@@ -35,12 +39,12 @@ class CudoConfigurator(Configurator):
35
39
  auth=CudoCreds.parse_obj(config.creds).json(),
36
40
  )
37
41
 
38
- def get_backend_config(
39
- self, record: BackendRecord, include_creds: bool
40
- ) -> AnyCudoBackendConfig:
42
+ def get_backend_config_with_creds(self, record: BackendRecord) -> CudoBackendConfigWithCreds:
43
+ config = self._get_config(record)
44
+ return CudoBackendConfigWithCreds.__response__.parse_obj(config)
45
+
46
+ def get_backend_config_without_creds(self, record: BackendRecord) -> CudoBackendConfig:
41
47
  config = self._get_config(record)
42
- if include_creds:
43
- return CudoBackendConfigWithCreds.__response__.parse_obj(config)
44
48
  return CudoBackendConfig.__response__.parse_obj(config)
45
49
 
46
50
  def get_backend(self, record: BackendRecord) -> CudoBackend:
@@ -161,7 +161,10 @@ class DataCrunchCompute(
161
161
  try:
162
162
  self.client.instances.action(id_list=[instance_id], action="delete")
163
163
  except APIException as e:
164
- if e.message == "Invalid instance id":
164
+ if e.message in [
165
+ "Invalid instance id",
166
+ "Can't discontinue a discontinued instance",
167
+ ]:
165
168
  logger.debug("Skipping instance %s termination. Instance not found.", instance_id)
166
169
  return
167
170
  raise
@@ -243,6 +246,7 @@ def _deploy_instance(
243
246
  hostname=hostname,
244
247
  description=description,
245
248
  startup_script_id=startup_script_id,
249
+ pricing="FIXED_PRICE",
246
250
  is_spot=is_spot,
247
251
  location=location,
248
252
  os_volume={"name": "OS volume", "size": disk_size},
@@ -10,7 +10,6 @@ from dstack._internal.core.backends.base.configurator import (
10
10
  )
11
11
  from dstack._internal.core.backends.datacrunch.backend import DataCrunchBackend
12
12
  from dstack._internal.core.backends.datacrunch.models import (
13
- AnyDataCrunchBackendConfig,
14
13
  DataCrunchBackendConfig,
15
14
  DataCrunchBackendConfigWithCreds,
16
15
  DataCrunchConfig,
@@ -22,7 +21,12 @@ from dstack._internal.core.models.backends.base import (
22
21
  )
23
22
 
24
23
 
25
- class DataCrunchConfigurator(Configurator):
24
+ class DataCrunchConfigurator(
25
+ Configurator[
26
+ DataCrunchBackendConfig,
27
+ DataCrunchBackendConfigWithCreds,
28
+ ]
29
+ ):
26
30
  TYPE = BackendType.DATACRUNCH
27
31
  BACKEND_CLASS = DataCrunchBackend
28
32
 
@@ -41,12 +45,14 @@ class DataCrunchConfigurator(Configurator):
41
45
  auth=DataCrunchCreds.parse_obj(config.creds).json(),
42
46
  )
43
47
 
44
- def get_backend_config(
45
- self, record: BackendRecord, include_creds: bool
46
- ) -> AnyDataCrunchBackendConfig:
48
+ def get_backend_config_with_creds(
49
+ self, record: BackendRecord
50
+ ) -> DataCrunchBackendConfigWithCreds:
51
+ config = self._get_config(record)
52
+ return DataCrunchBackendConfigWithCreds.__response__.parse_obj(config)
53
+
54
+ def get_backend_config_without_creds(self, record: BackendRecord) -> DataCrunchBackendConfig:
47
55
  config = self._get_config(record)
48
- if include_creds:
49
- return DataCrunchBackendConfigWithCreds.__response__.parse_obj(config)
50
56
  return DataCrunchBackendConfig.__response__.parse_obj(config)
51
57
 
52
58
  def get_backend(self, record: BackendRecord) -> DataCrunchBackend:
@@ -11,7 +11,6 @@ from dstack._internal.core.backends.base.configurator import (
11
11
  from dstack._internal.core.backends.gcp import auth, resources
12
12
  from dstack._internal.core.backends.gcp.backend import GCPBackend
13
13
  from dstack._internal.core.backends.gcp.models import (
14
- AnyGCPBackendConfig,
15
14
  GCPBackendConfig,
16
15
  GCPBackendConfigWithCreds,
17
16
  GCPConfig,
@@ -109,7 +108,12 @@ DEFAULT_REGIONS = REGIONS
109
108
  MAIN_REGION = "us-east1"
110
109
 
111
110
 
112
- class GCPConfigurator(Configurator):
111
+ class GCPConfigurator(
112
+ Configurator[
113
+ GCPBackendConfig,
114
+ GCPBackendConfigWithCreds,
115
+ ]
116
+ ):
113
117
  TYPE = BackendType.GCP
114
118
  BACKEND_CLASS = GCPBackend
115
119
 
@@ -147,12 +151,12 @@ class GCPConfigurator(Configurator):
147
151
  auth=GCPCreds.parse_obj(config.creds).json(),
148
152
  )
149
153
 
150
- def get_backend_config(
151
- self, record: BackendRecord, include_creds: bool
152
- ) -> AnyGCPBackendConfig:
154
+ def get_backend_config_with_creds(self, record: BackendRecord) -> GCPBackendConfigWithCreds:
155
+ config = self._get_config(record)
156
+ return GCPBackendConfigWithCreds.__response__.parse_obj(config)
157
+
158
+ def get_backend_config_without_creds(self, record: BackendRecord) -> GCPBackendConfig:
153
159
  config = self._get_config(record)
154
- if include_creds:
155
- return GCPBackendConfigWithCreds.__response__.parse_obj(config)
156
160
  return GCPBackendConfig.__response__.parse_obj(config)
157
161
 
158
162
  def get_backend(self, record: BackendRecord) -> GCPBackend:
@@ -7,7 +7,6 @@ from dstack._internal.core.backends.base.configurator import (
7
7
  from dstack._internal.core.backends.hotaisle.api_client import HotAisleAPIClient
8
8
  from dstack._internal.core.backends.hotaisle.backend import HotAisleBackend
9
9
  from dstack._internal.core.backends.hotaisle.models import (
10
- AnyHotAisleBackendConfig,
11
10
  AnyHotAisleCreds,
12
11
  HotAisleBackendConfig,
13
12
  HotAisleBackendConfigWithCreds,
@@ -20,7 +19,12 @@ from dstack._internal.core.models.backends.base import (
20
19
  )
21
20
 
22
21
 
23
- class HotAisleConfigurator(Configurator):
22
+ class HotAisleConfigurator(
23
+ Configurator[
24
+ HotAisleBackendConfig,
25
+ HotAisleBackendConfigWithCreds,
26
+ ]
27
+ ):
24
28
  TYPE = BackendType.HOTAISLE
25
29
  BACKEND_CLASS = HotAisleBackend
26
30
 
@@ -37,12 +41,14 @@ class HotAisleConfigurator(Configurator):
37
41
  auth=HotAisleCreds.parse_obj(config.creds).json(),
38
42
  )
39
43
 
40
- def get_backend_config(
41
- self, record: BackendRecord, include_creds: bool
42
- ) -> AnyHotAisleBackendConfig:
44
+ def get_backend_config_with_creds(
45
+ self, record: BackendRecord
46
+ ) -> HotAisleBackendConfigWithCreds:
47
+ config = self._get_config(record)
48
+ return HotAisleBackendConfigWithCreds.__response__.parse_obj(config)
49
+
50
+ def get_backend_config_without_creds(self, record: BackendRecord) -> HotAisleBackendConfig:
43
51
  config = self._get_config(record)
44
- if include_creds:
45
- return HotAisleBackendConfigWithCreds.__response__.parse_obj(config)
46
52
  return HotAisleBackendConfig.__response__.parse_obj(config)
47
53
 
48
54
  def get_backend(self, record: BackendRecord) -> HotAisleBackend:
@@ -6,7 +6,6 @@ from dstack._internal.core.backends.base.configurator import (
6
6
  from dstack._internal.core.backends.kubernetes import utils as kubernetes_utils
7
7
  from dstack._internal.core.backends.kubernetes.backend import KubernetesBackend
8
8
  from dstack._internal.core.backends.kubernetes.models import (
9
- AnyKubernetesBackendConfig,
10
9
  KubernetesBackendConfig,
11
10
  KubernetesBackendConfigWithCreds,
12
11
  KubernetesConfig,
@@ -18,7 +17,12 @@ from dstack._internal.utils.logging import get_logger
18
17
  logger = get_logger(__name__)
19
18
 
20
19
 
21
- class KubernetesConfigurator(Configurator):
20
+ class KubernetesConfigurator(
21
+ Configurator[
22
+ KubernetesBackendConfig,
23
+ KubernetesBackendConfigWithCreds,
24
+ ]
25
+ ):
22
26
  TYPE = BackendType.KUBERNETES
23
27
  BACKEND_CLASS = KubernetesBackend
24
28
 
@@ -40,12 +44,14 @@ class KubernetesConfigurator(Configurator):
40
44
  auth="",
41
45
  )
42
46
 
43
- def get_backend_config(
44
- self, record: BackendRecord, include_creds: bool
45
- ) -> AnyKubernetesBackendConfig:
47
+ def get_backend_config_with_creds(
48
+ self, record: BackendRecord
49
+ ) -> KubernetesBackendConfigWithCreds:
50
+ config = self._get_config(record)
51
+ return KubernetesBackendConfigWithCreds.__response__.parse_obj(config)
52
+
53
+ def get_backend_config_without_creds(self, record: BackendRecord) -> KubernetesBackendConfig:
46
54
  config = self._get_config(record)
47
- if include_creds:
48
- return KubernetesBackendConfigWithCreds.__response__.parse_obj(config)
49
55
  return KubernetesBackendConfig.__response__.parse_obj(config)
50
56
 
51
57
  def get_backend(self, record: BackendRecord) -> KubernetesBackend:
@@ -8,7 +8,6 @@ from dstack._internal.core.backends.base.configurator import (
8
8
  from dstack._internal.core.backends.lambdalabs import api_client
9
9
  from dstack._internal.core.backends.lambdalabs.backend import LambdaBackend
10
10
  from dstack._internal.core.backends.lambdalabs.models import (
11
- AnyLambdaBackendConfig,
12
11
  LambdaBackendConfig,
13
12
  LambdaBackendConfigWithCreds,
14
13
  LambdaConfig,
@@ -20,7 +19,12 @@ from dstack._internal.core.models.backends.base import (
20
19
  )
21
20
 
22
21
 
23
- class LambdaConfigurator(Configurator):
22
+ class LambdaConfigurator(
23
+ Configurator[
24
+ LambdaBackendConfig,
25
+ LambdaBackendConfigWithCreds,
26
+ ]
27
+ ):
24
28
  TYPE = BackendType.LAMBDA
25
29
  BACKEND_CLASS = LambdaBackend
26
30
 
@@ -37,12 +41,12 @@ class LambdaConfigurator(Configurator):
37
41
  auth=LambdaCreds.parse_obj(config.creds).json(),
38
42
  )
39
43
 
40
- def get_backend_config(
41
- self, record: BackendRecord, include_creds: bool
42
- ) -> AnyLambdaBackendConfig:
44
+ def get_backend_config_with_creds(self, record: BackendRecord) -> LambdaBackendConfigWithCreds:
45
+ config = self._get_config(record)
46
+ return LambdaBackendConfigWithCreds.__response__.parse_obj(config)
47
+
48
+ def get_backend_config_without_creds(self, record: BackendRecord) -> LambdaBackendConfig:
43
49
  config = self._get_config(record)
44
- if include_creds:
45
- return LambdaBackendConfigWithCreds.__response__.parse_obj(config)
46
50
  return LambdaBackendConfig.__response__.parse_obj(config)
47
51
 
48
52
  def get_backend(self, record: BackendRecord) -> LambdaBackend:
@@ -364,7 +364,7 @@ def _wait_for_instance(sdk: SDK, op: SDKOperation[Operation]) -> None:
364
364
  )
365
365
  time.sleep(WAIT_FOR_INSTANCE_UPDATE_INTERVAL)
366
366
  resources.LOOP.await_(
367
- op.update(timeout=resources.REQUEST_TIMEOUT, metadata=resources.REQUEST_MD)
367
+ op.update(per_retry_timeout=resources.REQUEST_TIMEOUT, metadata=resources.REQUEST_MD)
368
368
  )
369
369
 
370
370
 
@@ -11,7 +11,6 @@ from dstack._internal.core.backends.nebius import resources
11
11
  from dstack._internal.core.backends.nebius.backend import NebiusBackend
12
12
  from dstack._internal.core.backends.nebius.fabrics import get_all_infiniband_fabrics
13
13
  from dstack._internal.core.backends.nebius.models import (
14
- AnyNebiusBackendConfig,
15
14
  NebiusBackendConfig,
16
15
  NebiusBackendConfigWithCreds,
17
16
  NebiusConfig,
@@ -22,7 +21,12 @@ from dstack._internal.core.backends.nebius.models import (
22
21
  from dstack._internal.core.models.backends.base import BackendType
23
22
 
24
23
 
25
- class NebiusConfigurator(Configurator):
24
+ class NebiusConfigurator(
25
+ Configurator[
26
+ NebiusBackendConfig,
27
+ NebiusBackendConfigWithCreds,
28
+ ]
29
+ ):
26
30
  TYPE = BackendType.NEBIUS
27
31
  BACKEND_CLASS = NebiusBackend
28
32
 
@@ -60,12 +64,12 @@ class NebiusConfigurator(Configurator):
60
64
  auth=NebiusCreds.parse_obj(config.creds).json(),
61
65
  )
62
66
 
63
- def get_backend_config(
64
- self, record: BackendRecord, include_creds: bool
65
- ) -> AnyNebiusBackendConfig:
67
+ def get_backend_config_with_creds(self, record: BackendRecord) -> NebiusBackendConfigWithCreds:
68
+ config = self._get_config(record)
69
+ return NebiusBackendConfigWithCreds.__response__.parse_obj(config)
70
+
71
+ def get_backend_config_without_creds(self, record: BackendRecord) -> NebiusBackendConfig:
66
72
  config = self._get_config(record)
67
- if include_creds:
68
- return NebiusBackendConfigWithCreds.__response__.parse_obj(config)
69
73
  return NebiusBackendConfig.__response__.parse_obj(config)
70
74
 
71
75
  def get_backend(self, record: BackendRecord) -> NebiusBackend:
@@ -119,7 +119,7 @@ def wait_for_operation(
119
119
  if time.monotonic() + interval > deadline:
120
120
  raise TimeoutError(f"Operation {op.id} wait timeout")
121
121
  time.sleep(interval)
122
- LOOP.await_(op.update(timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
122
+ LOOP.await_(op.update(per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
123
123
 
124
124
 
125
125
  def get_region_to_project_id_map(
@@ -155,7 +155,7 @@ def validate_regions(configured: set[str], available: set[str]) -> None:
155
155
  def list_tenant_projects(sdk: SDK) -> Sequence[Container]:
156
156
  tenants = LOOP.await_(
157
157
  TenantServiceClient(sdk).list(
158
- ListTenantsRequest(), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
158
+ ListTenantsRequest(), per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
159
159
  )
160
160
  )
161
161
  if len(tenants.items) != 1:
@@ -164,7 +164,7 @@ def list_tenant_projects(sdk: SDK) -> Sequence[Container]:
164
164
  projects = LOOP.await_(
165
165
  ProjectServiceClient(sdk).list(
166
166
  ListProjectsRequest(parent_id=tenant_id, page_size=999),
167
- timeout=REQUEST_TIMEOUT,
167
+ per_retry_timeout=REQUEST_TIMEOUT,
168
168
  metadata=REQUEST_MD,
169
169
  )
170
170
  )
@@ -238,7 +238,7 @@ def get_default_subnet(sdk: SDK, project_id: str) -> Subnet:
238
238
  subnets = LOOP.await_(
239
239
  SubnetServiceClient(sdk).list(
240
240
  ListSubnetsRequest(parent_id=project_id, page_size=999),
241
- timeout=REQUEST_TIMEOUT,
241
+ per_retry_timeout=REQUEST_TIMEOUT,
242
242
  metadata=REQUEST_MD,
243
243
  )
244
244
  )
@@ -264,13 +264,15 @@ def create_disk(
264
264
  ),
265
265
  )
266
266
  with wrap_capacity_errors():
267
- return LOOP.await_(client.create(request, timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
267
+ return LOOP.await_(
268
+ client.create(request, per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD)
269
+ )
268
270
 
269
271
 
270
272
  def delete_disk(sdk: SDK, disk_id: str) -> None:
271
273
  LOOP.await_(
272
274
  DiskServiceClient(sdk).delete(
273
- DeleteDiskRequest(id=disk_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
275
+ DeleteDiskRequest(id=disk_id), per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
274
276
  )
275
277
  )
276
278
 
@@ -318,13 +320,17 @@ def create_instance(
318
320
  ),
319
321
  )
320
322
  with wrap_capacity_errors():
321
- return LOOP.await_(client.create(request, timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
323
+ return LOOP.await_(
324
+ client.create(request, per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD)
325
+ )
322
326
 
323
327
 
324
328
  def get_instance(sdk: SDK, instance_id: str) -> Instance:
325
329
  return LOOP.await_(
326
330
  InstanceServiceClient(sdk).get(
327
- GetInstanceRequest(id=instance_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
331
+ GetInstanceRequest(id=instance_id),
332
+ per_retry_timeout=REQUEST_TIMEOUT,
333
+ metadata=REQUEST_MD,
328
334
  )
329
335
  )
330
336
 
@@ -332,7 +338,9 @@ def get_instance(sdk: SDK, instance_id: str) -> Instance:
332
338
  def delete_instance(sdk: SDK, instance_id: str) -> SDKOperation[Operation]:
333
339
  return LOOP.await_(
334
340
  InstanceServiceClient(sdk).delete(
335
- DeleteInstanceRequest(id=instance_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
341
+ DeleteInstanceRequest(id=instance_id),
342
+ per_retry_timeout=REQUEST_TIMEOUT,
343
+ metadata=REQUEST_MD,
336
344
  )
337
345
  )
338
346
 
@@ -345,7 +353,7 @@ def create_cluster(sdk: SDK, name: str, project_id: str, fabric: str) -> SDKOper
345
353
  metadata=ResourceMetadata(name=name, parent_id=project_id),
346
354
  spec=GpuClusterSpec(infiniband_fabric=fabric),
347
355
  ),
348
- timeout=REQUEST_TIMEOUT,
356
+ per_retry_timeout=REQUEST_TIMEOUT,
349
357
  metadata=REQUEST_MD,
350
358
  )
351
359
  )
@@ -354,6 +362,8 @@ def create_cluster(sdk: SDK, name: str, project_id: str, fabric: str) -> SDKOper
354
362
  def delete_cluster(sdk: SDK, cluster_id: str) -> None:
355
363
  return LOOP.await_(
356
364
  GpuClusterServiceClient(sdk).delete(
357
- DeleteGpuClusterRequest(id=cluster_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
365
+ DeleteGpuClusterRequest(id=cluster_id),
366
+ per_retry_timeout=REQUEST_TIMEOUT,
367
+ metadata=REQUEST_MD,
358
368
  )
359
369
  )
@@ -10,7 +10,6 @@ from dstack._internal.core.backends.oci import resources
10
10
  from dstack._internal.core.backends.oci.backend import OCIBackend
11
11
  from dstack._internal.core.backends.oci.exceptions import any_oci_exception
12
12
  from dstack._internal.core.backends.oci.models import (
13
- AnyOCIBackendConfig,
14
13
  OCIBackendConfig,
15
14
  OCIBackendConfigWithCreds,
16
15
  OCIConfig,
@@ -42,7 +41,12 @@ SUPPORTED_REGIONS = frozenset(
42
41
  )
43
42
 
44
43
 
45
- class OCIConfigurator(Configurator):
44
+ class OCIConfigurator(
45
+ Configurator[
46
+ OCIBackendConfig,
47
+ OCIBackendConfigWithCreds,
48
+ ]
49
+ ):
46
50
  TYPE = BackendType.OCI
47
51
  BACKEND_CLASS = OCIBackend
48
52
 
@@ -83,12 +87,12 @@ class OCIConfigurator(Configurator):
83
87
  auth=OCICreds.parse_obj(config.creds).json(),
84
88
  )
85
89
 
86
- def get_backend_config(
87
- self, record: BackendRecord, include_creds: bool
88
- ) -> AnyOCIBackendConfig:
90
+ def get_backend_config_with_creds(self, record: BackendRecord) -> OCIBackendConfigWithCreds:
91
+ config = self._get_config(record)
92
+ return OCIBackendConfigWithCreds.__response__.parse_obj(config)
93
+
94
+ def get_backend_config_without_creds(self, record: BackendRecord) -> OCIBackendConfig:
89
95
  config = self._get_config(record)
90
- if include_creds:
91
- return OCIBackendConfigWithCreds.__response__.parse_obj(config)
92
96
  return OCIBackendConfig.__response__.parse_obj(config)
93
97
 
94
98
  def get_backend(self, record: BackendRecord) -> OCIBackend:
@@ -8,7 +8,6 @@ from dstack._internal.core.backends.base.configurator import (
8
8
  from dstack._internal.core.backends.runpod import api_client
9
9
  from dstack._internal.core.backends.runpod.backend import RunpodBackend
10
10
  from dstack._internal.core.backends.runpod.models import (
11
- AnyRunpodBackendConfig,
12
11
  RunpodBackendConfig,
13
12
  RunpodBackendConfigWithCreds,
14
13
  RunpodConfig,
@@ -18,7 +17,12 @@ from dstack._internal.core.backends.runpod.models import (
18
17
  from dstack._internal.core.models.backends.base import BackendType
19
18
 
20
19
 
21
- class RunpodConfigurator(Configurator):
20
+ class RunpodConfigurator(
21
+ Configurator[
22
+ RunpodBackendConfig,
23
+ RunpodBackendConfigWithCreds,
24
+ ]
25
+ ):
22
26
  TYPE = BackendType.RUNPOD
23
27
  BACKEND_CLASS = RunpodBackend
24
28
 
@@ -35,12 +39,12 @@ class RunpodConfigurator(Configurator):
35
39
  auth=RunpodCreds.parse_obj(config.creds).json(),
36
40
  )
37
41
 
38
- def get_backend_config(
39
- self, record: BackendRecord, include_creds: bool
40
- ) -> AnyRunpodBackendConfig:
42
+ def get_backend_config_with_creds(self, record: BackendRecord) -> RunpodBackendConfigWithCreds:
43
+ config = self._get_config(record)
44
+ return RunpodBackendConfigWithCreds.__response__.parse_obj(config)
45
+
46
+ def get_backend_config_without_creds(self, record: BackendRecord) -> RunpodBackendConfig:
41
47
  config = self._get_config(record)
42
- if include_creds:
43
- return RunpodBackendConfigWithCreds.__response__.parse_obj(config)
44
48
  return RunpodBackendConfig.__response__.parse_obj(config)
45
49
 
46
50
  def get_backend(self, record: BackendRecord) -> RunpodBackend:
@@ -7,7 +7,6 @@ from dstack._internal.core.backends.base.configurator import (
7
7
  )
8
8
  from dstack._internal.core.backends.{{ backend_name|lower }}.backend import {{ backend_name }}Backend
9
9
  from dstack._internal.core.backends.{{ backend_name|lower }}.models import (
10
- Any{{ backend_name }}BackendConfig,
11
10
  Any{{ backend_name }}Creds,
12
11
  {{ backend_name }}BackendConfig,
13
12
  {{ backend_name }}BackendConfigWithCreds,
@@ -20,7 +19,12 @@ from dstack._internal.core.models.backends.base import (
20
19
  )
21
20
 
22
21
 
23
- class {{ backend_name }}Configurator(Configurator):
22
+ class {{ backend_name }}Configurator(
23
+ Configurator[
24
+ {{ backend_name }}BackendConfig,
25
+ {{ backend_name }}BackendConfigWithCreds,
26
+ ]
27
+ ):
24
28
  TYPE = BackendType.{{ backend_name|upper }}
25
29
  BACKEND_CLASS = {{ backend_name }}Backend
26
30
 
@@ -40,12 +44,12 @@ class {{ backend_name }}Configurator(Configurator):
40
44
  auth={{ backend_name }}Creds.parse_obj(config.creds).json(),
41
45
  )
42
46
 
43
- def get_backend_config(
44
- self, record: BackendRecord, include_creds: bool
45
- ) -> Any{{ backend_name }}BackendConfig:
47
+ def get_backend_config_with_creds(self, record: BackendRecord) -> {{ backend_name }}BackendConfigWithCreds:
48
+ config = self._get_config(record)
49
+ return {{ backend_name }}BackendConfigWithCreds.__response__.parse_obj(config)
50
+
51
+ def get_backend_config_without_creds(self, record: BackendRecord) -> {{ backend_name }}BackendConfig:
46
52
  config = self._get_config(record)
47
- if include_creds:
48
- return {{ backend_name }}BackendConfigWithCreds.__response__.parse_obj(config)
49
53
  return {{ backend_name }}BackendConfig.__response__.parse_obj(config)
50
54
 
51
55
  def get_backend(self, record: BackendRecord) -> {{ backend_name }}Backend: