diff --git a/proto/worker_scheduler.proto b/proto/worker_scheduler.proto index 8a086cc9..d443eb1c 100644 --- a/proto/worker_scheduler.proto +++ b/proto/worker_scheduler.proto @@ -565,6 +565,12 @@ message JobMetrics { // th#913/gw#596: the CONCRETE lane that actually served this request // (e.g. "fp8-w8a8-dynamic+compiled"). "" = unmeasured (pre-lane worker). string lane = 13; + // th#1051: the declared runtime formula's term feature values, evaluated + // worker-side on the EXECUTED payload (defaults applied), keyed by + // canonical term key (e.g. "num_inference_steps", "1"). Empty = no + // declared formula (or a pre-formula worker) — the hub falls back to + // evaluating terms from the raw payload + schema defaults. + map runtime_terms = 14; } // Worker -> orchestrator streaming output chunk. Ordered per (request_id, diff --git a/pyproject.toml b/pyproject.toml index 6e69e4ee..a639a307 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "gen-worker" -version = "0.48.2" +version = "0.50.0" description = "A library used to build custom functions in Cozy Creator's serverless function platform." readme = "README.md" license = "MIT" diff --git a/src/gen_worker/__init__.py b/src/gen_worker/__init__.py index 045a2eba..3f34c9f1 100644 --- a/src/gen_worker/__init__.py +++ b/src/gen_worker/__init__.py @@ -16,6 +16,7 @@ def hello(ctx: RequestContext, data: In) -> Out: ... from . import io from .api.binding import Civitai, HF, Hub, ModelRef, ModelScope from .api.decorators import Compile, NoWarmup, Resources, endpoint, variant_of +from .api.formula import RuntimeFormula from .api.model import Model, ModelChoice, ModelDefaults from .api.slot import ResolvedSlot, Slot from .families import FamilyDefaults diff --git a/src/gen_worker/api/decorators.py b/src/gen_worker/api/decorators.py index 6ea5f799..fe34a45c 100644 --- a/src/gen_worker/api/decorators.py +++ b/src/gen_worker/api/decorators.py @@ -39,6 +39,7 @@ def generate(self, ctx, p: Input) -> Output: ... import msgspec from .binding import BINDING_TYPES, Binding +from .formula import RuntimeFormula from .slot import DEFAULT_REGIMES, REGIMES, Slot T = TypeVar("T") @@ -212,6 +213,47 @@ def _validate_function_regimes( return _validate_regime_tuple(owner, regimes) +def _validate_handles(owner: str, handles: Any) -> Tuple[str, ...]: + """th#1050 opt-in lane declaration: `handles=["fp8-w8a8-dynamic"]` marks + that this endpoint's code BRANCHES on the executing lane (ctx.lane) — + behavioral divergence only, never inventory. Tokens are concrete lane + BODIES (no `+eager|+compiled`: execution is platform-managed). Nothing + declared = fully platform-managed lanes, exactly as before.""" + from ..models import lanes as lanespec + + if handles is None: + return () + if isinstance(handles, str) or not isinstance(handles, (list, tuple)): + raise TypeError( + f"@endpoint {owner}: handles= must be a list/tuple of lane body " + f"strings, got {type(handles).__name__}" + ) + known = lanespec.known_lane_bodies() + out: list[str] = [] + for token in handles: + t = str(token or "").strip().lower() + if "+" in t: + raise ValueError( + f"@endpoint {owner}: handles= token {token!r} carries an " + "execution axis; declare the lane body (execution is " + "platform-managed)" + ) + if t in lanespec.FAMILIES: + raise ValueError( + f"@endpoint {owner}: handles= token {token!r} is a coarse " + f"family; declare a concrete lane body (one of {known})" + ) + if t not in known: + raise ValueError( + f"@endpoint {owner}: handles= token {token!r} is not a known " + f"lane body (known: {known})" + ) + if t in out: + raise ValueError(f"@endpoint {owner}: handles= repeats {token!r}") + out.append(t) + return tuple(out) + + class Compile(msgspec.Struct, frozen=True): """Opt-in torch.compile over pre-built per-SKU cache artifacts (#384). @@ -324,6 +366,14 @@ class EndpointDecl(msgspec.Struct, frozen=True, kw_only=True): # th#1017: per-handler declared regimes, attr_name -> (regime, ...). # Function-shaped endpoints key their single handler under "". regimes: Mapping[str, Tuple[str, ...]] = msgspec.field(default_factory=dict) + # th#1050: opt-in declared lane bodies this endpoint's code branches on + # (ctx.lane). Empty = platform-managed behavior only. + handles: Tuple[str, ...] = () + # th#1051: declared compute-time formulas, attr_name -> RuntimeFormula + # (function-shaped endpoints key their single handler under ""). + # Declared via the runtime= kwarg overload; payload-field validation + # happens at registry walk time when payload types are resolved. + runtime_formula: Mapping[str, RuntimeFormula] = msgspec.field(default_factory=dict) ATTR = "__gen_worker_endpoint__" @@ -565,6 +615,57 @@ def _validate_class_models( ) +def _split_runtime_kwarg( + owner: str, runtime: Any, +) -> Tuple[Optional[str], Dict[str, RuntimeFormula]]: + """th#1051 runtime= overload: a str selects an engine runtime (vllm / + llama-server, classes only); a RuntimeFormula declares the compute-time + formula for every handler ("*"); a mapping declares per-handler formulas + (classes only). Returns (engine_runtime, {attr_or_"*": formula}).""" + if runtime is None: + return None, {} + if isinstance(runtime, str): + return runtime, {} + if isinstance(runtime, RuntimeFormula): + return None, {"*": runtime} + if isinstance(runtime, Mapping): + out: Dict[str, RuntimeFormula] = {} + for attr, rf in runtime.items(): + if not isinstance(attr, str) or not isinstance(rf, RuntimeFormula): + raise TypeError( + f"@endpoint {owner}: runtime= mapping must be " + f"{{handler_name: RuntimeFormula}}" + ) + out[attr] = rf + return None, out + raise TypeError( + f"@endpoint {owner}: runtime= must be an engine runtime string, a " + f"RuntimeFormula, or a mapping of handler name -> RuntimeFormula, " + f"got {type(runtime).__name__}" + ) + + +def _expand_formula_map( + owner: str, formulas: Dict[str, RuntimeFormula], handler_attrs: "list[str]", +) -> Dict[str, RuntimeFormula]: + if not formulas: + return {} + if "*" in formulas: + if len(formulas) > 1: + raise ValueError( + f"@endpoint {owner}: runtime= cannot mix a bare RuntimeFormula " + f"with per-handler entries" + ) + return {attr: formulas["*"] for attr in handler_attrs} + unknown = [a for a in formulas if a not in handler_attrs] + if unknown: + raise ValueError( + f"@endpoint {owner}: runtime= names unknown handler(s) {unknown} " + f"(handlers: {sorted(handler_attrs)})" + ) + return dict(formulas) + + def _decorate_class( cls: type, *, @@ -573,10 +674,12 @@ def _decorate_class( models: Dict[str, Binding], slots: Dict[str, Slot], runtime: Optional[str], + runtime_formula: Optional[Dict[str, RuntimeFormula]] = None, compile: Optional[Compile] = None, child_calls: bool = False, warmup: Optional[WarmupDecl] = None, regimes: Optional[RegimesDecl] = None, + handles: Optional[Any] = None, ) -> type: handlers = _find_handler_methods(cls) for attr, member in handlers: @@ -597,6 +700,10 @@ def _decorate_class( regimes=_validate_class_regimes( cls.__name__, regimes, {attr for attr, _ in handlers} ), + handles=_validate_handles(cls.__name__, handles), + runtime_formula=_expand_formula_map( + cls.__name__, runtime_formula or {}, [attr for attr, _ in handlers] + ), ) setattr(cls, ATTR, decl) setattr(cls, "__gen_worker_handlers__", handlers) @@ -617,11 +724,13 @@ def _decorate_function( models: Dict[str, Binding], slots: Dict[str, Slot], runtime: Optional[str], + runtime_formula: Optional[Dict[str, RuntimeFormula]] = None, name: Optional[str], compile: Optional[Compile] = None, child_calls: bool = False, warmup: Optional[WarmupDecl] = None, regimes: Optional[RegimesDecl] = None, + handles: Optional[Any] = None, ) -> Callable[..., Any]: _validate_handler_shape(fn.__name__, fn, is_method=False) _reject_producer_generator(fn.__name__, fn, kind) @@ -640,19 +749,27 @@ def _decorate_function( slots[slot_name] = slots.pop("") if runtime is not None: raise ValueError( - f"@endpoint function {fn.__name__!r}: runtime= requires a class " - "with setup() (the engine server outlives single calls)." + f"@endpoint function {fn.__name__!r}: an engine runtime= requires " + "a class with setup() (the engine server outlives single calls)." ) if warmup is not None: raise ValueError( f"@endpoint function {fn.__name__!r}: warmup= requires a class " "with setup() (stateless functions hold nothing to warm)." ) + formulas = runtime_formula or {} + if set(formulas) - {"*"}: + raise ValueError( + f"@endpoint function {fn.__name__!r}: runtime= takes a bare " + "RuntimeFormula (per-handler mappings are for classes)" + ) decl = EndpointDecl( kind=kind, resources=resources, models=models, slots=slots, runtime=None, name=(name or fn.__name__), is_function=True, compile=compile, child_calls=child_calls, regimes={"": _validate_function_regimes(fn.__name__, regimes)}, + handles=_validate_handles(fn.__name__, handles), + runtime_formula={"": formulas["*"]} if "*" in formulas else {}, ) setattr(fn, ATTR, decl) return fn @@ -669,12 +786,13 @@ def endpoint( model: Optional[SlotLike] = ..., models: Optional[Mapping[str, SlotLike]] = ..., resources: Optional[Resources] = ..., - runtime: Optional[str] = ..., + runtime: Union[str, RuntimeFormula, Mapping[str, RuntimeFormula], None] = ..., name: Optional[str] = ..., compile: Optional[Compile] = ..., child_calls: bool = ..., warmup: Optional[WarmupDecl] = ..., regimes: Optional[RegimesDecl] = ..., + handles: Optional[Any] = ..., ) -> Callable[[T], T]: ... # configured @endpoint(...) form @@ -685,12 +803,13 @@ def endpoint( model: Optional[SlotLike] = None, models: Optional[Mapping[str, SlotLike]] = None, resources: Optional[Resources] = None, - runtime: Optional[str] = None, + runtime: Union[str, RuntimeFormula, Mapping[str, RuntimeFormula], None] = None, name: Optional[str] = None, compile: Optional[Compile] = None, child_calls: bool = False, warmup: Optional[WarmupDecl] = None, regimes: Optional[RegimesDecl] = None, + handles: Optional[Any] = None, ) -> Any: """The one endpoint decorator. See the module docstring for shapes. @@ -713,6 +832,8 @@ def endpoint( f"@endpoint compile= must be a Compile, got {type(compile).__name__}" ) model_map, slot_map = _normalize_models(model, models) + owner = getattr(target, "__name__", "") if target is not None else "" + engine_runtime, runtime_formulas = _split_runtime_kwarg(owner, runtime) def apply(obj: Any) -> Any: if inspect.isclass(obj): @@ -721,14 +842,18 @@ def apply(obj: Any) -> Any: "class handlers route by method name") return _decorate_class( obj, kind=kind, resources=resources_value, models=dict(model_map), - slots=dict(slot_map), runtime=runtime, compile=compile, + slots=dict(slot_map), runtime=engine_runtime, + runtime_formula=runtime_formulas, compile=compile, child_calls=child_calls, warmup=warmup, regimes=regimes, + handles=handles, ) if inspect.isfunction(obj): return _decorate_function( obj, kind=kind, resources=resources_value, models=dict(model_map), - slots=dict(slot_map), runtime=runtime, name=name, compile=compile, + slots=dict(slot_map), runtime=engine_runtime, + runtime_formula=runtime_formulas, name=name, compile=compile, child_calls=child_calls, warmup=warmup, regimes=regimes, + handles=handles, ) raise TypeError( f"@endpoint requires a function or class, got {type(obj).__name__}" diff --git a/src/gen_worker/api/formula.py b/src/gen_worker/api/formula.py new file mode 100644 index 00000000..405a897e --- /dev/null +++ b/src/gen_worker/api/formula.py @@ -0,0 +1,294 @@ +"""RuntimeFormula — declared compute-time formula terms (th#1051). + +``runtime=RuntimeFormula("a + b*num_inference_steps + c*num_inference_steps*megapixels")`` +on ``@endpoint``: the author declares the SHAPE of compute time as a sum of +terms, each a learned constant times an expression over payload field names. +The platform learns the constants per physics cell; the worker evaluates the +term expressions on the executed payload and reports them on the observation +back-channel (``JobMetrics.runtime_terms``). + +The grammar and term-key canonicalization mirror tensorhub's +``internal/formula`` package exactly — the formula STRING is the wire +contract and both sides must mint identical term keys. + +Grammar: ``+ - * /``, parentheses, numeric literals, identifiers. Top level +must be a '+'-joined sum; each term starts with a bare constant identifier +(the learned coefficient), optionally ``* ``. +""" + +from __future__ import annotations + +import ast +import math +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple + +import msgspec + +__all__ = ["RuntimeFormula"] + +_NUMERIC_TYPES = (int, float, bool) + + +class _Term: + __slots__ = ("constant", "key", "factor") + + def __init__(self, constant: str, key: str, factor: Optional[ast.expr]) -> None: + self.constant = constant + self.key = key # canonical feature key ("1" for the intercept) + self.factor = factor # None for the bare-constant intercept + + +class RuntimeFormula: + """Parsed, validated compute-time formula declaration.""" + + def __init__(self, source: str) -> None: + if not isinstance(source, str) or not source.strip(): + raise ValueError("RuntimeFormula: source must be a non-empty string") + self.source = source.strip() + self.terms: List[_Term] = _parse_terms(self.source) + self.fields: Tuple[str, ...] = tuple(sorted({ + name for t in self.terms if t.factor is not None + for name in _names(t.factor) + })) + + # -- declaration-time validation ------------------------------------ + + def validate_for_payload(self, payload_type: type, owner: str) -> None: + """Every payload identifier must be a numeric/bool payload field WITH + a declared default (the defaults are the reference payload); no + constant may collide with a field name.""" + try: + field_map = {f.name: f for f in msgspec.structs.fields(payload_type)} + except Exception as exc: # not a Struct — walker validates elsewhere + raise ValueError( + f"{owner}: runtime= formula needs a msgspec.Struct payload ({exc})" + ) from exc + for t in self.terms: + if t.constant in field_map: + raise ValueError( + f"{owner}: runtime formula constant {t.constant!r} collides " + f"with a payload field name" + ) + for name in self.fields: + f = field_map.get(name) + if f is None: + raise ValueError( + f"{owner}: runtime formula field {name!r} is not a payload field" + ) + default = f.default + if default is msgspec.NODEFAULT and f.default_factory is not msgspec.NODEFAULT: + default = f.default_factory() + if default is msgspec.NODEFAULT or not isinstance(default, _NUMERIC_TYPES): + raise ValueError( + f"{owner}: runtime formula field {name!r} needs a numeric/bool " + f"default (the defaults are the reference payload)" + ) + + # -- worker-side evaluation ------------------------------------------ + + def term_values(self, values: Mapping[str, Any]) -> Optional[Dict[str, float]]: + """Evaluate each term's factor -> {term key: value}. None when any + referenced field is missing/non-finite (the hub then falls back to + its own evaluation).""" + out: Dict[str, float] = {} + for t in self.terms: + if t.factor is None: + out[t.key] = 1.0 + continue + v = _eval(t.factor, values) + if v is None: + return None + out[t.key] = v + return out + + def term_values_from_struct(self, payload: Any) -> Optional[Dict[str, float]]: + values: Dict[str, float] = {} + for name in self.fields: + raw = getattr(payload, name, None) + if isinstance(raw, bool): + values[name] = 1.0 if raw else 0.0 + elif isinstance(raw, (int, float)): + values[name] = float(raw) + return self.term_values(values) + + def __repr__(self) -> str: # pragma: no cover + return f"RuntimeFormula({self.source!r})" + + +# --------------------------------------------------------------------------- +# Parsing (mirrors tensorhub internal/formula) +# --------------------------------------------------------------------------- + + +def _parse_terms(source: str) -> List[_Term]: + try: + tree = ast.parse(source, mode="eval") + except SyntaxError as exc: + raise ValueError(f"runtime formula: {exc}") from exc + _check_nodes(tree.body) + term_exprs = _split_sum(tree.body) + terms: List[_Term] = [] + seen: Set[str] = set() + for expr in term_exprs: + constant, factor = _split_leading_constant(expr) + if constant in seen: + raise ValueError( + f"runtime formula: constant {constant!r} used in more than one term" + ) + seen.add(constant) + key = "1" if factor is None else _canonical(factor) + if factor is not None and not _names(factor): + raise ValueError( + f"runtime formula term with constant {constant!r}: factor " + f"references no payload field" + ) + terms.append(_Term(constant, key, factor)) + return terms + + +def _check_nodes(node: ast.expr) -> None: + allowed = ( + ast.BinOp, ast.UnaryOp, ast.Name, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.USub, ast.Load, + ) + for n in ast.walk(node): + if not isinstance(n, allowed): + raise ValueError( + f"runtime formula: {type(n).__name__} is not allowed " + f"(arithmetic over payload fields only)" + ) + if isinstance(n, ast.BinOp) and not isinstance(n.op, (ast.Add, ast.Sub, ast.Mult, ast.Div)): + raise ValueError("runtime formula: only + - * / are allowed") + if isinstance(n, ast.UnaryOp) and not isinstance(n.op, ast.USub): + raise ValueError("runtime formula: only unary '-' is allowed") + if isinstance(n, ast.Constant) and ( + isinstance(n.value, bool) or not isinstance(n.value, (int, float)) + ): + raise ValueError("runtime formula: literals must be numeric") + + +def _split_sum(node: ast.expr) -> List[ast.expr]: + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): + return _split_sum(node.left) + _split_sum(node.right) + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Sub): + raise ValueError( + "runtime formula: top-level '-' not allowed (constants are " + "learned signed; write '+' terms only)" + ) + return [node] + + +def _split_leading_constant(term: ast.expr) -> Tuple[str, Optional[ast.expr]]: + """A term is ``constIdent`` or ``constIdent * factor [* factor...]``. + The leftmost factor of the product chain must be a bare Name.""" + if isinstance(term, ast.Name): + return term.id, None + if isinstance(term, ast.BinOp) and isinstance(term.op, (ast.Mult, ast.Div)): + # Walk to the leftmost leaf of the product chain. + chain: List[Tuple[ast.BinOp, str]] = [] + node: ast.expr = term + while isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Mult, ast.Div)): + chain.append((node, "l")) + node = node.left + if not isinstance(node, ast.Name): + raise ValueError( + "runtime formula: each term must start with a constant " + "identifier (e.g. \"b*steps\")" + ) + # The constant is multiplied (never divided) into the term: the + # innermost BinOp holding the Name on its left must be Mult. + innermost = chain[-1][0] + if not isinstance(innermost.op, ast.Mult): + raise ValueError( + f"runtime formula: constant {node.id!r} must be followed by '*'" + ) + factor = _rebuild_without_leading(term) + return node.id, factor + raise ValueError( + "runtime formula: each term must start with a constant identifier" + ) + + +def _rebuild_without_leading(term: ast.BinOp) -> ast.expr: + """Drop the leftmost leaf (the constant) from a product chain: + ``c*a*b/d`` -> ``a*b/d``.""" + if isinstance(term.left, ast.Name): + return term.right + assert isinstance(term.left, ast.BinOp) + new_left = _rebuild_without_leading(term.left) + return ast.BinOp(left=new_left, op=term.op, right=term.right) + + +def _names(node: ast.expr) -> Set[str]: + return {n.id for n in ast.walk(node) if isinstance(n, ast.Name)} + + +def _prec(op: ast.operator) -> int: + return 2 if isinstance(op, (ast.Mult, ast.Div)) else 1 + + +def _canonical(node: ast.expr) -> str: + """Whitespace-free canonical serialization — byte-identical to + tensorhub internal/formula's Expr.Canonical for the same source.""" + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): + v = float(node.value) + if v == math.trunc(v) and abs(v) < 1e15: + return str(int(v)) + return repr(v) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + inner = _canonical(node.operand) + if isinstance(node.operand, ast.BinOp) and _prec(node.operand.op) < 2: + return f"-({inner})" + return f"-{inner}" + if isinstance(node, ast.BinOp): + op_char = {ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/"}[type(node.op)] + left = _canonical(node.left) + right = _canonical(node.right) + if isinstance(node.left, ast.BinOp) and _prec(node.left.op) < _prec(node.op): + left = f"({left})" + if isinstance(node.right, ast.BinOp) and ( + _prec(node.right.op) < _prec(node.op) + or (_prec(node.right.op) == _prec(node.op) and op_char in "-/") + ): + right = f"({right})" + return f"{left}{op_char}{right}" + raise ValueError(f"runtime formula: cannot canonicalize {type(node).__name__}") + + +def _eval(node: ast.expr, values: Mapping[str, Any]) -> Optional[float]: + if isinstance(node, ast.Name): + v = values.get(node.id) + if v is None or not isinstance(v, (int, float)) or isinstance(v, bool): + if isinstance(v, bool): + return 1.0 if v else 0.0 + return None + f = float(v) + return None if (math.isnan(f) or math.isinf(f)) else f + if isinstance(node, ast.Constant): + if not isinstance(node.value, (int, float)): + return None + return float(node.value) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + v = _eval(node.operand, values) + return None if v is None else -v + if isinstance(node, ast.BinOp): + left = _eval(node.left, values) + right = _eval(node.right, values) + if left is None or right is None: + return None + if isinstance(node.op, ast.Add): + out = left + right + elif isinstance(node.op, ast.Sub): + out = left - right + elif isinstance(node.op, ast.Mult): + out = left * right + elif isinstance(node.op, ast.Div): + if right == 0: + return None + out = left / right + else: # pragma: no cover + return None + return None if (math.isnan(out) or math.isinf(out)) else out + return None # pragma: no cover diff --git a/src/gen_worker/cli/run.py b/src/gen_worker/cli/run.py index fd9e5283..1c205703 100644 --- a/src/gen_worker/cli/run.py +++ b/src/gen_worker/cli/run.py @@ -197,6 +197,7 @@ def __init__( bindings: Dict[str, Any], resources: Any = None, slots: Optional[Dict[str, Any]] = None, + handles: Tuple[str, ...] = (), ) -> None: self.cls = cls self.attr_name = attr_name @@ -209,6 +210,7 @@ def __init__( self.bindings = bindings self.resources = resources self.slots = slots or {} + self.handles = tuple(handles) def _collect_class_methods(mod: Any) -> List[_SelectedFunction]: @@ -230,6 +232,7 @@ def _collect_class_methods(mod: Any) -> List[_SelectedFunction]: bindings=dict(es.models), resources=es.resources, slots=dict(es.slots), + handles=es.handles, ) for es in collect_from_namespace(mod) ] @@ -672,6 +675,34 @@ def _resolve_ctx_slots(ctx: Any, selected: "_SelectedFunction") -> None: set_slots(resolved, errors) +def _local_executing_lane( + bindings: Dict[str, Any], lane_str: str, handles: Tuple[str, ...] +) -> str: + """th#1050 ctx.lane for local runs: a --lane the endpoint DECLARES wins + (author kernels execute it); otherwise the most-quantized binding's lane + (the local twin of Executor._served_lane, eager execution).""" + from ..models import lanes as lanespec + + if lane_str and handles: + try: + req = lanespec.parse_lane_spec(lane_str) + if req.lane is not None and lanespec.lane_body_id(req.lane) in handles: + return lanespec.lane_id(req.lane) + except ValueError: + pass + ranked = {b: i for i, b in enumerate(lanespec.known_lanes())} + best, best_key = None, (2, len(ranked) + 1) + for b in (bindings or {}).values(): + lane = lanespec.lane_of_binding( + getattr(b, "flavor", "") or "", + getattr(b, "storage_dtype", "") or "", False) + quant = 1 if lanespec.family_of(lane) == lanespec.FAMILY_BF16 else 0 + key = (quant, ranked.get(lanespec.lane_id(lane), len(ranked))) + if best is None or key < best_key: + best, best_key = lane, key + return lanespec.lane_id(best) if best is not None else "bf16-w16a16+eager" + + def _apply_lane_to_bindings(bindings: Dict[str, Any], lane_str: str) -> Dict[str, Any]: """th#913/gw#596: fold a --lane choice into the declared bindings. @@ -952,6 +983,8 @@ def _run_inner(args: argparse.Namespace) -> int: kind=selected.kind, allow_publish=bool(args.allow_publish), ) + ctx._set_lane(_local_executing_lane( + selected.bindings, getattr(args, "lane", ""), selected.handles)) if source_path is not None: set_source = getattr(ctx, "_set_source_path", None) if not callable(set_source): diff --git a/src/gen_worker/discovery/discover.py b/src/gen_worker/discovery/discover.py index 3971a4e0..e2855a60 100644 --- a/src/gen_worker/discovery/discover.py +++ b/src/gen_worker/discovery/discover.py @@ -761,6 +761,13 @@ def _extract_entries(obj: Any, module_name: str) -> List[Dict[str, Any]]: # default (absent = ["standard"]). if tuple(es.regimes) != DEFAULT_REGIMES: fn["regimes"] = list(es.regimes) + # th#1050: opt-in declared lane bodies (behavioral divergence marker). + if es.handles: + fn["handles"] = list(es.handles) + # th#1051: declared compute-time formula — the hub learns the + # constants per physics cell; the source string is the contract. + if es.runtime_formula is not None: + fn["runtime_formula"] = es.runtime_formula.source if model_key is not None: fn["model"] = model_key if slots_block: diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index e35052b9..a3e1976a 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -362,6 +362,19 @@ def _map_exception(exc: BaseException) -> Tuple["pb.JobStatus", str]: return pb.JOB_STATUS_FATAL, f"{label}: {detail}"[:512] if detail else label +def _runtime_term_values(spec: Any, payload: Any) -> "Optional[Dict[str, float]]": + """th#1051: evaluate the declared runtime formula's terms on the EXECUTED + payload (defaults already applied by msgspec decode). None = undeclared + or unevaluable — the hub then evaluates from the raw payload itself.""" + rf = getattr(spec, "runtime_formula", None) + if rf is None: + return None + try: + return rf.term_values_from_struct(payload) + except Exception: + return None + + def _scan_output_assets(output: Any) -> Tuple[float, int]: """One walk over the job output: (summed MEDIA seconds, count of output ``Asset``s). Billing sources for ``per_output_second`` (th#572) and @@ -3183,7 +3196,16 @@ def _lane_effective_spec(self, spec: EndpointSpec, lane_str: str) -> EndpointSpe raise ValidationError(str(exc)) from None if req.is_zero: return spec - if (req.lane is not None and req.lane.weights == lanespec.WEIGHTS_FP8 + # th#1050: a lane the endpoint DECLARES (handles=) is served by the + # author's own code branching on ctx.lane — satisfiable with no + # laddered rebind (custom loaders/kernels have nothing to rebind), + # and exempt from the platform-kernel compiled-only rule. + declared_lane = ( + req.lane is not None + and lanespec.lane_body_id(req.lane) in getattr(spec, "handles", ()) + ) + if not declared_lane and ( + req.lane is not None and req.lane.weights == lanespec.WEIGHTS_FP8 and req.lane.activation == lanespec.ACT_W8A8 and req.lane.execution == lanespec.EXEC_EAGER): # gw#586: w8a8 serves compiled-only; an eager w8a8 request would @@ -3206,8 +3228,9 @@ def pick_lane_of(pick: "Optional[Tuple[str, str, str]]") -> "Optional[Any]": effective = dict(spec.models) changed = False # bf16 is trivially serveable (the declared base IS bf16); quantized - # lanes must find at least one laddered ref that can serve them. - satisfied = req.family == lanespec.FAMILY_BF16 + # lanes must find at least one laddered ref that can serve them — + # unless the endpoint declares the lane (author code serves it). + satisfied = req.family == lanespec.FAMILY_BF16 or declared_lane want_w8a8 = req.lane is not None and req.lane.activation == lanespec.ACT_W8A8 for slot, base_binding in declared.items(): if slot not in effective: @@ -3271,10 +3294,28 @@ def pick_lane_of(pick: "Optional[Tuple[str, str, str]]") -> "Optional[Any]": if wire_ref(spec.models[s]) != wire_ref(b)}) return derived - def _served_lane(self, spec: EndpointSpec) -> str: + def _handled_lane_body(self, spec: EndpointSpec, instructed: str) -> str: + """th#1050: the instructed lane's body when the endpoint DECLARES it + (handles=) — the author's code, not binding surgery, serves it.""" + from .models import lanes as lanespec + + if not instructed or not getattr(spec, "handles", ()): + return "" + try: + req = lanespec.parse_lane_spec(instructed) + except ValueError: + return "" + if req.lane is None: + return "" + body = lanespec.lane_body_id(req.lane) + return body if body in spec.handles else "" + + def _served_lane(self, spec: EndpointSpec, instructed: str = "") -> str: """The CONCRETE lane this spec's instance executes as, for - JobMetrics.lane reporting: the most-quantized pipeline binding's lane - (table rank), execution axis from the record's live compile state.""" + JobMetrics.lane and ctx.lane reporting: the most-quantized pipeline + binding's lane (table rank), execution axis from the record's live + compile state. A declared (handles=) instructed lane wins outright — + the author's kernels execute it regardless of binding surgery.""" from .models import lanes as lanespec compiled = False @@ -3284,6 +3325,10 @@ def _served_lane(self, spec: EndpointSpec) -> str: compiled = any( getattr(t, "active_compile_ref", "") for t in rec.compile_targets.values()) + handled = self._handled_lane_body(spec, instructed) + if handled: + exec_axis = lanespec.EXEC_COMPILED if compiled else lanespec.EXEC_EAGER + return f"{handled}+{exec_axis}" # Report the most-quantized binding's lane: quantized lanes always # outrank bf16 (a bf16 VAE riding a w8a16 pipeline is still the # w8a16 lane), ties by table rank. @@ -3737,6 +3782,7 @@ async def _one(wj: Any, build: Any, mode: str, *, variant: bool) -> bool: **_resolve_slots_kwargs(wj.spec, None), boot_warmup=True, ) + ctx._set_lane(self._served_lane(wj.spec)) try: await self._invoke_warmup(wj.spec, instance, ctx, payload, handler_kwargs) except Exception as exc: @@ -6594,7 +6640,10 @@ async def _run_job_pinned( await self._materialize_datasets(ctx, payload) instance = await self.ensure_setup(spec, snapshots, promote_slots=routed) # th#913/gw#596: the concrete lane actually serving this job. - job.lane = self._served_lane(spec) + # th#1050: ctx.lane exposes the same post-degrade truth to the + # handler (declared-lane endpoints branch on it). + job.lane = self._served_lane(spec, instructed=run.lane) + ctx._set_lane(job.lane) kwargs = await self._handler_kwargs(spec, snapshots) adapters = await self._prepare_adapters(run, spec, snapshots) ctx.raise_if_cancelled("canceled") @@ -6695,7 +6744,8 @@ async def _run_job_pinned( self._adapters.deactivate, ref, pipe, run.request_id ) metrics = self._metrics(queue_ms, started, concurrency_at_start, gpu_index, - output=output, lane=job.lane) + output=output, lane=job.lane, + runtime_terms=_runtime_term_values(spec, payload)) handler_done = time.monotonic() # Handler GPU work is done — free the slot before result-blob # upload and result send so the next job's compute starts now. @@ -7173,6 +7223,7 @@ async def _serialize_output( def _metrics( self, queue_ms: int, started: float, concurrency_at_start: int, gpu_index: int, output: Any = None, lane: str = "", + runtime_terms: "Optional[Dict[str, float]]" = None, ) -> pb.JobMetrics: runtime_ms = int((time.monotonic() - started) * 1000) # rss_at_end_bytes (pgw#513): instantaneous RSS, honestly named — the @@ -7201,6 +7252,9 @@ def _metrics( input_cached_tokens=usage.cached_tokens if usage is not None else 0, output_tokens=usage.completion_tokens if usage is not None else 0, lane=lane, + # th#1051: declared-formula term features from the EXECUTED + # payload (defaults applied); empty = no declared formula. + runtime_terms=runtime_terms or {}, ) async def _send_result( diff --git a/src/gen_worker/models/lanes.py b/src/gen_worker/models/lanes.py index eba46815..da8ef3db 100644 --- a/src/gen_worker/models/lanes.py +++ b/src/gen_worker/models/lanes.py @@ -85,6 +85,24 @@ def known_lanes() -> list[str]: return out +def lane_body_id(lane: Lane) -> str: + """The lane id without the execution axis (verdict/declaration token).""" + body = f"{lane.weights}-{lane.activation}" + if lane.scale: + body += f"-{lane.scale}" + return body + + +def known_lane_bodies() -> list[str]: + """Every concrete lane BODY token, ranked (table order). These are the + valid `handles=` declaration tokens (th#1050) — execution axis excluded: + author kernels declare the quant scheme, the platform owns eager/compiled.""" + return [ + lane_body_id(Lane(weights=w, activation=a, scale=s, execution=EXEC_EAGER)) + for w, a, s in _KNOWN_BODIES + ] + + def valid_lane(lane: Lane) -> bool: if lane.execution not in (EXEC_EAGER, EXEC_COMPILED): return False @@ -219,7 +237,9 @@ def __init__(self, lane: str, detail: str) -> None: "WEIGHTS_SVDQ_INT4", "family_of", "is_w8a8_flavor", + "known_lane_bodies", "known_lanes", + "lane_body_id", "lane_id", "lane_of_binding", "parse_lane", diff --git a/src/gen_worker/pb/worker_scheduler_pb2.py b/src/gen_worker/pb/worker_scheduler_pb2.py index a3d36fb7..f9edfa0e 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.py +++ b/src/gen_worker/pb/worker_scheduler_pb2.py @@ -24,7 +24,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xa9\x04\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x12\x39\n\x0f\x61\x63tivity_update\x18\t \x01(\x0b\x32\x1e.cozy.scheduler.ActivityUpdateH\x00\x12\x41\n\x13hardware_unsuitable\x18\n \x01(\x0b\x32\".cozy.scheduler.HardwareUnsuitableH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xc7\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\x12\x1d\n\x15heartbeat_interval_ms\x18\x08 \x01(\x04\"\x9b\x02\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xa2\x02\n\x12HardwareUnsuitable\x12\x11\n\tworker_id\x18\x01 \x01(\t\x12\x12\n\nrelease_id\x18\x02 \x01(\t\x12\x14\n\x0creason_class\x18\x03 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x04 \x01(\t\x12\x16\n\x0e\x64river_version\x18\x05 \x01(\t\x12\x10\n\x08gpu_name\x18\x06 \x01(\t\x12\x15\n\rtorch_version\x18\x07 \x01(\t\x12\x1a\n\x12torch_cuda_version\x18\x08 \x01(\t\x12\x1a\n\x12gen_worker_version\x18\t \x01(\t\x12\x14\n\x0cimage_digest\x18\n \x01(\t\x12\x13\n\x0binstance_id\x18\x0b \x01(\t\x12\x1b\n\x13reported_at_unix_ms\x18\x0c \x01(\x03\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xdd\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\x12;\n\x11\x64\x65sired_residency\x18\x05 \x01(\x0b\x32 .cozy.scheduler.DesiredResidency\"\xf7\x01\n\x10\x44\x65siredResidency\x12\x12\n\ngeneration\x18\x01 \x01(\x04\x12\x11\n\tdisk_refs\x18\x02 \x03(\t\x12,\n\x03hot\x18\x03 \x03(\x0b\x32\x1f.cozy.scheduler.DesiredInstance\x12\x42\n\tsnapshots\x18\x04 \x03(\x0b\x32/.cozy.scheduler.DesiredResidency.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"V\n\x0f\x44\x65siredInstance\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\x06models\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\"P\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\x12\x0c\n\x04lane\x18\x04 \x01(\t\"\xe8\x02\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\x12%\n\x1dobserved_residency_generation\x18\x06 \x01(\x04\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\x12\x33\n\ndisk_usage\x18\t \x01(\x0b\x32\x1f.cozy.scheduler.DiskUsageReport\"\xa9\x01\n\x10StorageTierUsage\x12)\n\x04tier\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.StorageTier\x12\x12\n\nmount_path\x18\x02 \x01(\t\x12\x13\n\x0btotal_bytes\x18\x03 \x01(\x03\x12\x12\n\nfree_bytes\x18\x04 \x01(\x03\x12\x12\n\nused_bytes\x18\x05 \x01(\x03\x12\x19\n\x11reclaimable_bytes\x18\x06 \x01(\x03\"_\n\x0f\x44iskUsageReport\x12/\n\x05tiers\x18\x01 \x03(\x0b\x32 .cozy.scheduler.StorageTierUsage\x12\x1b\n\x13\x63\x61pacity_generation\x18\x02 \x01(\x04\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\xc8\x04\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\x12\x0c\n\x04lane\x18\x0e \x01(\t\x12\x30\n\x0cinput_assets\x18\x0f \x03(\x0b\x32\x1a.cozy.scheduler.InputAsset\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"w\n\nInputAsset\x12\x10\n\x08\x61sset_id\x18\x01 \x01(\t\x12\x12\n\nsource_ref\x18\x02 \x01(\t\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x12\n\nsize_bytes\x18\x04 \x01(\x03\x12\x0c\n\x04kind\x18\x05 \x01(\t\x12\x11\n\tmime_type\x18\x06 \x01(\t\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"\xe6\x01\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\x12@\n\ncomponents\x18\x05 \x03(\x0b\x32,.cozy.scheduler.ModelBinding.ComponentsEntry\x1a\x31\n\x0f\x43omponentsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xc2\x02\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\x12\x0c\n\x04lane\x18\r \x01(\t\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xa0\x01\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x9b\x04\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\x12\x15\n\rnetwork_bytes\x18\x14 \x01(\x03\"\xdc\x02\n\x0e\x41\x63tivityUpdate\x12\x0c\n\x04kind\x18\x01 \x01(\t\x12\r\n\x05phase\x18\x02 \x01(\t\x12\x0c\n\x04step\x18\x03 \x01(\x03\x12\x13\n\x0btotal_steps\x18\x04 \x01(\x03\x12\x0b\n\x03seq\x18\x05 \x01(\x04\x12,\n\x05state\x18\x06 \x01(\x0e\x32\x1d.cozy.scheduler.ActivityState\x12\r\n\x05\x65rror\x18\x07 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x08 \x01(\t\x12\x1a\n\x12updated_at_unix_ms\x18\t \x01(\x03\x12\x0f\n\x07\x63ounter\x18\n \x01(\t\x12\x14\n\x0c\x63ounter_unit\x18\x0b \x01(\t\x12\x14\n\x0c\x63ounter_done\x18\x0c \x01(\x01\x12\x15\n\rcounter_total\x18\r \x01(\x01\x12\x12\n\nrate_per_s\x18\x0e \x01(\x01\x12\x14\n\x0cself_stalled\x18\x0f \x01(\x08\x12\x16\n\x0estalled_for_ms\x18\x10 \x01(\x03\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x04*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*v\n\x0bStorageTier\x12\x1c\n\x18STORAGE_TIER_UNSPECIFIED\x10\x00\x12\x1a\n\x16STORAGE_TIER_CONTAINER\x10\x01\x12\x17\n\x13STORAGE_TIER_VOLUME\x10\x02\x12\x14\n\x10STORAGE_TIER_NFS\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa7\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04\"\x04\x08\x01\x10\x01\"\x04\x08\x02\x10\x02\"\x04\x08\x03\x10\x03*\x16MODEL_OP_KIND_DOWNLOAD*\x12MODEL_OP_KIND_LOAD*\x14MODEL_OP_KIND_UNLOAD*\x82\x02\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08*\x84\x01\n\rActivityState\x12\x1e\n\x1a\x41\x43TIVITY_STATE_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x41\x43TIVITY_STATE_RUNNING\x10\x01\x12\x1c\n\x18\x41\x43TIVITY_STATE_COMPLETED\x10\x02\x12\x19\n\x15\x41\x43TIVITY_STATE_FAILED\x10\x03\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xa9\x04\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x12\x39\n\x0f\x61\x63tivity_update\x18\t \x01(\x0b\x32\x1e.cozy.scheduler.ActivityUpdateH\x00\x12\x41\n\x13hardware_unsuitable\x18\n \x01(\x0b\x32\".cozy.scheduler.HardwareUnsuitableH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xc7\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\x12\x1d\n\x15heartbeat_interval_ms\x18\x08 \x01(\x04\"\x9b\x02\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xa2\x02\n\x12HardwareUnsuitable\x12\x11\n\tworker_id\x18\x01 \x01(\t\x12\x12\n\nrelease_id\x18\x02 \x01(\t\x12\x14\n\x0creason_class\x18\x03 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x04 \x01(\t\x12\x16\n\x0e\x64river_version\x18\x05 \x01(\t\x12\x10\n\x08gpu_name\x18\x06 \x01(\t\x12\x15\n\rtorch_version\x18\x07 \x01(\t\x12\x1a\n\x12torch_cuda_version\x18\x08 \x01(\t\x12\x1a\n\x12gen_worker_version\x18\t \x01(\t\x12\x14\n\x0cimage_digest\x18\n \x01(\t\x12\x13\n\x0binstance_id\x18\x0b \x01(\t\x12\x1b\n\x13reported_at_unix_ms\x18\x0c \x01(\x03\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xdd\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\x12;\n\x11\x64\x65sired_residency\x18\x05 \x01(\x0b\x32 .cozy.scheduler.DesiredResidency\"\xf7\x01\n\x10\x44\x65siredResidency\x12\x12\n\ngeneration\x18\x01 \x01(\x04\x12\x11\n\tdisk_refs\x18\x02 \x03(\t\x12,\n\x03hot\x18\x03 \x03(\x0b\x32\x1f.cozy.scheduler.DesiredInstance\x12\x42\n\tsnapshots\x18\x04 \x03(\x0b\x32/.cozy.scheduler.DesiredResidency.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"V\n\x0f\x44\x65siredInstance\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\x06models\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\"P\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\x12\x0c\n\x04lane\x18\x04 \x01(\t\"\xe8\x02\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\x12%\n\x1dobserved_residency_generation\x18\x06 \x01(\x04\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\x12\x33\n\ndisk_usage\x18\t \x01(\x0b\x32\x1f.cozy.scheduler.DiskUsageReport\"\xa9\x01\n\x10StorageTierUsage\x12)\n\x04tier\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.StorageTier\x12\x12\n\nmount_path\x18\x02 \x01(\t\x12\x13\n\x0btotal_bytes\x18\x03 \x01(\x03\x12\x12\n\nfree_bytes\x18\x04 \x01(\x03\x12\x12\n\nused_bytes\x18\x05 \x01(\x03\x12\x19\n\x11reclaimable_bytes\x18\x06 \x01(\x03\"_\n\x0f\x44iskUsageReport\x12/\n\x05tiers\x18\x01 \x03(\x0b\x32 .cozy.scheduler.StorageTierUsage\x12\x1b\n\x13\x63\x61pacity_generation\x18\x02 \x01(\x04\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\xc8\x04\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\x12\x0c\n\x04lane\x18\x0e \x01(\t\x12\x30\n\x0cinput_assets\x18\x0f \x03(\x0b\x32\x1a.cozy.scheduler.InputAsset\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"w\n\nInputAsset\x12\x10\n\x08\x61sset_id\x18\x01 \x01(\t\x12\x12\n\nsource_ref\x18\x02 \x01(\t\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x12\n\nsize_bytes\x18\x04 \x01(\x03\x12\x0c\n\x04kind\x18\x05 \x01(\t\x12\x11\n\tmime_type\x18\x06 \x01(\t\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"\xe6\x01\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\x12@\n\ncomponents\x18\x05 \x03(\x0b\x32,.cozy.scheduler.ModelBinding.ComponentsEntry\x1a\x31\n\x0f\x43omponentsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xbc\x03\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\x12\x0c\n\x04lane\x18\r \x01(\t\x12\x43\n\rruntime_terms\x18\x0e \x03(\x0b\x32,.cozy.scheduler.JobMetrics.RuntimeTermsEntry\x1a\x33\n\x11RuntimeTermsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01:\x02\x38\x01\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xa0\x01\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x9b\x04\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\x12\x15\n\rnetwork_bytes\x18\x14 \x01(\x03\"\xdc\x02\n\x0e\x41\x63tivityUpdate\x12\x0c\n\x04kind\x18\x01 \x01(\t\x12\r\n\x05phase\x18\x02 \x01(\t\x12\x0c\n\x04step\x18\x03 \x01(\x03\x12\x13\n\x0btotal_steps\x18\x04 \x01(\x03\x12\x0b\n\x03seq\x18\x05 \x01(\x04\x12,\n\x05state\x18\x06 \x01(\x0e\x32\x1d.cozy.scheduler.ActivityState\x12\r\n\x05\x65rror\x18\x07 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x08 \x01(\t\x12\x1a\n\x12updated_at_unix_ms\x18\t \x01(\x03\x12\x0f\n\x07\x63ounter\x18\n \x01(\t\x12\x14\n\x0c\x63ounter_unit\x18\x0b \x01(\t\x12\x14\n\x0c\x63ounter_done\x18\x0c \x01(\x01\x12\x15\n\rcounter_total\x18\r \x01(\x01\x12\x12\n\nrate_per_s\x18\x0e \x01(\x01\x12\x14\n\x0cself_stalled\x18\x0f \x01(\x08\x12\x16\n\x0estalled_for_ms\x18\x10 \x01(\x03\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x04*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*v\n\x0bStorageTier\x12\x1c\n\x18STORAGE_TIER_UNSPECIFIED\x10\x00\x12\x1a\n\x16STORAGE_TIER_CONTAINER\x10\x01\x12\x17\n\x13STORAGE_TIER_VOLUME\x10\x02\x12\x14\n\x10STORAGE_TIER_NFS\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa7\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04\"\x04\x08\x01\x10\x01\"\x04\x08\x02\x10\x02\"\x04\x08\x03\x10\x03*\x16MODEL_OP_KIND_DOWNLOAD*\x12MODEL_OP_KIND_LOAD*\x14MODEL_OP_KIND_UNLOAD*\x82\x02\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08*\x84\x01\n\rActivityState\x12\x1e\n\x1a\x41\x43TIVITY_STATE_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x41\x43TIVITY_STATE_RUNNING\x10\x01\x12\x1c\n\x18\x41\x43TIVITY_STATE_COMPLETED\x10\x02\x12\x19\n\x15\x41\x43TIVITY_STATE_FAILED\x10\x03\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -40,26 +40,28 @@ _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_options = b'8\001' _globals['_MODELBINDING_COMPONENTSENTRY']._loaded_options = None _globals['_MODELBINDING_COMPONENTSENTRY']._serialized_options = b'8\001' + _globals['_JOBMETRICS_RUNTIMETERMSENTRY']._loaded_options = None + _globals['_JOBMETRICS_RUNTIMETERMSENTRY']._serialized_options = b'8\001' _globals['_FNUNAVAILABLE_AXESENTRY']._loaded_options = None _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_options = b'8\001' - _globals['_PROTOCOLVERSION']._serialized_start=7664 - _globals['_PROTOCOLVERSION']._serialized_end=7745 - _globals['_RESIDENCYTIER']._serialized_start=7747 - _globals['_RESIDENCYTIER']._serialized_end=7868 - _globals['_STORAGETIER']._serialized_start=7870 - _globals['_STORAGETIER']._serialized_end=7988 - _globals['_WORKERPHASE']._serialized_start=7991 - _globals['_WORKERPHASE']._serialized_end=8207 - _globals['_OUTPUTMODE']._serialized_start=8209 - _globals['_OUTPUTMODE']._serialized_end=8295 - _globals['_JOBSTATUS']._serialized_start=8298 - _globals['_JOBSTATUS']._serialized_end=8453 - _globals['_MODELOPKIND']._serialized_start=8456 - _globals['_MODELOPKIND']._serialized_end=8623 - _globals['_MODELSTATE']._serialized_start=8626 - _globals['_MODELSTATE']._serialized_end=8884 - _globals['_ACTIVITYSTATE']._serialized_start=8887 - _globals['_ACTIVITYSTATE']._serialized_end=9019 + _globals['_PROTOCOLVERSION']._serialized_start=7786 + _globals['_PROTOCOLVERSION']._serialized_end=7867 + _globals['_RESIDENCYTIER']._serialized_start=7869 + _globals['_RESIDENCYTIER']._serialized_end=7990 + _globals['_STORAGETIER']._serialized_start=7992 + _globals['_STORAGETIER']._serialized_end=8110 + _globals['_WORKERPHASE']._serialized_start=8113 + _globals['_WORKERPHASE']._serialized_end=8329 + _globals['_OUTPUTMODE']._serialized_start=8331 + _globals['_OUTPUTMODE']._serialized_end=8417 + _globals['_JOBSTATUS']._serialized_start=8420 + _globals['_JOBSTATUS']._serialized_end=8575 + _globals['_MODELOPKIND']._serialized_start=8578 + _globals['_MODELOPKIND']._serialized_end=8745 + _globals['_MODELSTATE']._serialized_start=8748 + _globals['_MODELSTATE']._serialized_end=9006 + _globals['_ACTIVITYSTATE']._serialized_start=9009 + _globals['_ACTIVITYSTATE']._serialized_end=9141 _globals['_WORKERMESSAGE']._serialized_start=43 _globals['_WORKERMESSAGE']._serialized_end=596 _globals['_SCHEDULERMESSAGE']._serialized_start=599 @@ -125,27 +127,29 @@ _globals['_JOBRESULT']._serialized_start=5521 _globals['_JOBRESULT']._serialized_end=5727 _globals['_JOBMETRICS']._serialized_start=5730 - _globals['_JOBMETRICS']._serialized_end=6052 - _globals['_JOBPROGRESS']._serialized_start=6054 - _globals['_JOBPROGRESS']._serialized_end=6153 - _globals['_CANCELJOB']._serialized_start=6155 - _globals['_CANCELJOB']._serialized_end=6203 - _globals['_MODELOP']._serialized_start=6206 - _globals['_MODELOP']._serialized_end=6366 - _globals['_MODELEVENT']._serialized_start=6369 - _globals['_MODELEVENT']._serialized_end=6908 - _globals['_ACTIVITYUPDATE']._serialized_start=6911 - _globals['_ACTIVITYUPDATE']._serialized_end=7259 - _globals['_FNUNAVAILABLE']._serialized_start=7262 - _globals['_FNUNAVAILABLE']._serialized_end=7432 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=7389 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=7432 - _globals['_FNDEGRADED']._serialized_start=7435 - _globals['_FNDEGRADED']._serialized_end=7576 - _globals['_DRAIN']._serialized_start=7578 - _globals['_DRAIN']._serialized_end=7606 - _globals['_TOKENREFRESH']._serialized_start=7608 - _globals['_TOKENREFRESH']._serialized_end=7662 - _globals['_WORKERSCHEDULER']._serialized_start=9021 - _globals['_WORKERSCHEDULER']._serialized_end=9118 + _globals['_JOBMETRICS']._serialized_end=6174 + _globals['_JOBMETRICS_RUNTIMETERMSENTRY']._serialized_start=6123 + _globals['_JOBMETRICS_RUNTIMETERMSENTRY']._serialized_end=6174 + _globals['_JOBPROGRESS']._serialized_start=6176 + _globals['_JOBPROGRESS']._serialized_end=6275 + _globals['_CANCELJOB']._serialized_start=6277 + _globals['_CANCELJOB']._serialized_end=6325 + _globals['_MODELOP']._serialized_start=6328 + _globals['_MODELOP']._serialized_end=6488 + _globals['_MODELEVENT']._serialized_start=6491 + _globals['_MODELEVENT']._serialized_end=7030 + _globals['_ACTIVITYUPDATE']._serialized_start=7033 + _globals['_ACTIVITYUPDATE']._serialized_end=7381 + _globals['_FNUNAVAILABLE']._serialized_start=7384 + _globals['_FNUNAVAILABLE']._serialized_end=7554 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=7511 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=7554 + _globals['_FNDEGRADED']._serialized_start=7557 + _globals['_FNDEGRADED']._serialized_end=7698 + _globals['_DRAIN']._serialized_start=7700 + _globals['_DRAIN']._serialized_end=7728 + _globals['_TOKENREFRESH']._serialized_start=7730 + _globals['_TOKENREFRESH']._serialized_end=7784 + _globals['_WORKERSCHEDULER']._serialized_start=9143 + _globals['_WORKERSCHEDULER']._serialized_end=9240 # @@protoc_insertion_point(module_scope) diff --git a/src/gen_worker/pb/worker_scheduler_pb2.pyi b/src/gen_worker/pb/worker_scheduler_pb2.pyi index 4f979c97..6c8adea6 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.pyi +++ b/src/gen_worker/pb/worker_scheduler_pb2.pyi @@ -579,7 +579,14 @@ class JobResult(_message.Message): def __init__(self, request_id: _Optional[str] = ..., attempt: _Optional[int] = ..., status: _Optional[_Union[JobStatus, str]] = ..., inline: _Optional[bytes] = ..., blob_ref: _Optional[str] = ..., safe_message: _Optional[str] = ..., metrics: _Optional[_Union[JobMetrics, _Mapping]] = ...) -> None: ... class JobMetrics(_message.Message): - __slots__ = ("runtime_ms", "queue_ms", "rss_at_end_bytes", "peak_vram_bytes", "concurrency_at_start", "output_media_duration_s", "input_tokens", "input_cached_tokens", "output_tokens", "output_count", "slot_held_ms", "finalize_wall_ms", "lane") + __slots__ = ("runtime_ms", "queue_ms", "rss_at_end_bytes", "peak_vram_bytes", "concurrency_at_start", "output_media_duration_s", "input_tokens", "input_cached_tokens", "output_tokens", "output_count", "slot_held_ms", "finalize_wall_ms", "lane", "runtime_terms") + class RuntimeTermsEntry(_message.Message): + __slots__ = ("key", "value") + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: float + def __init__(self, key: _Optional[str] = ..., value: _Optional[float] = ...) -> None: ... RUNTIME_MS_FIELD_NUMBER: _ClassVar[int] QUEUE_MS_FIELD_NUMBER: _ClassVar[int] RSS_AT_END_BYTES_FIELD_NUMBER: _ClassVar[int] @@ -593,6 +600,7 @@ class JobMetrics(_message.Message): SLOT_HELD_MS_FIELD_NUMBER: _ClassVar[int] FINALIZE_WALL_MS_FIELD_NUMBER: _ClassVar[int] LANE_FIELD_NUMBER: _ClassVar[int] + RUNTIME_TERMS_FIELD_NUMBER: _ClassVar[int] runtime_ms: int queue_ms: int rss_at_end_bytes: int @@ -606,7 +614,8 @@ class JobMetrics(_message.Message): slot_held_ms: int finalize_wall_ms: int lane: str - def __init__(self, runtime_ms: _Optional[int] = ..., queue_ms: _Optional[int] = ..., rss_at_end_bytes: _Optional[int] = ..., peak_vram_bytes: _Optional[int] = ..., concurrency_at_start: _Optional[int] = ..., output_media_duration_s: _Optional[float] = ..., input_tokens: _Optional[int] = ..., input_cached_tokens: _Optional[int] = ..., output_tokens: _Optional[int] = ..., output_count: _Optional[int] = ..., slot_held_ms: _Optional[int] = ..., finalize_wall_ms: _Optional[int] = ..., lane: _Optional[str] = ...) -> None: ... + runtime_terms: _containers.ScalarMap[str, float] + def __init__(self, runtime_ms: _Optional[int] = ..., queue_ms: _Optional[int] = ..., rss_at_end_bytes: _Optional[int] = ..., peak_vram_bytes: _Optional[int] = ..., concurrency_at_start: _Optional[int] = ..., output_media_duration_s: _Optional[float] = ..., input_tokens: _Optional[int] = ..., input_cached_tokens: _Optional[int] = ..., output_tokens: _Optional[int] = ..., output_count: _Optional[int] = ..., slot_held_ms: _Optional[int] = ..., finalize_wall_ms: _Optional[int] = ..., lane: _Optional[str] = ..., runtime_terms: _Optional[_Mapping[str, float]] = ...) -> None: ... class JobProgress(_message.Message): __slots__ = ("request_id", "attempt", "seq", "data", "content_type") diff --git a/src/gen_worker/registry.py b/src/gen_worker/registry.py index 227465fa..df69917f 100644 --- a/src/gen_worker/registry.py +++ b/src/gen_worker/registry.py @@ -17,6 +17,7 @@ from .api.binding import Binding, ModelRef from .api.decorators import ATTR, VARIANT_ATTR, Compile, EndpointDecl, Resources, VariantDecl +from .api.formula import RuntimeFormula from .api.slot import DEFAULT_REGIMES, Slot from .discovery.names import slugify_name from .discovery.walk import find_endpoints @@ -78,6 +79,12 @@ class EndpointSpec: # th#1017: this handler's declared inference regimes (from @endpoint's # regimes=); ("standard",) when undeclared. regimes: tuple = DEFAULT_REGIMES + # th#1050: declared lane bodies this endpoint's code branches on + # (ctx.lane). Empty = platform-managed only. + handles: tuple = () + # th#1051: declared compute-time formula for this handler; None = + # undeclared (scalar EWMA fall-through hub-side). + runtime_formula: Optional["RuntimeFormula"] = None module: str = "" # declaring module walked_module: str = "" # top-level package the object was found under @@ -283,6 +290,12 @@ def _spec_for_handler( if variant_of_slug == slug: raise ValueError(f"{owner}: @variant_of cannot target itself") + # th#1051: resolve + validate this handler's declared compute-time + # formula now that the payload type is known. + runtime_formula = decl.runtime_formula.get(attr_name) + if runtime_formula is not None: + runtime_formula.validate_for_payload(payload_type, owner) + return EndpointSpec( name=slug, method=method, @@ -309,6 +322,8 @@ def _spec_for_handler( variant_of=variant_of_slug, variant_kind=variant_kind, regimes=decl.regimes.get(attr_name, DEFAULT_REGIMES), + handles=tuple(decl.handles), + runtime_formula=runtime_formula, module=getattr(cls or method, "__module__", "") or "", walked_module=walked_module, ) diff --git a/src/gen_worker/request_context/__init__.py b/src/gen_worker/request_context/__init__.py index b812ed7d..b254a32a 100644 --- a/src/gen_worker/request_context/__init__.py +++ b/src/gen_worker/request_context/__init__.py @@ -173,6 +173,7 @@ def __init__( self._deadline = self._started_at + (timeout_ms / 1000.0) self._canceled = False self._boot_warmup = bool(boot_warmup) + self._lane = "" # th#1050: executing lane, set by the executor self._cancel_event = threading.Event() self._emitter = emitter self._cached_repo_job_scope: Optional[tuple[str, str, str]] = None @@ -216,6 +217,18 @@ def deadline(self) -> Optional[float]: """Absolute unix-time deadline, or None when the request is unbounded.""" return self._deadline + @property + def lane(self) -> str: + """The EXECUTING precision lane of this call (th#1050), a full + descriptor id like ``"fp8-w8a8-dynamic+compiled"`` — post-degrade + truth, the same value JobMetrics.lane reports (th#1043 consistent). + Read-only; always available. Handlers that declared + ``handles=[...]`` branch on it; everyone else may ignore it.""" + return self._lane or "bf16-w16a16+eager" + + def _set_lane(self, lane: str) -> None: + self._lane = str(lane or "").strip() + @property def models(self) -> Dict[str, str]: """Resolved model refs for this invocation, keyed by slot name.""" diff --git a/tests/harness/toy_endpoints.py b/tests/harness/toy_endpoints.py index 665a7e8f..14e1e72f 100644 --- a/tests/harness/toy_endpoints.py +++ b/tests/harness/toy_endpoints.py @@ -282,6 +282,20 @@ def composed_echo(self, ctx: RequestContext, data: EchoIn) -> EchoOut: # --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# th#1050: opt-in declared lane — the endpoint's code BRANCHES on ctx.lane. +# --------------------------------------------------------------------------- + + +@endpoint(handles=["fp8-w8a8-dynamic"]) +class LaneAwareEndpoint: + def lane_echo(self, ctx: RequestContext, data: EchoIn) -> EchoOut: + # The divergence the declaration marks: an author kernel branch. + if ctx.lane.startswith("fp8-w8a8-dynamic"): + return EchoOut(response=f"author-kernel:{ctx.lane}") + return EchoOut(response=f"reference:{ctx.lane}") + + @endpoint class BillableEndpoint: def small_usage(self, ctx: RequestContext, data: EchoIn) -> StreamResult: diff --git a/tests/test_declared_lane_th1050.py b/tests/test_declared_lane_th1050.py new file mode 100644 index 00000000..05da88bb --- /dev/null +++ b/tests/test_declared_lane_th1050.py @@ -0,0 +1,114 @@ +"""th#1050 SDK lane contract, real codepaths over the hub-double gRPC wire. + +A toy endpoint DECLARES `handles=["fp8-w8a8-dynamic"]` and branches on +ctx.lane. Proven here: + 1. Discovery manifest emits the `handles` block (hub parse input). + 2. A dispatched declared lane executes with NO laddered binding rebind — + ctx.lane reports it, the handler branches on it, and JobMetrics.lane + reports the SAME value (post-degrade consistency). + 3. Undeclared dispatch (no RunJob.lane) = today's behavior exactly: + ctx.lane reports bf16-w16a16+eager and the reference branch runs. + 4. An instructed lane the endpoint does NOT declare still refuses TYPED + (declaration is divergence marking, never a blanket bypass). + 5. Decorator validation rejects coarse/unknown/execution-suffixed tokens. +""" + +from __future__ import annotations + +import msgspec +import pytest + +from gen_worker.api.decorators import endpoint +from gen_worker.pb import worker_scheduler_pb2 as pb + +from harness.hub_double import hub_double, is_ready, is_result_for +from harness.toy_endpoints import EchoIn + + +def _payload() -> bytes: + return msgspec.msgpack.encode(EchoIn(text="x")) + + +def _decode(res: "pb.JobResult") -> dict: + return msgspec.msgpack.decode(res.inline) + + +def test_manifest_emits_handles_block() -> None: + from gen_worker.discovery.discover import _extract_entries + from harness import toy_endpoints + + (fn,) = _extract_entries(toy_endpoints.LaneAwareEndpoint, "harness.toy_endpoints") + assert fn["handles"] == ["fp8-w8a8-dynamic"] + plain = _extract_entries(toy_endpoints.Basics, "harness.toy_endpoints") + assert all("handles" not in f for f in plain) + + +def test_declared_lane_executes_and_ctx_lane_reports_it() -> None: + with hub_double() as (scheduler, _harness): + conn = scheduler.wait_connection(0) + conn.wait_for(is_ready) + conn.send(run_job=pb.RunJob( + request_id="r-lane", attempt=1, function_name="lane-echo", + input_payload=_payload(), lane="fp8-w8a8-dynamic+eager")) + res = conn.wait_for(is_result_for("r-lane")).job_result + assert res.status == pb.JOB_STATUS_OK, res.safe_message + body = _decode(res) + # ctx.lane carried the executing lane into the author branch... + assert body["response"] == "author-kernel:fp8-w8a8-dynamic+eager" + # ...and JobMetrics.lane reports the SAME executing truth. + assert res.metrics.lane == "fp8-w8a8-dynamic+eager" + + +def test_undeclared_dispatch_is_todays_behavior() -> None: + with hub_double() as (scheduler, _harness): + conn = scheduler.wait_connection(0) + conn.wait_for(is_ready) + conn.send(run_job=pb.RunJob( + request_id="r-base", attempt=1, function_name="lane-echo", + input_payload=_payload())) + res = conn.wait_for(is_result_for("r-base")).job_result + assert res.status == pb.JOB_STATUS_OK, res.safe_message + assert _decode(res)["response"] == "reference:bf16-w16a16+eager" + assert res.metrics.lane == "bf16-w16a16+eager" + + +def test_unhandled_lane_still_refuses_typed() -> None: + with hub_double() as (scheduler, _harness): + conn = scheduler.wait_connection(0) + conn.wait_for(is_ready) + conn.send(run_job=pb.RunJob( + request_id="r-nope", attempt=1, function_name="lane-echo", + input_payload=_payload(), lane="nvfp4-w4a4-static+eager")) + res = conn.wait_for(is_result_for("r-nope")).job_result + assert res.status == pb.JOB_STATUS_INVALID + assert "lane_unavailable" in res.safe_message + assert "nvfp4-w4a4-static" in res.safe_message + + +def test_decorator_rejects_bad_handles_tokens() -> None: + class _In(msgspec.Struct): + text: str = "" + + class _Out(msgspec.Struct): + response: str = "" + + def _make(tokens): + @endpoint(handles=tokens) + class _E: + def go(self, ctx, data: _In) -> _Out: # pragma: no cover + return _Out() + return _E + + for bad, needle in ( + (["fp8"], "coarse family"), + (["fp8-w8a8-dynamic+compiled"], "execution axis"), + (["int8-w8a8"], "not a known lane body"), + (["fp8-w8a8-dynamic", "fp8-w8a8-dynamic"], "repeats"), + ): + with pytest.raises(ValueError, match=needle): + _make(bad) + # A valid declaration normalizes to a tuple of bodies. + cls = _make(["FP8-W8A8-DYNAMIC"]) + from gen_worker.api.decorators import ATTR + + assert getattr(cls, ATTR).handles == ("fp8-w8a8-dynamic",) diff --git a/tests/test_runtime_formula_th1051.py b/tests/test_runtime_formula_th1051.py new file mode 100644 index 00000000..b419ab88 --- /dev/null +++ b/tests/test_runtime_formula_th1051.py @@ -0,0 +1,110 @@ +"""th#1051 RuntimeFormula: parse, canonical term keys (must byte-match the Go +side), payload validation, worker-side term evaluation.""" + +import msgspec +import pytest + +from gen_worker.api.formula import RuntimeFormula + + +class Payload(msgspec.Struct): + prompt: str = "" + num_inference_steps: int = 28 + megapixels: float = 1.0 + width: int = 1024 + height: int = 1024 + hires: bool = False + + +def test_terms_and_keys(): + rf = RuntimeFormula( + "a + b*num_inference_steps + c*num_inference_steps*megapixels" + ) + assert [t.key for t in rf.terms] == [ + "1", "num_inference_steps", "num_inference_steps*megapixels", + ] + assert rf.fields == ("megapixels", "num_inference_steps") + + +def test_composite_key_matches_go_canonical(): + # Go: "steps*(width*height/1000000)" -> "steps*width*height/1000000" + rf = RuntimeFormula("a + b*num_inference_steps*(width*height/1000000)") + assert rf.terms[1].key == "num_inference_steps*width*height/1000000" + + +def test_validate_for_payload(): + rf = RuntimeFormula("a + b*num_inference_steps + c*num_inference_steps*megapixels") + rf.validate_for_payload(Payload, "gen") # ok + + with pytest.raises(ValueError, match="not a payload field"): + RuntimeFormula("a + b*frames").validate_for_payload(Payload, "gen") + with pytest.raises(ValueError, match="collides"): + RuntimeFormula("megapixels + b*num_inference_steps").validate_for_payload( + Payload, "gen" + ) + + class NoDefault(msgspec.Struct): + steps: int + + with pytest.raises(ValueError, match="default"): + RuntimeFormula("a + b*steps").validate_for_payload(NoDefault, "gen") + + +def test_rejects(): + for src in [ + "a - b*num_inference_steps", # top-level minus + "2 + b*num_inference_steps", # term must start with a constant + "a + a*num_inference_steps", # duplicate constant + "a + b/num_inference_steps", # constant must multiply + "a + b*max(x, y)", # calls + "a + b*3", # factor with no payload field + ]: + with pytest.raises(ValueError): + RuntimeFormula(src) + + +def test_term_values_from_struct(): + rf = RuntimeFormula("a + b*num_inference_steps + c*num_inference_steps*megapixels") + got = rf.term_values_from_struct(Payload(num_inference_steps=50, megapixels=2.0)) + assert got == { + "1": 1.0, + "num_inference_steps": 50.0, + "num_inference_steps*megapixels": 100.0, + } + # Bool fields coerce 0/1. + rf2 = RuntimeFormula("a + b*hires") + assert rf2.term_values_from_struct(Payload(hires=True)) == {"1": 1.0, "hires": 1.0} + + +def test_manifest_emits_runtime_formula(tmp_path, monkeypatch): + import textwrap + + monkeypatch.syspath_prepend(str(tmp_path)) + pkg = tmp_path / "ep_th1051" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "main.py").write_text(textwrap.dedent(""" + import msgspec + from gen_worker import RequestContext, RuntimeFormula, endpoint + + class In_(msgspec.Struct): + prompt: str = "" + num_inference_steps: int = 28 + megapixels: float = 1.0 + + class Out_(msgspec.Struct): + y: str + + @endpoint(runtime=RuntimeFormula( + "a + b*num_inference_steps + c*num_inference_steps*megapixels" + )) + def generate(ctx: RequestContext, data: In_) -> Out_: + return Out_(y="ok") + """)) + + from gen_worker.discovery.discover import discover_functions + + (fn,) = discover_functions(tmp_path, main_module="ep_th1051.main") + assert fn["runtime_formula"] == ( + "a + b*num_inference_steps + c*num_inference_steps*megapixels" + ) diff --git a/uv.lock b/uv.lock index a10931bf..7a58ae85 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "gen-worker" -version = "0.48.2" +version = "0.50.0" source = { editable = "." } dependencies = [ { name = "blake3" },