zenml-nightly 0.61.0.dev20240712__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/zen_stores/sql_zen_store.py

@@ -33,6 +33,7 @@ from typing import (
     NoReturn,
     Optional,
     Sequence,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -107,6 +108,7 @@ from zenml.enums import (
     ExecutionStatus,
     LoggingLevels,
     ModelStages,
+    OnboardingStep,
     SecretScope,
     SecretsStoreType,
     SorterOps,
@@ -284,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 (
@@ -797,6 +800,7 @@ class SqlZenStore(BaseZenStore):
     _secrets_store: Optional[BaseSecretsStore] = None
     _backup_secrets_store: Optional[BaseSecretsStore] = None
     _should_send_user_enriched_events: bool = False
+    _cached_onboarding_state: Optional[Set[str]] = None
 
     @property
     def secrets_store(self) -> "BaseSecretsStore":
@@ -1688,6 +1692,60 @@ class SqlZenStore(BaseZenStore):
 
             return settings.to_model(include_metadata=True)
 
+    def get_onboarding_state(self) -> List[str]:
+        """Get the server onboarding state.
+
+        Returns:
+            The server onboarding state.
+        """
+        with Session(self.engine) as session:
+            settings = self._get_server_settings(session=session)
+            if settings.onboarding_state:
+                self._cached_onboarding_state = set(
+                    json.loads(settings.onboarding_state)
+                )
+                return list(self._cached_onboarding_state)
+            else:
+                return []
+
+    def _update_onboarding_state(
+        self, completed_steps: Set[str], session: Session
+    ) -> None:
+        """Update the server onboarding state.
+
+        Args:
+            completed_steps: Newly completed onboarding steps.
+            session: DB session.
+        """
+        if self._cached_onboarding_state and completed_steps.issubset(
+            self._cached_onboarding_state
+        ):
+            # All the onboarding steps are already completed, no need to query
+            # the DB
+            return
+
+        settings = self._get_server_settings(session=session)
+        settings.update_onboarding_state(completed_steps=completed_steps)
+        session.add(settings)
+        session.commit()
+        session.refresh(settings)
+
+        assert settings.onboarding_state
+        self._cached_onboarding_state = set(
+            json.loads(settings.onboarding_state)
+        )
+
+    def update_onboarding_state(self, completed_steps: Set[str]) -> None:
+        """Update the server onboarding state.
+
+        Args:
+            completed_steps: Newly completed onboarding steps.
+        """
+        with Session(self.engine) as session:
+            self._update_onboarding_state(
+                completed_steps=completed_steps, session=session
+            )
+
     def activate_server(
         self, request: ServerActivationRequest
     ) -> Optional[UserResponse]:
@@ -3128,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,
@@ -6156,6 +6237,7 @@ class SqlZenStore(BaseZenStore):
 
             connector = new_service_connector.to_model(include_metadata=True)
             self._populate_connector_type(connector)
+
             return connector
 
     def get_service_connector(
@@ -6931,6 +7013,16 @@ class SqlZenStore(BaseZenStore):
             session.commit()
             session.refresh(new_stack_schema)
 
+            for component in defined_components:
+                if component.type == StackComponentType.ORCHESTRATOR:
+                    if component.flavor not in {"local", "local_docker"}:
+                        self._update_onboarding_state(
+                            completed_steps={
+                                OnboardingStep.STACK_WITH_REMOTE_ORCHESTRATOR_CREATED
+                            },
+                            session=session,
+                        )
+
             return new_stack_schema.to_model(include_metadata=True)
 
     def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse:
@@ -6964,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)
                     )
@@ -6979,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={
@@ -7965,6 +8096,25 @@ class SqlZenStore(BaseZenStore):
                         "duration_seconds": duration_seconds,
                         **stack_metadata,
                     }
+
+                completed_onboarding_steps: Set[str] = {
+                    OnboardingStep.PIPELINE_RUN,
+                    OnboardingStep.STARTER_SETUP_COMPLETED,
+                }
+                if stack_metadata["orchestrator"] not in {
+                    "local",
+                    "local_docker",
+                }:
+                    completed_onboarding_steps.update(
+                        {
+                            OnboardingStep.PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR,
+                            OnboardingStep.PRODUCTION_SETUP_COMPLETED,
+                        }
+                    )
+
+                self._update_onboarding_state(
+                    completed_steps=completed_onboarding_steps, session=session
+                )
             pipeline_run.update(run_update)
             session.add(pipeline_run)
 

{zenml_nightly-0.61.0.dev20240712.dist-info → zenml_nightly-0.62.0.dev20240726.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.61.0.dev20240712
+Version: 0.62.0.dev20240726
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -59,7 +59,7 @@ Requires-Dist: copier (>=8.1.0) ; extra == "templates"
 Requires-Dist: coverage[toml] (>=5.5,<6.0) ; extra == "dev"
 Requires-Dist: darglint (>=1.8.1,<2.0.0) ; extra == "dev"
 Requires-Dist: distro (>=1.6.0,<2.0.0)
-Requires-Dist: docker (>=6.1.0,<6.2.0)
+Requires-Dist: docker (>=7.1.0,<7.2.0)
 Requires-Dist: fastapi (>=0.100,<=0.110) ; extra == "server"
 Requires-Dist: gcsfs (>=2022.11.0) ; extra == "gcsfs"
 Requires-Dist: gitpython (>=3.1.18,<4.0.0)
@@ -83,6 +83,7 @@ Requires-Dist: mlstacks (==0.9.0) ; extra == "mlstacks"
 Requires-Dist: mypy (==1.7.1) ; extra == "dev"
 Requires-Dist: numpy (<2.0.0)
 Requires-Dist: orjson (>=3.10.0,<3.11.0) ; extra == "server"
+Requires-Dist: packaging (>=24.1)
 Requires-Dist: pandas (>=1.1.5)
 Requires-Dist: passlib[bcrypt] (>=1.7.4,<1.8.0)
 Requires-Dist: psutil (>=5.0.0)
@@ -104,7 +105,6 @@ Requires-Dist: python-multipart (>=0.0.9,<0.1.0) ; extra == "server"
 Requires-Dist: python-terraform (>=0.10.1,<0.11.0) ; extra == "terraform"
 Requires-Dist: pyyaml (>=6.0.1)
 Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
-Requires-Dist: requests (<2.32.0)
 Requires-Dist: rich[jupyter] (>=12.0.0)
 Requires-Dist: ruff (>=0.1.7) ; extra == "templates" or extra == "dev"
 Requires-Dist: s3fs (>=2022.11.0) ; extra == "s3fs"
@@ -458,7 +458,7 @@ the Apache License Version 2.0.
     ·
     <a href="https://github.com/zenml-io/zenml/issues">Report Bug</a>
     ·
-    <a href="https://zenml.io/cloud">Sign up for ZenML Pro</a>
+    <a href="https://zenml.io/pro">Sign up for ZenML Pro</a>
     ·
     <a href="https://www.zenml.io/blog">Read Blog</a>
     ·
@@ -467,7 +467,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.61.0 is out. Check out the release notes
+    🎉 Version 0.62.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.