zenml-nightly 0.61.0.dev20240711__py3-none-any.whl → 0.61.0.dev20240713__py3-none-any.whl

zenml/zen_stores/migrations/versions/b4fca5241eea_migrate_onboarding_state.py

@@ -0,0 +1,167 @@
+"""Migrate onboarding state [b4fca5241eea].
+
+Revision ID: b4fca5241eea
+Revises: 0.61.0
+Create Date: 2024-06-20 15:01:22.414801
+
+"""
+
+import json
+from typing import Dict, List
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b4fca5241eea"
+down_revision = "0.61.0"
+branch_labels = None
+depends_on = None
+
+
+ONBOARDING_KEY_MAPPING = {
+    "connect_zenml": ["device_verified"],
+    "run_first_pipeline": ["pipeline_run", "starter_setup_completed"],
+    "run_remote_pipeline": [
+        "production_setup_completed",
+    ],
+}
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    connection = op.get_bind()
+
+    meta = sa.MetaData()
+    meta.reflect(only=("server_settings",), bind=connection)
+
+    server_settings_table = sa.Table("server_settings", meta)
+
+    existing_onboarding_state = connection.execute(
+        sa.select(server_settings_table.c.onboarding_state)
+    ).scalar_one_or_none()
+
+    new_state = []
+
+    if existing_onboarding_state:
+        # There was already an existing onboarding state in the DB
+        # -> Migrate to the new server keys
+        state = json.loads(existing_onboarding_state)
+
+        if isinstance(state, Dict):
+            for key in state.keys():
+                if key in ONBOARDING_KEY_MAPPING:
+                    new_state.extend(ONBOARDING_KEY_MAPPING[key])
+        elif isinstance(state, List):
+            # Somehow the state is already converted, probably shouldn't happen
+            return
+
+    # We now query the DB and complete all onboarding steps that we can detect
+    # from the database
+    meta = sa.MetaData()
+    meta.reflect(
+        only=(
+            "pipeline_run",
+            "stack_component",
+            "stack",
+            "stack_composition",
+            "pipeline_deployment",
+        ),
+        bind=connection,
+    )
+
+    pipeline_run_table = sa.Table("pipeline_run", meta)
+    stack_component_table = sa.Table("stack_component", meta)
+    stack_table = sa.Table("stack", meta)
+    stack_composition_table = sa.Table("stack_composition", meta)
+    pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
+
+    pipeline_run_count = connection.execute(
+        sa.select(sa.func.count(pipeline_run_table.c.id))
+    ).scalar()
+    if pipeline_run_count and pipeline_run_count > 0:
+        new_state.extend(ONBOARDING_KEY_MAPPING["run_first_pipeline"])
+
+    stack_with_remote_orchestrator_count = connection.execute(
+        sa.select(sa.func.count(stack_table.c.id))
+        .where(stack_composition_table.c.stack_id == stack_table.c.id)
+        .where(
+            stack_composition_table.c.component_id
+            == stack_component_table.c.id
+        )
+        .where(
+            stack_component_table.c.flavor.not_in(["local", "local_docker"])
+        )
+        .where(stack_component_table.c.type == "orchestrator")
+    ).scalar()
+    if (
+        stack_with_remote_orchestrator_count
+        and stack_with_remote_orchestrator_count > 0
+    ):
+        new_state.append("stack_with_remote_orchestrator_created")
+
+    pipeline_run_with_remote_artifact_store_count = connection.execute(
+        sa.select(sa.func.count(pipeline_run_table.c.id))
+        .where(
+            pipeline_run_table.c.deployment_id
+            == pipeline_deployment_table.c.id
+        )
+        .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
+        .where(stack_composition_table.c.stack_id == stack_table.c.id)
+        .where(
+            stack_composition_table.c.component_id
+            == stack_component_table.c.id
+        )
+        .where(stack_component_table.c.flavor != "local")
+        .where(stack_component_table.c.type == "artifact_store")
+    ).scalar()
+    if (
+        pipeline_run_with_remote_artifact_store_count
+        and pipeline_run_with_remote_artifact_store_count > 0
+    ):
+        new_state.append("production_setup_completed")
+
+    pipeline_run_with_remote_orchestrator_count = connection.execute(
+        sa.select(sa.func.count(pipeline_run_table.c.id))
+        .where(
+            pipeline_run_table.c.deployment_id
+            == pipeline_deployment_table.c.id
+        )
+        .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
+        .where(stack_composition_table.c.stack_id == stack_table.c.id)
+        .where(
+            stack_composition_table.c.component_id
+            == stack_component_table.c.id
+        )
+        .where(
+            stack_component_table.c.flavor.not_in(["local", "local_docker"])
+        )
+        .where(stack_component_table.c.type == "orchestrator")
+    ).scalar()
+    if (
+        pipeline_run_with_remote_orchestrator_count
+        and pipeline_run_with_remote_orchestrator_count > 0
+    ):
+        new_state.append("pipeline_run_with_remote_orchestrator")
+        new_state.append("production_setup_completed")
+
+    if new_state:
+        # If any of the items are finished, we also complete the initial
+        # onboarding step which is not explicitly tracked in the database
+        new_state.append("device_verified")
+
+        # Remove duplicate keys
+        new_state = list(set(new_state))
+
+    connection.execute(
+        sa.update(server_settings_table).values(
+            onboarding_state=json.dumps(new_state)
+        )
+    )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###

zenml/zen_stores/rest_zen_store.py

@@ -58,6 +58,7 @@ from zenml.constants import (
     ARTIFACTS,
     CODE_REFERENCES,
     CODE_REPOSITORIES,
+    CONFIG,
     CURRENT_USER,
     DEACTIVATE,
     DEFAULT_HTTP_TIMEOUT,
@@ -91,6 +92,7 @@ from zenml.constants import (
     SERVICE_CONNECTOR_RESOURCES,
     SERVICE_CONNECTOR_TYPES,
     SERVICE_CONNECTOR_VERIFY,
+    SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
     SERVICE_CONNECTORS,
     SERVICES,
     STACK,
@@ -101,7 +103,6 @@ from zenml.constants import (
     TAGS,
     TRIGGER_EXECUTIONS,
     TRIGGERS,
-    URL,
     USERS,
     VERSION_1,
     WORKSPACES,
@@ -216,6 +217,7 @@ from zenml.models import (
     ServiceRequest,
     ServiceResponse,
     ServiceUpdate,
+    StackDeploymentConfig,
     StackDeploymentInfo,
     StackFilter,
     StackRequest,
@@ -2476,6 +2478,10 @@ class RestZenStore(BaseZenStore):
             f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}",
             body=service_connector,
             params={"list_resources": list_resources},
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         resources = ServiceConnectorResourcesModel.model_validate(
@@ -2513,6 +2519,10 @@ class RestZenStore(BaseZenStore):
         response_body = self.put(
             f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}",
             params=params,
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         resources = ServiceConnectorResourcesModel.model_validate(
@@ -2859,13 +2869,13 @@ class RestZenStore(BaseZenStore):
         )
         return StackDeploymentInfo.model_validate(body)
 
-    def get_stack_deployment_url(
+    def get_stack_deployment_config(
         self,
         provider: StackDeploymentProvider,
         stack_name: str,
         location: Optional[str] = None,
-    ) -> Tuple[str, str]:
-        """Return the URL to deploy the ZenML stack to the specified cloud provider.
+    ) -> StackDeploymentConfig:
+        """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.
 
         Args:
             provider: The stack deployment provider.
@@ -2873,11 +2883,8 @@ class RestZenStore(BaseZenStore):
             location: The location where the stack should be deployed.
 
         Returns:
-            The URL to deploy the ZenML stack to the specified cloud provider
-            and a text description of the URL.
-
-        Raises:
-            ValueError: If the response body is not as expected.
+            The cloud provider console URL and configuration needed to deploy
+            the ZenML stack to the specified cloud provider.
         """
         params = {
             "provider": provider.value,
@@ -2885,14 +2892,8 @@ class RestZenStore(BaseZenStore):
         }
         if location:
             params["location"] = location
-        body = self.get(f"{STACK_DEPLOYMENT}{URL}", params=params)
-
-        if not isinstance(body, list) or len(body) != 2:
-            raise ValueError(
-                "Bad response body received from the stack deployment URL "
-                "endpoint."
-            )
-        return body[0], body[1]
+        body = self.get(f"{STACK_DEPLOYMENT}{CONFIG}", params=params)
+        return StackDeploymentConfig.model_validate(body)
 
     def get_stack_deployment_stack(
         self,
@@ -4081,6 +4082,7 @@ class RestZenStore(BaseZenStore):
         method: str,
         url: str,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a request to the REST API.
@@ -4089,6 +4091,7 @@ class RestZenStore(BaseZenStore):
             method: The HTTP method to use.
             url: The URL to request.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4111,7 +4114,7 @@ class RestZenStore(BaseZenStore):
                     url,
                     params=params,
                     verify=self.config.verify_ssl,
-                    timeout=self.config.http_timeout,
+                    timeout=timeout or self.config.http_timeout,
                     **kwargs,
                 )
             )
@@ -4140,13 +4143,18 @@ class RestZenStore(BaseZenStore):
             raise
 
     def get(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a GET request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4154,17 +4162,26 @@ class RestZenStore(BaseZenStore):
         """
         logger.debug(f"Sending GET request to {path}...")
         return self._request(
-            "GET", self.url + API + VERSION_1 + path, params=params, **kwargs
+            "GET",
+            self.url + API + VERSION_1 + path,
+            params=params,
+            timeout=timeout,
+            **kwargs,
         )
 
     def delete(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a DELETE request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4175,6 +4192,7 @@ class RestZenStore(BaseZenStore):
             "DELETE",
             self.url + API + VERSION_1 + path,
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -4183,6 +4201,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: BaseModel,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a POST request to the given endpoint path.
@@ -4191,6 +4210,7 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4202,6 +4222,7 @@ class RestZenStore(BaseZenStore):
             self.url + API + VERSION_1 + path,
             data=body.model_dump_json(),
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -4210,6 +4231,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: Optional[BaseModel] = None,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a PUT request to the given endpoint path.
@@ -4218,6 +4240,7 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4230,6 +4253,7 @@ class RestZenStore(BaseZenStore):
             self.url + API + VERSION_1 + path,
             data=data,
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 

zenml/zen_stores/schemas/server_settings_schemas.py

@@ -15,7 +15,7 @@
 
 import json
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any, Optional, Set
 from uuid import UUID
 
 from sqlmodel import Field, SQLModel
@@ -59,16 +59,33 @@ class ServerSettingsSchema(SQLModel, table=True):
         for field, value in settings_update.model_dump(
             exclude_unset=True
         ).items():
-            if field == "onboarding_state":
-                if value is not None:
-                    self.onboarding_state = json.dumps(value)
-            elif hasattr(self, field):
+            if hasattr(self, field):
                 setattr(self, field, value)
 
         self.updated = datetime.utcnow()
 
         return self
 
+    def update_onboarding_state(
+        self, completed_steps: Set[str]
+    ) -> "ServerSettingsSchema":
+        """Update the onboarding state.
+
+        Args:
+            completed_steps: Newly completed onboarding steps.
+
+        Returns:
+            The updated schema.
+        """
+        old_state = set(
+            json.loads(self.onboarding_state) if self.onboarding_state else []
+        )
+        new_state = old_state.union(completed_steps)
+        self.onboarding_state = json.dumps(list(new_state))
+        self.updated = datetime.utcnow()
+
+        return self
+
     def to_model(
         self,
         include_metadata: bool = False,
@@ -82,7 +99,6 @@ class ServerSettingsSchema(SQLModel, table=True):
             include_resources: Whether the resources will be filled.
             **kwargs: Keyword arguments to allow schema specific logic
 
-
         Returns:
             The created `SettingsResponse`.
         """
@@ -101,11 +117,7 @@ class ServerSettingsSchema(SQLModel, table=True):
         resources = None
 
         if include_metadata:
-            metadata = ServerSettingsResponseMetadata(
-                onboarding_state=json.loads(self.onboarding_state)
-                if self.onboarding_state
-                else {},
-            )
+            metadata = ServerSettingsResponseMetadata()
 
         if include_resources:
             resources = ServerSettingsResponseResources()

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,
@@ -245,6 +247,7 @@ from zenml.models import (
     ServiceRequest,
     ServiceResponse,
     ServiceUpdate,
+    StackDeploymentConfig,
     StackDeploymentInfo,
     StackFilter,
     StackRequest,
@@ -796,6 +799,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":
@@ -1552,12 +1556,18 @@ class SqlZenStore(BaseZenStore):
         revisions_afterwards = self.alembic.current_revisions()
 
         if current_revisions != revisions_afterwards:
-            if current_revisions and version.parse(
-                current_revisions[0]
-            ) < version.parse("0.57.1"):
-                # We want to send the missing user enriched events for users
-                # which were created pre 0.57.1 and only on one upgrade
-                self._should_send_user_enriched_events = True
+            try:
+                if current_revisions and version.parse(
+                    current_revisions[0]
+                ) < version.parse("0.57.1"):
+                    # We want to send the missing user enriched events for users
+                    # which were created pre 0.57.1 and only on one upgrade
+                    self._should_send_user_enriched_events = True
+            except version.InvalidVersion:
+                # This can happen if the database is not currently
+                # stamped with an official ZenML version (e.g. in
+                # development environments).
+                pass
 
             self._sync_flavors()
 
@@ -1681,6 +1691,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]:
@@ -6149,6 +6213,7 @@ class SqlZenStore(BaseZenStore):
 
             connector = new_service_connector.to_model(include_metadata=True)
             self._populate_connector_type(connector)
+
             return connector
 
     def get_service_connector(
@@ -6924,6 +6989,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:
@@ -7000,7 +7075,6 @@ class SqlZenStore(BaseZenStore):
 
             # Stack Components
             components_mapping: Dict[StackComponentType, List[UUID]] = {}
-
             for (
                 component_type,
                 component_info,
@@ -7012,6 +7086,20 @@ class SqlZenStore(BaseZenStore):
                     )
                 # Create a new component
                 else:
+                    flavor_list = self.list_flavors(
+                        flavor_filter_model=FlavorFilter(
+                            name=component_info.flavor,
+                            type=component_type,
+                        )
+                    )
+                    if not len(flavor_list):
+                        raise ValueError(
+                            f"Flavor '{component_info.flavor}' not found "
+                            f"for component type '{component_type}'."
+                        )
+
+                    flavor_model = flavor_list[0]
+
                     component_name = full_stack.name
                     while True:
                         try:
@@ -7039,15 +7127,6 @@ class SqlZenStore(BaseZenStore):
                         service_connector = service_connectors[
                             component_info.service_connector_index
                         ]
-                        flavor_list = self.list_flavors(
-                            flavor_filter_model=FlavorFilter(
-                                name=component_info.flavor,
-                                type=component_type,
-                            )
-                        )
-                        assert len(flavor_list) == 1
-
-                        flavor_model = flavor_list[0]
 
                         requirements = flavor_model.connector_requirements
 
@@ -7434,13 +7513,13 @@ class SqlZenStore(BaseZenStore):
             "Stack deployments are not supported by local ZenML deployments."
         )
 
-    def get_stack_deployment_url(
+    def get_stack_deployment_config(
         self,
         provider: StackDeploymentProvider,
         stack_name: str,
         location: Optional[str] = None,
-    ) -> Tuple[str, str]:
-        """Return the URL to deploy the ZenML stack to the specified cloud provider.
+    ) -> StackDeploymentConfig:
+        """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.
 
         Args:
             provider: The stack deployment provider.
@@ -7954,6 +8033,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)