Improve metrics

docker-compose.lite.yml

@@ -27,6 +27,31 @@ services:
     network_mode: host
     volumes:
         - workdir:/workdir
+  prometheus:
+    image: prom/prometheus:latest
+    network_mode: host
+    user: "0:0"
+    volumes:
+        - ./prometheus.yml:/etc/prometheus/prometheus.yml:ro
+        - prometheus_data:/prometheus
+    command:
+        - --config.file=/etc/prometheus/prometheus.yml
+        - --storage.tsdb.path=/prometheus
+        - --storage.tsdb.retention.time=15d
+        - --web.console.libraries=/usr/share/prometheus/console_libraries
+        - --web.console.templates=/usr/share/prometheus/consoles
+  grafana:
+    image: grafana/grafana:latest
+    network_mode: host
+    user: "0:0"
+    environment:
+        - GF_SECURITY_ADMIN_USER=admin
+        - GF_SECURITY_ADMIN_PASSWORD=admin
+        - GF_SERVER_HTTP_PORT=3500
+    volumes:
+        - grafana_data:/var/lib/grafana
 
 volumes:
   workdir:
+  prometheus_data:
+  grafana_data:

file-tracker/file_tracker/main.py

@@ -4,13 +4,16 @@
 
 from cleanup import TerminationHandler, setup_cleanup_handlers
 from connection_manager import ConnectionManager
+from metrics import start_metrics_server
 from task_listener import TaskListener
 
 
 async def main():
     workdir = os.getenv("WORKDIR", "/workdir")
     os.chdir(workdir)
 
+    await start_metrics_server()
+
     connection_manager = ConnectionManager.from_env()
     file_tracker_host = os.getenv("FILE_TRACKER_HOST", "0.0.0.0")
     file_tracker_port = int(os.getenv("FILE_TRACKER_PORT", "5000"))

file-tracker/file_tracker/metrics.py

@@ -0,0 +1,24 @@
+"""Prometheus metrics server for file-tracker service."""
+import logging
+import os
+
+from aiohttp import web
+from prometheus_client import generate_latest
+
+
+async def metrics_handler(request):
+    """Handler for /metrics endpoint exposing Prometheus client metrics."""
+    return web.Response(body=generate_latest(), content_type="text/plain")
+
+
+async def start_metrics_server(host='0.0.0.0', port=None):
+    """Start a basic Prometheus metrics server exposing client metrics."""
+    port = port or int(os.getenv('FILE_TRACKER_METRICS_PORT', '9091'))
+    app = web.Application()
+    app.router.add_get('/metrics', metrics_handler)
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, host, port)
+    await site.start()
+    logging.info("Metrics server started on %s:%s", host, port)
+    return runner

file-tracker/requirements.txt

@@ -1,2 +1,3 @@
 aiortc
-aiohttp
+aiohttp
+prometheus-client

prometheus.yml

@@ -0,0 +1,16 @@
+global:
+  scrape_interval: 15s
+  evaluation_interval: 15s
+
+scrape_configs:
+  - job_name: "file-tracker"
+    static_configs:
+      - targets: ["localhost:9091"]
+    metrics_path: "/metrics"
+
+  - job_name: "task-runner-lite"
+    static_configs:
+      - targets: ["localhost:8000"]
+    metrics_path: "/metrics"
+
+

task-runner/requirements.txt

@@ -10,3 +10,4 @@ gputil
 pydantic~=2.11.7
 tenacity
 pytest
+prometheus-client

task-runner/task_runner/main.py

@@ -29,6 +29,7 @@
     task_execution_loop,
     utils,
 )
+from task_runner.metrics import start_metrics_server
 from task_runner.register_task_runner import register_task_runner
 from task_runner.task_request_handler import TaskRequestHandler
 from task_runner.task_status import TaskRunnerTerminationReason
@@ -66,6 +67,9 @@ def _set_socks_proxy():
 
 def main(_):
     _set_socks_proxy()
+
+    start_metrics_server()
+
     workdir = os.getenv("WORKDIR", "/workdir")
     executer_images_dir = os.getenv("EXECUTER_IMAGES_DIR", "/apptainer")
     if not executer_images_dir:

task-runner/task_runner/metrics.py

@@ -0,0 +1,18 @@
+"""Prometheus metrics server for task-runner service."""
+import logging
+import os
+
+from prometheus_client import Counter, Gauge, Histogram, start_http_server
+
+# Task metrics
+tasks_active = Gauge('tasks_active', 'Number of currently active tasks')
+tasks_total = Counter('tasks_total', 'Total number of tasks', ['status'])
+task_duration = Histogram('task_duration_seconds',
+                          'Task execution duration in seconds')
+
+
+def start_metrics_server(port=None):
+    """Start a basic Prometheus metrics server exposing client metrics."""
+    port = port or int(os.getenv('TASK_RUNNER_METRICS_PORT', '8000'))
+    start_http_server(port)
+    logging.info("Metrics server started on port %s", port)

task-runner/task_runner/task_request_handler.py

@@ -30,10 +30,11 @@
     executers,
     observers,
     task_message_listener,
+    task_status,
     utils,
 )
+from task_runner.metrics import task_duration, tasks_active, tasks_total
 from task_runner.operations_logger import OperationName, OperationsLogger
-from task_runner.task_status import task_status
 from task_runner.utils import files
 
 KILL_MESSAGE = "kill"
@@ -287,6 +288,10 @@ def __call__(self, request: dict[str, str]) -> None:
         Args:
             request: Request describing the task to be executed.
         """
+        task_start_time = time.time()
+        task_status_str = 'failed'
+        tasks_active.inc()
+
         # Save the task request to use during output recovery
         with open(self.request_path, "w", encoding="utf-8") as request_file:
             json.dump(request, request_file, default=str)
@@ -416,12 +421,15 @@ def __call__(self, request: dict[str, str]) -> None:
 
             if exit_reason == TaskExitReason.KILLED:
                 new_status = task_status.TaskStatusCode.KILLED.value
+                task_status_str = 'killed'
             elif exit_reason == TaskExitReason.TTL_EXCEEDED:
                 new_status = task_status.TaskStatusCode.TTL_EXCEEDED.value
+                task_status_str = 'ttl_exceeded'
             else:
                 new_status = (task_status.TaskStatusCode.SUCCESS.value
                               if exit_code == 0 else
                               task_status.TaskStatusCode.FAILED.value)
+                task_status_str = 'success' if exit_code == 0 else 'failed'
 
             safely_delete = self.save_output(new_task_status=new_status)
 
@@ -436,6 +444,7 @@ def __call__(self, request: dict[str, str]) -> None:
 
         # Catch all exceptions to ensure that we log the error message
         except Exception as e:  # noqa: BLE001
+            task_status_str = 'error'
             message = utils.get_exception_root_cause_message(e)
             try:
                 self._publish_event(
@@ -461,6 +470,11 @@ def __call__(self, request: dict[str, str]) -> None:
                 safely_delete = False
 
         finally:
+            # Track metrics
+            task_duration.observe(time.time() - task_start_time)
+            tasks_total.labels(status=task_status_str).inc()
+            tasks_active.dec()
+
             self.cleaning_up = True
             self._cleanup(safely_delete)