fix: fixes heartbeat errors that were getting hidden (#5471)

Author: ravenac95

warehouse/oso_dagster/assets/default/heartbeat.py

@@ -3,6 +3,7 @@
 """
 
 import dagster as dg
+from oso_dagster.config import DagsterConfig
 from oso_dagster.factories.common import (
     AssetFactoryResponse,
     early_resources_asset_factory,
@@ -15,7 +16,7 @@ class HeartbeatConfig(dg.Config):
 
 
 @early_resources_asset_factory()
-def heartbeat_factory() -> AssetFactoryResponse:
+def heartbeat_factory(global_config: DagsterConfig) -> AssetFactoryResponse:
     @dg.op
     async def heartbeat_fake_beat(
         config: HeartbeatConfig, heartbeat: HeartBeatResource
@@ -44,7 +45,33 @@ async def heartbeat_get_last_heartbeat(
     def heartbeat_get_last_heartbeat_job():
         heartbeat_get_last_heartbeat()
 
+    @dg.asset
+    async def heartbeat_test_asset(
+        context: dg.AssetExecutionContext, heartbeat: HeartBeatResource
+    ) -> dg.MaterializeResult:
+        # Return a basic dataframe with a row of data with 3 columns
+
+        async with heartbeat.heartbeat(
+            "heartbeat_noop_asset", interval_seconds=5, log_override=context.log
+        ):
+            # Intentionally do some work here to simulate a long running asset
+            # that consumes a lot of cpu
+            def fibonacci(n):
+                if n <= 1:
+                    return n
+                else:
+                    return fibonacci(n - 1) + fibonacci(n - 2)
+
+            for i in range(40):
+                context.log.info(f"Fibonacci({i}) = {fibonacci(i)}")
+            context.log.info("Heartbeat asset completed work")
+            return dg.MaterializeResult(metadata={"info": "Heartbeat asset completed"})
+
+    assets = []
+    if global_config.test_assets_enabled:
+        assets.append(heartbeat_test_asset)
+
     return AssetFactoryResponse(
-        assets=[],
+        assets=assets,
         jobs=[heartbeat_fake_beat_job, heartbeat_get_last_heartbeat_job],
     )

warehouse/oso_dagster/definitions/resources.py

@@ -415,15 +415,17 @@ def nessie_factory(global_config: DagsterConfig):
 @time_function(logger)
 def heartbeat_factory(global_config: DagsterConfig) -> HeartBeatResource:
     """Factory function to create a heartbeat resource."""
-    if global_config.k8s_enabled:
+    if global_config.k8s_enabled or global_config.redis_host:
         assert global_config.redis_host is not None, (
             "Redis host must be set for Redis heartbeat."
         )
+        logger.info("Using RedisHeartBeatResource for heartbeat.")
         return RedisHeartBeatResource(
             host=global_config.redis_host,
             port=global_config.redis_port,
         )
     else:
+        logger.info("Using FilebasedHeartBeatResource for heartbeat.")
         return FilebasedHeartBeatResource(
             directory=global_config.dagster_home,
         )

warehouse/oso_dagster/resources/heartbeat.py

@@ -8,6 +8,7 @@
 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
@@ -16,8 +17,34 @@
 
 logger = logging.getLogger(__name__)
 
+BeatLoopFunc = t.Callable[..., t.Coroutine[None, None, None]]
+
+
+def run_beat_loop(
+    interval_seconds: int,
+    queue: Queue[bool],
+    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
+
 
 class HeartBeatResource(dg.ConfigurableResource):
+    def beat_loop_func(self) -> BeatLoopFunc:
+        raise NotImplementedError()
+
+    def beat_loop_kwargs(self) -> dict[str, t.Any]:
+        return {}
+
     async def get_last_heartbeat_for(self, job_name: str) -> datetime | None:
         raise NotImplementedError()
 
@@ -30,51 +57,58 @@ async def heartbeat(
         job_name: str,
         interval_seconds: int = 120,
         log_override: logging.Logger | None = None,
-    ):
-        """Asynchronously run a heartbeat that updates every `interval_seconds`. We
-        use a separate process which should only live as long as the entire pod
-        for the dagster job is alive.
-        """
-
+    ) -> t.AsyncIterator[None]:
         log_override = log_override or logger
-
-        async def _beat_loop():
-            log_override.info(
-                f"Starting heartbeat for job {job_name} every {interval_seconds} seconds"
+        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]()
+            beat_task = loop.run_in_executor(
+                executor,
+                run_beat_loop,
+                interval_seconds,
+                queue,
+                self.beat_loop_func(),
+                kwargs,
             )
-            while True:
-                log_override.info(f"Beating heartbeat for job {job_name}")
-                await self.beat(job_name)
-                await asyncio.sleep(interval_seconds)
+            try:
+                yield
+            finally:
+                queue.put(True)
+                beat_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await beat_task
 
-        async def _beat_process():
-            loop = asyncio.get_running_loop()
-            with concurrent.futures.ProcessPoolExecutor(max_workers=1) as pool:
-                await loop.run_in_executor(pool, asyncio.run, _beat_loop())
 
-        task = asyncio.create_task(_beat_process())
-        try:
-            yield
-        finally:
-            task.cancel()
-            with suppress(asyncio.CancelledError):
-                await task
+@asynccontextmanager
+async def async_redis_client(host: str, port: int) -> t.AsyncIterator[Redis]:
+    client = Redis(host=host, port=port)
+    try:
+        yield client
+    finally:
+        await client.aclose()
+
+
+async def redis_send_heartbeat(*, host: str, port: int, job_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()
+        )
 
 
 class RedisHeartBeatResource(HeartBeatResource):
     host: str = Field(description="Redis host for heartbeat storage.")
     port: int = Field(default=6379, description="Redis port for heartbeat storage.")
 
-    @asynccontextmanager
-    async def redis_client(self) -> t.AsyncIterator[Redis]:
-        client = Redis(host=self.host, port=self.port)
-        try:
-            yield client
-        finally:
-            await client.aclose()
+    def beat_loop_func(self) -> BeatLoopFunc:
+        return redis_send_heartbeat
+
+    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 with self.redis_client() as redis_client:
+        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}")
             if isinstance(timestamp, str):
@@ -85,11 +119,19 @@ async def get_last_heartbeat_for(self, job_name: str) -> datetime | None:
                 return None
 
     async def beat(self, job_name: str) -> None:
-        async with self.redis_client() as redis_client:
-            logger.info(f"Setting heartbeat for job {job_name}")
-            await redis_client.set(
-                f"heartbeat:{job_name}", datetime.now(timezone.utc).isoformat()
-            )
+        return await redis_send_heartbeat(
+            host=self.host, port=self.port, job_name=job_name
+        )
+
+
+async def filebased_send_heartbeat(*, directory: str, job_name: str) -> None:
+    from pathlib import Path
+
+    import aiofiles
+
+    filepath = Path(directory) / f"{job_name}_heartbeat.txt"
+    async with aiofiles.open(filepath, mode="w") as f:
+        await f.write(datetime.now(timezone.utc).isoformat())
 
 
 class FilebasedHeartBeatResource(HeartBeatResource):
@@ -115,3 +157,33 @@ async def beat(self, job_name: str) -> None:
         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())
+
+    @asynccontextmanager
+    async def heartbeat(
+        self,
+        job_name: str,
+        interval_seconds: int = 120,
+        log_override: logging.Logger | None = None,
+    ) -> t.AsyncIterator[None]:
+        logger_to_use = log_override or logger
+
+        async def beat_loop():
+            while True:
+                try:
+                    await self.beat(job_name)
+                    logger_to_use.info(
+                        f"Heartbeat sent for job {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}"
+                    )
+                await asyncio.sleep(interval_seconds)
+
+        beat_task = asyncio.create_task(beat_loop())
+        try:
+            yield
+        finally:
+            beat_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await beat_task