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