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

zenml/zen_server/deploy/helm/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v2
 name: zenml
-version: "0.61.0"
+version: "0.62.0"
 description: Open source MLOps framework for portable production ready ML pipelines
 keywords:
 - mlops

zenml/zen_server/deploy/helm/README.md

@@ -20,8 +20,8 @@ ZenML is an open-source MLOps framework designed to help you create robust, main
 To install the ZenML chart directly from Amazon ECR, use the following command:
 
 ```bash
-# example command for version 0.61.0
-helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.61.0
+# example command for version 0.62.0
+helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.62.0
 ```
 
 Note: Ensure you have OCI support enabled in your Helm client and that you are authenticated with Amazon ECR.

zenml/zen_server/rbac/utils.py

@@ -103,7 +103,9 @@ def dehydrate_response_model(
         )
 
     dehydrated_values = {}
-    for key, value in dict(model).items():
+    # See `get_subresources_for_model(...)` for a detailed explanation why we
+    # need to use `model.__iter__()` here
+    for key, value in model.__iter__():
         dehydrated_values[key] = _dehydrate_value(
             value, permissions=permissions
         )
@@ -484,7 +486,13 @@ def get_subresources_for_model(
     """
     resources = set()
 
-    for value in dict(model).values():
+    # We don't want to use `model.model_dump()` here as that recursively
+    # converts models to dicts, but we want to preserve those classes for
+    # the recursive `_get_subresources_for_value` calls.
+    # We previously used `dict(model)` here, but that lead to issues with
+    # models overwriting `__getattr__`, this `model.__iter__()` has the same
+    # results though.
+    for _, value in model.__iter__():
         resources.update(_get_subresources_for_value(value))
 
     return resources

zenml/zen_server/routers/devices_endpoints.py

@@ -27,7 +27,7 @@ from zenml.constants import (
     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
 
 

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

zenml/zen_server/routers/service_connectors_endpoints.py

@@ -21,6 +21,7 @@ from fastapi import APIRouter, Depends, Security
 from zenml.constants import (
     API,
     SERVICE_CONNECTOR_CLIENT,
+    SERVICE_CONNECTOR_FULL_STACK,
     SERVICE_CONNECTOR_TYPES,
     SERVICE_CONNECTOR_VERIFY,
     SERVICE_CONNECTORS,
@@ -35,6 +36,13 @@ from zenml.models import (
     ServiceConnectorTypeModel,
     ServiceConnectorUpdate,
 )
+from zenml.models.v2.misc.full_stack import (
+    ServiceConnectorInfo,
+    ServiceConnectorResourcesInfo,
+)
+from zenml.service_connectors.service_connector_utils import (
+    get_resources_options_from_resource_model_for_full_stack,
+)
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
@@ -393,3 +401,52 @@ def get_service_connector_type(
         The requested service connector type.
     """
     return zen_store().get_service_connector_type(connector_type)
+
+
+@router.post(
+    SERVICE_CONNECTOR_FULL_STACK,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_resources_based_on_service_connector_info(
+    connector_info: Optional[ServiceConnectorInfo] = None,
+    connector_uuid: Optional[UUID] = None,
+    _: AuthContext = Security(authorize),
+) -> ServiceConnectorResourcesInfo:
+    """Gets the list of resources that a service connector can access.
+
+    Args:
+        connector_info: The service connector info.
+        connector_uuid: The service connector uuid.
+
+    Returns:
+        The list of resources that the service connector configuration has
+        access to and consumable from UI/CLI.
+
+    Raises:
+        ValueError: If both connector_info and connector_uuid are provided.
+        ValueError: If neither connector_info nor connector_uuid are provided.
+    """
+    if connector_info is not None and connector_uuid is not None:
+        raise ValueError(
+            "Only one of connector_info or connector_uuid must be provided."
+        )
+    if connector_info is None and connector_uuid is None:
+        raise ValueError(
+            "Either connector_info or connector_uuid must be provided."
+        )
+
+    if connector_info is not None:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+        )
+    elif connector_uuid is not None:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR,
+            action=Action.READ,
+            resource_id=connector_uuid,
+        )
+
+    return get_resources_options_from_resource_model_for_full_stack(
+        connector_details=connector_info or connector_uuid  # type: ignore[arg-type]
+    )

zenml/zen_server/routers/steps_endpoints.py

@@ -117,7 +117,8 @@ def create_run_step(
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
     verify_permission_for_model(pipeline_run, action=Action.UPDATE)
 
-    return zen_store().create_run_step(step_run=step)
+    step_response = zen_store().create_run_step(step_run=step)
+    return dehydrate_response_model(step_response)
 
 
 @router.get(

zenml/zen_stores/migrations/versions/0.62.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.62.0].
+
+Revision ID: 0.62.0
+Revises: b4fca5241eea
+Create Date: 2024-07-15 14:15:45.347033
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.62.0"
+down_revision = "b4fca5241eea"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

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

@@ -2591,6 +2591,10 @@ class RestZenStore(BaseZenStore):
         response_body = self.get(
             f"{WORKSPACES}/{workspace_name_or_id}{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_RESOURCES}",
             params=params,
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         assert isinstance(response_body, list)

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/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()