fix(oso_dagster): add heartbeat to any usage of producer trino (#5475)

ravenac95 · web-flow · commit 10d957fe56af · 2025-10-23T18:55:55.000Z
* fix(dagster): adds the heartbeat to the trino exporter

* fix: make heartbeat naming more generic

Update warehouse/oso_dagster/assets/default/trino_automation.py

* fix: move all trino heartbeat logic into a single place

* fix

* fix: further clean up

* clean up
diff --git a/warehouse/oso_dagster/assets/default/trino_automation.py b/warehouse/oso_dagster/assets/default/trino_automation.py
@@ -0,0 +1,65 @@
+import dagster as dg
+from oso_dagster.config import DagsterConfig
+from oso_dagster.factories.common import (
+    AssetFactoryResponse,
+    early_resources_asset_factory,
+)
+from oso_dagster.resources.heartbeat import HeartBeatResource
+from oso_dagster.resources.trino import TrinoResource
+
+
+@early_resources_asset_factory()
+def trino_automation_assets() -> AssetFactoryResponse:
+    # Define a job that checks the heartbeat of "producer_trino". If the
+    # heartbeat is older than 30 minutes, we scale trino down to zero.
+    @dg.op
+    async def trino_heartbeat_checker(
+        context: dg.OpExecutionContext,
+        global_config: dg.ResourceParam[DagsterConfig],
+        trino: dg.ResourceParam[TrinoResource],
+        heartbeat: dg.ResourceParam[HeartBeatResource],
+    ) -> None:
+        from datetime import datetime, timedelta, timezone
+
+        now = datetime.now(timezone.utc)
+
+        last_heartbeat = await heartbeat.get_last_heartbeat_for("producer_trino")
+        if last_heartbeat is None:
+            context.log.info(
+                "No heartbeat found for sqlmesh now ensuring trino shutdown."
+            )
+            last_heartbeat = now - timedelta(
+                minutes=global_config.sqlmesh_trino_ttl_minutes + 1
+            )
+
+        # Only scale down trino if we're in a k8s environment
+        if not global_config.k8s_enabled:
+            return
+
+        if now - last_heartbeat > timedelta(
+            minutes=global_config.sqlmesh_trino_ttl_minutes
+        ):
+            context.log.info(
+                f"No heartbeat detected for sqlmesh in the last {global_config.sqlmesh_trino_ttl_minutes} minutes. Ensuring that producer trino is scaled down."
+            )
+            await trino.ensure_shutdown()
+        else:
+            context.log.info("Heartbeat detected for sqlmesh. No action needed.")
+
+    # Use the in-process executor for the heartbeat monitor job to avoid
+    # the overhead of spinning up a new k8s pod in addition to the run launcher
+    @dg.job(executor_def=dg.in_process_executor)
+    def trino_heartbeat_monitor_job():
+        trino_heartbeat_checker()
+
+    return AssetFactoryResponse(
+        assets=[],
+        jobs=[trino_heartbeat_monitor_job],
+        schedules=[
+            dg.ScheduleDefinition(
+                name="trino_heartbeat_monitor_schedule",
+                job=trino_heartbeat_monitor_job,
+                cron_schedule="*/15 * * * *",
+            )
+        ],
+    )
diff --git a/warehouse/oso_dagster/assets/sqlmesh/sqlmesh.py b/warehouse/oso_dagster/assets/sqlmesh/sqlmesh.py
@@ -24,7 +24,6 @@
 from oso_dagster.config import DagsterConfig
 from oso_dagster.factories import AssetFactoryResponse, cacheable_asset_factory
 from oso_dagster.factories.common import CacheableDagsterContext
-from oso_dagster.resources.heartbeat import HeartBeatResource
 from oso_dagster.resources.trino import TrinoResource
 from oso_dagster.utils.asynctools import multiple_async_contexts
 from pydantic import BaseModel
@@ -164,7 +163,6 @@ async def sqlmesh_project(
             global_config: ResourceParam[DagsterConfig],
             sqlmesh: SQLMeshResource,
             trino: TrinoResource,
-            heartbeat: HeartBeatResource,
             config: SQLMeshRunConfig,
         ):
             restate_models = config.restate_models[:] if config.restate_models else []
@@ -234,68 +232,21 @@ def run_sqlmesh(
 
             # Trino can either be `local-trino` or `trino`
             if "trino" in global_config.sqlmesh_gateway:
-                # Start a heartbeat to indicate that sqlmesh is running
-                async with heartbeat.heartbeat(
-                    job_name="sqlmesh", interval_seconds=300, log_override=context.log
+                async with multiple_async_contexts(
+                    trino=trino.ensure_available(log_override=context.log),
                 ):
-                    async with multiple_async_contexts(
-                        trino=trino.ensure_available(log_override=context.log),
-                    ):
-                        for result in run_sqlmesh(context, sqlmesh, config):
-                            yield result
+                    for result in run_sqlmesh(context, sqlmesh, config):
+                        yield result
             else:
                 # If we are not running trino we are using duckdb
                 for result in run_sqlmesh(context, sqlmesh, config):
                     yield result
 
-        # Define a job that checks the heartbeat of sqlmesh runs. If the
-        # heartbeat is older than 30 minutes, we scale trino down to zero.
-        @dg.op
-        async def sqlmesh_heartbeat_checker(
-            context: dg.OpExecutionContext,
-            global_config: ResourceParam[DagsterConfig],
-            trino: TrinoResource,
-            heartbeat: HeartBeatResource,
-        ) -> None:
-            from datetime import datetime, timedelta, timezone
-
-            now = datetime.now(timezone.utc)
-
-            last_heartbeat = await heartbeat.get_last_heartbeat_for("sqlmesh")
-            if last_heartbeat is None:
-                context.log.info(
-                    "No heartbeat found for sqlmesh now ensuring trino shutdown."
-                )
-                last_heartbeat = now - timedelta(
-                    minutes=global_config.sqlmesh_trino_ttl_minutes + 1
-                )
-
-            # Only scale down trino if we're in a k8s environment
-            if not global_config.k8s_enabled:
-                return
-
-            if now - last_heartbeat > timedelta(
-                minutes=global_config.sqlmesh_trino_ttl_minutes
-            ):
-                context.log.info(
-                    f"No heartbeat detected for sqlmesh in the last {global_config.sqlmesh_trino_ttl_minutes} minutes. Ensuring that producer trino is scaled down."
-                )
-                await trino.ensure_shutdown()
-            else:
-                context.log.info("Heartbeat detected for sqlmesh. No action needed.")
-
-        # Use the in-process executor for the heartbeat monitor job to avoid
-        # the overhead of spinning up a new k8s pod in addition to the run launcher
-        @dg.job(executor_def=dg.in_process_executor)
-        def sqlmesh_heartbeat_monitor_job():
-            sqlmesh_heartbeat_checker()
-
         all_assets_selection = AssetSelection.assets(sqlmesh_project)
 
         return AssetFactoryResponse(
             assets=[sqlmesh_project],
             jobs=[
-                sqlmesh_heartbeat_monitor_job,
                 define_asset_job(
                     name="sqlmesh_all_assets",
                     selection=all_assets_selection,
@@ -356,13 +307,6 @@ def sqlmesh_heartbeat_monitor_job():
                     ),
                 ),
             ],
-            schedules=[
-                dg.ScheduleDefinition(
-                    name="sqlmesh_heartbeat_monitor_schedule",
-                    job=sqlmesh_heartbeat_monitor_job,
-                    cron_schedule="*/15 * * * *",
-                )
-            ],
         )
 
     return cache_context
diff --git a/warehouse/oso_dagster/definitions/resources.py b/warehouse/oso_dagster/definitions/resources.py
@@ -188,7 +188,9 @@ def k8s_resource_factory(global_config: DagsterConfig) -> K8sResource | K8sApiRe
 @resource_factory("trino")
 @time_function(logger)
 def trino_resource_factory(
-    global_config: DagsterConfig, k8s: K8sResource | K8sApiResource
+    global_config: DagsterConfig,
+    k8s: K8sResource | K8sApiResource,
+    heartbeat: HeartBeatResource,
 ) -> TrinoResource:
     if not global_config.k8s_enabled:
         return TrinoRemoteResource()
@@ -199,6 +201,7 @@ def trino_resource_factory(
         coordinator_deployment_name=global_config.trino_k8s_coordinator_deployment_name,
         worker_deployment_name=global_config.trino_k8s_worker_deployment_name,
         use_port_forward=global_config.k8s_use_port_forward,
+        heartbeat=heartbeat,
     )
 
 
diff --git a/warehouse/oso_dagster/resources/heartbeat.py b/warehouse/oso_dagster/resources/heartbeat.py
@@ -5,10 +5,10 @@
 import asyncio
 import concurrent.futures
 import logging
+import threading
 import typing as t
 from contextlib import asynccontextmanager, suppress
 from datetime import datetime, timezone
-from queue import Empty, Queue
 
 import aiofiles
 import dagster as dg
@@ -22,60 +22,72 @@
 
 def run_beat_loop(
     interval_seconds: int,
-    queue: Queue[bool],
+    stop: threading.Event,
     beat_loop_func: BeatLoopFunc,
     beat_loop_kwargs: dict[str, t.Any],
 ) -> None:
     logger.info("Starting heartbeat beat loop")
     while True:
         logger.info("Running heartbeat beat loop function")
         asyncio.run(beat_loop_func(**beat_loop_kwargs))
-        try:
-            if queue.get(timeout=interval_seconds):
-                logger.info("Stopping heartbeat beat loop")
-                break
-        except Empty:
-            continue
+        if stop.wait(timeout=float(interval_seconds)):
+            logger.info("Stopping heartbeat beat loop")
+            break
 
 
 class HeartBeatResource(dg.ConfigurableResource):
     def beat_loop_func(self) -> BeatLoopFunc:
+        """Return the function to be called in the heartbeat loop"""
         raise NotImplementedError()
 
     def beat_loop_kwargs(self) -> dict[str, t.Any]:
+        """Return the kwargs to be passed to the heartbeat loop function"""
         return {}
 
-    async def get_last_heartbeat_for(self, job_name: str) -> datetime | None:
+    async def get_last_heartbeat_for(self, name: str) -> datetime | None:
         raise NotImplementedError()
 
-    async def beat(self, job_name: str) -> None:
+    async def beat(self, name: str) -> None:
         raise NotImplementedError()
 
     @asynccontextmanager
     async def heartbeat(
         self,
-        job_name: str,
+        name: str,
         interval_seconds: int = 120,
         log_override: logging.Logger | None = None,
     ) -> t.AsyncIterator[None]:
         log_override = log_override or logger
         loop = asyncio.get_running_loop()
         with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
             kwargs = self.beat_loop_kwargs().copy()
-            kwargs.update({"job_name": job_name})
-            queue = Queue[bool]()
+            kwargs.update({"heartbeat_name": name})
+            stop = threading.Event()
+
+            # The beat loop must run in a separate thread because dagster's
+            # async event loop tends to block despite being async. Using a
+            # thread ensures that the heartbeat will run assuming the process is
+            # alive and not completely blocked.
+            #
+            # So in order to make this work we must pass all of the context we
+            # need for the beat loop function as the function running in a
+            # separate thread might not have access to the same context.
+            # Additionally, this means changing this to be separate processes
+            # can be done in the future if needed. This is also why the
+            # functions used for the beat loop are not methods of the resource
+            # class implementations.
             beat_task = loop.run_in_executor(
                 executor,
                 run_beat_loop,
                 interval_seconds,
-                queue,
+                stop,
                 self.beat_loop_func(),
                 kwargs,
             )
             try:
                 yield
             finally:
-                queue.put(True)
+                stop.set()
                 beat_task.cancel()
                 with suppress(asyncio.CancelledError):
                     await beat_task
@@ -90,10 +102,10 @@ async def async_redis_client(host: str, port: int) -> t.AsyncIterator[Redis]:
         await client.aclose()
 
 
-async def redis_send_heartbeat(*, host: str, port: int, job_name: str) -> None:
+async def redis_send_heartbeat(*, host: str, port: int, heartbeat_name: str) -> None:
     async with async_redis_client(host, port) as redis_client:
         await redis_client.set(
-            f"heartbeat:{job_name}", datetime.now(timezone.utc).isoformat()
+            f"heartbeat:{heartbeat_name}", datetime.now(timezone.utc).isoformat()
         )
 
 
@@ -107,29 +119,29 @@ def beat_loop_func(self) -> BeatLoopFunc:
     def beat_loop_kwargs(self) -> dict[str, t.Any]:
         return {"host": self.host, "port": self.port}
 
-    async def get_last_heartbeat_for(self, job_name: str) -> datetime | None:
+    async def get_last_heartbeat_for(self, name: str) -> datetime | None:
         async with async_redis_client(self.host, self.port) as redis_client:
-            timestamp = await redis_client.get(f"heartbeat:{job_name}")
-            logger.info(f"Fetched heartbeat for job {job_name}: {timestamp}")
+            timestamp = await redis_client.get(f"heartbeat:{name}")
+            logger.info(f"Fetched heartbeat `{name}`: {timestamp}")
             if isinstance(timestamp, str):
                 return datetime.fromisoformat(timestamp)
             elif isinstance(timestamp, bytes):
                 return datetime.fromisoformat(timestamp.decode("utf-8"))
             else:
                 return None
 
-    async def beat(self, job_name: str) -> None:
+    async def beat(self, name: str) -> None:
         return await redis_send_heartbeat(
-            host=self.host, port=self.port, job_name=job_name
+            host=self.host, port=self.port, heartbeat_name=name
         )
 
 
-async def filebased_send_heartbeat(*, directory: str, job_name: str) -> None:
+async def filebased_send_heartbeat(*, directory: str, heartbeat_name: str) -> None:
     from pathlib import Path
 
     import aiofiles
 
-    filepath = Path(directory) / f"{job_name}_heartbeat.txt"
+    filepath = Path(directory) / f"{heartbeat_name}_heartbeat.txt"
     async with aiofiles.open(filepath, mode="w") as f:
         await f.write(datetime.now(timezone.utc).isoformat())
 
@@ -139,29 +151,25 @@ class FilebasedHeartBeatResource(HeartBeatResource):
 
     directory: str = Field(description="Directory to store heartbeat files.")
 
-    async def get_last_heartbeat_for(self, job_name: str) -> datetime | None:
+    async def get_last_heartbeat_for(self, name: str) -> datetime | None:
         from pathlib import Path
 
-        filepath = Path(self.directory) / f"{job_name}_heartbeat.txt"
+        filepath = Path(self.directory) / f"{name}_heartbeat.txt"
         if not filepath.exists():
             return None
         async with aiofiles.open(filepath, mode="r") as f:
             timestamp = await f.read()
             return datetime.fromisoformat(timestamp)
 
-    async def beat(self, job_name: str) -> None:
-        from pathlib import Path
-
-        import aiofiles
-
-        filepath = Path(self.directory) / f"{job_name}_heartbeat.txt"
-        async with aiofiles.open(filepath, mode="w") as f:
-            await f.write(datetime.now(timezone.utc).isoformat())
+    async def beat(self, name: str) -> None:
+        return await filebased_send_heartbeat(
+            directory=self.directory, heartbeat_name=name
+        )
 
     @asynccontextmanager
     async def heartbeat(
         self,
-        job_name: str,
+        name: str,
         interval_seconds: int = 120,
         log_override: logging.Logger | None = None,
     ) -> t.AsyncIterator[None]:
@@ -170,14 +178,12 @@ async def heartbeat(
         async def beat_loop():
             while True:
                 try:
-                    await self.beat(job_name)
+                    await self.beat(name)
                     logger_to_use.info(
-                        f"Heartbeat sent for job {job_name} at {datetime.now(timezone.utc).isoformat()}"
+                        f"Heartbeat sent for job {name} at {datetime.now(timezone.utc).isoformat()}"
                     )
                 except Exception as e:
-                    logger_to_use.error(
-                        f"Error sending heartbeat for job {job_name}: {e}"
-                    )
+                    logger_to_use.error(f"Error sending heartbeat for job {name}: {e}")
                 await asyncio.sleep(interval_seconds)
 
         beat_task = asyncio.create_task(beat_loop())
diff --git a/warehouse/oso_dagster/resources/trino.py b/warehouse/oso_dagster/resources/trino.py
