Fix staging cold-start timeout

Dockerfile

@@ -44,4 +44,13 @@ ENV FUNCTION_COMMAND="python"
 ENV FUNCTION_ARGS="-m,evaluation_function.main"
 
 ENV FUNCTION_INTERFACE="rpc"
-ENV LOG_LEVEL="debug"
+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"

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()")
 

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"]
+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"]

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