zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

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/schemas/stack_schemas.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Stacks."""
 
+import base64
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
@@ -75,6 +77,7 @@ class StackSchema(NamedSchema, table=True):
 
     __tablename__ = "stack"
     stack_spec_path: Optional[str]
+    labels: Optional[bytes]
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -124,6 +127,10 @@ class StackSchema(NamedSchema, table=True):
         ).items():
             if field == "components":
                 self.components = components
+            elif field == "labels":
+                self.labels = base64.b64encode(
+                    json.dumps(stack_update.labels).encode("utf-8")
+                )
             else:
                 setattr(self, field, value)
 
@@ -158,6 +165,9 @@ class StackSchema(NamedSchema, table=True):
                 workspace=self.workspace.to_model(),
                 components={c.type: [c.to_model()] for c in self.components},
                 stack_spec_path=self.stack_spec_path,
+                labels=json.loads(base64.b64decode(self.labels).decode())
+                if self.labels
+                else None,
             )
 
         return StackResponse(

zenml/zen_stores/schemas/step_run_schemas.py

@@ -192,6 +192,7 @@ class StepRunSchema(NamedSchema, table=True):
             The created StepRunResponse.
 
         Raises:
+            ValueError: In case the step run configuration can not be loaded.
             RuntimeError: If the step run schema does not have a deployment_id
                 or a step_configuration.
         """
@@ -212,19 +213,34 @@ class StepRunSchema(NamedSchema, table=True):
             for artifact in self.output_artifacts
         }
 
+        full_step_config = None
         if self.deployment is not None:
-            full_step_config = Step.model_validate(
-                json.loads(self.deployment.step_configurations)[self.name]
-            )
-        elif self.step_configuration is not None:
-            full_step_config = Step.model_validate_json(
-                self.step_configuration
-            )
-        else:
-            raise RuntimeError(
-                "Step run model creation has failed. Each step run entry "
-                "should either have a deployment_id or step_configuration."
+            step_configuration = json.loads(
+                self.deployment.step_configurations
             )
+            if self.name in step_configuration:
+                full_step_config = Step.model_validate(
+                    step_configuration[self.name]
+                )
+            elif not self.step_configuration:
+                raise ValueError(
+                    f"Unable to load the configuration for step `{self.name}` from the"
+                    f"database. To solve this please delete the pipeline run that this"
+                    f"step run belongs to. Pipeline Run ID: `{self.pipeline_run_id}`."
+                )
+
+        # the step configuration moved into the deployment - the following case is to ensure
+        # backwards compatibility
+        if full_step_config is None:
+            if self.step_configuration:
+                full_step_config = Step.model_validate_json(
+                    self.step_configuration
+                )
+            else:
+                raise RuntimeError(
+                    "Step run model creation has failed. Each step run entry "
+                    "should either have a deployment_id or step_configuration."
+                )
 
         body = StepRunResponseBody(
             user=self.user.to_model() if self.user else None,