pytensorforge 0.1.0__py3-none-any.whl

src/serving/metrics.py

@@ -0,0 +1,293 @@
+import bisect
+import math
+import os
+import resource
+import threading
+import time
+
+LATENCY_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 300)
+RATE_BUCKETS = (1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000)
+TOKEN_BUCKETS = (16, 64, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768)
+
+
+def _escape(value):
+    return str(value).replace("\\", "\\\\").replace("\n", "\\n").replace('"', '\\"')
+
+
+def _labels(names, values):
+    if not names:
+        return ""
+
+    return "{" + ",".join(f'{n}="{_escape(v)}"' for n, v in zip(names, values)) + "}"
+
+
+def _fmt(v):
+    if v == math.inf:
+        return "+Inf"
+
+    if isinstance(v, float) and v.is_integer() and abs(v) < 1e15:
+        return str(int(v))
+
+    return repr(float(v)) if isinstance(v, float) else str(v)
+
+
+class _Metric:
+    kind = ""
+
+    def __init__(self, name, help_text, labels=()):
+        self.name = name
+        self.help = help_text
+        self.label_names = tuple(labels)
+        self._lock = threading.Lock()
+
+    def _key(self, labels):
+        if set(labels) != set(self.label_names):
+            raise ValueError(f"{self.name} expects labels {self.label_names}, got {tuple(labels)}")
+
+        return tuple(str(labels[n]) for n in self.label_names)
+
+
+class Counter(_Metric):
+    kind = "counter"
+
+    def __init__(self, *a, **k):
+        super().__init__(*a, **k)
+        self._values = {}
+
+    def inc(self, amount=1, **labels):
+        key = self._key(labels)
+
+        with self._lock:
+            self._values[key] = self._values.get(key, 0) + amount
+
+    def get(self, **labels):
+        return self._values.get(self._key(labels), 0)
+
+    def samples(self):
+        with self._lock:
+            return [(self.name + "_total", k, v) for k, v in sorted(self._values.items())]
+
+    def snapshot(self):
+        with self._lock:
+            return {",".join(k) or "": v for k, v in self._values.items()}
+
+
+class Gauge(_Metric):
+    kind = "gauge"
+
+    def __init__(self, *a, fn=None, **k):
+        super().__init__(*a, **k)
+        self._values = {}
+        self._fn = fn
+
+    def set(self, value, **labels):
+        key = self._key(labels)
+
+        with self._lock:
+            self._values[key] = value
+
+    def inc(self, amount=1, **labels):
+        key = self._key(labels)
+
+        with self._lock:
+            self._values[key] = self._values.get(key, 0) + amount
+
+    def dec(self, amount=1, **labels):
+        self.inc(-amount, **labels)
+
+    def remove(self, **labels):
+        with self._lock:
+            self._values.pop(self._key(labels), None)
+
+    def _current(self):
+        if self._fn is not None:
+            return {self._key(lbl): v for lbl, v in self._fn()}
+
+        with self._lock:
+            return dict(self._values)
+
+    def samples(self):
+        return [(self.name, k, v) for k, v in sorted(self._current().items())]
+
+    def snapshot(self):
+        return {",".join(k) or "": v for k, v in self._current().items()}
+
+
+class Histogram(_Metric):
+    kind = "histogram"
+
+    def __init__(self, *a, buckets=LATENCY_BUCKETS, **k):
+        super().__init__(*a, **k)
+        self.buckets = tuple(sorted(buckets))
+        self._data = {}
+
+    def observe(self, value, **labels):
+        if value is None or not math.isfinite(value):
+            return
+
+        key = self._key(labels)
+        idx = bisect.bisect_left(self.buckets, value)
+
+        with self._lock:
+            d = self._data.get(key)
+
+            if d is None:
+                d = self._data[key] = [[0] * (len(self.buckets) + 1), 0.0, 0]
+
+            d[0][idx] += 1
+            d[1] += value
+            d[2] += 1
+
+    def samples(self):
+        out = []
+
+        with self._lock:
+            items = sorted((k, (list(v[0]), v[1], v[2])) for k, v in self._data.items())
+
+        for key, (counts, total, n) in items:
+            running = 0
+
+            for bound, c in zip(self.buckets + (math.inf,), counts):
+                running += c
+                out.append((self.name + "_bucket", key + (_fmt(bound),), running))
+
+            out.append((self.name + "_sum", key, total))
+            out.append((self.name + "_count", key, n))
+
+        return out
+
+    def label_names_for(self, sample_name):
+        if sample_name.endswith("_bucket"):
+            return self.label_names + ("le",)
+
+        return self.label_names
+
+    def snapshot(self):
+        with self._lock:
+            return {
+                ",".join(k) or "": {"count": v[2], "sum": v[1], "mean": (v[1] / v[2]) if v[2] else None}
+                for k, v in self._data.items()
+            }
+
+
+class Registry:
+
+    def __init__(self, namespace="ptf"):
+        self.namespace = namespace
+        self._metrics = []
+
+    def _add(self, metric):
+        metric.name = f"{self.namespace}_{metric.name}"
+        self._metrics.append(metric)
+        return metric
+
+    def counter(self, name, help_text, labels=()):
+        return self._add(Counter(name, help_text, labels))
+
+    def gauge(self, name, help_text, labels=(), fn=None):
+        return self._add(Gauge(name, help_text, labels, fn=fn))
+
+    def histogram(self, name, help_text, labels=(), buckets=LATENCY_BUCKETS):
+        return self._add(Histogram(name, help_text, labels, buckets=buckets))
+
+    def render_prometheus(self):
+        lines = []
+
+        for m in self._metrics:
+            lines.append(f"# HELP {m.name} {m.help}")
+            lines.append(f"# TYPE {m.name} {m.kind}")
+
+            for sample_name, key, value in m.samples():
+                names = m.label_names_for(sample_name) if isinstance(m, Histogram) else m.label_names
+                lines.append(f"{sample_name}{_labels(names, key)} {_fmt(value)}")
+
+        return "\n".join(lines) + "\n"
+
+    def snapshot(self):
+        return {m.name: m.snapshot() for m in self._metrics}
+
+
+def _rss_bytes():
+    try:
+        with open("/proc/self/statm") as f:
+            return int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE")
+    except (OSError, ValueError, IndexError):
+        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024
+
+
+class ServerMetrics:
+
+    def __init__(self, model_server_ref):
+        r = self.registry = Registry("ptf")
+        self._server = model_server_ref
+        self.started = time.time()
+
+        self.http_requests = r.counter("http_requests", "HTTP requests by route and status", ("route", "status"))
+        self.http_duration = r.histogram("http_request_duration_seconds", "Full HTTP request duration", ("route",))
+        self.open_connections = r.gauge("open_connections", "Open client TCP connections")
+
+        self.requests = r.counter("completion_requests", "Completion requests accepted", ("model", "endpoint", "stream"))
+        self.finished = r.counter("completion_finished", "Completions by finish reason", ("model", "reason"))
+        self.cancellations = r.counter("cancellations", "Generations cancelled before finishing", ("model", "cause"))
+        self.errors = r.counter("errors", "Errors by kind", ("kind",))
+        self.auth_failures = r.counter("auth_failures", "Rejected authentication attempts")
+        self.rate_limited = r.counter("rate_limited", "Requests rejected by rate limits", ("reason",))
+
+        self.prompt_tokens = r.counter("prompt_tokens", "Prompt tokens processed", ("model",))
+        self.completion_tokens = r.counter("completion_tokens", "Completion tokens generated", ("model",))
+        self.active = r.gauge("active_requests", "Requests in flight (queued or generating)", ("model",))
+
+        self.tokenize_s = r.histogram("tokenize_seconds", "Chat templating + prompt tokenization time", ("model",))
+        self.queue_s = r.histogram("queue_wait_seconds", "Time from submit to admission into a batch", ("model",))
+        self.prefill_s = r.histogram("prefill_seconds", "Prompt prefill compute time", ("model",))
+        self.ttft_s = r.histogram("time_to_first_token_seconds", "Submit to first generated token", ("model",))
+        self.generation_s = r.histogram("generation_duration_seconds", "Submit to final token", ("model",))
+        self.decode_rate = r.histogram("decode_tokens_per_second", "Per-request generation rate after first token",
+                                       ("model",), buckets=RATE_BUCKETS)
+        self.network_s = r.histogram("network_write_seconds", "Time blocked writing response bytes to the client",
+                                     ("route",))
+        self.prompt_len = r.histogram("prompt_length_tokens", "Prompt length", ("model",), buckets=TOKEN_BUCKETS)
+
+        r.gauge("process_resident_memory_bytes", "Resident set size", fn=lambda: [({}, _rss_bytes())])
+        r.gauge("process_cpu_seconds", "CPU seconds used by the server process",
+                fn=lambda: [({}, time.process_time())])
+        r.gauge("uptime_seconds", "Seconds since server start", fn=lambda: [({}, time.time() - self.started)])
+
+        self._scheduler_gauges(r)
+
+    def _each_model(self):
+        server = self._server()
+
+        if server is None:
+            return []
+
+        return server.loaded_models()
+
+    def _scheduler_gauges(self, r):
+        def per_model(fn):
+            return lambda: [({"model": m.name}, fn(m)) for m in self._each_model()]
+
+        def sched(key):
+            return per_model(lambda m: m.generator.scheduler.stats[key])
+
+        r.gauge("scheduler_waiting", "Requests waiting for a batch slot or cache", ("model",),
+                fn=per_model(lambda m: m.generator.scheduler.num_waiting))
+        r.gauge("scheduler_running", "Sequences in the running batch", ("model",),
+                fn=per_model(lambda m: m.generator.scheduler.num_active))
+        r.gauge("kv_cache_used_bytes", "KV cache bytes allocated", ("model",),
+                fn=per_model(lambda m: m.generator.scheduler.cache_manager.used_bytes))
+        r.gauge("kv_cache_budget_bytes", "KV cache byte budget", ("model",),
+                fn=per_model(lambda m: m.generator.scheduler.cache_manager.max_bytes or 0))
+        r.gauge("model_weights_bytes", "Weights resident in memory", ("model",),
+                fn=per_model(lambda m: m.generator.weights_nbytes))
+        r.gauge("scheduler_decode_seconds", "Cumulative batched decode compute", ("model",),
+                fn=sched("decode_seconds"))
+        r.gauge("scheduler_prefill_seconds", "Cumulative prefill compute", ("model",),
+                fn=sched("prefill_seconds"))
+        r.gauge("scheduler_decode_steps", "Batched decode steps executed", ("model",),
+                fn=sched("decode_steps"))
+        r.gauge("scheduler_mean_batch_size", "Generated tokens per decode step", ("model",),
+                fn=per_model(lambda m: (m.generator.scheduler.stats["generated_tokens"]
+                                        / max(1, m.generator.scheduler.stats["decode_steps"]))))
+        r.gauge("scheduler_max_batch_seen", "Largest decode batch so far", ("model",),
+                fn=sched("max_batch_seen"))

src/serving/model_server.py

@@ -0,0 +1,287 @@
+import logging
+import threading
+import time
+from dataclasses import dataclass
+
+from src.inference.runtime import load_model
+from src.serving import errors
+from src.models.gpt.context import describe_context
+
+log = logging.getLogger("ptf.server")
+
+SUPPORTED_DEVICES = ("cpu",)
+
+
+def resolve_device(requested):
+    requested = (requested or "auto").lower()
+
+    if requested in ("auto", "cpu"):
+        return "cpu"
+
+    raise ValueError(
+        f"device '{requested}' is not available: this build's inference engine runs on {SUPPORTED_DEVICES} only"
+    )
+
+
+@dataclass
+class ModelLimits:
+    context_length: int
+    max_generation_tokens: int
+    max_prompt_tokens: int
+
+
+class ServedModel:
+
+    def __init__(self, entry, generator, template, template_source, limits, device):
+        self.entry = entry
+        self.name = entry.name
+        self.generator = generator
+        self.template = template
+        self.template_source = template_source
+        self.limits = limits
+        self.device = device
+        self.loaded_at = int(time.time())
+        self.in_flight = 0
+        self.accepting = True
+        self._idle = threading.Condition()
+
+    @property
+    def fingerprint(self):
+        return "ptf-" + self.generator.metadata.get("weights_sha256", "")[:12]
+
+    @property
+    def cache_budget_bytes(self):
+        return self.generator.scheduler.cache_manager.max_bytes
+
+    @property
+    def memory_bytes(self):
+        return self.generator.weights_nbytes + (self.cache_budget_bytes or 0)
+
+    def enter(self):
+        with self._idle:
+            if not self.accepting:
+                raise errors.overloaded(f"model '{self.name}' is being unloaded")
+
+            self.in_flight += 1
+
+    def exit(self):
+        with self._idle:
+            self.in_flight -= 1
+
+            if self.in_flight <= 0:
+                self._idle.notify_all()
+
+    def wait_idle(self, timeout):
+        deadline = time.monotonic() + timeout
+
+        with self._idle:
+            while self.in_flight > 0:
+                left = deadline - time.monotonic()
+
+                if left <= 0:
+                    return False
+
+                self._idle.wait(left)
+
+        return True
+
+    def describe(self):
+        cfg = self.generator.config
+
+        return {
+            "id": self.name,
+            "object": "model",
+            "created": self.loaded_at,
+            "owned_by": "pytensorforge",
+            "context_length": self.limits.context_length,
+            "context": describe_context(cfg),
+            "max_generation_tokens": self.limits.max_generation_tokens,
+            "chat_template": self.template.name,
+            "chat_template_source": self.template_source,
+            "architecture": cfg.architecture,
+            "vocab_size": cfg.vocab_size,
+            "n_layers": cfg.n_layers,
+            "d_model": cfg.d_model,
+            "n_heads": cfg.n_heads,
+            "device": self.device,
+            "fingerprint": self.fingerprint,
+        }
+
+
+class ModelServer:
+
+    def __init__(self, config):
+        self.config = config
+        self.device = resolve_device(config.runtime.device)
+        self._entries = {m.name: m for m in config.models}
+        self._models = {}
+        self._lock = threading.RLock()
+        self._loading = set()
+        self._global_in_flight = 0
+        self._global_lock = threading.Lock()
+
+    def declared(self):
+        return list(self._entries.values())
+
+    def loaded_models(self):
+        with self._lock:
+            return list(self._models.values())
+
+    def get(self, name):
+        with self._lock:
+            model = self._models.get(name)
+
+        if model is None or not model.accepting:
+            if name in self._entries:
+                raise errors.APIError(503, f"model '{name}' is not loaded", "server_error", code="model_not_loaded")
+
+            raise errors.not_found(f"the model '{name}' does not exist", code="model_not_found")
+
+        return model
+
+    def _default_budget(self, generator, entry):
+        ctx = generator.config.context_length
+
+        if entry.max_context_tokens:
+            ctx = min(ctx, entry.max_context_tokens)
+
+        return generator.scheduler.cache_manager.bytes_for(ctx) * entry.max_batch_size
+
+    def _committed_bytes(self, exclude=None):
+        return sum(m.memory_bytes for n, m in self._models.items() if n != exclude)
+
+    def load(self, name):
+        entry = self._entries.get(name)
+
+        if entry is None:
+            raise errors.not_found(f"model '{name}' is not declared in the server configuration", code="model_not_found")
+
+        with self._lock:
+            if name in self._models:
+                return self._models[name]
+
+            if name in self._loading:
+                raise errors.APIError(409, f"model '{name}' is already loading", "invalid_request_error", code="conflict")
+
+            self._loading.add(name)
+
+        try:
+            return self._load(entry)
+        finally:
+            with self._lock:
+                self._loading.discard(name)
+
+    def _load(self, entry):
+        started = time.perf_counter()
+        budget = int(entry.kv_cache_budget_mb * (1 << 20)) if entry.kv_cache_budget_mb else None
+
+        generator = load_model(
+            entry.path,
+            max_batch_size=entry.max_batch_size,
+            cache_budget_bytes=budget,
+            prefill_chunk=entry.prefill_chunk,
+            chat_template=entry.chat_template,
+            context_length=entry.extend_context_to,
+            context_extension=entry.context_extension,
+        )
+
+        if budget is None:
+            generator.scheduler.cache_manager.max_bytes = self._default_budget(generator, entry)
+
+        if generator.chat_template is not None:
+            source = "config" if entry.chat_template else "model"
+            template = generator.bound_chat_template()
+        else:
+            source = "fallback"
+            template = generator.bound_chat_template(self.config.runtime.fallback_chat_template)
+            log.warning("model '%s' has no chat template; using fallback '%s'", entry.name, template.name)
+
+        ctx = generator.config.context_length
+
+        if entry.max_context_tokens:
+            ctx = min(ctx, entry.max_context_tokens)
+
+        max_gen = min(entry.max_generation_tokens or self.config.limits.max_generation_tokens, ctx - 1)
+        max_prompt = min(self.config.limits.max_prompt_tokens or ctx, ctx - 1)
+        limits = ModelLimits(context_length=ctx, max_generation_tokens=max_gen, max_prompt_tokens=max_prompt)
+
+        model = ServedModel(entry, generator, template, source, limits, self.device)
+
+        with self._lock:
+            limit_mb = self.config.runtime.memory_limit_mb
+
+            if limit_mb is not None:
+                needed = self._committed_bytes() + model.memory_bytes
+
+                if needed > limit_mb * (1 << 20):
+                    raise errors.APIError(
+                        507,
+                        f"loading '{entry.name}' needs {model.memory_bytes / (1 << 20):.1f} MiB "
+                        f"(weights + KV cache budget); memory limit of {limit_mb:.0f} MiB would be exceeded",
+                        "server_error",
+                        code="insufficient_memory",
+                    )
+
+            generator.start()
+            self._models[entry.name] = model
+
+        log.info("loaded model '%s' in %.2fs (%.1f MiB weights, %.1f MiB kv budget, template %s/%s)",
+                 entry.name, time.perf_counter() - started, generator.weights_nbytes / (1 << 20),
+                 (model.cache_budget_bytes or 0) / (1 << 20), template.name, source)
+
+        return model
+
+    def unload(self, name, drain_timeout_s=30.0):
+        with self._lock:
+            model = self._models.get(name)
+
+            if model is None:
+                raise errors.not_found(f"model '{name}' is not loaded", code="model_not_loaded")
+
+            model.accepting = False
+
+        drained = model.wait_idle(drain_timeout_s)
+
+        if not drained:
+            sched = model.generator.scheduler
+
+            for rid in list(sched._active.keys()) + [h.request_id for h in list(sched._waiting)]:
+                sched.request_cancel(rid)
+
+            model.wait_idle(5.0)
+
+        model.generator.stop()
+
+        with self._lock:
+            self._models.pop(name, None)
+
+        log.info("unloaded model '%s' (%s)", name, "drained" if drained else "cancelled in-flight requests")
+
+        return drained
+
+    def preload(self):
+        for entry in self._entries.values():
+            if entry.preload:
+                self.load(entry.name)
+
+    def acquire_global(self):
+        with self._global_lock:
+            if self._global_in_flight >= self.config.limits.max_concurrent_requests:
+                raise errors.overloaded()
+
+            self._global_in_flight += 1
+
+    def release_global(self):
+        with self._global_lock:
+            self._global_in_flight -= 1
+
+    @property
+    def global_in_flight(self):
+        return self._global_in_flight
+
+    def shutdown(self, drain_timeout_s):
+        for model in self.loaded_models():
+            model.accepting = False
+
+        for model in self.loaded_models():
+            self.unload(model.name, drain_timeout_s)