fix(deployments): harden workload_calculator against bad return values

A user-supplied workload_calculator now runs on the live request path
(count_workload at backend.py:443) and the benchmark path. Without a
guard, a calculator that raises or returns nan/inf/negative would 500
every request and poison the autoscaler's load sums (workload_pending +=
... never recovers from a nan). Fall back to the default 100 and warn.

Also rename the kwargs test (it never exercised positional args) and add
coverage for the raise / non-finite / negative fallbacks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: vastzuby

tests/serverless/test_remote_workload_calculator.py

@@ -48,7 +48,7 @@ async def mul(a, b):
     assert entry.workload_calculator is calc
 
 
-def test_workload_calculator_receives_deserialized_args() -> None:
+def test_workload_calculator_receives_deserialized_kwargs() -> None:
     d = Deployment(name="wl-args")
 
     @d.remote(workload_calculator=lambda a, b: float(len(a) * len(b)))
@@ -76,6 +76,39 @@ async def mul(a, b):
     assert payload.count_workload() == 100.0
 
 
+def test_workload_calculator_falls_back_when_it_raises() -> None:
+    d = Deployment(name="wl-raises")
+
+    def boom(a, b):
+        raise ValueError("bad input")
+
+    @d.remote(workload_calculator=boom)
+    async def mul(a, b):
+        return a
+
+    handler = _handler_for(d, "/remote/mul", boom)
+    payload = handler.payload_cls().from_json_msg(
+        _client_payload(d, a=[1, 2, 3], b=[4, 5])
+    )
+    assert payload.count_workload() == 100.0
+
+
+def test_workload_calculator_falls_back_on_negative_or_nan() -> None:
+    d = Deployment(name="wl-bad-value")
+
+    @d.remote()
+    async def mul(a, b):
+        return a
+
+    for bad in (-1.0, float("nan"), float("inf")):
+        calc = (lambda v: lambda a, b: v)(bad)
+        handler = _handler_for(d, "/remote/mul", calc)
+        payload = handler.payload_cls().from_json_msg(
+            _client_payload(d, a=[1, 2, 3], b=[4, 5])
+        )
+        assert payload.count_workload() == 100.0
+
+
 def test_into_worker_wires_workload_calculator(monkeypatch) -> None:
     d = Deployment(name="wl-into-worker")
 

vastai/serverless/server/worker.py

@@ -9,6 +9,7 @@
 from aiohttp import web, ClientResponse
 import logging
 import json
+import math
 import random
 import inspect
 import logging
@@ -30,6 +31,10 @@
 ]
 WorkloadCalculator = Callable[[Dict[str, Any]], float]
 
+log = logging.getLogger(__name__)
+
+DEFAULT_WORKLOAD = 100.0
+
 
 @dataclass
 class LogActionConfig:
@@ -172,11 +177,26 @@ def from_dict(cls, input: Dict[str, Any]) -> "GenericApiPayload":
                     return cls(input=input)
 
                 def count_workload(self) -> float:
-                    # Use custom workload calculator if provided
-                    if user_workload_calculator:
-                        return user_workload_calculator(self.input)
-                    # Default to 100 unless overridden
-                    return 100.0
+                    if not user_workload_calculator:
+                        return DEFAULT_WORKLOAD
+                    # Fall back rather than let a bad calculator 500 requests or
+                    # poison the autoscaler's load sums.
+                    try:
+                        workload = float(user_workload_calculator(self.input))
+                    except Exception:
+                        log.warning(
+                            f'workload_calculator for "{route_path}" raised; '
+                            f"falling back to {DEFAULT_WORKLOAD}",
+                            exc_info=True,
+                        )
+                        return DEFAULT_WORKLOAD
+                    if not math.isfinite(workload) or workload < 0:
+                        log.warning(
+                            f'workload_calculator for "{route_path}" returned '
+                            f"{workload!r}; falling back to {DEFAULT_WORKLOAD}"
+                        )
+                        return DEFAULT_WORKLOAD
+                    return workload
 
                 @classmethod
                 def from_json_msg(cls, json_msg: Dict[str, Any]) -> "GenericApiPayload":