Make runs obey event log entry timestamp

Author: dpeng817

python_modules/dagster/dagster/_core/execution/context_creation_job.py

@@ -44,7 +44,6 @@
 from dagster._core.execution.retries import RetryMode
 from dagster._core.executor.init import InitExecutorContext
 from dagster._core.instance import DagsterInstance
-from dagster._core.log_manager import DagsterLogManager
 from dagster._core.storage.dagster_run import DagsterRun
 from dagster._core.system_config.objects import ResolvedRunConfig
 from dagster._loggers import default_loggers, default_system_loggers
@@ -55,10 +54,16 @@
     from dagster._core.execution.plan.outputs import StepOutputHandle
     from dagster._core.executor.base import Executor
 
+    # Import within functions so that we can mock the class in tests using freeze_time.
+    # Essentially, we want to be able to control the timestamp of the log records.
+    from dagster._core.log_manager import DagsterLogManager
+
 
 def initialize_console_manager(
     dagster_run: Optional[DagsterRun], instance: Optional[DagsterInstance] = None
-) -> DagsterLogManager:
+) -> "DagsterLogManager":
+    from dagster._core.log_manager import DagsterLogManager
+
     # initialize default colored console logger
     loggers = []
     for logger_def, logger_config in default_system_loggers(instance):
@@ -458,7 +463,9 @@ def scoped_job_context(
 
 def create_log_manager(
     context_creation_data: ContextCreationData,
-) -> DagsterLogManager:
+) -> "DagsterLogManager":
+    from dagster._core.log_manager import DagsterLogManager
+
     check.inst_param(context_creation_data, "context_creation_data", ContextCreationData)
 
     job_def, resolved_run_config, dagster_run = (
@@ -506,14 +513,16 @@ def create_log_manager(
 
 def create_context_free_log_manager(
     instance: DagsterInstance, dagster_run: DagsterRun
-) -> DagsterLogManager:
+) -> "DagsterLogManager":
     """In the event of pipeline initialization failure, we want to be able to log the failure
     without a dependency on the PlanExecutionContext to initialize DagsterLogManager.
 
     Args:
         dagster_run (PipelineRun)
         pipeline_def (JobDefinition)
     """
+    from dagster._core.log_manager import DagsterLogManager
+
     check.inst_param(instance, "instance", DagsterInstance)
     check.inst_param(dagster_run, "dagster_run", DagsterRun)
 

python_modules/dagster/dagster/_core/execution/host_mode.py

@@ -26,7 +26,6 @@
 from dagster._core.executor.base import Executor
 from dagster._core.executor.init import InitExecutorContext
 from dagster._core.instance import DagsterInstance
-from dagster._core.log_manager import DagsterLogManager
 from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatus
 from dagster._loggers import default_system_loggers
 from dagster._utils import ensure_single_item
@@ -80,6 +79,8 @@ def host_mode_execution_context_event_generator(
     output_capture: None,
     resume_from_failure: bool = False,
 ) -> Iterator[Union[PlanOrchestrationContext, DagsterEvent]]:
+    from dagster._core.log_manager import DagsterLogManager
+
     check.inst_param(execution_plan, "execution_plan", ExecutionPlan)
     check.inst_param(pipeline, "pipeline", ReconstructableJob)
 

python_modules/dagster/dagster/_core/instance/__init__.py

@@ -84,7 +84,7 @@
 )
 from dagster._core.types.pagination import PaginatedResults
 from dagster._serdes import ConfigurableClass
-from dagster._time import get_current_datetime, get_current_timestamp
+from dagster._time import datetime_from_timestamp, get_current_datetime, get_current_timestamp
 from dagster._utils import PrintFn, is_uuid, traced
 from dagster._utils.error import serializable_error_info_from_exc_info
 from dagster._utils.merger import merge_dicts
@@ -2667,7 +2667,9 @@ def handle_new_event(
                 and event.is_dagster_event
                 and event.get_dagster_event().is_job_event
             ):
-                self._run_storage.handle_run_event(run_id, event.get_dagster_event())
+                self._run_storage.handle_run_event(
+                    run_id, event.get_dagster_event(), datetime_from_timestamp(event.timestamp)
+                )
                 run = self.get_run_by_id(run_id)
                 if run and event.get_dagster_event().is_run_failure and self.run_retries_enabled:
                     # Note that this tag is only applied to runs that fail. Successful runs will not

python_modules/dagster/dagster/_core/storage/runs/sql_run_storage.py

@@ -169,7 +169,9 @@ def add_historical_run(
         )
         return dagster_run
 
-    def handle_run_event(self, run_id: str, event: DagsterEvent) -> None:
+    def handle_run_event(
+        self, run_id: str, event: DagsterEvent, update_timestamp: Optional[datetime] = None
+    ) -> None:
         from dagster._core.events import JobFailureData
 
         check.str_param(run_id, "run_id")
@@ -189,19 +191,17 @@ def handle_run_event(self, run_id: str, event: DagsterEvent) -> None:
 
         kwargs = {}
 
-        # consider changing the `handle_run_event` signature to get timestamp off of the
-        # EventLogEntry instead of the DagsterEvent, for consistency
-        now = get_current_datetime()
+        update_timestamp = update_timestamp or get_current_datetime()
 
         if run_stats_cols_in_index and event.event_type == DagsterEventType.PIPELINE_START:
-            kwargs["start_time"] = now.timestamp()
+            kwargs["start_time"] = update_timestamp.timestamp()
 
         if run_stats_cols_in_index and event.event_type in {
             DagsterEventType.PIPELINE_CANCELED,
             DagsterEventType.PIPELINE_FAILURE,
             DagsterEventType.PIPELINE_SUCCESS,
         }:
-            kwargs["end_time"] = now.timestamp()
+            kwargs["end_time"] = update_timestamp.timestamp()
 
         with self.connect() as conn:
             conn.execute(
@@ -210,7 +210,7 @@ def handle_run_event(self, run_id: str, event: DagsterEvent) -> None:
                 .values(
                     run_body=serialize_value(run.with_status(new_job_status)),
                     status=new_job_status.value,
-                    update_timestamp=now,
+                    update_timestamp=update_timestamp,
                     **kwargs,
                 )
             )

python_modules/dagster/dagster/_core/test_utils.py

@@ -1,5 +1,6 @@
 import asyncio
 import datetime
+import logging
 import os
 import re
 import sys
@@ -14,6 +15,7 @@
 from pathlib import Path
 from signal import Signals
 from threading import Event
+from types import TracebackType
 from typing import (  # noqa: UP035
     AbstractSet,
     Any,
@@ -78,7 +80,7 @@
 from dagster._core.workspace.workspace import CodeLocationEntry, CurrentWorkspace
 from dagster._serdes import ConfigurableClass
 from dagster._serdes.config_class import ConfigurableClassData
-from dagster._time import create_datetime, get_timezone
+from dagster._time import create_datetime, get_current_timestamp, get_timezone
 from dagster._utils import Counter, get_terminate_signal, traced, traced_counter
 from dagster._utils.log import configure_loggers
 
@@ -724,6 +726,37 @@ def create_test_asset_job(
     ).get_job_def(name)
 
 
+def get_freezable_log_manager():
+    # The log manager usually sets its own timestamp in the guts of python internals, but we want to be able to control it in test scenarios.
+    from dagster._core.log_manager import DagsterLogManager
+
+    class FreezableLogManager(DagsterLogManager):
+        def makeRecord(
+            self,
+            name: str,
+            level: int,
+            fn: str,
+            lno: int,
+            msg: object,
+            args: tuple[object, ...] | Mapping[str, object],
+            exc_info: tuple[type[BaseException], BaseException, TracebackType | None]
+            | tuple[None, None, None]
+            | None,
+            func: str | None = None,
+            extra: Mapping[str, object] | None = None,
+            sinfo: str | None = None,
+        ) -> logging.LogRecord:
+            record = super().makeRecord(
+                name, level, fn, lno, msg, args, exc_info, func, extra, sinfo
+            )
+            record.created = get_current_timestamp()
+            record.msecs = (record.created - int(record.created)) * 1000
+            record.relativeCreated = (record.created - logging._startTime) * 1000  # noqa: SLF001
+            return record
+
+    return FreezableLogManager
+
+
 @contextmanager
 def freeze_time(new_now: Union[datetime.datetime, float]):
     new_dt = (
@@ -737,6 +770,9 @@ def freeze_time(new_now: Union[datetime.datetime, float]):
         unittest.mock.patch(
             "dagster._time._mockable_get_current_timestamp", return_value=new_dt.timestamp()
         ),
+        unittest.mock.patch(
+            "dagster._core.log_manager.DagsterLogManager", new=get_freezable_log_manager()
+        ),
     ):
         yield