[multi-lora v2 5/7] Tinker gateway protocol face

miles/tinker/core/future.py

@@ -0,0 +1,84 @@
+"""In-memory future results. Request deduplication lasts only for the current process."""
+
+import asyncio
+import time
+import uuid
+from dataclasses import dataclass, field
+
+from miles.tinker.core.types import OwnershipError
+
+PENDING = "pending"
+DONE = "done"
+FAILED = "failed"
+
+_FINISHED_TTL_S = 3600.0
+
+
+@dataclass
+class RequestFuture:
+    request_id: str
+    model_id: str
+    tenant: str
+    state: str = PENDING
+    result: dict | None = None
+    error: str | None = None
+    error_category: str | None = None
+    finished_at: float | None = None
+    # long-poll wakeup: set when the future settles
+    settled: asyncio.Event = field(default_factory=asyncio.Event)
+
+
+class RequestFutureStore:
+    def __init__(self) -> None:
+        self._futures: dict[str, RequestFuture] = {}
+
+    def create(self, model_id: str, tenant: str) -> RequestFuture:
+        future = RequestFuture(request_id=f"req-{uuid.uuid4().hex}", model_id=model_id, tenant=tenant)
+        self._futures[future.request_id] = future
+        return future
+
+    def resolve(self, request_id: str, result: dict) -> None:
+        future = self._futures[request_id]
+        if future.state != PENDING:  # e.g. already failed by lease expiry
+            return
+        future.state = DONE
+        future.result = result
+        future.finished_at = time.monotonic()
+        future.settled.set()
+
+    def fail(self, request_id: str, error: str, category: str) -> None:
+        future = self._futures[request_id]
+        if future.state != PENDING:
+            return
+        future.state = FAILED
+        future.error = error
+        future.error_category = category
+        future.finished_at = time.monotonic()
+        future.settled.set()
+
+    def get(self, request_id: str, tenant: str) -> RequestFuture | None:
+        """Return None for unknown or expired futures; enforce tenant ownership otherwise."""
+        self._sweep()
+        future = self._futures.get(request_id)
+        if future is None:
+            return None
+        if future.tenant != tenant:
+            raise OwnershipError(f"request {request_id} does not belong to this tenant")
+        return future
+
+    def request_id_for_retry(self, request_id: str, model_id: str, tenant: str) -> str:
+        if self.get(request_id, tenant) is not None:
+            return request_id
+        replacement = self.create(model_id, tenant)
+        self.fail(replacement.request_id, "result expired after retention", "user")
+        return replacement.request_id
+
+    def _sweep(self) -> None:
+        now = time.monotonic()
+        expired = [
+            request_id
+            for request_id, future in self._futures.items()
+            if future.finished_at is not None and now - future.finished_at > _FINISHED_TTL_S
+        ]
+        for request_id in expired:
+            del self._futures[request_id]

miles/tinker/core/input_validation.py

@@ -0,0 +1,156 @@
+import math
+import re
+
+from miles.tinker.core.types import (
+    LOSS_FN_INPUTS,
+    LOSS_INPUT_KEYS,
+    CommandOp,
+    GatewayConfig,
+    ModelRecord,
+    UserInputError,
+)
+
+
+def validate_model_config(lora_config: dict, config: GatewayConfig) -> None:
+    """Reject per-model settings that conflict with the fixed server adapter layout."""
+    if lora_config.get("seed") is not None:
+        raise UserInputError("lora_config.seed is not supported: adapter initialization is not per-model seedable")
+    layout = {
+        "train_attn": config.trains_attn,
+        "train_mlp": config.trains_mlp,
+        "train_unembed": config.trains_unembed,
+    }
+    for field, layout_trains in layout.items():
+        requested = lora_config.get(field)
+        if requested is not None and requested != layout_trains:
+            raise UserInputError(
+                f"lora_config.{field}={requested} conflicts with this gateway's adapter layout "
+                f"({field}={layout_trains}); the layout is fixed by --target-modules at server start"
+            )
+    rank = lora_config.get("rank", 32)
+    if type(rank) is not int or not 1 <= rank <= config.max_lora_rank:
+        raise UserInputError(f"lora_config.rank must be an integer in [1, {config.max_lora_rank}], got {rank!r}")
+
+
+def validate_batch_payload(op: CommandOp, payload: dict, config: GatewayConfig) -> None:
+    if not op.is_batch():
+        return
+    datums = payload["datums"]
+    if not datums:
+        raise UserInputError("forward_backward with no data")
+    if len(datums) > config.max_datums_per_request:
+        raise UserInputError(f"{len(datums)} datums exceeds max_datums_per_request={config.max_datums_per_request}")
+    required_inputs = LOSS_FN_INPUTS.get(payload["loss_fn"])
+    if required_inputs is None:
+        raise UserInputError(f"unknown loss_fn {payload['loss_fn']!r}; known: {sorted(LOSS_FN_INPUTS)}")
+    loss_fn_config = payload.get("loss_fn_config")
+    if loss_fn_config is not None:
+        if not isinstance(loss_fn_config, dict):
+            raise UserInputError("loss_fn_config must be an object")
+        config_keys = {
+            "ppo": ("clip_low_threshold", "clip_high_threshold"),
+            "cispo": ("clip_low_threshold", "clip_high_threshold"),
+            "dro": ("beta",),
+        }.get(payload["loss_fn"], ())
+        for key in config_keys:
+            if key in loss_fn_config:
+                value = loss_fn_config[key]
+                if type(value) not in (int, float) or not math.isfinite(value):
+                    raise UserInputError(f"loss_fn_config[{key!r}] must be a finite number")
+    total_tokens = 0
+    for index, datum in enumerate(datums):
+        validate_token_ids(datum["tokens"], config.vocab_size, f"datum {index}: model_input")
+        validate_token_ids(datum["target_tokens"], config.vocab_size, f"datum {index}: target_tokens")
+        if len(datum["tokens"]) > config.max_tokens_per_datum:
+            raise UserInputError(f"datum {index}: {len(datum['tokens'])} tokens exceeds {config.max_tokens_per_datum}")
+        total_tokens += len(datum["tokens"])
+        for wire_key in required_inputs:
+            values = datum.get(LOSS_INPUT_KEYS[wire_key])
+            if values is None:
+                raise UserInputError(
+                    f"datum {index}: loss_fn {payload['loss_fn']!r} needs loss_fn_inputs[{wire_key!r}]"
+                )
+            if len(values) != datum["target_len"]:
+                raise UserInputError(
+                    f"datum {index}: loss_fn_inputs[{wire_key!r}] has {len(values)} values "
+                    f"for {datum['target_len']} target tokens"
+                )
+        unread = [
+            wire_key
+            for wire_key, datum_key in LOSS_INPUT_KEYS.items()
+            if wire_key not in required_inputs and datum_key in datum
+        ]
+        if unread:
+            raise UserInputError(
+                f"datum {index}: loss_fn {payload['loss_fn']!r} does not read loss_fn_inputs {unread}"
+            )
+    if total_tokens > config.max_tokens_per_request:
+        raise UserInputError(f"{total_tokens} tokens exceeds max_tokens_per_request={config.max_tokens_per_request}")
+
+
+def validate_sample_payload(payload: dict, config: GatewayConfig) -> None:
+    validate_token_ids(payload["prompt_tokens"], config.vocab_size, "prompt")
+    num_samples = payload.get("num_samples", 1)
+    if type(num_samples) is not int or not 1 <= num_samples <= config.max_samples_per_request:
+        raise UserInputError(f"num_samples must be an integer in [1, {config.max_samples_per_request}]")
+    topk = payload.get("topk_prompt_logprobs", 0)
+    if type(topk) is not int or topk < 0:
+        raise UserInputError("topk_prompt_logprobs must be a nonnegative integer")
+
+
+def validate_seq_id(value, name: str, minimum: int = 1) -> int:
+    # request seq_ids are 1-based (last_enqueued_seq_id starts at 0); idempotency keys are 0-based
+    if not isinstance(value, int) or value < minimum:
+        raise UserInputError(f"{name} must be an integer >= {minimum}, got {value!r}")
+    return value
+
+
+def validate_checkpoint_segment(segment: str) -> None:
+    """Reject client path segments that could escape the checkpoint root."""
+    if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}", segment) is None:
+        raise UserInputError(f"invalid checkpoint path segment {segment!r}")
+
+
+def validate_checkpoint_compatibility(meta: dict, record: ModelRecord, config: GatewayConfig, shown_path: str) -> None:
+    """Reject settings that would change the saved tensors' meaning."""
+    expected = {
+        "base_model": record.base_model,
+        "lora_rank": record.lora_rank,
+        "lora_alpha": record.lora_alpha,
+        "train_attn": config.trains_attn,
+        "train_mlp": config.trains_mlp,
+        "train_unembed": config.trains_unembed,
+    }
+    for key, value in expected.items():
+        if meta[key] != value:
+            raise UserInputError(
+                f"checkpoint {shown_path!r} was saved with {key}={meta[key]!r}; this model expects {key}={value!r}"
+            )
+
+
+def validate_save_options(payload: dict) -> None:
+    if payload.get("ttl_seconds") is not None:
+        raise UserInputError("ttl_seconds is not supported: checkpoints on this gateway do not expire")
+    if payload.get("user_metadata") is not None:
+        raise UserInputError("user_metadata is not supported by this gateway")
+
+
+def validate_token_ids(tokens: list, vocab_size: int, field: str) -> None:
+    if any(type(token) is not int or not 0 <= token < vocab_size for token in tokens):
+        raise UserInputError(f"{field} must contain integer token IDs in [0, {vocab_size})")
+
+
+def validate_checkpoint_metadata(meta, shown_path: str) -> None:
+    fields = {
+        "tenant_digest": (str,),
+        "base_model": (str,),
+        "lora_rank": (int,),
+        "lora_alpha": (int, float),
+        "train_attn": (bool,),
+        "train_mlp": (bool,),
+        "train_unembed": (bool,),
+    }
+    if not isinstance(meta, dict) or any(type(meta.get(key)) not in types for key, types in fields.items()):
+        raise UserInputError(f"checkpoint {shown_path!r} has invalid or unsupported metadata")
+    if meta["lora_rank"] <= 0 or not math.isfinite(meta["lora_alpha"]):
+        raise UserInputError(f"checkpoint {shown_path!r} has invalid LoRA metadata")

miles/tinker/core/model_queue.py

@@ -0,0 +1,77 @@
+"""Per-model command ordering and batch runs separated by barriers."""
+
+from collections import deque
+from dataclasses import dataclass, field
+
+from miles.tinker.core.types import Command
+
+
+@dataclass
+class PendingRequest:
+    """One submitted command and its completion accounting."""
+
+    command: Command
+    datums: list[dict] = field(default_factory=list)  # batch ops only
+    num_issued_datums: int = 0
+    outputs: list[dict | None] = field(default_factory=list)
+    remaining: int = 0
+
+    @property
+    def is_batch_op(self) -> bool:
+        return self.command.op.is_batch()
+
+    def record_output(self, local_index: int, output: dict) -> bool:
+        """Store one datum's result; True once every datum has reported."""
+        self.outputs[local_index] = output
+        self.remaining -= 1
+        return self.remaining == 0
+
+    def pack_key(self) -> tuple:
+        """Datums pack into one BatchUnit only within the same (op, loss_fn, config)."""
+        config = self.command.payload.get("loss_fn_config") or {}
+        return (self.command.op, self.command.payload["loss_fn"], tuple(sorted(config.items())))
+
+
+class ModelRequestQueue:
+    def __init__(self, model_id: str, tenant: str, slot: int) -> None:
+        self.model_id = model_id
+        self.tenant = tenant
+        self.slot = slot
+        # seq_ids are 1-based
+        self.last_enqueued_seq_id = 0
+        self.pending_by_seq: dict[int, Command] = {}
+        self.queue: deque[PendingRequest] = deque()
+
+    def submit(self, command: Command) -> None:
+        """Accept one deduplicated command; feed the queue in seq order."""
+        self.pending_by_seq[command.seq_id] = command
+        self._enqueue_contiguous_commands()
+
+    def _enqueue_contiguous_commands(self) -> None:
+        while self.last_enqueued_seq_id + 1 in self.pending_by_seq:
+            self.last_enqueued_seq_id += 1
+            next_command = self.pending_by_seq.pop(self.last_enqueued_seq_id)
+            pending = PendingRequest(command=next_command)
+            if pending.is_batch_op and pending.command.validation_error is None:
+                pending.datums = next_command.payload["datums"]
+                pending.remaining = len(pending.datums)
+                pending.outputs = [None] * len(pending.datums)
+            self.queue.append(pending)
+
+    def open_batch_run(self) -> list[PendingRequest]:
+        """The leading run of batch-op commands; their datums are all issuable."""
+        run = []
+        for pending in self.queue:
+            if not pending.is_batch_op or pending.command.validation_error is not None:
+                break
+            run.append(pending)
+        return run
+
+    def ready_barrier(self) -> PendingRequest | None:
+        """The head barrier, executable once every batch op ahead of it completed."""
+        if self.queue and not self.queue[0].is_batch_op and self.queue[0].command.validation_error is None:
+            return self.queue[0]
+        return None
+
+    def finish(self, pending: PendingRequest) -> None:
+        self.queue.remove(pending)