zenml-nightly 0.72.0.dev20250121__py3-none-any.whl → 0.73.0.dev20250124__py3-none-any.whl

zenml/zen_stores/schemas/device_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation for authorized OAuth2 devices."""
 
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from secrets import token_hex
 from typing import Any, Optional, Tuple
 from uuid import UUID
@@ -115,7 +115,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
         device_code = cls._generate_device_code()
         hashed_user_code = cls._get_hashed_code(user_code)
         hashed_device_code = cls._get_hashed_code(device_code)
-        now = datetime.utcnow()
+        now = datetime.now(timezone.utc)
         return (
             cls(
                 client_id=request.client_id,
@@ -159,7 +159,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
         elif device_update.locked is False:
             self.status = OAuthDeviceStatus.ACTIVE.value
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def internal_update(
@@ -174,7 +174,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             The updated `OAuthDeviceSchema` and the new user code and device
             code, if they were generated.
         """
-        now = datetime.utcnow()
+        now = datetime.now(timezone.utc)
         user_code: Optional[str] = None
         device_code: Optional[str] = None
 

zenml/zen_stores/schemas/event_source_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional, cast
 from uuid import UUID
 
@@ -184,5 +184,5 @@ class EventSourceSchema(NamedSchema, table=True):
                 )
             else:
                 setattr(self, field, value)
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self

zenml/zen_stores/schemas/flavor_schemas.py

@@ -14,7 +14,7 @@
 """SQL Model Implementations for Flavors."""
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -103,7 +103,7 @@ class FlavorSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/model_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID, uuid4
 
@@ -224,7 +224,7 @@ class ModelSchema(NamedSchema, table=True):
             exclude_unset=True, exclude_none=True
         ).items():
             setattr(self, field, value)
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
 
@@ -504,7 +504,7 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
             self.name = target_name
         if target_description is not None:
             self.description = target_description
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
 

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation of pipeline run tables."""
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
@@ -269,6 +269,13 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
             for k, v in step_metadata.items():
                 metadata_collection[f"{s.name}::{k}"] = v
 
+        # Fetch the metadata related to the schedule of this run
+        if self.deployment is not None:
+            if schedule := self.deployment.schedule:
+                schedule_metadata = schedule.fetch_metadata_collection()
+                for k, v in schedule_metadata.items():
+                    metadata_collection[f"schedule:{k}"] = v
+
         return metadata_collection
 
     def to_model(
@@ -430,7 +437,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         if run_update.model_version_id and self.model_version_id is None:
             self.model_version_id = run_update.model_version_id
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def update_placeholder(
@@ -471,7 +478,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         self.orchestrator_environment = orchestrator_environment
         self.status = request.status.value
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
 
         return self
 

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipelines and Pipeline Runs."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -185,5 +185,5 @@ class PipelineSchema(NamedSchema, table=True):
             The updated `PipelineSchema`.
         """
         self.description = pipeline_update.description
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self

zenml/zen_stores/schemas/run_template_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of run template tables."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -156,7 +156,7 @@ class RunTemplateSchema(BaseSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/schedule_schema.py

@@ -13,12 +13,13 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipeline Schedules."""
 
-from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Optional
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
 from sqlmodel import Field, Relationship
 
+from zenml.enums import MetadataResourceTypes
 from zenml.models import (
     ScheduleRequest,
     ScheduleResponse,
@@ -31,15 +32,19 @@ from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import RunMetadataInterface
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
         PipelineDeploymentSchema,
     )
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataSchema,
+    )
 
 
-class ScheduleSchema(NamedSchema, table=True):
+class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for schedules."""
 
     __tablename__ = "schedule"
@@ -89,6 +94,15 @@ class ScheduleSchema(NamedSchema, table=True):
         back_populates="schedules"
     )
 
+    run_metadata: List["RunMetadataSchema"] = Relationship(
+        sa_relationship_kwargs=dict(
+            secondary="run_metadata_resource",
+            primaryjoin=f"and_(foreign(RunMetadataResourceSchema.resource_type)=='{MetadataResourceTypes.SCHEDULE.value}', foreign(RunMetadataResourceSchema.resource_id)==ScheduleSchema.id)",
+            secondaryjoin="RunMetadataSchema.id==foreign(RunMetadataResourceSchema.run_metadata_id)",
+            overlaps="run_metadata",
+        ),
+    )
+
     active: bool
     cron_expression: Optional[str] = Field(nullable=True)
     start_time: Optional[datetime] = Field(nullable=True)
@@ -153,7 +167,7 @@ class ScheduleSchema(NamedSchema, table=True):
             )
         if schedule_update.catchup is not None:
             self.catchup = schedule_update.catchup
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(
@@ -196,6 +210,7 @@ class ScheduleSchema(NamedSchema, table=True):
                 workspace=self.workspace.to_model(),
                 pipeline_id=self.pipeline_id,
                 orchestrator_id=self.orchestrator_id,
+                run_metadata=self.fetch_metadata(),
             )
 
         return ScheduleResponse(

zenml/zen_stores/schemas/secret_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Dict, Optional, cast
 from uuid import UUID
 
@@ -209,7 +209,7 @@ class SecretSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/server_settings_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation for the server settings table."""
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Optional, Set
 from uuid import UUID
 
@@ -42,8 +42,12 @@ class ServerSettingsSchema(SQLModel, table=True):
     display_announcements: Optional[bool] = Field(nullable=True)
     display_updates: Optional[bool] = Field(nullable=True)
     onboarding_state: Optional[str] = Field(nullable=True)
-    last_user_activity: datetime = Field(default_factory=datetime.utcnow)
-    updated: datetime = Field(default_factory=datetime.utcnow)
+    last_user_activity: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc)
+    )
+    updated: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc)
+    )
 
     def update(
         self, settings_update: ServerSettingsUpdate
@@ -63,7 +67,7 @@ class ServerSettingsSchema(SQLModel, table=True):
             if hasattr(self, field):
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
 
         return self
 
@@ -83,7 +87,7 @@ class ServerSettingsSchema(SQLModel, table=True):
         )
         new_state = old_state.union(completed_steps)
         self.onboarding_state = json.dumps(list(new_state))
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
 
         return self
 

zenml/zen_stores/schemas/service_connector_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
 
@@ -227,7 +227,7 @@ class ServiceConnectorSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
         self.secret_id = secret_id
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/service_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -210,7 +210,7 @@ class ServiceSchema(NamedSchema, table=True):
                 )
             else:
                 setattr(self, field, value)
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     @classmethod

zenml/zen_stores/schemas/stack_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -134,7 +134,7 @@ class StackSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/step_run_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation of step run tables."""
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
@@ -358,7 +358,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
                 if value and self.model_version_id is None:
                     self.model_version_id = value
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
 
         return self
 

zenml/zen_stores/schemas/tag_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of tag tables."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, List
 from uuid import UUID
 
@@ -103,7 +103,7 @@ class TagSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
 

zenml/zen_stores/schemas/trigger_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, List, Optional, cast
 from uuid import UUID
 
@@ -133,7 +133,7 @@ class TriggerSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     @classmethod

zenml/zen_stores/schemas/user_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation of user tables."""
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -225,7 +225,7 @@ class UserSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def update_service_account(
@@ -245,7 +245,7 @@ class UserSchema(NamedSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/schemas/workspace_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Workspaces."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List
 
 from sqlmodel import Relationship
@@ -168,7 +168,7 @@ class WorkspaceSchema(NamedSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.utcnow()
+        self.updated = datetime.now(timezone.utc)
         return self
 
     def to_model(

zenml/zen_stores/sql_zen_store.py

@@ -8634,10 +8634,19 @@ class SqlZenStore(BaseZenStore):
                 ExecutionStatus.COMPLETED,
                 ExecutionStatus.FAILED,
             }:
-                run_update.end_time = datetime.utcnow()
+                run_update.end_time = datetime.now(timezone.utc)
                 if pipeline_run.start_time and isinstance(
                     pipeline_run.start_time, datetime
                 ):
+                    # We need to ensure both datetimes are timezone-aware to avoid TypeError when subtracting
+                    # Now calculate the duration time
+                    if pipeline_run.start_time.tzinfo is None:
+                        pipeline_run.start_time = (
+                            pipeline_run.start_time.replace(
+                                tzinfo=timezone.utc
+                            )
+                        )
+
                     duration_time = (
                         run_update.end_time - pipeline_run.start_time
                     )

{zenml_nightly-0.72.0.dev20250121.dist-info → zenml_nightly-0.73.0.dev20250124.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.72.0.dev20250121
+Version: 0.73.0.dev20250124
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
@@ -107,7 +107,7 @@ Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
 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"
-Requires-Dist: sagemaker (>=2.117.0) ; extra == "sagemaker"
+Requires-Dist: sagemaker (>=2.199.0) ; extra == "sagemaker"
 Requires-Dist: secure (>=0.3.0,<0.4.0) ; extra == "server"
 Requires-Dist: setuptools
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)
@@ -476,7 +476,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.72.0 is out. Check out the release notes
+    🎉 Version 0.73.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>.