diff --git a/runpod/serverless/modules/rp_logger.py b/runpod/serverless/modules/rp_logger.py index 47ee46712..6ef4c5f73 100644 --- a/runpod/serverless/modules/rp_logger.py +++ b/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 diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 4546162d9..4cbf94ffb 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/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)