From 7e8a71c7362cbb78192f41d9fccaf18f0e38469c Mon Sep 17 00:00:00 2001 From: jhcipar Date: Tue, 4 Aug 2026 09:31:42 -0700 Subject: [PATCH 1/2] feat(serverless): add batch ID to log messages --- runpod/serverless/modules/rp_logger.py | 15 +++++++++++++++ runpod/serverless/modules/rp_scale.py | 18 +++++++++++------- 2 files changed, 26 insertions(+), 7 deletions(-) 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..11083c4f1 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"]) @@ -359,11 +360,14 @@ async def handle_job(self, session: ClientSession, job: dict): raise finally: - # Inform Queue of a task completion - self.jobs_queue.task_done() + try: + # Inform Queue of a task completion + self.jobs_queue.task_done() - # Job is no longer in progress - self.job_progress.remove(job) - self.jobs_tasks.pop(job["id"], None) + # Job is no longer in progress + self.job_progress.remove(job) + self.jobs_tasks.pop(job["id"], None) - log.debug("Finished Job", job["id"]) + log.debug("Finished Job", job["id"]) + finally: + _reset_batch_id(batch_id_token) From 34a87bddf92881d0930d7e6701acd7ac0171007a Mon Sep 17 00:00:00 2001 From: jhcipar Date: Tue, 4 Aug 2026 10:27:01 -0700 Subject: [PATCH 2/2] refactor(serverless): preserve existing job cleanup flow --- runpod/serverless/modules/rp_scale.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 11083c4f1..4cbf94ffb 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -360,14 +360,12 @@ async def handle_job(self, session: ClientSession, job: dict): raise finally: - try: - # Inform Queue of a task completion - self.jobs_queue.task_done() + # Inform Queue of a task completion + self.jobs_queue.task_done() - # Job is no longer in progress - self.job_progress.remove(job) - self.jobs_tasks.pop(job["id"], None) + # Job is no longer in progress + self.job_progress.remove(job) + self.jobs_tasks.pop(job["id"], None) - log.debug("Finished Job", job["id"]) - finally: - _reset_batch_id(batch_id_token) + log.debug("Finished Job", job["id"]) + _reset_batch_id(batch_id_token)