zenml-nightly 0.62.0.dev20240719__py3-none-any.whl → 0.62.0.dev20240721__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.
zenml/VERSION CHANGED
@@ -1 +1 @@
1
- 0.62.0.dev20240719
1
+ 0.62.0.dev20240721
zenml/cli/stack.py CHANGED
@@ -491,18 +491,6 @@ def register_stack(
491
491
  artifact_store = component_name
492
492
  if component_type == StackComponentType.ORCHESTRATOR:
493
493
  orchestrator = component_name
494
- if not isinstance(
495
- component_info, UUID
496
- ) and component_info.flavor.startswith("vm"):
497
- if isinstance(
498
- service_connector, ServiceConnectorInfo
499
- ) and service_connector.auth_method in {
500
- "service-account",
501
- "external-account",
502
- }:
503
- service_connector.configuration[
504
- "generate_temporary_tokens"
505
- ] = False
506
494
  if component_type == StackComponentType.CONTAINER_REGISTRY:
507
495
  container_registry = component_name
508
496
 
@@ -2316,7 +2304,7 @@ def _get_service_connector_info(
2316
2304
  """
2317
2305
  from rich.prompt import Prompt
2318
2306
 
2319
- if cloud_provider not in {"aws", "gcp"}:
2307
+ if cloud_provider not in {"aws", "gcp", "azure"}:
2320
2308
  raise ValueError(f"Unknown cloud provider {cloud_provider}")
2321
2309
 
2322
2310
  client = Client()
@@ -2431,19 +2419,25 @@ def _get_stack_component_info(
2431
2419
 
2432
2420
  selected_storage = choices[selected_storage_idx]
2433
2421
 
2422
+ flavor = selected_storage[0]
2434
2423
  config = {"path": selected_storage[1]}
2435
2424
  service_connector_resource_id = selected_storage[1]
2436
2425
  elif component_type == "orchestrator":
2437
2426
 
2438
- def query_gcp_region(compute_type: str) -> str:
2427
+ def query_region(
2428
+ provider: StackDeploymentProvider,
2429
+ compute_type: str,
2430
+ is_skypilot: bool = False,
2431
+ ) -> str:
2432
+ deployment_info = Client().zen_store.get_stack_deployment_info(
2433
+ provider
2434
+ )
2439
2435
  region = Prompt.ask(
2440
2436
  f"Select the location for your {compute_type}:",
2441
2437
  choices=sorted(
2442
- Client()
2443
- .zen_store.get_stack_deployment_info(
2444
- StackDeploymentProvider.GCP
2445
- )
2446
- .locations.values()
2438
+ deployment_info.skypilot_default_regions.values()
2439
+ if is_skypilot
2440
+ else deployment_info.locations.values()
2447
2441
  ),
2448
2442
  show_choices=True,
2449
2443
  )
@@ -2467,10 +2461,22 @@ def _get_stack_component_info(
2467
2461
  config["execution_role"] = execution_role
2468
2462
  elif flavor == "vm_aws":
2469
2463
  config["region"] = selected_orchestrator[1]
2464
+ elif flavor == "vm_gcp":
2465
+ config["region"] = query_region(
2466
+ StackDeploymentProvider.GCP,
2467
+ "Skypilot cluster",
2468
+ is_skypilot=True,
2469
+ )
2470
2470
  elif flavor == "vm_azure":
2471
- config["region"] = query_gcp_region("Skypilot cluster")
2471
+ config["region"] = query_region(
2472
+ StackDeploymentProvider.AZURE,
2473
+ "Skypilot cluster",
2474
+ is_skypilot=True,
2475
+ )
2472
2476
  elif flavor == "vertex":
2473
- config["location"] = query_gcp_region("Vertex AI job")
2477
+ config["location"] = query_region(
2478
+ StackDeploymentProvider.GCP, "Vertex AI job"
2479
+ )
2474
2480
  service_connector_resource_id = selected_orchestrator[1]
2475
2481
  elif component_type == "container_registry":
2476
2482
  selected_registry_idx = cli_utils.multi_choice_prompt(
@@ -2482,6 +2488,7 @@ def _get_stack_component_info(
2482
2488
  if selected_registry_idx is None:
2483
2489
  cli_utils.error("No container registry selected.")
2484
2490
  selected_registry = choices[selected_registry_idx]
2491
+ flavor = selected_registry[0]
2485
2492
  config = {"uri": selected_registry[1]}
2486
2493
  service_connector_resource_id = selected_registry[1]
2487
2494
  else:
zenml/enums.py CHANGED
@@ -406,3 +406,4 @@ class StackDeploymentProvider(StrEnum):
406
406
 
407
407
  AWS = "aws"
408
408
  GCP = "gcp"
409
+ AZURE = "azure"
@@ -42,7 +42,7 @@ class AzureIntegration(Integration):
42
42
  "adlfs>=2021.10.0",
43
43
  "azure-keyvault-keys",
44
44
  "azure-keyvault-secrets",
45
- "azure-identity==1.10.0",
45
+ "azure-identity",
46
46
  "azureml-core==1.54.0.post1",
47
47
  "azure-mgmt-containerservice>=20.0.0",
48
48
  "azure-storage-blob==12.17.0", # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056
@@ -19,13 +19,11 @@ orchestrator for a remote orchestration of ZenML pipelines on VMs.
19
19
  from typing import List, Type
20
20
 
21
21
  from zenml.integrations.constants import (
22
-
23
22
  SKYPILOT_AZURE,
24
23
  )
25
24
  from zenml.integrations.integration import Integration
26
25
  from zenml.stack import Flavor
27
26
 
28
-
29
27
  SKYPILOT_AZURE_ORCHESTRATOR_FLAVOR = "vm_azure"
30
28
 
31
29
 
@@ -33,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
33
31
  """Definition of Skypilot (Azure) Integration for ZenML."""
34
32
 
35
33
  NAME = SKYPILOT_AZURE
36
- REQUIREMENTS = ["skypilot[azure]~=0.6.0"]
34
+ REQUIREMENTS = ["skypilot-nightly[azure]==1.0.0.dev20240716"]
37
35
  APT_PACKAGES = ["openssh-client", "rsync"]
38
36
 
39
37
  @classmethod
@@ -55,6 +55,11 @@ class StackDeploymentInfo(BaseModel):
55
55
  description="The locations where the stack can be deployed, as a "
56
56
  "dictionary mapping location names to descriptions.",
57
57
  )
58
+ skypilot_default_regions: Dict[str, str] = Field(
59
+ title="The locations where the Skypilot clusters can be deployed by default.",
60
+ description="The locations where the Skypilot clusters can be deployed by default, as a "
61
+ "dictionary mapping location names to descriptions.",
62
+ )
58
63
 
59
64
 
60
65
  class StackDeploymentConfig(BaseModel):
@@ -66,6 +66,9 @@ def _raise_specific_cloud_exception_if_needed(
66
66
  GCP_DOCS = (
67
67
  "https://docs.zenml.io/how-to/auth-management/gcp-service-connector"
68
68
  )
69
+ AZURE_DOCS = (
70
+ "https://docs.zenml.io/how-to/auth-management/azure-service-connector"
71
+ )
69
72
 
70
73
  if not artifact_stores:
71
74
  error_msg = (
@@ -89,7 +92,12 @@ def _raise_specific_cloud_exception_if_needed(
89
92
  )
90
93
  )
91
94
  elif cloud_provider == "azure":
92
- pass
95
+ raise ValueError(
96
+ error_msg.format(
97
+ obj_name="Blob Container",
98
+ docs=f"{AZURE_DOCS}#azure-blob-storage-container",
99
+ )
100
+ )
93
101
  if not orchestrators:
94
102
  error_msg = (
95
103
  "We were unable to find any orchestrator engines "
@@ -105,7 +113,7 @@ def _raise_specific_cloud_exception_if_needed(
105
113
  error_msg.format(
106
114
  cloud_name="AWS",
107
115
  gen_docs=f"{AWS_DOCS}#generic-aws-resource",
108
- k8s_name="Kubernetes",
116
+ k8s_name="EKS",
109
117
  k8s_docs=f"{AWS_DOCS}#eks-kubernetes-cluster",
110
118
  )
111
119
  )
@@ -115,12 +123,19 @@ def _raise_specific_cloud_exception_if_needed(
115
123
  error_msg.format(
116
124
  cloud_name="GCP",
117
125
  gen_docs=f"{GCP_DOCS}#generic-gcp-resource",
118
- k8s_name="GKE Kubernetes",
126
+ k8s_name="GKE",
119
127
  k8s_docs=f"{GCP_DOCS}#gke-kubernetes-cluster",
120
128
  )
121
129
  )
122
130
  elif cloud_provider == "azure":
123
- pass
131
+ raise ValueError(
132
+ error_msg.format(
133
+ cloud_name="Azure",
134
+ gen_docs=f"{AZURE_DOCS}#generic-azure-resource",
135
+ k8s_name="AKS",
136
+ k8s_docs=f"{AZURE_DOCS}#aks-kubernetes-cluster",
137
+ )
138
+ )
124
139
  if not container_registries:
125
140
  error_msg = (
126
141
  "We were unable to find any container registries "
@@ -145,7 +160,12 @@ def _raise_specific_cloud_exception_if_needed(
145
160
  )
146
161
  )
147
162
  elif cloud_provider == "azure":
148
- pass
163
+ raise ValueError(
164
+ error_msg.format(
165
+ registry_name="ACR",
166
+ docs_link=f"{AZURE_DOCS}#acr-container-registry",
167
+ )
168
+ )
149
169
 
150
170
 
151
171
  def get_resources_options_from_resource_model_for_full_stack(
@@ -162,15 +182,11 @@ def get_resources_options_from_resource_model_for_full_stack(
162
182
  client = Client()
163
183
  zen_store = client.zen_store
164
184
 
165
- can_generate_long_tokens = False
166
185
  if isinstance(connector_details, UUID):
167
186
  resource_model = zen_store.verify_service_connector(
168
187
  connector_details,
169
188
  list_resources=True,
170
189
  )
171
- can_generate_long_tokens = not zen_store.get_service_connector(
172
- connector_details
173
- ).configuration.get("generate_temporary_tokens", True)
174
190
  else:
175
191
  resource_model = zen_store.verify_service_connector_config(
176
192
  service_connector=ServiceConnectorRequest(
@@ -185,7 +201,6 @@ def get_resources_options_from_resource_model_for_full_stack(
185
201
  ),
186
202
  list_resources=True,
187
203
  )
188
- can_generate_long_tokens = True
189
204
 
190
205
  resources = resource_model.resources
191
206
 
@@ -229,18 +244,17 @@ def get_resources_options_from_resource_model_for_full_stack(
229
244
  flavor_display_name="AWS Sagemaker",
230
245
  )
231
246
  )
232
- if can_generate_long_tokens:
233
- orchestrators.append(
234
- _prepare_resource_info(
235
- connector_details=connector_details,
236
- resource_ids=each.resource_ids,
237
- stack_component_type=StackComponentType.ORCHESTRATOR,
238
- flavor="vm_aws",
239
- required_configuration={"region": "region"},
240
- use_resource_value_as_fixed_config=True,
241
- flavor_display_name="Skypilot (EC2)",
242
- )
247
+ orchestrators.append(
248
+ _prepare_resource_info(
249
+ connector_details=connector_details,
250
+ resource_ids=each.resource_ids,
251
+ stack_component_type=StackComponentType.ORCHESTRATOR,
252
+ flavor="vm_aws",
253
+ required_configuration={"region": "region"},
254
+ use_resource_value_as_fixed_config=True,
255
+ flavor_display_name="Skypilot (EC2)",
243
256
  )
257
+ )
244
258
 
245
259
  if each.resource_type == "kubernetes-cluster":
246
260
  orchestrators.append(
@@ -291,19 +305,16 @@ def get_resources_options_from_resource_model_for_full_stack(
291
305
  flavor_display_name="Vertex AI",
292
306
  )
293
307
  )
294
- if can_generate_long_tokens:
295
- orchestrators.append(
296
- _prepare_resource_info(
297
- connector_details=connector_details,
298
- resource_ids=each.resource_ids,
299
- stack_component_type=StackComponentType.ORCHESTRATOR,
300
- flavor="vm_gcp",
301
- required_configuration={
302
- "region": "region name"
303
- },
304
- flavor_display_name="Skypilot (Compute)",
305
- )
308
+ orchestrators.append(
309
+ _prepare_resource_info(
310
+ connector_details=connector_details,
311
+ resource_ids=each.resource_ids,
312
+ stack_component_type=StackComponentType.ORCHESTRATOR,
313
+ flavor="vm_gcp",
314
+ required_configuration={"region": "region name"},
315
+ flavor_display_name="Skypilot (Compute)",
306
316
  )
317
+ )
307
318
 
308
319
  if each.resource_type == "kubernetes-cluster":
309
320
  orchestrators.append(
@@ -330,7 +341,55 @@ def get_resources_options_from_resource_model_for_full_stack(
330
341
  )
331
342
 
332
343
  elif connector_type == "azure":
333
- pass
344
+ for each in resources:
345
+ if each.resource_ids:
346
+ if each.resource_type == "blob-container":
347
+ artifact_stores.append(
348
+ _prepare_resource_info(
349
+ connector_details=connector_details,
350
+ resource_ids=each.resource_ids,
351
+ stack_component_type=StackComponentType.ARTIFACT_STORE,
352
+ flavor="azure",
353
+ required_configuration={},
354
+ flavor_display_name="Blob container",
355
+ )
356
+ )
357
+ if each.resource_type == "azure-generic":
358
+ # No native orchestrator ATM
359
+ orchestrators.append(
360
+ _prepare_resource_info(
361
+ connector_details=connector_details,
362
+ resource_ids=each.resource_ids,
363
+ stack_component_type=StackComponentType.ORCHESTRATOR,
364
+ flavor="vm_azure",
365
+ required_configuration={"region": "region name"},
366
+ flavor_display_name="Skypilot (VM)",
367
+ )
368
+ )
369
+
370
+ if each.resource_type == "kubernetes-cluster":
371
+ orchestrators.append(
372
+ _prepare_resource_info(
373
+ connector_details=connector_details,
374
+ resource_ids=each.resource_ids,
375
+ stack_component_type=StackComponentType.ORCHESTRATOR,
376
+ flavor="kubernetes",
377
+ required_configuration={},
378
+ flavor_display_name="Kubernetes",
379
+ )
380
+ )
381
+ if each.resource_type == "docker-registry":
382
+ container_registries.append(
383
+ _prepare_resource_info(
384
+ connector_details=connector_details,
385
+ resource_ids=each.resource_ids,
386
+ stack_component_type=StackComponentType.CONTAINER_REGISTRY,
387
+ flavor="azure",
388
+ required_configuration={"uri": "URI"},
389
+ use_resource_value_as_fixed_config=True,
390
+ flavor_display_name="ACR",
391
+ )
392
+ )
334
393
 
335
394
  _raise_specific_cloud_exception_if_needed(
336
395
  cloud_provider=connector_type,
@@ -0,0 +1,179 @@
1
+ # Copyright (c) ZenML GmbH 2024. All Rights Reserved.
2
+
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at:
6
+
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
12
+ # or implied. See the License for the specific language governing
13
+ # permissions and limitations under the License.
14
+ """Functionality to deploy a ZenML stack to Azure."""
15
+
16
+ import re
17
+ from typing import ClassVar, Dict, List
18
+
19
+ from zenml.enums import StackDeploymentProvider
20
+ from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
21
+
22
+
23
+ # TODO: this class just implements the regions list, and is not suitable for other
24
+ # deployment tasks.
25
+ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
26
+ """Azure ZenML Cloud Stack Deployment."""
27
+
28
+ provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AZURE
29
+
30
+ @classmethod
31
+ def description(cls) -> str:
32
+ """Return a description of the ZenML Cloud Stack Deployment.
33
+
34
+ This will be displayed when the user is prompted to deploy
35
+ the ZenML stack.
36
+
37
+ Returns:
38
+ A MarkDown description of the ZenML Cloud Stack Deployment.
39
+ """
40
+ # TODO: Implement this
41
+ return ""
42
+
43
+ @classmethod
44
+ def instructions(cls) -> str:
45
+ """Return instructions on how to deploy the ZenML stack to the specified cloud provider.
46
+
47
+ This will be displayed before the user is prompted to deploy the ZenML
48
+ stack.
49
+
50
+ Returns:
51
+ MarkDown instructions on how to deploy the ZenML stack to the
52
+ specified cloud provider.
53
+ """
54
+ # TODO: Implement this
55
+ return ""
56
+
57
+ @classmethod
58
+ def post_deploy_instructions(cls) -> str:
59
+ """Return instructions on what to do after the deployment is complete.
60
+
61
+ This will be displayed after the deployment is complete.
62
+
63
+ Returns:
64
+ MarkDown instructions on what to do after the deployment is
65
+ complete.
66
+ """
67
+ # TODO: Implement this
68
+ return ""
69
+
70
+ @classmethod
71
+ def integrations(cls) -> List[str]:
72
+ """Return the ZenML integrations required for the stack.
73
+
74
+ Returns:
75
+ The list of ZenML integrations that need to be installed for the
76
+ stack to be usable.
77
+ """
78
+ return ["azure"]
79
+
80
+ @classmethod
81
+ def permissions(cls) -> Dict[str, List[str]]:
82
+ """Return the permissions granted to ZenML to access the cloud resources.
83
+
84
+ Returns:
85
+ The permissions granted to ZenML to access the cloud resources, as
86
+ a dictionary grouping permissions by resource.
87
+ """
88
+ # TODO: Implement this
89
+ return {}
90
+
91
+ @classmethod
92
+ def locations(cls) -> Dict[str, str]:
93
+ """Return the locations where the ZenML stack can be deployed.
94
+
95
+ Returns:
96
+ The regions where the ZenML stack can be deployed as a map of region
97
+ names to region descriptions.
98
+ """
99
+ # Based on `az account list-locations -o table` on 16.07.2024
100
+ return {
101
+ "(US) East US": "eastus",
102
+ "(US) South Central US": "southcentralus",
103
+ "(US) West US 2": "westus2",
104
+ "(US) West US 3": "westus3",
105
+ "(Asia Pacific) Australia East": "australiaeast",
106
+ "(Asia Pacific) Southeast Asia": "southeastasia",
107
+ "(Europe) North Europe": "northeurope",
108
+ "(Europe) Sweden Central": "swedencentral",
109
+ "(Europe) UK South": "uksouth",
110
+ "(Europe) West Europe": "westeurope",
111
+ "(US) Central US": "centralus",
112
+ "(Africa) South Africa North": "southafricanorth",
113
+ "(Asia Pacific) Central India": "centralindia",
114
+ "(Asia Pacific) East Asia": "eastasia",
115
+ "(Asia Pacific) Japan East": "japaneast",
116
+ "(Asia Pacific) Korea Central": "koreacentral",
117
+ "(Canada) Canada Central": "canadacentral",
118
+ "(Europe) France Central": "francecentral",
119
+ "(Europe) Germany West Central": "germanywestcentral",
120
+ "(Europe) Italy North": "italynorth",
121
+ "(Europe) Norway East": "norwayeast",
122
+ "(Europe) Poland Central": "polandcentral",
123
+ "(Europe) Spain Central": "spaincentral",
124
+ "(Europe) Switzerland North": "switzerlandnorth",
125
+ "(Mexico) Mexico Central": "mexicocentral",
126
+ "(Middle East) UAE North": "uaenorth",
127
+ "(South America) Brazil South": "brazilsouth",
128
+ "(Middle East) Israel Central": "israelcentral",
129
+ "(Middle East) Qatar Central": "qatarcentral",
130
+ "(US) Central US (Stage)": "centralusstage",
131
+ "(US) East US (Stage)": "eastusstage",
132
+ "(US) East US 2 (Stage)": "eastus2stage",
133
+ "(US) North Central US (Stage)": "northcentralusstage",
134
+ "(US) South Central US (Stage)": "southcentralusstage",
135
+ "(US) West US (Stage)": "westusstage",
136
+ "(US) West US 2 (Stage)": "westus2stage",
137
+ "(Asia Pacific) East Asia (Stage)": "eastasiastage",
138
+ "(Asia Pacific) Southeast Asia (Stage)": "southeastasiastage",
139
+ "(South America) Brazil US": "brazilus",
140
+ "(US) East US 2": "eastus2",
141
+ "(US) East US STG": "eastusstg",
142
+ "(US) North Central US": "northcentralus",
143
+ "(US) West US": "westus",
144
+ "(Asia Pacific) Japan West": "japanwest",
145
+ "(Asia Pacific) Jio India West": "jioindiawest",
146
+ "(US) Central US EUAP": "centraluseuap",
147
+ "(US) East US 2 EUAP": "eastus2euap",
148
+ "(US) West Central US": "westcentralus",
149
+ "(Africa) South Africa West": "southafricawest",
150
+ "(Asia Pacific) Australia Central": "australiacentral",
151
+ "(Asia Pacific) Australia Central 2": "australiacentral2",
152
+ "(Asia Pacific) Australia Southeast": "australiasoutheast",
153
+ "(Asia Pacific) Jio India Central": "jioindiacentral",
154
+ "(Asia Pacific) Korea South": "koreasouth",
155
+ "(Asia Pacific) South India": "southindia",
156
+ "(Asia Pacific) West India": "westindia",
157
+ "(Canada) Canada East": "canadaeast",
158
+ "(Europe) France South": "francesouth",
159
+ "(Europe) Germany North": "germanynorth",
160
+ "(Europe) Norway West": "norwaywest",
161
+ "(Europe) Switzerland West": "switzerlandwest",
162
+ "(Europe) UK West": "ukwest",
163
+ "(Middle East) UAE Central": "uaecentral",
164
+ "(South America) Brazil Southeast": "brazilsoutheast",
165
+ }
166
+
167
+ @classmethod
168
+ def skypilot_default_regions(cls) -> Dict[str, str]:
169
+ """Returns the regions supported by default for the Skypilot.
170
+
171
+ Returns:
172
+ The regions supported by default for the Skypilot.
173
+ """
174
+ matcher = re.compile(r".*us\d*( |$)")
175
+ return {
176
+ k: v
177
+ for k, v in cls.locations().items()
178
+ if "(US)" in k and matcher.match(v)
179
+ }
@@ -13,6 +13,7 @@
13
13
  # permissions and limitations under the License.
14
14
  """Functionality to deploy a ZenML stack to GCP."""
15
15
 
16
+ import re
16
17
  from typing import ClassVar, Dict, List
17
18
 
18
19
  from zenml.enums import StackDeploymentProvider
@@ -208,6 +209,16 @@ GCP project and to clean up the resources created by the stack by using
208
209
  "US West (Las Vegas)": "us-west4",
209
210
  }
210
211
 
212
+ @classmethod
213
+ def skypilot_default_regions(cls) -> Dict[str, str]:
214
+ """Returns the regions supported by default for the Skypilot.
215
+
216
+ Returns:
217
+ The regions supported by default for the Skypilot.
218
+ """
219
+ matcher = re.compile(r"us-.*")
220
+ return {k: v for k, v in cls.locations().items() if matcher.match(v)}
221
+
211
222
  def get_deployment_config(
212
223
  self,
213
224
  ) -> StackDeploymentConfig:
@@ -105,6 +105,15 @@ class ZenMLCloudStackDeployment(BaseModel):
105
105
  names to region descriptions.
106
106
  """
107
107
 
108
+ @classmethod
109
+ def skypilot_default_regions(cls) -> Dict[str, str]:
110
+ """Returns the regions supported by default for the Skypilot.
111
+
112
+ Returns:
113
+ The regions supported by default for the Skypilot.
114
+ """
115
+ return cls.locations()
116
+
108
117
  @classmethod
109
118
  def get_deployment_info(cls) -> StackDeploymentInfo:
110
119
  """Return information about the ZenML Cloud Stack Deployment.
@@ -120,6 +129,7 @@ class ZenMLCloudStackDeployment(BaseModel):
120
129
  integrations=cls.integrations(),
121
130
  permissions=cls.permissions(),
122
131
  locations=cls.locations(),
132
+ skypilot_default_regions=cls.skypilot_default_regions(),
123
133
  )
124
134
 
125
135
  @abstractmethod
@@ -19,6 +19,9 @@ from zenml.enums import StackDeploymentProvider
19
19
  from zenml.stack_deployments.aws_stack_deployment import (
20
20
  AWSZenMLCloudStackDeployment,
21
21
  )
22
+ from zenml.stack_deployments.azure_stack_deployment import (
23
+ AZUREZenMLCloudStackDeployment,
24
+ )
22
25
  from zenml.stack_deployments.gcp_stack_deployment import (
23
26
  GCPZenMLCloudStackDeployment,
24
27
  )
@@ -27,6 +30,7 @@ from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
27
30
  STACK_DEPLOYMENT_PROVIDERS = {
28
31
  StackDeploymentProvider.AWS: AWSZenMLCloudStackDeployment,
29
32
  StackDeploymentProvider.GCP: GCPZenMLCloudStackDeployment,
33
+ StackDeploymentProvider.AZURE: AZUREZenMLCloudStackDeployment,
30
34
  }
31
35
 
32
36
 
@@ -286,6 +286,7 @@ from zenml.service_connectors.service_connector_registry import (
286
286
  service_connector_registry,
287
287
  )
288
288
  from zenml.stack.flavor_registry import FlavorRegistry
289
+ from zenml.stack_deployments.utils import get_stack_deployment_class
289
290
  from zenml.utils import uuid_utils
290
291
  from zenml.utils.enum_utils import StrEnum
291
292
  from zenml.utils.networking_utils import (
@@ -3185,6 +3186,29 @@ class SqlZenStore(BaseZenStore):
3185
3186
  "found."
3186
3187
  )
3187
3188
 
3189
+ # warn about skypilot regions, if needed
3190
+ if component.flavor in {"vm_gcp", "vm_azure"}:
3191
+ stack_deployment_class = get_stack_deployment_class(
3192
+ StackDeploymentProvider.GCP
3193
+ if component.flavor == "vm_gcp"
3194
+ else StackDeploymentProvider.AZURE
3195
+ )
3196
+ skypilot_regions = (
3197
+ stack_deployment_class.skypilot_default_regions().values()
3198
+ )
3199
+ if (
3200
+ component.configuration.get("region", None)
3201
+ and component.configuration["region"]
3202
+ not in skypilot_regions
3203
+ ):
3204
+ logger.warning(
3205
+ f"Region `{component.configuration['region']}` is not enabled in Skypilot "
3206
+ f"by default. Supported regions by default are: {skypilot_regions}. "
3207
+ "Check the Skypilot documentation to learn how to enable regions rather "
3208
+ "than default ones. (If you have already extended your configuration - "
3209
+ "simply ignore this warning)"
3210
+ )
3211
+
3188
3212
  # Create the component
3189
3213
  new_component = StackComponentSchema(
3190
3214
  name=component.name,
@@ -7032,9 +7056,44 @@ class SqlZenStore(BaseZenStore):
7032
7056
  # Service Connectors
7033
7057
  service_connectors: List[ServiceConnectorResponse] = []
7034
7058
 
7059
+ need_to_generate_permanent_tokens = False
7060
+ orchestrator_component = full_stack.components[
7061
+ StackComponentType.ORCHESTRATOR
7062
+ ]
7063
+ if isinstance(orchestrator_component, UUID):
7064
+ orchestrator = self.get_stack_component(
7065
+ orchestrator_component,
7066
+ hydrate=False,
7067
+ )
7068
+ need_to_generate_permanent_tokens = (
7069
+ orchestrator.flavor.startswith("vm_")
7070
+ )
7071
+ else:
7072
+ need_to_generate_permanent_tokens = (
7073
+ orchestrator_component.flavor.startswith("vm_")
7074
+ )
7075
+
7035
7076
  for connector_id_or_info in full_stack.service_connectors:
7036
7077
  # Fetch an existing service connector
7037
7078
  if isinstance(connector_id_or_info, UUID):
7079
+ existing_service_connector = self.get_service_connector(
7080
+ connector_id_or_info
7081
+ )
7082
+ if need_to_generate_permanent_tokens:
7083
+ if (
7084
+ existing_service_connector.configuration.get(
7085
+ "generate_temporary_tokens", None
7086
+ )
7087
+ is not False
7088
+ ):
7089
+ self.update_service_connector(
7090
+ existing_service_connector.id,
7091
+ ServiceConnectorUpdate(
7092
+ configuration=existing_service_connector.configuration.update(
7093
+ {"generate_temporary_tokens": False}
7094
+ )
7095
+ ),
7096
+ )
7038
7097
  service_connectors.append(
7039
7098
  self.get_service_connector(connector_id_or_info)
7040
7099
  )
@@ -7047,7 +7106,11 @@ class SqlZenStore(BaseZenStore):
7047
7106
  name=connector_name,
7048
7107
  connector_type=connector_id_or_info.type,
7049
7108
  auth_method=connector_id_or_info.auth_method,
7050
- configuration=connector_id_or_info.configuration,
7109
+ configuration=connector_id_or_info.configuration.update(
7110
+ {
7111
+ "generate_temporary_tokens": not need_to_generate_permanent_tokens
7112
+ }
7113
+ ),
7051
7114
  user=full_stack.user,
7052
7115
  workspace=full_stack.workspace,
7053
7116
  labels={
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: zenml-nightly
3
- Version: 0.62.0.dev20240719
3
+ Version: 0.62.0.dev20240721
4
4
  Summary: ZenML: Write production-ready ML code.
5
5
  Home-page: https://zenml.io
6
6
  License: Apache-2.0
@@ -6,7 +6,7 @@ RELEASE_NOTES.md,sha256=rF05H0U7U4lCYWomDqKhWdaWtOfCFSOf10LNBoFnl3Y,349994
6
6
  ROADMAP.md,sha256=hiLSmr16BH8Dfx7SaQM4JcXCGCVl6mFZPFAwJeDTrJU,407
7
7
  SECURITY.md,sha256=9DepA8y03yvCZLHEfcXLTDH4lUyKHquAdukBsccNN7c,682
8
8
  zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
9
- zenml/VERSION,sha256=Gbq4CdB9Z7CjJ5RLBoMark6aI20BmsoTqXQU17hRTPw,19
9
+ zenml/VERSION,sha256=3JOMcHTOmS-pZYA-nJSR5miiyVyMd1DEC3YvDHa-HDo,19
10
10
  zenml/__init__.py,sha256=6zRqB0FeBWX5E4Np9k9jzJgKaEHKvLFbtmOZQYrGpD8,2453
11
11
  zenml/_hub/__init__.py,sha256=6qDzpQAAZa__Aiiz0mC1qM-9dw9_jk_v_aXeJknxbDE,644
12
12
  zenml/_hub/client.py,sha256=Um2df1JO7-BmNm65efHSPpV5e0GITuoQJmod_wkdvbw,9136
@@ -58,7 +58,7 @@ zenml/cli/served_model.py,sha256=lIPd6kOS9kfWFOsd3DSILtiYomINpmXrdpyP5vtQlj8,146
58
58
  zenml/cli/server.py,sha256=N2cvu9C_z8CzwGu3a7P6fE2X8Se_rHEx30iq9nYMy9o,28682
59
59
  zenml/cli/service_accounts.py,sha256=MO8pVVjppdvIULFxpvQKlhfEDmRqvlH8aIbPTdNSHXU,16489
60
60
  zenml/cli/service_connectors.py,sha256=v6o-FpTEnDXYcQXczBELwfWPCcYBi9z7b1GC90Fvp30,74397
61
- zenml/cli/stack.py,sha256=M3V6TUCqWQLnHPtVUNsdZyZcid0DYes9JPYyBl5--_k,82437
61
+ zenml/cli/stack.py,sha256=xNB2sXXFJe3VC_Gow-xWm5s3ae3NKFS2pXAiYeviUlg,82520
62
62
  zenml/cli/stack_components.py,sha256=BxWsqeu5e2DtP5TMjzE98d8Lqfwe5o45TC2awjil5wY,74144
63
63
  zenml/cli/stack_recipes.py,sha256=ZhJE2HdytilaybhJgPkNp9DvSa5LiVgxS19jR0UxHuk,15716
64
64
  zenml/cli/tag.py,sha256=Co-AL1eNZMf_AsqlcnJpoZ_k7UKa0l5X_g_TLpJdGes,4884
@@ -113,7 +113,7 @@ zenml/entrypoints/base_entrypoint_configuration.py,sha256=8dD_TRMUAXbRj6XNboQjyz
113
113
  zenml/entrypoints/entrypoint.py,sha256=Cxh4S_F4URQQ1WF2fGuJEIfLcUWcxVDeDiKYJA4Kxps,2174
114
114
  zenml/entrypoints/pipeline_entrypoint_configuration.py,sha256=To-vTP29qAE36ndJDF1fRw9wL2Nk2bsBuO-ayAwvSmo,1646
115
115
  zenml/entrypoints/step_entrypoint_configuration.py,sha256=NskQF6OFFCIoD7v0JxR_cZRVXBRaKl3R0fwqpaQCPLM,6571
116
- zenml/enums.py,sha256=3_SpffPCzk0MRPUiqKl20qGyq6UOBHzK6F58we7pEDo,10033
116
+ zenml/enums.py,sha256=tyte4wjKGy3bmTydYMkJaXGaxPKED21XhNXWOqYpM5c,10053
117
117
  zenml/environment.py,sha256=xqsj8G-qr_W2KyyWnJABMsNQ7Tl32MUhFir7kUDPpKc,20276
118
118
  zenml/event_hub/__init__.py,sha256=-fD9mPOslf4J-_XFBPp5gYlPz7-6ZaAKHa5jxf_nyAo,757
119
119
  zenml/event_hub/base_event_hub.py,sha256=jqVdy2ix5JyaVV08DHTqyFn1NP03GWtBQXgbjQdm_fo,6574
@@ -163,7 +163,7 @@ zenml/integrations/aws/service_connectors/aws_service_connector.py,sha256=YssWip
163
163
  zenml/integrations/aws/step_operators/__init__.py,sha256=HK5oIqLixYKUKoiN4LnqTyGjAZJUdUqWYGqJhFSWyo0,817
164
164
  zenml/integrations/aws/step_operators/sagemaker_step_operator.py,sha256=-y1zqfIDJm2wwy_KETJ1nHdKuSbIHPFK3bT-0i9tpoA,10113
165
165
  zenml/integrations/aws/step_operators/sagemaker_step_operator_entrypoint_config.py,sha256=2cXroe6-bXyHVkO5yPnZagNpqx6MrMDcvuvNr8J2j-A,1581
166
- zenml/integrations/azure/__init__.py,sha256=TlRmc8-FNjRA84OW6ikJgSKu_N5assjB9I2bc7VAjU0,2570
166
+ zenml/integrations/azure/__init__.py,sha256=Ttt7jgnpIFu-AWiVBKFbPmuLX6vdMZjTEV1IVwfZ2yI,2562
167
167
  zenml/integrations/azure/artifact_stores/__init__.py,sha256=dlIwbpgjE0Hy4rhMbelNJHVKm4t8tj_hRu9mQ_cEIAg,820
168
168
  zenml/integrations/azure/artifact_stores/azure_artifact_store.py,sha256=uVMpsWDIyeytn8aHxBAma5pdPwiWuK5fwew2Tvxsk3w,11955
169
169
  zenml/integrations/azure/flavors/__init__.py,sha256=KEwsaLeKsvdI5tqYFzji_VXBLlZuGoQXXmOT8vZ8PKE,1069
@@ -462,7 +462,7 @@ zenml/integrations/skypilot_aws/flavors/__init__.py,sha256=UGL7NuZ2BkB5IxL9WUHmG
462
462
  zenml/integrations/skypilot_aws/flavors/skypilot_orchestrator_aws_vm_flavor.py,sha256=PqcUCXRyJoT7jmR15A5d7N5NJxXkgLKWm8BC3YW6Omw,3789
463
463
  zenml/integrations/skypilot_aws/orchestrators/__init__.py,sha256=xxFZ2iM9H9LVKfSB8N-yKGxuB-lJE3Y368lcj7MfhI0,1006
464
464
  zenml/integrations/skypilot_aws/orchestrators/skypilot_aws_vm_orchestrator.py,sha256=0uFfRa8lGsc7D_zdrGGMqKoL6eKUwuf2ImtDw_RRg6Q,3094
465
- zenml/integrations/skypilot_azure/__init__.py,sha256=n8ODpRF2blirwCMUL5SKdioDSja0135pubcFeI4rUcA,1740
465
+ zenml/integrations/skypilot_azure/__init__.py,sha256=HuOab2Bv8Fuz5n6_nNiR7_Brdzd8UjgRggosRtql8r0,1758
466
466
  zenml/integrations/skypilot_azure/flavors/__init__.py,sha256=9JB07IZbIZhY6yL9TjkYFTCUUfxajMajsMbbjC1BZ_s,1024
467
467
  zenml/integrations/skypilot_azure/flavors/skypilot_orchestrator_azure_vm_flavor.py,sha256=aL8ibfQGo2IcyckaesjAFk6m8kazy5pq6ycArYRD5w0,3858
468
468
  zenml/integrations/skypilot_azure/orchestrators/__init__.py,sha256=gW7JOTLdGB542Cv1ydSqF4kVBAo7LUnfkrlj45kCWm8,1017
@@ -625,7 +625,7 @@ zenml/models/v2/misc/hub_plugin_models.py,sha256=ogWYrotKq__2vR2LsrInlOSR-DJqcqd
625
625
  zenml/models/v2/misc/loaded_visualization.py,sha256=u6lapDNZDtU9eS-_EMzl00urj0yPNhiqhZcvjIz7DSg,946
626
626
  zenml/models/v2/misc/server_models.py,sha256=wd14abzVBnv4lBIb_dbBSTkqGcdQjI-LcdoZC_2ucSw,3613
627
627
  zenml/models/v2/misc/service_connector_type.py,sha256=jGJLvIsBB87ZIEeZCLLueUf28H1QJFqXOJYtiRVvZ4M,28748
628
- zenml/models/v2/misc/stack_deployment.py,sha256=tSqb-E9ErvsNpeM0WzXxpTcOlXYLQ9-x3ioaLzDyOcs,3315
628
+ zenml/models/v2/misc/stack_deployment.py,sha256=D_eIGhM6zZysiNOQvDXbjQN11pguVJyD9jJIMECwNhk,3622
629
629
  zenml/models/v2/misc/user_auth.py,sha256=XLhagiLN27wlS70LVVvF5PVBDBO4PXBRgXo2VGtVbnI,7198
630
630
  zenml/new/__init__.py,sha256=knhroJ2h0uHBCGzAiBBGJEiuhEA3cwI6XYBRIyXdbkQ,613
631
631
  zenml/new/pipelines/__init__.py,sha256=knhroJ2h0uHBCGzAiBBGJEiuhEA3cwI6XYBRIyXdbkQ,613
@@ -678,7 +678,7 @@ zenml/service_connectors/__init__.py,sha256=gIBAVUPBZtdO6JtEBCixy9YG4wZRA1mnaxaG
678
678
  zenml/service_connectors/docker_service_connector.py,sha256=vplyZNUxZhAyzXRdtAiSGMFZT87h9Wq8izE0HWh7OSQ,12441
679
679
  zenml/service_connectors/service_connector.py,sha256=HW8dTE826ZSD1vWQlqQRm8XYq08v7tDZQ8ORnJ5rTcE,55351
680
680
  zenml/service_connectors/service_connector_registry.py,sha256=lP0s0k7RIp9WZEW36Q0ygHD5urzF-wUwh41rbXWS5ss,9186
681
- zenml/service_connectors/service_connector_utils.py,sha256=TGk9IIw4zhxvcNqZDQHUGs4i1Edrh3yeM0wTETDekdQ,14428
681
+ zenml/service_connectors/service_connector_utils.py,sha256=pUen7wvsxNf12VmpGp5PvSiPWllNXbgrKH559p30U34,17195
682
682
  zenml/services/__init__.py,sha256=8aSLtYlX9EVLB1Se1I4dbhZdTqrjNrAfsn1d_RCc_nk,2800
683
683
  zenml/services/container/__init__.py,sha256=dFHcmlXgNXjUaCR18oGQJ19-I_6f3UeAUURHjqldjTs,668
684
684
  zenml/services/container/container_service.py,sha256=DJIK7DD53Q_mgZByM0oeIyAgkrwN12T5E3QCWjubkLs,18383
@@ -705,9 +705,10 @@ zenml/stack/stack_validator.py,sha256=hWbvvGIeWLj6NwSsF4GCc6RAxAWvxHXTcBZL9nJvca
705
705
  zenml/stack/utils.py,sha256=W_ZKttALLGABYSzqapkIdHerbOWf6gRL8Tl9Bgj_q5s,5460
706
706
  zenml/stack_deployments/__init__.py,sha256=-7593cQ_ZgRn774Ol-8AKXXQquIU4DSiaThVEr6TfWM,644
707
707
  zenml/stack_deployments/aws_stack_deployment.py,sha256=gRjG7aCK8cl-AxFOgI00eTMIpaIGXfsd9UfGmbQXLJ4,10239
708
- zenml/stack_deployments/gcp_stack_deployment.py,sha256=gLo1tLJaK8sW1N2C04Pdk0LDv_Y-wtm33NYqlIdUX_k,11000
709
- zenml/stack_deployments/stack_deployment.py,sha256=XTURbeOUMdByRO0jMN1tXY1ryrDZ2Z83rOzzwLtTylo,6871
710
- zenml/stack_deployments/utils.py,sha256=Vpt5hWJLvaHyPfYJ7vx0CEqrTrGCVXy2DEB1uTLr2EA,1572
708
+ zenml/stack_deployments/azure_stack_deployment.py,sha256=qFPcTfubZWsbXPP-xuTLKKXfOIqA7o36eZ2mWTriRfM,7332
709
+ zenml/stack_deployments/gcp_stack_deployment.py,sha256=XRwAB8tq1hAB2jlr3c7TqvNH-HWjCgs-_TP4kKK7y-U,11365
710
+ zenml/stack_deployments/stack_deployment.py,sha256=R57UqafZshcs0DXLKX2OWll_65UzyHK8w4kf8tS4Oqw,7209
711
+ zenml/stack_deployments/utils.py,sha256=4rU0JM6wAmSLLlH7f9xk3g3EsacMHQYDTgN1OvU9pRs,1738
711
712
  zenml/step_operators/__init__.py,sha256=tqj7fgnQyZubLjwUu4ITwkA-70KMQz4g37agbfF7ZgE,1228
712
713
  zenml/step_operators/base_step_operator.py,sha256=ZRnY6lcEHY8xZskrKKdPhgKg2BlEoh2_kb8xCaM2D1g,3522
713
714
  zenml/step_operators/step_operator_entrypoint_configuration.py,sha256=vG0Z9D2C72PYOZ_P3eXYTmm2TvL-dS_kHhSMq-coAKE,3690
@@ -1334,10 +1335,10 @@ zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py,sha256=NfW1EHIA99lseb
1334
1335
  zenml/zen_stores/secrets_stores/secrets_store_interface.py,sha256=Q2Jbnt2Pp7NGlR-u1YBfRZV2g8su2Fd0ArBMdksAE-Q,2819
1335
1336
  zenml/zen_stores/secrets_stores/service_connector_secrets_store.py,sha256=kPYX-Z_OOhZCI1CP77ncfV7IsV4e8brklnTXmKxZYNc,7078
1336
1337
  zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=Bq1djrUP9saoD7vECjS7-TlA_7sjJGgw1talri4spjU,8656
1337
- zenml/zen_stores/sql_zen_store.py,sha256=K4FAR764dyEKFwDZTCHgnfseTx_tBtvHNSdnTWywGx8,378688
1338
+ zenml/zen_stores/sql_zen_store.py,sha256=j1D_vqLzg8lK2wT-a8moDOXJfTVGAa_oA0bb-RHWEVo,381780
1338
1339
  zenml/zen_stores/zen_store_interface.py,sha256=ZcgPtlAMdO5die42Hwd6l8glxHYsxFjzuCBewIu5Hrs,91984
1339
- zenml_nightly-0.62.0.dev20240719.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
1340
- zenml_nightly-0.62.0.dev20240719.dist-info/METADATA,sha256=8RdicRE2A9zcjwf1hGBMeEd2gzQ_JSZvWLz-CvAMpuI,21026
1341
- zenml_nightly-0.62.0.dev20240719.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1342
- zenml_nightly-0.62.0.dev20240719.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
1343
- zenml_nightly-0.62.0.dev20240719.dist-info/RECORD,,
1340
+ zenml_nightly-0.62.0.dev20240721.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
1341
+ zenml_nightly-0.62.0.dev20240721.dist-info/METADATA,sha256=TkcxCHbj5x1YzReviLAju1XhSV3Jvfk-Qm5DU4bMw3w,21026
1342
+ zenml_nightly-0.62.0.dev20240721.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1343
+ zenml_nightly-0.62.0.dev20240721.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
1344
+ zenml_nightly-0.62.0.dev20240721.dist-info/RECORD,,