zenml-nightly 0.62.0.dev20240717__py3-none-any.whl → 0.62.0.dev20240720__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.62.0.dev20240717
+0.62.0.dev20240720

zenml/cli/stack.py

@@ -491,18 +491,6 @@ def register_stack(
                 artifact_store = component_name
             if component_type == StackComponentType.ORCHESTRATOR:
                 orchestrator = component_name
-                if not isinstance(
-                    component_info, UUID
-                ) and component_info.flavor.startswith("vm"):
-                    if isinstance(
-                        service_connector, ServiceConnectorInfo
-                    ) and service_connector.auth_method in {
-                        "service-account",
-                        "external-account",
-                    }:
-                        service_connector.configuration[
-                            "generate_temporary_tokens"
-                        ] = False
             if component_type == StackComponentType.CONTAINER_REGISTRY:
                 container_registry = component_name
 
@@ -2316,7 +2304,7 @@ def _get_service_connector_info(
     """
     from rich.prompt import Prompt
 
-    if cloud_provider not in {"aws", "gcp"}:
+    if cloud_provider not in {"aws", "gcp", "azure"}:
         raise ValueError(f"Unknown cloud provider {cloud_provider}")
 
     client = Client()
@@ -2431,19 +2419,25 @@ def _get_stack_component_info(
 
         selected_storage = choices[selected_storage_idx]
 
+        flavor = selected_storage[0]
         config = {"path": selected_storage[1]}
         service_connector_resource_id = selected_storage[1]
     elif component_type == "orchestrator":
 
-        def query_gcp_region(compute_type: str) -> str:
+        def query_region(
+            provider: StackDeploymentProvider,
+            compute_type: str,
+            is_skypilot: bool = False,
+        ) -> str:
+            deployment_info = Client().zen_store.get_stack_deployment_info(
+                provider
+            )
             region = Prompt.ask(
                 f"Select the location for your {compute_type}:",
                 choices=sorted(
-                    Client()
-                    .zen_store.get_stack_deployment_info(
-                        StackDeploymentProvider.GCP
-                    )
-                    .locations.values()
+                    deployment_info.skypilot_default_regions.values()
+                    if is_skypilot
+                    else deployment_info.locations.values()
                 ),
                 show_choices=True,
             )
@@ -2467,10 +2461,22 @@ def _get_stack_component_info(
             config["execution_role"] = execution_role
         elif flavor == "vm_aws":
             config["region"] = selected_orchestrator[1]
+        elif flavor == "vm_gcp":
+            config["region"] = query_region(
+                StackDeploymentProvider.GCP,
+                "Skypilot cluster",
+                is_skypilot=True,
+            )
         elif flavor == "vm_azure":
-            config["region"] = query_gcp_region("Skypilot cluster")
+            config["region"] = query_region(
+                StackDeploymentProvider.AZURE,
+                "Skypilot cluster",
+                is_skypilot=True,
+            )
         elif flavor == "vertex":
-            config["location"] = query_gcp_region("Vertex AI job")
+            config["location"] = query_region(
+                StackDeploymentProvider.GCP, "Vertex AI job"
+            )
         service_connector_resource_id = selected_orchestrator[1]
     elif component_type == "container_registry":
         selected_registry_idx = cli_utils.multi_choice_prompt(
@@ -2482,6 +2488,7 @@ def _get_stack_component_info(
         if selected_registry_idx is None:
             cli_utils.error("No container registry selected.")
         selected_registry = choices[selected_registry_idx]
+        flavor = selected_registry[0]
         config = {"uri": selected_registry[1]}
         service_connector_resource_id = selected_registry[1]
     else:

zenml/enums.py

@@ -406,3 +406,4 @@ class StackDeploymentProvider(StrEnum):
 
     AWS = "aws"
     GCP = "gcp"
+    AZURE = "azure"

zenml/integrations/azure/__init__.py

@@ -42,7 +42,7 @@ class AzureIntegration(Integration):
         "adlfs>=2021.10.0",
         "azure-keyvault-keys",
         "azure-keyvault-secrets",
-        "azure-identity==1.10.0",
+        "azure-identity",
         "azureml-core==1.54.0.post1",
         "azure-mgmt-containerservice>=20.0.0",
         "azure-storage-blob==12.17.0",  # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056

zenml/integrations/skypilot_azure/__init__.py

@@ -19,13 +19,11 @@ orchestrator for a remote orchestration of ZenML pipelines on VMs.
 from typing import List, Type
 
 from zenml.integrations.constants import (
-
     SKYPILOT_AZURE,
 )
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
-
 SKYPILOT_AZURE_ORCHESTRATOR_FLAVOR = "vm_azure"
 
 
@@ -33,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]~=0.6.0"]
+    REQUIREMENTS = ["skypilot-nightly[azure]==1.0.0.dev20240716"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/models/v2/core/component.py

@@ -191,6 +191,17 @@ class ComponentResponseBody(WorkspaceScopedResponseBody):
         title="The flavor of the stack component.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
+    integration: Optional[str] = Field(
+        default=None,
+        title="The name of the integration that the component's flavor "
+        "belongs to.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    logo_url: Optional[str] = Field(
+        default=None,
+        title="Optionally, a url pointing to a png,"
+        "svg or jpg can be attached.",
+    )
 
 
 class ComponentResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -285,6 +296,24 @@ class ComponentResponse(
         """
         return self.get_body().flavor
 
+    @property
+    def integration(self) -> Optional[str]:
+        """The `integration` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().integration
+
+    @property
+    def logo_url(self) -> Optional[str]:
+        """The `logo_url` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().logo_url
+
     @property
     def configuration(self) -> Dict[str, Any]:
         """The `configuration` property.

zenml/models/v2/misc/stack_deployment.py

@@ -55,6 +55,11 @@ class StackDeploymentInfo(BaseModel):
         description="The locations where the stack can be deployed, as a "
         "dictionary mapping location names to descriptions.",
     )
+    skypilot_default_regions: Dict[str, str] = Field(
+        title="The locations where the Skypilot clusters can be deployed by default.",
+        description="The locations where the Skypilot clusters can be deployed by default, as a "
+        "dictionary mapping location names to descriptions.",
+    )
 
 
 class StackDeploymentConfig(BaseModel):

zenml/service_connectors/service_connector_utils.py

@@ -66,6 +66,9 @@ def _raise_specific_cloud_exception_if_needed(
     GCP_DOCS = (
         "https://docs.zenml.io/how-to/auth-management/gcp-service-connector"
     )
+    AZURE_DOCS = (
+        "https://docs.zenml.io/how-to/auth-management/azure-service-connector"
+    )
 
     if not artifact_stores:
         error_msg = (
@@ -89,7 +92,12 @@ def _raise_specific_cloud_exception_if_needed(
                 )
             )
         elif cloud_provider == "azure":
-            pass
+            raise ValueError(
+                error_msg.format(
+                    obj_name="Blob Container",
+                    docs=f"{AZURE_DOCS}#azure-blob-storage-container",
+                )
+            )
     if not orchestrators:
         error_msg = (
             "We were unable to find any orchestrator engines "
@@ -105,7 +113,7 @@ def _raise_specific_cloud_exception_if_needed(
                 error_msg.format(
                     cloud_name="AWS",
                     gen_docs=f"{AWS_DOCS}#generic-aws-resource",
-                    k8s_name="Kubernetes",
+                    k8s_name="EKS",
                     k8s_docs=f"{AWS_DOCS}#eks-kubernetes-cluster",
                 )
             )
@@ -115,12 +123,19 @@ def _raise_specific_cloud_exception_if_needed(
                 error_msg.format(
                     cloud_name="GCP",
                     gen_docs=f"{GCP_DOCS}#generic-gcp-resource",
-                    k8s_name="GKE Kubernetes",
+                    k8s_name="GKE",
                     k8s_docs=f"{GCP_DOCS}#gke-kubernetes-cluster",
                 )
             )
         elif cloud_provider == "azure":
-            pass
+            raise ValueError(
+                error_msg.format(
+                    cloud_name="Azure",
+                    gen_docs=f"{AZURE_DOCS}#generic-azure-resource",
+                    k8s_name="AKS",
+                    k8s_docs=f"{AZURE_DOCS}#aks-kubernetes-cluster",
+                )
+            )
     if not container_registries:
         error_msg = (
             "We were unable to find any container registries "
@@ -145,7 +160,12 @@ def _raise_specific_cloud_exception_if_needed(
                 )
             )
         elif cloud_provider == "azure":
-            pass
+            raise ValueError(
+                error_msg.format(
+                    registry_name="ACR",
+                    docs_link=f"{AZURE_DOCS}#acr-container-registry",
+                )
+            )
 
 
 def get_resources_options_from_resource_model_for_full_stack(
@@ -162,15 +182,11 @@ def get_resources_options_from_resource_model_for_full_stack(
     client = Client()
     zen_store = client.zen_store
 
-    can_generate_long_tokens = False
     if isinstance(connector_details, UUID):
         resource_model = zen_store.verify_service_connector(
             connector_details,
             list_resources=True,
         )
-        can_generate_long_tokens = not zen_store.get_service_connector(
-            connector_details
-        ).configuration.get("generate_temporary_tokens", True)
     else:
         resource_model = zen_store.verify_service_connector_config(
             service_connector=ServiceConnectorRequest(
@@ -185,7 +201,6 @@ def get_resources_options_from_resource_model_for_full_stack(
             ),
             list_resources=True,
         )
-        can_generate_long_tokens = True
 
     resources = resource_model.resources
 
@@ -229,18 +244,17 @@ def get_resources_options_from_resource_model_for_full_stack(
                             flavor_display_name="AWS Sagemaker",
                         )
                     )
-                    if can_generate_long_tokens:
-                        orchestrators.append(
-                            _prepare_resource_info(
-                                connector_details=connector_details,
-                                resource_ids=each.resource_ids,
-                                stack_component_type=StackComponentType.ORCHESTRATOR,
-                                flavor="vm_aws",
-                                required_configuration={"region": "region"},
-                                use_resource_value_as_fixed_config=True,
-                                flavor_display_name="Skypilot (EC2)",
-                            )
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="vm_aws",
+                            required_configuration={"region": "region"},
+                            use_resource_value_as_fixed_config=True,
+                            flavor_display_name="Skypilot (EC2)",
                         )
+                    )
 
                 if each.resource_type == "kubernetes-cluster":
                     orchestrators.append(
@@ -291,19 +305,16 @@ def get_resources_options_from_resource_model_for_full_stack(
                             flavor_display_name="Vertex AI",
                         )
                     )
-                    if can_generate_long_tokens:
-                        orchestrators.append(
-                            _prepare_resource_info(
-                                connector_details=connector_details,
-                                resource_ids=each.resource_ids,
-                                stack_component_type=StackComponentType.ORCHESTRATOR,
-                                flavor="vm_gcp",
-                                required_configuration={
-                                    "region": "region name"
-                                },
-                                flavor_display_name="Skypilot (Compute)",
-                            )
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="vm_gcp",
+                            required_configuration={"region": "region name"},
+                            flavor_display_name="Skypilot (Compute)",
                         )
+                    )
 
                 if each.resource_type == "kubernetes-cluster":
                     orchestrators.append(
@@ -330,7 +341,55 @@ def get_resources_options_from_resource_model_for_full_stack(
                     )
 
     elif connector_type == "azure":
-        pass
+        for each in resources:
+            if each.resource_ids:
+                if each.resource_type == "blob-container":
+                    artifact_stores.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ARTIFACT_STORE,
+                            flavor="azure",
+                            required_configuration={},
+                            flavor_display_name="Blob container",
+                        )
+                    )
+                if each.resource_type == "azure-generic":
+                    # No native orchestrator ATM
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="vm_azure",
+                            required_configuration={"region": "region name"},
+                            flavor_display_name="Skypilot (VM)",
+                        )
+                    )
+
+                if each.resource_type == "kubernetes-cluster":
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="kubernetes",
+                            required_configuration={},
+                            flavor_display_name="Kubernetes",
+                        )
+                    )
+                if each.resource_type == "docker-registry":
+                    container_registries.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.CONTAINER_REGISTRY,
+                            flavor="azure",
+                            required_configuration={"uri": "URI"},
+                            use_resource_value_as_fixed_config=True,
+                            flavor_display_name="ACR",
+                        )
+                    )
 
     _raise_specific_cloud_exception_if_needed(
         cloud_provider=connector_type,

zenml/stack_deployments/azure_stack_deployment.py

@@ -0,0 +1,179 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+
+#       https://www.apache.org/licenses/LICENSE-2.0
+
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Functionality to deploy a ZenML stack to Azure."""
+
+import re
+from typing import ClassVar, Dict, List
+
+from zenml.enums import StackDeploymentProvider
+from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+
+
+# TODO: this class just implements the regions list, and is not suitable for other
+# deployment tasks.
+class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
+    """Azure ZenML Cloud Stack Deployment."""
+
+    provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AZURE
+
+    @classmethod
+    def description(cls) -> str:
+        """Return a description of the ZenML Cloud Stack Deployment.
+
+        This will be displayed when the user is prompted to deploy
+        the ZenML stack.
+
+        Returns:
+            A MarkDown description of the ZenML Cloud Stack Deployment.
+        """
+        # TODO: Implement this
+        return ""
+
+    @classmethod
+    def instructions(cls) -> str:
+        """Return instructions on how to deploy the ZenML stack to the specified cloud provider.
+
+        This will be displayed before the user is prompted to deploy the ZenML
+        stack.
+
+        Returns:
+            MarkDown instructions on how to deploy the ZenML stack to the
+            specified cloud provider.
+        """
+        # TODO: Implement this
+        return ""
+
+    @classmethod
+    def post_deploy_instructions(cls) -> str:
+        """Return instructions on what to do after the deployment is complete.
+
+        This will be displayed after the deployment is complete.
+
+        Returns:
+            MarkDown instructions on what to do after the deployment is
+            complete.
+        """
+        # TODO: Implement this
+        return ""
+
+    @classmethod
+    def integrations(cls) -> List[str]:
+        """Return the ZenML integrations required for the stack.
+
+        Returns:
+            The list of ZenML integrations that need to be installed for the
+            stack to be usable.
+        """
+        return ["azure"]
+
+    @classmethod
+    def permissions(cls) -> Dict[str, List[str]]:
+        """Return the permissions granted to ZenML to access the cloud resources.
+
+        Returns:
+            The permissions granted to ZenML to access the cloud resources, as
+            a dictionary grouping permissions by resource.
+        """
+        # TODO: Implement this
+        return {}
+
+    @classmethod
+    def locations(cls) -> Dict[str, str]:
+        """Return the locations where the ZenML stack can be deployed.
+
+        Returns:
+            The regions where the ZenML stack can be deployed as a map of region
+            names to region descriptions.
+        """
+        # Based on `az account list-locations -o table` on 16.07.2024
+        return {
+            "(US) East US": "eastus",
+            "(US) South Central US": "southcentralus",
+            "(US) West US 2": "westus2",
+            "(US) West US 3": "westus3",
+            "(Asia Pacific) Australia East": "australiaeast",
+            "(Asia Pacific) Southeast Asia": "southeastasia",
+            "(Europe) North Europe": "northeurope",
+            "(Europe) Sweden Central": "swedencentral",
+            "(Europe) UK South": "uksouth",
+            "(Europe) West Europe": "westeurope",
+            "(US) Central US": "centralus",
+            "(Africa) South Africa North": "southafricanorth",
+            "(Asia Pacific) Central India": "centralindia",
+            "(Asia Pacific) East Asia": "eastasia",
+            "(Asia Pacific) Japan East": "japaneast",
+            "(Asia Pacific) Korea Central": "koreacentral",
+            "(Canada) Canada Central": "canadacentral",
+            "(Europe) France Central": "francecentral",
+            "(Europe) Germany West Central": "germanywestcentral",
+            "(Europe) Italy North": "italynorth",
+            "(Europe) Norway East": "norwayeast",
+            "(Europe) Poland Central": "polandcentral",
+            "(Europe) Spain Central": "spaincentral",
+            "(Europe) Switzerland North": "switzerlandnorth",
+            "(Mexico) Mexico Central": "mexicocentral",
+            "(Middle East) UAE North": "uaenorth",
+            "(South America) Brazil South": "brazilsouth",
+            "(Middle East) Israel Central": "israelcentral",
+            "(Middle East) Qatar Central": "qatarcentral",
+            "(US) Central US (Stage)": "centralusstage",
+            "(US) East US (Stage)": "eastusstage",
+            "(US) East US 2 (Stage)": "eastus2stage",
+            "(US) North Central US (Stage)": "northcentralusstage",
+            "(US) South Central US (Stage)": "southcentralusstage",
+            "(US) West US (Stage)": "westusstage",
+            "(US) West US 2 (Stage)": "westus2stage",
+            "(Asia Pacific) East Asia (Stage)": "eastasiastage",
+            "(Asia Pacific) Southeast Asia (Stage)": "southeastasiastage",
+            "(South America) Brazil US": "brazilus",
+            "(US) East US 2": "eastus2",
+            "(US) East US STG": "eastusstg",
+            "(US) North Central US": "northcentralus",
+            "(US) West US": "westus",
+            "(Asia Pacific) Japan West": "japanwest",
+            "(Asia Pacific) Jio India West": "jioindiawest",
+            "(US) Central US EUAP": "centraluseuap",
+            "(US) East US 2 EUAP": "eastus2euap",
+            "(US) West Central US": "westcentralus",
+            "(Africa) South Africa West": "southafricawest",
+            "(Asia Pacific) Australia Central": "australiacentral",
+            "(Asia Pacific) Australia Central 2": "australiacentral2",
+            "(Asia Pacific) Australia Southeast": "australiasoutheast",
+            "(Asia Pacific) Jio India Central": "jioindiacentral",
+            "(Asia Pacific) Korea South": "koreasouth",
+            "(Asia Pacific) South India": "southindia",
+            "(Asia Pacific) West India": "westindia",
+            "(Canada) Canada East": "canadaeast",
+            "(Europe) France South": "francesouth",
+            "(Europe) Germany North": "germanynorth",
+            "(Europe) Norway West": "norwaywest",
+            "(Europe) Switzerland West": "switzerlandwest",
+            "(Europe) UK West": "ukwest",
+            "(Middle East) UAE Central": "uaecentral",
+            "(South America) Brazil Southeast": "brazilsoutheast",
+        }
+
+    @classmethod
+    def skypilot_default_regions(cls) -> Dict[str, str]:
+        """Returns the regions supported by default for the Skypilot.
+
+        Returns:
+            The regions supported by default for the Skypilot.
+        """
+        matcher = re.compile(r".*us\d*( |$)")
+        return {
+            k: v
+            for k, v in cls.locations().items()
+            if "(US)" in k and matcher.match(v)
+        }

zenml/stack_deployments/gcp_stack_deployment.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Functionality to deploy a ZenML stack to GCP."""
 
+import re
 from typing import ClassVar, Dict, List
 
 from zenml.enums import StackDeploymentProvider
@@ -208,6 +209,16 @@ GCP project and to clean up the resources created by the stack by using
             "US West (Las Vegas)": "us-west4",
         }
 
+    @classmethod
+    def skypilot_default_regions(cls) -> Dict[str, str]:
+        """Returns the regions supported by default for the Skypilot.
+
+        Returns:
+            The regions supported by default for the Skypilot.
+        """
+        matcher = re.compile(r"us-.*")
+        return {k: v for k, v in cls.locations().items() if matcher.match(v)}
+
     def get_deployment_config(
         self,
     ) -> StackDeploymentConfig:

zenml/stack_deployments/stack_deployment.py

@@ -105,6 +105,15 @@ class ZenMLCloudStackDeployment(BaseModel):
             names to region descriptions.
         """
 
+    @classmethod
+    def skypilot_default_regions(cls) -> Dict[str, str]:
+        """Returns the regions supported by default for the Skypilot.
+
+        Returns:
+            The regions supported by default for the Skypilot.
+        """
+        return cls.locations()
+
     @classmethod
     def get_deployment_info(cls) -> StackDeploymentInfo:
         """Return information about the ZenML Cloud Stack Deployment.
@@ -120,6 +129,7 @@ class ZenMLCloudStackDeployment(BaseModel):
             integrations=cls.integrations(),
             permissions=cls.permissions(),
             locations=cls.locations(),
+            skypilot_default_regions=cls.skypilot_default_regions(),
         )
 
     @abstractmethod

zenml/stack_deployments/utils.py

@@ -19,6 +19,9 @@ from zenml.enums import StackDeploymentProvider
 from zenml.stack_deployments.aws_stack_deployment import (
     AWSZenMLCloudStackDeployment,
 )
+from zenml.stack_deployments.azure_stack_deployment import (
+    AZUREZenMLCloudStackDeployment,
+)
 from zenml.stack_deployments.gcp_stack_deployment import (
     GCPZenMLCloudStackDeployment,
 )
@@ -27,6 +30,7 @@ from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
 STACK_DEPLOYMENT_PROVIDERS = {
     StackDeploymentProvider.AWS: AWSZenMLCloudStackDeployment,
     StackDeploymentProvider.GCP: GCPZenMLCloudStackDeployment,
+    StackDeploymentProvider.AZURE: AZUREZenMLCloudStackDeployment,
 }
 
 

zenml/zen_stores/schemas/component_schemas.py

@@ -38,6 +38,7 @@ from zenml.zen_stores.schemas.user_schemas import UserSchema
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
+    from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema
     from zenml.zen_stores.schemas.logs_schemas import LogsSchema
     from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
     from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
@@ -85,6 +86,13 @@ class StackComponentSchema(NamedSchema, table=True):
     run_metadata: List["RunMetadataSchema"] = Relationship(
         back_populates="stack_component",
     )
+    flavor_schema: Optional["FlavorSchema"] = Relationship(
+        sa_relationship_kwargs={
+            "primaryjoin": "and_(foreign(StackComponentSchema.flavor) == FlavorSchema.name, foreign(StackComponentSchema.type) == FlavorSchema.type)",
+            "lazy": "joined",
+            "uselist": False,
+        },
+    )
 
     run_or_step_logs: List["LogsSchema"] = Relationship(
         back_populates="artifact_store",
@@ -161,6 +169,12 @@ class StackComponentSchema(NamedSchema, table=True):
             user=self.user.to_model() if self.user else None,
             created=self.created,
             updated=self.updated,
+            logo_url=self.flavor_schema.logo_url
+            if self.flavor_schema
+            else None,
+            integration=self.flavor_schema.integration
+            if self.flavor_schema
+            else None,
         )
         metadata = None
         if include_metadata:

zenml/zen_stores/sql_zen_store.py

@@ -286,6 +286,7 @@ from zenml.service_connectors.service_connector_registry import (
     service_connector_registry,
 )
 from zenml.stack.flavor_registry import FlavorRegistry
+from zenml.stack_deployments.utils import get_stack_deployment_class
 from zenml.utils import uuid_utils
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.networking_utils import (
@@ -3185,6 +3186,29 @@ class SqlZenStore(BaseZenStore):
                         "found."
                     )
 
+            # warn about skypilot regions, if needed
+            if component.flavor in {"vm_gcp", "vm_azure"}:
+                stack_deployment_class = get_stack_deployment_class(
+                    StackDeploymentProvider.GCP
+                    if component.flavor == "vm_gcp"
+                    else StackDeploymentProvider.AZURE
+                )
+                skypilot_regions = (
+                    stack_deployment_class.skypilot_default_regions().values()
+                )
+                if (
+                    component.configuration.get("region", None)
+                    and component.configuration["region"]
+                    not in skypilot_regions
+                ):
+                    logger.warning(
+                        f"Region `{component.configuration['region']}` is not enabled in Skypilot "
+                        f"by default. Supported regions by default are: {skypilot_regions}. "
+                        "Check the Skypilot documentation to learn how to enable regions rather "
+                        "than default ones. (If you have already extended your configuration - "
+                        "simply ignore this warning)"
+                    )
+
             # Create the component
             new_component = StackComponentSchema(
                 name=component.name,
@@ -7032,9 +7056,44 @@ class SqlZenStore(BaseZenStore):
             # Service Connectors
             service_connectors: List[ServiceConnectorResponse] = []
 
+            need_to_generate_permanent_tokens = False
+            orchestrator_component = full_stack.components[
+                StackComponentType.ORCHESTRATOR
+            ]
+            if isinstance(orchestrator_component, UUID):
+                orchestrator = self.get_stack_component(
+                    orchestrator_component,
+                    hydrate=False,
+                )
+                need_to_generate_permanent_tokens = (
+                    orchestrator.flavor.startswith("vm_")
+                )
+            else:
+                need_to_generate_permanent_tokens = (
+                    orchestrator_component.flavor.startswith("vm_")
+                )
+
             for connector_id_or_info in full_stack.service_connectors:
                 # Fetch an existing service connector
                 if isinstance(connector_id_or_info, UUID):
+                    existing_service_connector = self.get_service_connector(
+                        connector_id_or_info
+                    )
+                    if need_to_generate_permanent_tokens:
+                        if (
+                            existing_service_connector.configuration.get(
+                                "generate_temporary_tokens", None
+                            )
+                            is not False
+                        ):
+                            self.update_service_connector(
+                                existing_service_connector.id,
+                                ServiceConnectorUpdate(
+                                    configuration=existing_service_connector.configuration.update(
+                                        {"generate_temporary_tokens": False}
+                                    )
+                                ),
+                            )
                     service_connectors.append(
                         self.get_service_connector(connector_id_or_info)
                     )
@@ -7047,7 +7106,11 @@ class SqlZenStore(BaseZenStore):
                                 name=connector_name,
                                 connector_type=connector_id_or_info.type,
                                 auth_method=connector_id_or_info.auth_method,
-                                configuration=connector_id_or_info.configuration,
+                                configuration=connector_id_or_info.configuration.update(
+                                    {
+                                        "generate_temporary_tokens": not need_to_generate_permanent_tokens
+                                    }
+                                ),
                                 user=full_stack.user,
                                 workspace=full_stack.workspace,
                                 labels={

{zenml_nightly-0.62.0.dev20240717.dist-info → zenml_nightly-0.62.0.dev20240720.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.62.0.dev20240717
+Version: 0.62.0.dev20240720
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.62.0.dev20240717.dist-info → zenml_nightly-0.62.0.dev20240720.dist-info}/RECORD

@@ -6,7 +6,7 @@ RELEASE_NOTES.md,sha256=rF05H0U7U4lCYWomDqKhWdaWtOfCFSOf10LNBoFnl3Y,349994
 ROADMAP.md,sha256=hiLSmr16BH8Dfx7SaQM4JcXCGCVl6mFZPFAwJeDTrJU,407
 SECURITY.md,sha256=9DepA8y03yvCZLHEfcXLTDH4lUyKHquAdukBsccNN7c,682
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=5f5pNofWBEtoppIGNX90R6zf6K2wZJhhFWS8AdpPMjc,19
+zenml/VERSION,sha256=K_-fpypC0EJWSkrK7_ymOkOfKqLH1wClnk7HXn4OH6o,19
 zenml/__init__.py,sha256=6zRqB0FeBWX5E4Np9k9jzJgKaEHKvLFbtmOZQYrGpD8,2453
 zenml/_hub/__init__.py,sha256=6qDzpQAAZa__Aiiz0mC1qM-9dw9_jk_v_aXeJknxbDE,644
 zenml/_hub/client.py,sha256=Um2df1JO7-BmNm65efHSPpV5e0GITuoQJmod_wkdvbw,9136
@@ -58,7 +58,7 @@ zenml/cli/served_model.py,sha256=lIPd6kOS9kfWFOsd3DSILtiYomINpmXrdpyP5vtQlj8,146
 zenml/cli/server.py,sha256=N2cvu9C_z8CzwGu3a7P6fE2X8Se_rHEx30iq9nYMy9o,28682
 zenml/cli/service_accounts.py,sha256=MO8pVVjppdvIULFxpvQKlhfEDmRqvlH8aIbPTdNSHXU,16489
 zenml/cli/service_connectors.py,sha256=v6o-FpTEnDXYcQXczBELwfWPCcYBi9z7b1GC90Fvp30,74397
-zenml/cli/stack.py,sha256=M3V6TUCqWQLnHPtVUNsdZyZcid0DYes9JPYyBl5--_k,82437
+zenml/cli/stack.py,sha256=xNB2sXXFJe3VC_Gow-xWm5s3ae3NKFS2pXAiYeviUlg,82520
 zenml/cli/stack_components.py,sha256=BxWsqeu5e2DtP5TMjzE98d8Lqfwe5o45TC2awjil5wY,74144
 zenml/cli/stack_recipes.py,sha256=ZhJE2HdytilaybhJgPkNp9DvSa5LiVgxS19jR0UxHuk,15716
 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
 zenml/entrypoints/entrypoint.py,sha256=Cxh4S_F4URQQ1WF2fGuJEIfLcUWcxVDeDiKYJA4Kxps,2174
 zenml/entrypoints/pipeline_entrypoint_configuration.py,sha256=To-vTP29qAE36ndJDF1fRw9wL2Nk2bsBuO-ayAwvSmo,1646
 zenml/entrypoints/step_entrypoint_configuration.py,sha256=NskQF6OFFCIoD7v0JxR_cZRVXBRaKl3R0fwqpaQCPLM,6571
-zenml/enums.py,sha256=3_SpffPCzk0MRPUiqKl20qGyq6UOBHzK6F58we7pEDo,10033
+zenml/enums.py,sha256=tyte4wjKGy3bmTydYMkJaXGaxPKED21XhNXWOqYpM5c,10053
 zenml/environment.py,sha256=xqsj8G-qr_W2KyyWnJABMsNQ7Tl32MUhFir7kUDPpKc,20276
 zenml/event_hub/__init__.py,sha256=-fD9mPOslf4J-_XFBPp5gYlPz7-6ZaAKHa5jxf_nyAo,757
 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
 zenml/integrations/aws/step_operators/__init__.py,sha256=HK5oIqLixYKUKoiN4LnqTyGjAZJUdUqWYGqJhFSWyo0,817
 zenml/integrations/aws/step_operators/sagemaker_step_operator.py,sha256=-y1zqfIDJm2wwy_KETJ1nHdKuSbIHPFK3bT-0i9tpoA,10113
 zenml/integrations/aws/step_operators/sagemaker_step_operator_entrypoint_config.py,sha256=2cXroe6-bXyHVkO5yPnZagNpqx6MrMDcvuvNr8J2j-A,1581
-zenml/integrations/azure/__init__.py,sha256=TlRmc8-FNjRA84OW6ikJgSKu_N5assjB9I2bc7VAjU0,2570
+zenml/integrations/azure/__init__.py,sha256=Ttt7jgnpIFu-AWiVBKFbPmuLX6vdMZjTEV1IVwfZ2yI,2562
 zenml/integrations/azure/artifact_stores/__init__.py,sha256=dlIwbpgjE0Hy4rhMbelNJHVKm4t8tj_hRu9mQ_cEIAg,820
 zenml/integrations/azure/artifact_stores/azure_artifact_store.py,sha256=uVMpsWDIyeytn8aHxBAma5pdPwiWuK5fwew2Tvxsk3w,11955
 zenml/integrations/azure/flavors/__init__.py,sha256=KEwsaLeKsvdI5tqYFzji_VXBLlZuGoQXXmOT8vZ8PKE,1069
@@ -462,7 +462,7 @@ zenml/integrations/skypilot_aws/flavors/__init__.py,sha256=UGL7NuZ2BkB5IxL9WUHmG
 zenml/integrations/skypilot_aws/flavors/skypilot_orchestrator_aws_vm_flavor.py,sha256=PqcUCXRyJoT7jmR15A5d7N5NJxXkgLKWm8BC3YW6Omw,3789
 zenml/integrations/skypilot_aws/orchestrators/__init__.py,sha256=xxFZ2iM9H9LVKfSB8N-yKGxuB-lJE3Y368lcj7MfhI0,1006
 zenml/integrations/skypilot_aws/orchestrators/skypilot_aws_vm_orchestrator.py,sha256=0uFfRa8lGsc7D_zdrGGMqKoL6eKUwuf2ImtDw_RRg6Q,3094
-zenml/integrations/skypilot_azure/__init__.py,sha256=n8ODpRF2blirwCMUL5SKdioDSja0135pubcFeI4rUcA,1740
+zenml/integrations/skypilot_azure/__init__.py,sha256=HuOab2Bv8Fuz5n6_nNiR7_Brdzd8UjgRggosRtql8r0,1758
 zenml/integrations/skypilot_azure/flavors/__init__.py,sha256=9JB07IZbIZhY6yL9TjkYFTCUUfxajMajsMbbjC1BZ_s,1024
 zenml/integrations/skypilot_azure/flavors/skypilot_orchestrator_azure_vm_flavor.py,sha256=aL8ibfQGo2IcyckaesjAFk6m8kazy5pq6ycArYRD5w0,3858
 zenml/integrations/skypilot_azure/orchestrators/__init__.py,sha256=gW7JOTLdGB542Cv1ydSqF4kVBAo7LUnfkrlj45kCWm8,1017
@@ -586,7 +586,7 @@ zenml/models/v2/core/artifact_version.py,sha256=kmMmYWnN6DO6jPWb-Ct89d0p9Vj4c32_
 zenml/models/v2/core/artifact_visualization.py,sha256=milAHTpfuYd8Dx_oB4hcPYJL-TMoiU0FwUI0UhB3F-E,3115
 zenml/models/v2/core/code_reference.py,sha256=9xy3Gzychq16lLUs01va3gGfwdu-AMrFgys_ua89uFY,3557
 zenml/models/v2/core/code_repository.py,sha256=S_p-cC99g0IxwWc8ux1enveM6pXhkyKMF8sQiUnR7qc,5925
-zenml/models/v2/core/component.py,sha256=zdXWZcVCSKDBjEfTQq44wGMwl6E3TKx3IapXzYCryhI,12619
+zenml/models/v2/core/component.py,sha256=7isCreRBVolPxMp1VrkoNsyOA4LdSY5zMaXXVHbY_xo,13392
 zenml/models/v2/core/device.py,sha256=Vlz4-6VwglKL3F-k7DObH-2LeHc8oRqKbYxSImHvDJk,13929
 zenml/models/v2/core/event_source.py,sha256=kOexFvxj3iqsQGiGHPe549v9cSTo2GfsPCK05eDiNCw,6899
 zenml/models/v2/core/event_source_flavor.py,sha256=TgrGkypnJUutMvvHE7Y9MbUiEeNXZxitc3eYlIbbsas,2077
@@ -625,7 +625,7 @@ zenml/models/v2/misc/hub_plugin_models.py,sha256=ogWYrotKq__2vR2LsrInlOSR-DJqcqd
 zenml/models/v2/misc/loaded_visualization.py,sha256=u6lapDNZDtU9eS-_EMzl00urj0yPNhiqhZcvjIz7DSg,946
 zenml/models/v2/misc/server_models.py,sha256=wd14abzVBnv4lBIb_dbBSTkqGcdQjI-LcdoZC_2ucSw,3613
 zenml/models/v2/misc/service_connector_type.py,sha256=jGJLvIsBB87ZIEeZCLLueUf28H1QJFqXOJYtiRVvZ4M,28748
-zenml/models/v2/misc/stack_deployment.py,sha256=tSqb-E9ErvsNpeM0WzXxpTcOlXYLQ9-x3ioaLzDyOcs,3315
+zenml/models/v2/misc/stack_deployment.py,sha256=D_eIGhM6zZysiNOQvDXbjQN11pguVJyD9jJIMECwNhk,3622
 zenml/models/v2/misc/user_auth.py,sha256=XLhagiLN27wlS70LVVvF5PVBDBO4PXBRgXo2VGtVbnI,7198
 zenml/new/__init__.py,sha256=knhroJ2h0uHBCGzAiBBGJEiuhEA3cwI6XYBRIyXdbkQ,613
 zenml/new/pipelines/__init__.py,sha256=knhroJ2h0uHBCGzAiBBGJEiuhEA3cwI6XYBRIyXdbkQ,613
@@ -678,7 +678,7 @@ zenml/service_connectors/__init__.py,sha256=gIBAVUPBZtdO6JtEBCixy9YG4wZRA1mnaxaG
 zenml/service_connectors/docker_service_connector.py,sha256=vplyZNUxZhAyzXRdtAiSGMFZT87h9Wq8izE0HWh7OSQ,12441
 zenml/service_connectors/service_connector.py,sha256=HW8dTE826ZSD1vWQlqQRm8XYq08v7tDZQ8ORnJ5rTcE,55351
 zenml/service_connectors/service_connector_registry.py,sha256=lP0s0k7RIp9WZEW36Q0ygHD5urzF-wUwh41rbXWS5ss,9186
-zenml/service_connectors/service_connector_utils.py,sha256=TGk9IIw4zhxvcNqZDQHUGs4i1Edrh3yeM0wTETDekdQ,14428
+zenml/service_connectors/service_connector_utils.py,sha256=pUen7wvsxNf12VmpGp5PvSiPWllNXbgrKH559p30U34,17195
 zenml/services/__init__.py,sha256=8aSLtYlX9EVLB1Se1I4dbhZdTqrjNrAfsn1d_RCc_nk,2800
 zenml/services/container/__init__.py,sha256=dFHcmlXgNXjUaCR18oGQJ19-I_6f3UeAUURHjqldjTs,668
 zenml/services/container/container_service.py,sha256=DJIK7DD53Q_mgZByM0oeIyAgkrwN12T5E3QCWjubkLs,18383
@@ -705,9 +705,10 @@ zenml/stack/stack_validator.py,sha256=hWbvvGIeWLj6NwSsF4GCc6RAxAWvxHXTcBZL9nJvca
 zenml/stack/utils.py,sha256=W_ZKttALLGABYSzqapkIdHerbOWf6gRL8Tl9Bgj_q5s,5460
 zenml/stack_deployments/__init__.py,sha256=-7593cQ_ZgRn774Ol-8AKXXQquIU4DSiaThVEr6TfWM,644
 zenml/stack_deployments/aws_stack_deployment.py,sha256=gRjG7aCK8cl-AxFOgI00eTMIpaIGXfsd9UfGmbQXLJ4,10239
-zenml/stack_deployments/gcp_stack_deployment.py,sha256=gLo1tLJaK8sW1N2C04Pdk0LDv_Y-wtm33NYqlIdUX_k,11000
-zenml/stack_deployments/stack_deployment.py,sha256=XTURbeOUMdByRO0jMN1tXY1ryrDZ2Z83rOzzwLtTylo,6871
-zenml/stack_deployments/utils.py,sha256=Vpt5hWJLvaHyPfYJ7vx0CEqrTrGCVXy2DEB1uTLr2EA,1572
+zenml/stack_deployments/azure_stack_deployment.py,sha256=qFPcTfubZWsbXPP-xuTLKKXfOIqA7o36eZ2mWTriRfM,7332
+zenml/stack_deployments/gcp_stack_deployment.py,sha256=XRwAB8tq1hAB2jlr3c7TqvNH-HWjCgs-_TP4kKK7y-U,11365
+zenml/stack_deployments/stack_deployment.py,sha256=R57UqafZshcs0DXLKX2OWll_65UzyHK8w4kf8tS4Oqw,7209
+zenml/stack_deployments/utils.py,sha256=4rU0JM6wAmSLLlH7f9xk3g3EsacMHQYDTgN1OvU9pRs,1738
 zenml/step_operators/__init__.py,sha256=tqj7fgnQyZubLjwUu4ITwkA-70KMQz4g37agbfF7ZgE,1228
 zenml/step_operators/base_step_operator.py,sha256=ZRnY6lcEHY8xZskrKKdPhgKg2BlEoh2_kb8xCaM2D1g,3522
 zenml/step_operators/step_operator_entrypoint_configuration.py,sha256=vG0Z9D2C72PYOZ_P3eXYTmm2TvL-dS_kHhSMq-coAKE,3690
@@ -1301,7 +1302,7 @@ zenml/zen_stores/schemas/artifact_schemas.py,sha256=TEOD23Mt8y5BsOdDbLnsUR-KrxgA
 zenml/zen_stores/schemas/artifact_visualization_schemas.py,sha256=zv_Mg--WMTtAPEy5PXIkZlxq-oLrXZNLHXb_YgZbm0g,3667
 zenml/zen_stores/schemas/base_schemas.py,sha256=T3d2IioA2XcCoeaHfEb_mp51ai1djegAi9BIHRE18s0,2015
 zenml/zen_stores/schemas/code_repository_schemas.py,sha256=_v-HozD6UbyQgnQ_u4jo2vFsbWzUnQCMARTz4XEyhq0,7405
-zenml/zen_stores/schemas/component_schemas.py,sha256=8lzoUj2MBVDn-GMtYUV9Qoy1ps2sMSvmJjiyfcGPCKA,6358
+zenml/zen_stores/schemas/component_schemas.py,sha256=KD5KdjG6s6LPSqVgtuX9IXUvrb2NyLrO46V9uSifb34,6966
 zenml/zen_stores/schemas/device_schemas.py,sha256=8OOnfVYG9GFmYC8_vysHrL2TGQ4iu4-7r_BivITWsxw,8934
 zenml/zen_stores/schemas/event_source_schemas.py,sha256=3CLt-TfJFQF-xJ6j5UnBFhNXMQTM-pLdIBazghyGmQw,6386
 zenml/zen_stores/schemas/flavor_schemas.py,sha256=rQl5_qSXlvTXwyzq4QPtKCMBeT4hRJmPPNsh9FGx38U,5053
@@ -1334,10 +1335,10 @@ zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py,sha256=NfW1EHIA99lseb
 zenml/zen_stores/secrets_stores/secrets_store_interface.py,sha256=Q2Jbnt2Pp7NGlR-u1YBfRZV2g8su2Fd0ArBMdksAE-Q,2819
 zenml/zen_stores/secrets_stores/service_connector_secrets_store.py,sha256=kPYX-Z_OOhZCI1CP77ncfV7IsV4e8brklnTXmKxZYNc,7078
 zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=Bq1djrUP9saoD7vECjS7-TlA_7sjJGgw1talri4spjU,8656
-zenml/zen_stores/sql_zen_store.py,sha256=K4FAR764dyEKFwDZTCHgnfseTx_tBtvHNSdnTWywGx8,378688
+zenml/zen_stores/sql_zen_store.py,sha256=j1D_vqLzg8lK2wT-a8moDOXJfTVGAa_oA0bb-RHWEVo,381780
 zenml/zen_stores/zen_store_interface.py,sha256=ZcgPtlAMdO5die42Hwd6l8glxHYsxFjzuCBewIu5Hrs,91984
-zenml_nightly-0.62.0.dev20240717.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.62.0.dev20240717.dist-info/METADATA,sha256=pozCSjQQPPCIOrAqtmaXTgPZ4j_ihTSkKXic2qHAuNs,21026
-zenml_nightly-0.62.0.dev20240717.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zenml_nightly-0.62.0.dev20240717.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.62.0.dev20240717.dist-info/RECORD,,
+zenml_nightly-0.62.0.dev20240720.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.62.0.dev20240720.dist-info/METADATA,sha256=LMGux7_P9Wjx-G4w8QhzCuEjg0Zk1sbgZ4kIiwtZYhI,21026
+zenml_nightly-0.62.0.dev20240720.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zenml_nightly-0.62.0.dev20240720.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.62.0.dev20240720.dist-info/RECORD,,