From 45b719a4230a5c857b72a764be72e6d87f1ef3df Mon Sep 17 00:00:00 2001 From: Dhruv Maniya Date: Wed, 15 Jul 2026 22:44:54 +0530 Subject: [PATCH] feat: log container execution heartbeats Signed-off-by: Dhruv Maniya --- src/forge/sandbox/runner.py | 68 ++++++++++++----- .../sandbox/test_runner_timeout_cleanup.py | 73 +++++++++++++++++++ 2 files changed, 123 insertions(+), 18 deletions(-) diff --git a/src/forge/sandbox/runner.py b/src/forge/sandbox/runner.py index a422ef91..106ba79d 100644 --- a/src/forge/sandbox/runner.py +++ b/src/forge/sandbox/runner.py @@ -17,6 +17,7 @@ import logging import os import shutil +from contextlib import suppress from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -36,6 +37,7 @@ EXIT_TASK_FAILED = 1 EXIT_TESTS_FAILED = 2 EXIT_CONFIG_ERROR = 3 +CONTAINER_HEARTBEAT_INTERVAL_SECONDS = 60 @dataclass @@ -366,6 +368,25 @@ async def _stop_timed_out_container( process.kill() await process.wait() + @staticmethod + async def _log_container_heartbeat( + container_name: str, + timeout_seconds: int, + ) -> None: + """Log periodic progress while a container process is running.""" + loop = asyncio.get_running_loop() + started_at = loop.time() + + while True: + await asyncio.sleep(CONTAINER_HEARTBEAT_INTERVAL_SECONDS) + elapsed = int(loop.time() - started_at) + remaining = max(0, timeout_seconds - elapsed) + logger.info( + f"Container {container_name} still running " + f"({elapsed}s elapsed, {remaining}s remaining of " + f"{timeout_seconds}s timeout)" + ) + async def run( self, workspace_path: Path, @@ -425,26 +446,37 @@ async def run( stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) + heartbeat_task = asyncio.create_task( + self._log_container_heartbeat( + container_name, + config.timeout_seconds, + ) + ) try: - stdout, stderr = await asyncio.wait_for( - process.communicate(), - timeout=config.timeout_seconds + 60, # Extra buffer - ) - except TimeoutError: - logger.error(f"Container execution timed out, stopping {container_name}") - await self._stop_timed_out_container(container_name, process) - return ContainerResult( - success=False, - exit_code=-1, - stdout="", - stderr="Container execution timed out", - error_message="Timeout exceeded", - ) - except asyncio.CancelledError: - logger.warning(f"Container execution cancelled, stopping {container_name}") - await self._stop_timed_out_container(container_name, process) - raise # Re-raise CancelledError + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), + timeout=config.timeout_seconds + 60, # Extra buffer + ) + except TimeoutError: + logger.error(f"Container execution timed out, stopping {container_name}") + await self._stop_timed_out_container(container_name, process) + return ContainerResult( + success=False, + exit_code=-1, + stdout="", + stderr="Container execution timed out", + error_message="Timeout exceeded", + ) + except asyncio.CancelledError: + logger.warning(f"Container execution cancelled, stopping {container_name}") + await self._stop_timed_out_container(container_name, process) + raise # Re-raise CancelledError + finally: + heartbeat_task.cancel() + with suppress(asyncio.CancelledError): + await heartbeat_task exit_code = process.returncode or 0 stdout_str = stdout.decode("utf-8", errors="replace") diff --git a/tests/unit/sandbox/test_runner_timeout_cleanup.py b/tests/unit/sandbox/test_runner_timeout_cleanup.py index de07e6b9..8ffba608 100644 --- a/tests/unit/sandbox/test_runner_timeout_cleanup.py +++ b/tests/unit/sandbox/test_runner_timeout_cleanup.py @@ -1,6 +1,8 @@ """Tests for container runner timeout cleanup.""" +import asyncio import json +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,6 +14,77 @@ def _runner_without_init() -> ContainerRunner: return object.__new__(ContainerRunner) +@pytest.mark.asyncio +async def test_container_heartbeat_logs_elapsed_and_remaining_time( + caplog: pytest.LogCaptureFixture, +) -> None: + runner = _runner_without_init() + loop = MagicMock() + loop.time.side_effect = [100.0, 160.0] + sleep_calls = 0 + + async def fake_sleep(delay: float) -> None: + nonlocal sleep_calls + assert delay == 60 + sleep_calls += 1 + if sleep_calls == 2: + raise asyncio.CancelledError + + with ( + patch("forge.sandbox.runner.asyncio.get_running_loop", return_value=loop), + patch("forge.sandbox.runner.asyncio.sleep", side_effect=fake_sleep), + caplog.at_level(logging.INFO, logger="forge.sandbox.runner"), + pytest.raises(asyncio.CancelledError), + ): + await runner._log_container_heartbeat("forge-ticket-abc123", 180) + + assert ( + "Container forge-ticket-abc123 still running (60s elapsed, 120s remaining of 180s timeout)" + ) in caplog.messages + + +@pytest.mark.asyncio +async def test_run_cancels_heartbeat_after_process_exits(tmp_path) -> None: + runner = _runner_without_init() + runner.settings = MagicMock() + runner.settings.container_keep = False + runner._build_container_name = MagicMock(return_value="forge-ticket-abc123") + runner._build_podman_command = MagicMock(return_value=["podman", "run", "fake"]) + + heartbeat_started = asyncio.Event() + heartbeat_cancelled = asyncio.Event() + + async def heartbeat(_container_name: str, _timeout_seconds: int) -> None: + heartbeat_started.set() + try: + await asyncio.Event().wait() + finally: + heartbeat_cancelled.set() + + async def communicate() -> tuple[bytes, bytes]: + await heartbeat_started.wait() + return b"ok", b"" + + runner._log_container_heartbeat = heartbeat # type: ignore[method-assign] + process = MagicMock() + process.communicate = AsyncMock(side_effect=communicate) + process.returncode = 0 + + with patch( + "forge.sandbox.runner.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=process), + ): + result = await runner.run( + workspace_path=tmp_path, + task_summary="Do it", + task_description="Details", + config=ContainerConfig(timeout_seconds=180), + ) + + assert result.success is True + assert heartbeat_cancelled.is_set() + + @pytest.mark.asyncio async def test_stop_failure_kills_container_and_waits_for_run_process() -> None: runner = _runner_without_init()