From 6f2c12f5f791175ed258c846f006ba18c14302c4 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 1 Sep 2026 19:18:33 +0100 Subject: [PATCH] Fix staging cold-start timeout: lazy-load models, raise shimmy worker timeouts Staging returned 503/500 on every request: the worker's cold-start import (torch + now pydantic/pydantic-core from the lf_toolkit v1.1.1 bump) takes ~34s on the 1024MB Lambda, longer than shimmy's 30s worker-send-timeout, so shimmy killed the half-booted worker before it became ready and never recovered. - evaluation.py / models/__init__.py: dispatch models via importlib against an AVAILABLE_MODELS allowlist instead of eagerly importing every model module. torch is no longer pulled in at worker startup, only when a model that needs it actually runs. Cold import drops ~34s -> ~8s (measured CPU/mem-constrained to match the Lambda). - models/utils.py: build the torch NeuralLM class lazily so importing utils for csv_to_lists / shard_for (shannon_letters_ngram, shannon_words_ngram) does not pull in torch either. - Dockerfile: raise FUNCTION_WORKER_START_TIMEOUT / FUNCTION_WORKER_SEND_TIMEOUT to 150s (< the 175s Lambda timeout) so shimmy waits out a slow first request instead of discarding the worker; make FUNCTION_RPC_TRANSPORT=stdio explicit. Verified in AWS RIE Lambda mode, container constrained to 1024MB / 0.6 vCPU: cold torch-free model 13.9s -> 200; cold basic_nn (first torch load) 35s -> 200 on direct invoke; warm calls ~0.2s. Worker now survives the slow first call. pytest + flake8 pass. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01UsgxJUsq1vG176rbznHRqz --- Dockerfile | 11 +++++- evaluation_function/evaluation.py | 12 ++++-- evaluation_function/models/__init__.py | 17 ++++++++- evaluation_function/models/utils.py | 53 ++++++++++++++++++-------- 4 files changed, 70 insertions(+), 23 deletions(-) diff --git a/Dockerfile b/Dockerfile index c7d45ab..1846dea 100755 --- a/Dockerfile +++ b/Dockerfile @@ -44,4 +44,13 @@ ENV FUNCTION_COMMAND="python" ENV FUNCTION_ARGS="-m,evaluation_function.main" ENV FUNCTION_INTERFACE="rpc" -ENV LOG_LEVEL="debug" \ No newline at end of file +ENV FUNCTION_RPC_TRANSPORT="stdio" + +# The worker pulls in torch on its first request; on a small (1024 MB) Lambda +# that cold-start import runs ~30-40s. Give shimmy room to wait for it instead +# of killing the half-booted worker at the 30s default. Keep these below the +# Lambda function timeout (currently 175s). +ENV FUNCTION_WORKER_START_TIMEOUT="150s" +ENV FUNCTION_WORKER_SEND_TIMEOUT="150s" + +ENV LOG_LEVEL="debug" diff --git a/evaluation_function/evaluation.py b/evaluation_function/evaluation.py index e303fa3..03d70a7 100755 --- a/evaluation_function/evaluation.py +++ b/evaluation_function/evaluation.py @@ -1,7 +1,9 @@ from typing import Any +import importlib + from lf_toolkit.evaluation import Result, Params -from . import models +from .models import AVAILABLE_MODELS def evaluation_function( response: Any, @@ -37,11 +39,13 @@ def evaluation_function( print(f"#### Response: {str(response)} ####") print(f"#### Answer: {str(answer)} ####") - try: - model = getattr(models, model_name) # e.g. models.basic_nn - except AttributeError: + if model_name not in AVAILABLE_MODELS: raise ValueError(f"Unknown model: {model_name}") + # Imported here rather than at module load so torch (pulled in by some + # models) is only imported when a request actually needs it. + model = importlib.import_module(f"{__package__}.models.{model_name}") + if not hasattr(model, "run"): raise ValueError(f"Model {model_name} has no run()") diff --git a/evaluation_function/models/__init__.py b/evaluation_function/models/__init__.py index e14bec7..dd1c7c7 100644 --- a/evaluation_function/models/__init__.py +++ b/evaluation_function/models/__init__.py @@ -1,3 +1,16 @@ -from . import basic_nn, shannon_letters_single, shannon_letters_ngram, shannon_words_ngram, bengio_infer, utils +"""Model registry. -__all__ = ["basic_nn", "shannon_letters_single", "shannon_letters_ngram", "shannon_words_ngram", "bengio_infer", "utils"] \ No newline at end of file +Model modules are imported lazily by ``evaluation_function.evaluation`` (not +here) so that a cold start does not pay torch's multi-second import cost until a +request actually selects a model that needs it. +""" + +AVAILABLE_MODELS = [ + "basic_nn", + "shannon_letters_single", + "shannon_letters_ngram", + "shannon_words_ngram", + "bengio_infer", +] + +__all__ = ["AVAILABLE_MODELS"] \ No newline at end of file diff --git a/evaluation_function/models/utils.py b/evaluation_function/models/utils.py index 2b35947..22feb45 100644 --- a/evaluation_function/models/utils.py +++ b/evaluation_function/models/utils.py @@ -1,7 +1,6 @@ import csv, pickle, bz2 from pathlib import Path -import torch.nn as nn import hashlib def csv_to_lists(filename: Path) -> list: @@ -29,21 +28,43 @@ def shard_for(ctx, n_shards): # Deterministic shard assignment h = hashlib.sha1(str(ctx).encode("utf8")).hexdigest() return int(h, 16) % n_shards -class NeuralLM(nn.Module): - def __init__(self, vocab_size, n_ctx, embed_dim, hidden, dropout_p): - super().__init__() - self.emb = nn.Embedding(vocab_size, embed_dim) - self.fc1 = nn.Linear(n_ctx * embed_dim, hidden) - self.act = nn.Tanh() - self.drop = nn.Dropout(dropout_p) - self.fc2 = nn.Linear(hidden, vocab_size) - def forward(self, ctx_idx): # ctx_idx: (B, n) - e = self.emb(ctx_idx) # (B, n, d) - x = e.reshape(e.size(0), -1) # (B, n*d) - h = self.act(self.fc1(x)) # (B, H) - logits = self.fc2(h) # (B, V) - return logits - +_NeuralLM_cls = None + + +def _neural_lm_cls(): + """Build the torch ``NeuralLM`` class on first use. + + Deferred so that importing this module (e.g. for ``csv_to_lists`` or + ``shard_for``) does not pull in torch. + """ + global _NeuralLM_cls + if _NeuralLM_cls is None: + import torch.nn as nn + + class NeuralLM(nn.Module): + def __init__(self, vocab_size, n_ctx, embed_dim, hidden, dropout_p): + super().__init__() + self.emb = nn.Embedding(vocab_size, embed_dim) + self.fc1 = nn.Linear(n_ctx * embed_dim, hidden) + self.act = nn.Tanh() + self.drop = nn.Dropout(dropout_p) + self.fc2 = nn.Linear(hidden, vocab_size) + def forward(self, ctx_idx): # ctx_idx: (B, n) + e = self.emb(ctx_idx) # (B, n, d) + x = e.reshape(e.size(0), -1) # (B, n*d) + h = self.act(self.fc1(x)) # (B, H) + logits = self.fc2(h) # (B, V) + return logits + + _NeuralLM_cls = NeuralLM + return _NeuralLM_cls + + +def NeuralLM(*args, **kwargs): + """Instantiate the (lazily built) torch ``NeuralLM`` model.""" + return _neural_lm_cls()(*args, **kwargs) + + def encode(seq): import sentencepiece as spm from pathlib import Path