zenml-nightly 0.73.0.dev20250124__py3-none-any.whl → 0.73.0.dev20250126__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.73.0.dev20250124
+0.73.0.dev20250126

zenml/analytics/context.py

@@ -17,7 +17,6 @@ This module is based on the 'analytics-python' package created by Segment.
 The base functionalities are adapted to work with the ZenML analytics server.
 """
 
-import datetime
 import locale
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
@@ -32,6 +31,7 @@ from zenml.constants import (
 )
 from zenml.environment import Environment, get_environment
 from zenml.logger import get_logger
+from zenml.utils.time_utils import utc_now_tz_aware
 
 if TYPE_CHECKING:
     from zenml.analytics.enums import AnalyticsEvent
@@ -284,11 +284,7 @@ class AnalyticsContext:
 
         try:
             # Timezone as tzdata
-            tz = (
-                datetime.datetime.now(datetime.timezone.utc)
-                .astimezone()
-                .tzname()
-            )
+            tz = utc_now_tz_aware().astimezone().tzname()
             if tz is not None:
                 properties.update({"timezone": tz})
 

zenml/cli/annotator.py

@@ -182,7 +182,7 @@ def register_annotator_subcommands() -> None:
         kwargs_dict = {}
         for arg in kwargs:
             if arg.startswith("--"):
-                key, value = arg.lstrip("--").split("=", 1)
+                key, value = arg.removeprefix("--").split("=", 1)
                 kwargs_dict[key] = value
 
         if annotator.flavor == "prodigy":

zenml/cli/login.py

@@ -792,15 +792,16 @@ def login(
     )
 
     if server is not None:
-        if not re.match(r"^(https?|mysql)://", server):
-            # The server argument is a ZenML Pro server name or UUID
-            connect_to_pro_server(
-                pro_server=server,
+        if re.match(r"^mysql://", server):
+            # The server argument is a MySQL URL, we can directly connect to it
+
+            connect_to_server(
+                url=server,
                 api_key=api_key_value,
+                verify_ssl=verify_ssl,
                 refresh=refresh,
-                pro_api_url=pro_api_url,
             )
-        else:
+        elif re.match(r"^https?://", server):
             # The server argument is a server URL
 
             # First, try to discover if the server is a ZenML Pro server or not
@@ -821,6 +822,14 @@ def login(
                     verify_ssl=verify_ssl,
                     refresh=refresh,
                 )
+        else:
+            # The server argument is a ZenML Pro server name or UUID
+            connect_to_pro_server(
+                pro_server=server,
+                api_key=api_key_value,
+                refresh=refresh,
+                pro_api_url=pro_api_url,
+            )
 
     elif current_non_local_server:
         # The server argument is not provided, so we default to

zenml/cli/service_connectors.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Service connector CLI commands."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Any, Dict, List, Optional, Union, cast
 from uuid import UUID
 
@@ -25,7 +25,6 @@ from zenml.cli.utils import (
     is_sorted_or_filtered,
     list_options,
     print_page_info,
-    seconds_to_human_readable,
 )
 from zenml.client import Client
 from zenml.console import console
@@ -37,6 +36,7 @@ from zenml.models import (
     ServiceConnectorResourcesModel,
     ServiceConnectorResponse,
 )
+from zenml.utils.time_utils import seconds_to_human_readable, utc_now
 
 
 # Service connectors
@@ -292,7 +292,7 @@ def prompt_expires_at(
         default_str = ""
         if default is not None:
             seconds = int(
-                (default - datetime.now(timezone.utc)).total_seconds()
+                (default - utc_now(tz_aware=default)).total_seconds()
             )
             default_str = (
                 f" [{str(default)} i.e. in "
@@ -309,7 +309,7 @@ def prompt_expires_at(
 
         assert expires_at is not None
         assert isinstance(expires_at, datetime)
-        if expires_at < datetime.now(timezone.utc):
+        if expires_at < utc_now(tz_aware=expires_at):
             cli_utils.warning(
                 "The expiration time must be in the future. Please enter a "
                 "later date and time."
@@ -317,7 +317,7 @@ def prompt_expires_at(
             continue
 
         seconds = int(
-            (expires_at - datetime.now(timezone.utc)).total_seconds()
+            (expires_at - utc_now(tz_aware=expires_at)).total_seconds()
         )
 
         confirm = click.confirm(

zenml/cli/stack.py

@@ -17,7 +17,6 @@ import getpass
 import re
 import time
 import webbrowser
-from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -77,6 +76,7 @@ from zenml.service_connectors.service_connector_utils import (
 )
 from zenml.utils import requirements_utils
 from zenml.utils.dashboard_utils import get_component_url, get_stack_url
+from zenml.utils.time_utils import utc_now_tz_aware
 from zenml.utils.yaml_utils import read_yaml, write_yaml
 
 if TYPE_CHECKING:
@@ -1575,7 +1575,7 @@ def deploy(
         ):
             raise click.Abort()
 
-        date_start = datetime.now(timezone.utc)
+        date_start = utc_now_tz_aware()
 
         webbrowser.open(deployment_config.deployment_url)
         console.print(

zenml/cli/utils.py

@@ -14,7 +14,6 @@
 """Utility functions for the CLI."""
 
 import contextlib
-import datetime
 import json
 import os
 import platform
@@ -79,6 +78,7 @@ from zenml.services import BaseService, ServiceState
 from zenml.stack import StackComponent
 from zenml.stack.stack_component import StackComponentConfig
 from zenml.utils import secret_utils
+from zenml.utils.time_utils import expires_in
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -1581,58 +1581,6 @@ def print_components_table(
     print_table(configurations)
 
 
-def seconds_to_human_readable(time_seconds: int) -> str:
-    """Converts seconds to human-readable format.
-
-    Args:
-        time_seconds: Seconds to convert.
-
-    Returns:
-        Human readable string.
-    """
-    seconds = time_seconds % 60
-    minutes = (time_seconds // 60) % 60
-    hours = (time_seconds // 3600) % 24
-    days = time_seconds // 86400
-    tokens = []
-    if days:
-        tokens.append(f"{days}d")
-    if hours:
-        tokens.append(f"{hours}h")
-    if minutes:
-        tokens.append(f"{minutes}m")
-    if seconds:
-        tokens.append(f"{seconds}s")
-
-    return "".join(tokens)
-
-
-def expires_in(
-    expires_at: datetime.datetime,
-    expired_str: str,
-    skew_tolerance: Optional[int] = None,
-) -> str:
-    """Returns a human-readable string of the time until the token expires.
-
-    Args:
-        expires_at: Datetime object of the token expiration.
-        expired_str: String to return if the token is expired.
-        skew_tolerance: Seconds of skew tolerance to subtract from the
-            expiration time. If the token expires within this time, it will be
-            considered expired.
-
-    Returns:
-        Human readable string.
-    """
-    now = datetime.datetime.now(datetime.timezone.utc)
-    expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
-    if skew_tolerance:
-        expires_at -= datetime.timedelta(seconds=skew_tolerance)
-    if expires_at < now:
-        return expired_str
-    return seconds_to_human_readable(int((expires_at - now).total_seconds()))
-
-
 def print_service_connectors_table(
     client: "Client",
     connectors: Sequence["ServiceConnectorResponse"],
@@ -2660,7 +2608,7 @@ def print_model_url(url: Optional[str]) -> None:
         warning(
             "You can display various ZenML entities including pipelines, "
             "runs, stacks and much more on the ZenML Dashboard. "
-            "You can try it locally, by running `zenml local --local`, or "
+            "You can try it locally, by running `zenml login --local`, or "
             "remotely, by deploying ZenML on the infrastructure of your choice."
         )
 

zenml/config/pipeline_configurations.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from pydantic import SerializeAsAny, field_validator
@@ -23,6 +23,7 @@ from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import SourceWithValidator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings
@@ -61,7 +62,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
             The full substitutions dict including date and time.
         """
         if start_time is None:
-            start_time = datetime.now(timezone.utc)
+            start_time = utc_now()
         ret = self.substitutions.copy()
         ret.setdefault("date", start_time.strftime("%Y_%m_%d"))
         ret.setdefault("time", start_time.strftime("%H_%M_%S_%f"))

zenml/config/schedule.py

@@ -98,27 +98,3 @@ class Schedule(BaseModel):
                 "or a run once start time "
                 "need to be set for a valid schedule."
             )
-
-    @property
-    def utc_start_time(self) -> Optional[str]:
-        """Optional ISO-formatted string of the UTC start time.
-
-        Returns:
-            Optional ISO-formatted string of the UTC start time.
-        """
-        if not self.start_time:
-            return None
-
-        return self.start_time.astimezone(datetime.timezone.utc).isoformat()
-
-    @property
-    def utc_end_time(self) -> Optional[str]:
-        """Optional ISO-formatted string of the UTC end time.
-
-        Returns:
-            Optional ISO-formatted string of the UTC end time.
-        """
-        if not self.end_time:
-            return None
-
-        return self.end_time.astimezone(datetime.timezone.utc).isoformat()

zenml/event_hub/base_event_hub.py

@@ -14,7 +14,7 @@
 """Base class for event hub implementations."""
 
 from abc import ABC, abstractmethod
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 from zenml import EventSourceResponse
@@ -28,6 +28,7 @@ from zenml.models import (
     TriggerExecutionResponse,
     TriggerResponse,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.jwt import JWTToken
 
@@ -134,9 +135,7 @@ class BaseEventHub(ABC):
         )
         expires: Optional[datetime] = None
         if trigger.action.auth_window:
-            expires = datetime.now(timezone.utc) + timedelta(
-                minutes=trigger.action.auth_window
-            )
+            expires = utc_now() + timedelta(minutes=trigger.action.auth_window)
         encoded_token = token.encode(expires=expires)
         auth_context = AuthContext(
             user=trigger.action.service_account,

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -42,6 +42,7 @@ from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils import io_utils
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.config import ResourceSettings
@@ -408,8 +409,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         if schedule:
             if schedule.cron_expression:
                 start_time = schedule.start_time or (
-                    datetime.datetime.now(datetime.timezone.utc)
-                    - datetime.timedelta(7)
+                    utc_now() - datetime.timedelta(7)
                 )
                 return {
                     "schedule": schedule.cron_expression,
@@ -429,7 +429,6 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
             "schedule": "@once",
             # set a start time in the past and disable catchup so airflow
             # runs the dag immediately
-            "start_date": datetime.datetime.now(datetime.timezone.utc)
-            - datetime.timedelta(7),
+            "start_date": utc_now() - datetime.timedelta(7),
             "catchup": False,
         }

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -15,7 +15,6 @@
 
 import os
 import re
-from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -64,6 +63,7 @@ from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils.env_utils import split_environment_variables
+from zenml.utils.time_utils import to_utc_timezone, utc_now_tz_aware
 
 if TYPE_CHECKING:
     from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
@@ -522,6 +522,11 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
 
             schedule_name = orchestrator_run_name
             next_execution = None
+            start_date = (
+                to_utc_timezone(deployment.schedule.start_time)
+                if deployment.schedule.start_time
+                else None
+            )
 
             # Create PipelineSchedule based on schedule type
             if deployment.schedule.cron_expression:
@@ -531,7 +536,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 schedule = PipelineSchedule(
                     name=schedule_name,
                     cron=cron_exp,
-                    start_date=deployment.schedule.start_time,
+                    start_date=start_date,
                     enabled=True,
                 )
             elif deployment.schedule.interval_second:
@@ -549,12 +554,11 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 schedule = PipelineSchedule(
                     name=schedule_name,
                     rate=(minutes, "minutes"),
-                    start_date=deployment.schedule.start_time,
+                    start_date=start_date,
                     enabled=True,
                 )
                 next_execution = (
-                    deployment.schedule.start_time
-                    or datetime.now(timezone.utc)
+                    deployment.schedule.start_time or utc_now_tz_aware()
                 ) + deployment.schedule.interval_second
             else:
                 # One-time schedule
@@ -569,7 +573,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                     )
                 schedule = PipelineSchedule(
                     name=schedule_name,
-                    at=execution_time.astimezone(timezone.utc),
+                    at=to_utc_timezone(execution_time),
                     enabled=True,
                 )
                 next_execution = execution_time
@@ -722,15 +726,13 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         Returns:
             A dictionary of metadata.
         """
-        from zenml import get_step_context
-
         execution_arn = os.environ[ENV_ZENML_SAGEMAKER_RUN_ID]
 
         run_metadata: Dict[str, "MetadataType"] = {}
 
         settings = cast(
             SagemakerOrchestratorSettings,
-            self.get_settings(get_step_context().pipeline_run),
+            self.get_settings(Client().get_pipeline_run(run_id)),
         )
 
         for metadata in self.compute_metadata(

zenml/integrations/aws/service_connectors/aws_service_connector.py

@@ -66,6 +66,7 @@ from zenml.service_connectors.service_connector import (
 )
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.secret_utils import PlainSerializedSecretStr
+from zenml.utils.time_utils import utc_now_tz_aware
 
 logger = get_logger(__name__)
 
@@ -711,7 +712,7 @@ class AWSServiceConnector(ServiceConnector):
                 return session, None
 
             # Refresh expired sessions
-            now = datetime.datetime.now(datetime.timezone.utc)
+            now = utc_now_tz_aware()
             expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
             # check if the token expires in the near future
             if expires_at > now + datetime.timedelta(
@@ -959,9 +960,7 @@ class AWSServiceConnector(ServiceConnector):
                 # determine the expiration time of the temporary credentials
                 # from the boto3 session, so we assume the default IAM role
                 # expiration date is used
-                expiration_time = datetime.datetime.now(
-                    tz=datetime.timezone.utc
-                ) + datetime.timedelta(
+                expiration_time = utc_now_tz_aware() + datetime.timedelta(
                     seconds=DEFAULT_IAM_ROLE_TOKEN_EXPIRATION
                 )
                 return session, expiration_time
@@ -1673,9 +1672,7 @@ class AWSServiceConnector(ServiceConnector):
                     # expiration time of the temporary credentials from the
                     # boto3 session, so we assume the default IAM role
                     # expiration period is used
-                    expires_at = datetime.datetime.now(
-                        tz=datetime.timezone.utc
-                    ) + datetime.timedelta(
+                    expires_at = utc_now_tz_aware() + datetime.timedelta(
                         seconds=DEFAULT_IAM_ROLE_TOKEN_EXPIRATION
                     )
 
@@ -1720,9 +1717,7 @@ class AWSServiceConnector(ServiceConnector):
                         aws_secret_access_key=credentials["SecretAccessKey"],
                         aws_session_token=credentials["SessionToken"],
                     )
-                    expires_at = datetime.datetime.now(
-                        tz=datetime.timezone.utc
-                    ) + datetime.timedelta(
+                    expires_at = utc_now_tz_aware() + datetime.timedelta(
                         seconds=DEFAULT_STS_TOKEN_EXPIRATION
                     )
 
@@ -2130,9 +2125,9 @@ class AWSServiceConnector(ServiceConnector):
             # Kubernetes authentication tokens issued by AWS EKS have a fixed
             # expiration time of 15 minutes
             # source: https://aws.github.io/aws-eks-best-practices/security/docs/iam/#controlling-access-to-eks-clusters
-            expires_at = datetime.datetime.now(
-                tz=datetime.timezone.utc
-            ) + datetime.timedelta(minutes=EKS_KUBE_API_TOKEN_EXPIRATION)
+            expires_at = utc_now_tz_aware() + datetime.timedelta(
+                minutes=EKS_KUBE_API_TOKEN_EXPIRATION
+            )
 
             # get cluster details
             cluster_arn = cluster["cluster"]["arn"]

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -58,6 +58,7 @@ from zenml.service_connectors.service_connector import (
 )
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.secret_utils import PlainSerializedSecretStr
+from zenml.utils.time_utils import to_local_tz, to_utc_timezone, utc_now
 
 # Configure the logging level for azure.identity
 logging.getLogger("azure.identity").setLevel(logging.WARNING)
@@ -171,12 +172,7 @@ class ZenMLAzureTokenCredential(TokenCredential):
         self.token = token
 
         # Convert the expiration time from UTC to local time
-        expires_at.replace(tzinfo=datetime.timezone.utc)
-        expires_at = expires_at.astimezone(
-            datetime.datetime.now().astimezone().tzinfo
-        )
-
-        self.expires_on = int(expires_at.timestamp())
+        self.expires_on = int(to_local_tz(expires_at).timestamp())
 
     def get_token(self, *scopes: str, **kwargs: Any) -> Any:
         """Get token.
@@ -604,11 +600,9 @@ class AzureServiceConnector(ServiceConnector):
                 return session, None
 
             # Refresh expired sessions
-            now = datetime.datetime.now(datetime.timezone.utc)
-            expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
 
             # check if the token expires in the near future
-            if expires_at > now + datetime.timedelta(
+            if expires_at > utc_now(tz_aware=expires_at) + datetime.timedelta(
                 minutes=AZURE_SESSION_EXPIRATION_BUFFER
             ):
                 return session, expires_at
@@ -1137,7 +1131,7 @@ class AzureServiceConnector(ServiceConnector):
             # format.
             expires_at = datetime.datetime.fromtimestamp(token.expires_on)
             # Convert the expiration timestamp from local time to UTC time.
-            expires_at = expires_at.astimezone(datetime.timezone.utc)
+            expires_at = to_utc_timezone(expires_at)
 
             auth_config = AzureAccessTokenConfig(
                 token=token.token,

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -78,6 +78,7 @@ from zenml.service_connectors.service_connector import (
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.utils.secret_utils import PlainSerializedSecretStr
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -1124,10 +1125,9 @@ class GCPServiceConnector(ServiceConnector):
                 return session, None
 
             # Refresh expired sessions
-            now = datetime.datetime.now(datetime.timezone.utc)
-            expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
+
             # check if the token expires in the near future
-            if expires_at > now + datetime.timedelta(
+            if expires_at > utc_now(tz_aware=expires_at) + datetime.timedelta(
                 minutes=GCP_SESSION_EXPIRATION_BUFFER
             ):
                 return session, expires_at

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -31,7 +31,6 @@ Internal interface: no backwards compatibility guarantees.
 Adjusted from https://github.com/tensorflow/tfx/blob/master/tfx/utils/kube_utils.py.
 """
 
-import datetime
 import enum
 import re
 import time
@@ -47,6 +46,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -248,7 +248,7 @@ def wait_pod(
     Returns:
         The pod object which meets the exit condition.
     """
-    start_time = datetime.datetime.now(datetime.timezone.utc)
+    start_time = utc_now()
 
     # Link to exponential back-off algorithm used here:
     # https://cloud.google.com/storage/docs/exponential-backoff
@@ -288,7 +288,7 @@ def wait_pod(
             return resp
 
         # Check if wait timed out.
-        elapse_time = datetime.datetime.now(datetime.timezone.utc) - start_time
+        elapse_time = utc_now() - start_time
         if elapse_time.seconds >= timeout_sec and timeout_sec != 0:
             raise RuntimeError(
                 f"Waiting for pod `{namespace}:{pod_name}` timed out after "

zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -504,13 +504,17 @@ class KubernetesServiceConnector(ServiceConnector):
                 ).decode("utf-8")
                 if kube_config.ssl_ca_cert
                 else None,
-                cluster_name=kube_config.host.strip("https://").split(":")[0],
+                cluster_name=kube_config.host.removeprefix("https://").split(
+                    ":"
+                )[0],
                 insecure=kube_config.verify_ssl is False,
             )
         else:
             token: Optional[str] = None
             if kube_config.api_key:
-                token = kube_config.api_key["authorization"].strip("Bearer ")
+                token = kube_config.api_key["authorization"].removeprefix(
+                    "Bearer "
+                )
 
             auth_method = KubernetesAuthenticationMethods.TOKEN
             auth_config = KubernetesTokenConfig(

zenml/integrations/whylogs/data_validators/whylogs_data_validator.py

@@ -34,6 +34,7 @@ from zenml.integrations.whylogs.secret_schemas.whylabs_secret_schema import (
 )
 from zenml.logger import get_logger
 from zenml.stack.authentication_mixin import AuthenticationMixin
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -97,9 +98,7 @@ class WhylogsDataValidator(BaseDataValidator, AuthenticationMixin):
         """
         results = why.log(pandas=dataset)
         profile = results.profile()
-        dataset_timestamp = dataset_timestamp or datetime.datetime.now(
-            datetime.timezone.utc
-        )
+        dataset_timestamp = dataset_timestamp or utc_now()
         profile.set_dataset_timestamp(dataset_timestamp=dataset_timestamp)
         return profile.view()
 

zenml/logging/step_logging.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """ZenML logging handler."""
 
-import datetime
 import os
 import re
 import sys
@@ -35,6 +34,7 @@ from zenml.logging import (
     STEP_LOGS_STORAGE_MAX_MESSAGES,
     STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 # Get the logger
@@ -303,9 +303,9 @@ class StepLogsStorage:
                             "w",
                         ) as file:
                             for message in self.buffer:
-                                timestamp = datetime.datetime.now(
-                                    datetime.timezone.utc
-                                ).strftime("%Y-%m-%d %H:%M:%S")
+                                timestamp = utc_now().strftime(
+                                    "%Y-%m-%d %H:%M:%S"
+                                )
                                 file.write(
                                     f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n"
                                 )
@@ -314,9 +314,9 @@ class StepLogsStorage:
                             self.logs_uri, "a"
                         ) as file:
                             for message in self.buffer:
-                                timestamp = datetime.datetime.now(
-                                    datetime.timezone.utc
-                                ).strftime("%Y-%m-%d %H:%M:%S")
+                                timestamp = utc_now().strftime(
+                                    "%Y-%m-%d %H:%M:%S"
+                                )
                                 file.write(
                                     f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n"
                                 )