From d539dcb7dc34bfbcdd84be10cb6a0377b96f6717 Mon Sep 17 00:00:00 2001
From: Radu Swigler <radu.porumba@gmail.com>
Date: Thu, 3 Sep 2026 03:06:51 -0400
Subject: [PATCH] fix: guard set_scale against invalid concurrency_modifier
 returns (#458)

If a user-provided concurrency_modifier callback returns None, a
non-integer, or a value < 1, the worker crashes with:
  TypeError: '<' not supported between instances of 'int' and 'NoneType'

Validate the return value in set_scale(): reject None, non-int, and
values < 1 (default to 1 with a warning). Wrap the callback in
try/except so a raising modifier keeps the current concurrency instead
of crashing the worker.

Closes #458
---
 runpod/serverless/modules/rp_scale.py         | 18 +++++++-
 .../test_modules/test_scale.py                | 42 ++++++++++++++++++-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py
index 4cbf94ffb..6355380a3 100644
--- a/runpod/serverless/modules/rp_scale.py
+++ b/runpod/serverless/modules/rp_scale.py
@@ -85,7 +85,23 @@ def __init__(self, config: Dict[str, Any]):
             self.stop_signals_fetcher_timeout = stop_signals_fetcher_timeout
 
     async def set_scale(self):
-        self.current_concurrency = self.concurrency_modifier(self.current_concurrency)
+        try:
+            result = self.concurrency_modifier(self.current_concurrency)
+        except Exception as err:
+            log.warning(
+                "concurrency_modifier raised %s: %s — keeping concurrency at %d",
+                type(err).__name__, err, self.current_concurrency,
+            )
+            result = self.current_concurrency
+
+        if not isinstance(result, int) or result < 1:
+            log.warning(
+                "concurrency_modifier returned invalid value %r — defaulting to 1",
+                result,
+            )
+            result = 1
+
+        self.current_concurrency = result
 
         if self.jobs_queue and (self.current_concurrency == self.jobs_queue.maxsize):
             # no need to resize
diff --git a/tests/test_serverless/test_modules/test_scale.py b/tests/test_serverless/test_modules/test_scale.py
index 640694225..3f8727957 100644
--- a/tests/test_serverless/test_modules/test_scale.py
+++ b/tests/test_serverless/test_modules/test_scale.py
@@ -1,9 +1,10 @@
+import asyncio
 import sys
 import traceback
 from unittest import TestCase
 from unittest.mock import patch
 
-from runpod.serverless.modules.rp_scale import _handle_uncaught_exception
+from runpod.serverless.modules.rp_scale import JobScaler, _handle_uncaught_exception
 
 
 class TestHandleUncaughtException(TestCase):
@@ -52,3 +53,42 @@ def test_handle_uncaught_exception_with_no_exception(self, mock_logger):
     def test_excepthook_not_set_when_start_not_invoked(self):
         assert sys.excepthook == sys.__excepthook__
         assert sys.excepthook != _handle_uncaught_exception
+
+
+class TestSetScaleConcurrencyValidation(TestCase):
+    """Tests that set_scale guards against invalid concurrency_modifier returns. (Fixes #458)"""
+
+    def _make_scaler(self, modifier):
+        config = {"handler": lambda job: job, "concurrency_modifier": modifier}
+        return JobScaler(config)
+
+    @patch("runpod.serverless.modules.rp_scale.log")
+    def test_none_return_defaults_to_1(self, _mock_log):
+        scaler = self._make_scaler(lambda _: None)
+        asyncio.run(scaler.set_scale())
+        assert scaler.current_concurrency == 1
+
+    @patch("runpod.serverless.modules.rp_scale.log")
+    def test_negative_return_defaults_to_1(self, _mock_log):
+        scaler = self._make_scaler(lambda _: -5)
+        asyncio.run(scaler.set_scale())
+        assert scaler.current_concurrency == 1
+
+    @patch("runpod.serverless.modules.rp_scale.log")
+    def test_zero_return_defaults_to_1(self, _mock_log):
+        scaler = self._make_scaler(lambda _: 0)
+        asyncio.run(scaler.set_scale())
+        assert scaler.current_concurrency == 1
+
+    @patch("runpod.serverless.modules.rp_scale.log")
+    def test_exception_keeps_current(self, _mock_log):
+        scaler = self._make_scaler(lambda _: (_ for _ in ()).throw(RuntimeError("boom")))
+        scaler.current_concurrency = 4
+        asyncio.run(scaler.set_scale())
+        assert scaler.current_concurrency == 4
+
+    @patch("runpod.serverless.modules.rp_scale.log")
+    def test_valid_int_applied(self, _mock_log):
+        scaler = self._make_scaler(lambda _: 8)
+        asyncio.run(scaler.set_scale())
+        assert scaler.current_concurrency == 8