zenml-nightly 0.73.0.dev20250123__py3-none-any.whl → 0.73.0.dev20250125__py3-none-any.whl

zenml/zen_stores/schemas/component_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -30,6 +29,7 @@ from zenml.models import (
     ComponentResponseResources,
     ComponentUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.service_connector_schemas import (
@@ -171,7 +171,7 @@ class StackComponentSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

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, timezone
+from datetime import datetime, timedelta
 from secrets import token_hex
 from typing import Any, Optional, Tuple
 from uuid import UUID
@@ -31,6 +31,7 @@ from zenml.models import (
     OAuthDeviceResponseMetadata,
     OAuthDeviceUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -115,7 +116,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.now(timezone.utc)
+        now = utc_now()
         return (
             cls(
                 client_id=request.client_id,
@@ -159,7 +160,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
         elif device_update.locked is False:
             self.status = OAuthDeviceStatus.ACTIVE.value
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def internal_update(
@@ -174,7 +175,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             The updated `OAuthDeviceSchema` and the new user code and device
             code, if they were generated.
         """
-        now = datetime.now(timezone.utc)
+        now = utc_now()
         user_code: Optional[str] = None
         device_code: Optional[str] = None
 

zenml/zen_stores/schemas/event_source_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional, cast
 from uuid import UUID
 
@@ -32,6 +31,7 @@ from zenml.models import (
     Page,
 )
 from zenml.utils.json_utils import pydantic_encoder
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -184,5 +184,5 @@ class EventSourceSchema(NamedSchema, table=True):
                 )
             else:
                 setattr(self, field, value)
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self

zenml/zen_stores/schemas/flavor_schemas.py

@@ -14,7 +14,6 @@
 """SQL Model Implementations for Flavors."""
 
 import json
-from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -28,6 +27,7 @@ from zenml.models import (
     FlavorResponseMetadata,
     FlavorUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -103,7 +103,7 @@ class FlavorSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/schemas/model_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID, uuid4
 
@@ -52,6 +51,7 @@ from zenml.models import (
     ModelVersionResponseResources,
     Page,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
@@ -224,7 +224,7 @@ class ModelSchema(NamedSchema, table=True):
             exclude_unset=True, exclude_none=True
         ).items():
             setattr(self, field, value)
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         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.now(timezone.utc)
+        self.updated = utc_now()
         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, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
@@ -37,6 +37,7 @@ from zenml.models import (
     RunMetadataEntry,
 )
 from zenml.models.v2.core.pipeline_run import PipelineRunResponseResources
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema
@@ -269,6 +270,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 +438,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.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def update_placeholder(
@@ -471,7 +479,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         self.orchestrator_environment = orchestrator_environment
         self.status = request.status.value
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
 
         return self
 

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipelines and Pipeline Runs."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -29,6 +28,7 @@ from zenml.models import (
     PipelineResponseResources,
     PipelineUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -185,5 +185,5 @@ class PipelineSchema(NamedSchema, table=True):
             The updated `PipelineSchema`.
         """
         self.description = pipeline_update.description
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self

zenml/zen_stores/schemas/run_template_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of run template tables."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -31,6 +30,7 @@ from zenml.models import (
     RunTemplateResponseResources,
     RunTemplateUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -156,7 +156,7 @@ class RunTemplateSchema(BaseSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         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, timezone
-from typing import TYPE_CHECKING, Any, Optional
+from datetime import datetime, timedelta
+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,
@@ -26,20 +27,25 @@ from zenml.models import (
     ScheduleResponseMetadata,
     ScheduleUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 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 +95,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 +168,7 @@ class ScheduleSchema(NamedSchema, table=True):
             )
         if schedule_update.catchup is not None:
             self.catchup = schedule_update.catchup
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(
@@ -196,6 +211,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,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import Any, Dict, Optional, cast
 from uuid import UUID
 
@@ -35,6 +34,7 @@ from zenml.models import (
     SecretResponseMetadata,
     SecretUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -209,7 +209,7 @@ class SecretSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         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, timezone
+from datetime import datetime
 from typing import Any, Optional, Set
 from uuid import UUID
 
@@ -27,6 +27,7 @@ from zenml.models import (
     ServerSettingsResponseResources,
     ServerSettingsUpdate,
 )
+from zenml.utils.time_utils import utc_now
 
 
 class ServerSettingsSchema(SQLModel, table=True):
@@ -42,12 +43,8 @@ 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=lambda: datetime.now(timezone.utc)
-    )
-    updated: datetime = Field(
-        default_factory=lambda: datetime.now(timezone.utc)
-    )
+    last_user_activity: datetime = Field(default_factory=utc_now)
+    updated: datetime = Field(default_factory=utc_now)
 
     def update(
         self, settings_update: ServerSettingsUpdate
@@ -67,7 +64,7 @@ class ServerSettingsSchema(SQLModel, table=True):
             if hasattr(self, field):
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
 
         return self
 
@@ -87,7 +84,7 @@ class ServerSettingsSchema(SQLModel, table=True):
         )
         new_state = old_state.union(completed_steps)
         self.onboarding_state = json.dumps(list(new_state))
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
 
         return self
 

zenml/zen_stores/schemas/service_connector_schemas.py

@@ -15,7 +15,7 @@
 
 import base64
 import json
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
 
@@ -29,6 +29,7 @@ from zenml.models import (
     ServiceConnectorResponseMetadata,
     ServiceConnectorUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -227,7 +228,7 @@ class ServiceConnectorSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
         self.secret_id = secret_id
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/schemas/service_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -32,6 +31,7 @@ from zenml.models.v2.core.service import (
     ServiceUpdate,
 )
 from zenml.utils.dict_utils import dict_to_bytes
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
@@ -210,7 +210,7 @@ class ServiceSchema(NamedSchema, table=True):
                 )
             else:
                 setattr(self, field, value)
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     @classmethod

zenml/zen_stores/schemas/stack_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -27,6 +26,7 @@ from zenml.models import (
     StackResponseMetadata,
     StackUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -134,7 +134,7 @@ class StackSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         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, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
@@ -43,6 +43,7 @@ from zenml.models.v2.core.step_run import (
     StepRunInputResponse,
     StepRunResponseResources,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
@@ -358,7 +359,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
                 if value and self.model_version_id is None:
                     self.model_version_id = value
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
 
         return self
 

zenml/zen_stores/schemas/tag_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of tag tables."""
 
-from datetime import datetime, timezone
 from typing import Any, List
 from uuid import UUID
 
@@ -30,6 +29,7 @@ from zenml.models import (
     TagResponseBody,
     TagUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 
@@ -103,7 +103,7 @@ class TagSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
 

zenml/zen_stores/schemas/trigger_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import Any, List, Optional, cast
 from uuid import UUID
 
@@ -38,6 +37,7 @@ from zenml.models import (
     TriggerUpdate,
 )
 from zenml.utils.json_utils import pydantic_encoder
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.action_schemas import ActionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema
@@ -133,7 +133,7 @@ class TriggerSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     @classmethod

zenml/zen_stores/schemas/user_schemas.py

@@ -14,7 +14,6 @@
 """SQLModel implementation of user tables."""
 
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -33,6 +32,7 @@ from zenml.models import (
     UserResponseMetadata,
     UserUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
@@ -225,7 +225,7 @@ class UserSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def update_service_account(
@@ -245,7 +245,7 @@ class UserSchema(NamedSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/schemas/workspace_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Workspaces."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List
 
 from sqlmodel import Relationship
@@ -25,6 +24,7 @@ from zenml.models import (
     WorkspaceResponseMetadata,
     WorkspaceUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
@@ -168,7 +168,7 @@ class WorkspaceSchema(NamedSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/sql_zen_store.py

@@ -22,7 +22,7 @@ import random
 import re
 import sys
 import time
-from datetime import datetime, timezone
+from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
 from typing import (
@@ -309,6 +309,7 @@ from zenml.utils.string_utils import (
     random_str,
     validate_name,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores import template_utils
 from zenml.zen_stores.base_zen_store import (
     BaseZenStore,
@@ -795,7 +796,7 @@ class SqlZenStoreConfiguration(StoreConfiguration):
                     logger.warning(
                         f"Database SSL setting `{key}` is not a file. "
                     )
-                sqlalchemy_ssl_args[key.lstrip("ssl_")] = ssl_setting
+                sqlalchemy_ssl_args[key.removeprefix("ssl_")] = ssl_setting
             if len(sqlalchemy_ssl_args) > 0:
                 sqlalchemy_ssl_args["check_hostname"] = (
                     self.ssl_verify_server_cert
@@ -1743,7 +1744,7 @@ class SqlZenStore(BaseZenStore):
             settings = self._get_server_settings(session=session)
 
             if last_user_activity < settings.last_user_activity.replace(
-                tzinfo=timezone.utc
+                tzinfo=None
             ):
                 return
 
@@ -4048,7 +4049,7 @@ class SqlZenStore(BaseZenStore):
                 # Delete devices that have expired
                 if (
                     device.expires is not None
-                    and device.expires < datetime.now()
+                    and device.expires < utc_now()
                     and device.user_id is None
                 ):
                     session.delete(device)
@@ -8634,19 +8635,10 @@ class SqlZenStore(BaseZenStore):
                 ExecutionStatus.COMPLETED,
                 ExecutionStatus.FAILED,
             }:
-                run_update.end_time = datetime.now(timezone.utc)
+                run_update.end_time = utc_now()
                 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
                     )