Update onboarding

src/zenml/constants.py

@@ -390,6 +390,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 TAGS = "/tags"
 TRIGGERS = "/triggers"
 TRIGGER_EXECUTIONS = "/trigger_executions"
+ONBOARDING_STATE = "/onboarding_state"
 USERS = "/users"
 URL = "/url"
 VERSION_1 = "/v1"

src/zenml/enums.py

@@ -386,6 +386,21 @@ class PluginSubType(StrEnum):
     PIPELINE_RUN = "pipeline_run"
 
 
+class OnboardingStep(StrEnum):
+    """All onboarding steps."""
+
+    DEVICE_VERIFIED = "device_verified"
+    PIPELINE_RUN = "pipeline_run"
+    STARTER_SETUP_COMPLETED = "starter_setup_completed"
+    STACK_WITH_REMOTE_ORCHESTRATOR_CREATED = (
+        "stack_with_remote_orchestrator_created"
+    )
+    PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR = (
+        "pipeline_run_with_remote_orchestrator"
+    )
+    PRODUCTION_SETUP_COMPLETED = "production_setup_completed"
+
+
 class StackDeploymentProvider(StrEnum):
     """All possible stack deployment providers."""
 

src/zenml/models/v2/core/server_settings.py

@@ -15,8 +15,6 @@
 
 from datetime import datetime
 from typing import (
-    Any,
-    Dict,
     Optional,
 )
 from uuid import UUID
@@ -57,10 +55,6 @@ class ServerSettingsUpdate(BaseZenModel):
         default=None,
         title="Whether to display notifications about ZenML updates in the dashboard.",
     )
-    onboarding_state: Optional[Dict[str, Any]] = Field(
-        default=None,
-        title="The server's onboarding state.",
-    )
 
 
 # ------------------ Response Model ------------------
@@ -96,11 +90,6 @@ class ServerSettingsResponseBody(BaseResponseBody):
 class ServerSettingsResponseMetadata(BaseResponseMetadata):
     """Response metadata for server settings."""
 
-    onboarding_state: Dict[str, Any] = Field(
-        default={},
-        title="The server's onboarding state.",
-    )
-
 
 class ServerSettingsResponseResources(BaseResponseResources):
     """Response resources for server settings."""
@@ -199,15 +188,6 @@ def updated(self) -> datetime:
         """
         return self.get_body().updated
 
-    @property
-    def onboarding_state(self) -> Dict[str, Any]:
-        """The `onboarding_state` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().onboarding_state
-
 
 # ------------------ Filter Model ------------------
 

src/zenml/zen_server/routers/devices_endpoints.py

@@ -27,7 +27,7 @@
     DEVICES,
     VERSION_1,
 )
-from zenml.enums import OAuthDeviceStatus
+from zenml.enums import OAuthDeviceStatus, OnboardingStep
 from zenml.models import (
     OAuthDeviceFilter,
     OAuthDeviceInternalUpdate,
@@ -270,6 +270,9 @@ def verify_authorized_device(
             update=update,
         )
 
+    store.update_onboarding_state(
+        completed_steps={OnboardingStep.DEVICE_VERIFIED}
+    )
     return device_model
 
 

src/zenml/zen_server/routers/server_endpoints.py

@@ -13,12 +13,19 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from typing import Optional
+from typing import List, Optional
 
 from fastapi import APIRouter, Security
 
 import zenml
-from zenml.constants import ACTIVATE, API, INFO, SERVER_SETTINGS, VERSION_1
+from zenml.constants import (
+    ACTIVATE,
+    API,
+    INFO,
+    ONBOARDING_STATE,
+    SERVER_SETTINGS,
+    VERSION_1,
+)
 from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
@@ -64,6 +71,26 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    ONBOARDING_STATE,
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def get_onboarding_state(
+    _: AuthContext = Security(authorize),
+) -> List[str]:
+    """Get the onboarding state of the server.
+
+    Returns:
+        The onboarding state of the server.
+    """
+    return zen_store().get_onboarding_state()
+
+
 # We don't have any concrete value that tells us whether a server is a cloud
 # tenant, so we use `external_server_id` as the best proxy option.
 # For cloud tenants, we don't add these endpoints as the server settings don't

src/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 ###

src/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 @@ def update(
         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 @@ def to_model(
             include_resources: Whether the resources will be filled.
             **kwargs: Keyword arguments to allow schema specific logic
 
-
         Returns:
             The created `SettingsResponse`.
         """
@@ -101,11 +117,7 @@ def to_model(
         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()