feat: SLS-389: add batch ID to log messages

runpod/serverless/modules/rp_logger.py

@@ -10,12 +10,24 @@
 ERROR - 4 - Serious problem, the software has not been able to perform some function.
 """
 
+from contextvars import ContextVar, Token
 import json
 import os
 from typing import Optional
 
 MAX_MESSAGE_LENGTH = 4096
 LOG_LEVELS = ["NOTSET", "TRACE", "DEBUG", "INFO", "WARN", "ERROR"]
+_batch_id: ContextVar[Optional[str]] = ContextVar("runpod_batch_id", default=None)
+
+
+def _set_batch_id(batch_id: Optional[str]) -> Token:
+    """Set the batch ID associated with the current job task."""
+    return _batch_id.set(batch_id)
+
+
+def _reset_batch_id(token: Token):
+    """Restore the previous batch ID for the current job task."""
+    _batch_id.reset(token)
 
 
 def _validate_log_level(log_level):
@@ -74,6 +86,9 @@ def log(self, message, message_level="INFO", job_id=None):
             return
 
         message = str(message)
+        if batch_id := _batch_id.get():
+            message = f"[batchId={batch_id}] {message}"
+
         # Truncate message over 10MB, remove chunk from the middle
         if len(message) > MAX_MESSAGE_LENGTH:
             half_max_length = MAX_MESSAGE_LENGTH // 2

runpod/serverless/modules/rp_scale.py

@@ -11,7 +11,7 @@
 
 from ...http_client import AsyncClientSession, ClientSession, TooManyRequests
 from .rp_job import _job_stop_url, get_job, get_stop_signals, handle_job
-from .rp_logger import RunPodLogger
+from .rp_logger import RunPodLogger, _reset_batch_id, _set_batch_id
 from .worker_state import JobsProgress, IS_LOCAL_TEST
 
 log = RunPodLogger()
@@ -342,6 +342,7 @@ async def handle_job(self, session: ClientSession, job: dict):
         """
         Process an individual job. This function is run concurrently for multiple jobs.
         """
+        batch_id_token = _set_batch_id(job.get("batchId"))
         try:
             log.debug("Handling Job", job["id"])
 
@@ -367,3 +368,4 @@ async def handle_job(self, session: ClientSession, job: dict):
             self.jobs_tasks.pop(job["id"], None)
 
             log.debug("Finished Job", job["id"])
+            _reset_batch_id(batch_id_token)