diff --git a/src/livepeer_gateway/byoc.py b/src/livepeer_gateway/byoc.py index 4d27781..09961c3 100644 --- a/src/livepeer_gateway/byoc.py +++ b/src/livepeer_gateway/byoc.py @@ -144,6 +144,36 @@ def audio_url(self) -> Optional[str]: # Header building # --------------------------------------------------------------------------- +# Keys, in priority order, under which a BYOC request conventionally carries +# the specific provider model id (distinct from the coarse capability name). +_MODEL_ID_KEYS = ("model_id", "model") + + +def _extract_model_id( + payload: Optional[dict[str, Any]], + parameters: Optional[dict[str, Any]] = None, +) -> str: + """Best-effort extraction of the provider model id from a BYOC request. + + The *model id* is the specific provider model the worker runs (e.g. the + value billed/attributed per generation), as opposed to the *capability* + name (e.g. ``nano-banana``) which selects the orchestrator pipeline. It is + conventionally carried in the request ``payload`` under ``model_id`` or + ``model``; some capabilities place it in the job ``parameters`` instead. + + Returns an empty string when no model id can be determined so callers can + fall back to today's behavior (the signer then derives/defaults it). + """ + for source in (payload, parameters): + if not isinstance(source, dict): + continue + for key in _MODEL_ID_KEYS: + value = source.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + def _create_byoc_payment( *, orch_origin: str, @@ -151,6 +181,7 @@ def _create_byoc_payment( livepeer_hdr: str, signer_url: str, signer_headers: Optional[dict[str, str]] = None, + model_id: str = "", timeout: float = 30.0, ) -> dict[str, str]: """ @@ -194,11 +225,27 @@ def _create_byoc_payment( signer_origin = _http_origin(signer_url) payment_url = f"{signer_origin}/generate-live-payment" - payment_body = json.dumps({ + # `type` stays "lv2v" so the remote signer keeps reusing the live + # fee/pixel routing for BYOC jobs (BYOC has no per-job pixel count). The + # additive `capability` + `model_id` fields below carry the REAL usage + # attribution so the signer's `create_signed_ticket` metering event + # records the true pipeline + model instead of `live-video-to-video` / + # `unknown`. Both are backward-compatible: an older signer ignores the + # extra fields, and an empty `model_id` is simply omitted so behavior is + # byte-identical to today. + payment_payload: dict[str, Any] = { "orchestrator": orch_info_b64, "type": "lv2v", "capability": capability, - }).encode("utf-8") + } + # Defensively strip (and ignore non-strings): callers normally pass a value + # already cleaned by `_extract_model_id`, but a direct caller could pass a + # whitespace-only or non-string value, which must still be omitted so the + # wire body stays byte-identical to today and the signer falls back. + model_id = model_id.strip() if isinstance(model_id, str) else "" + if model_id: + payment_payload["model_id"] = model_id + payment_body = json.dumps(payment_payload).encode("utf-8") payment_headers = { "Content-Type": "application/json", "Livepeer-Capability": capability, @@ -390,6 +437,7 @@ def submit_byoc_job( livepeer_hdr=livepeer_hdr, signer_url=signer_url, signer_headers=signer_headers, + model_id=_extract_model_id(req.payload, req.parameters), timeout=http_timeout, ) headers.update(payment_headers) @@ -673,6 +721,11 @@ def submit_training_job( livepeer_hdr=livepeer_hdr, signer_url=signer_url, signer_headers=signer_headers, + # Attribute payment to the explicit top-level training + # model_id (matching the orchestrator body above), not the + # merged payload where a `model_id` inside req.params could + # otherwise override it and diverge from the real model. + model_id=_extract_model_id({"model_id": req.model_id}, req.params), timeout=http_timeout, ) headers.update(payment_headers) diff --git a/tests/test_byoc_usage_attribution.py b/tests/test_byoc_usage_attribution.py new file mode 100644 index 0000000..08f7ad3 --- /dev/null +++ b/tests/test_byoc_usage_attribution.py @@ -0,0 +1,274 @@ +""" +Unit tests for BYOC usage attribution (real capability + model_id). + +Root cause being fixed: ``_create_byoc_payment`` used to send only a bare +``capability`` field (which the remote signer dropped) and no model id, so the +signer's ``create_signed_ticket`` metering event recorded +``pipeline=live-video-to-video`` / ``model_id=unknown`` for every BYOC job. + +These tests prove the additive wire contract on the gateway side: +- the ``/generate-live-payment`` body now carries the REAL ``capability`` and + ``model_id`` for a representative BYOC capability (e.g. ``nano-banana``); +- ``type`` stays ``"lv2v"`` (fee/pixel routing is unchanged); +- when no model id can be determined the field is omitted (backward-compatible, + byte-identical to today so the signer falls back). +""" +from __future__ import annotations + +import json +from contextlib import contextmanager +from unittest.mock import MagicMock, patch +from urllib.request import Request + +from livepeer_gateway.byoc import ( + ByocJobRequest, + _create_byoc_payment, + _extract_model_id, +) + + +def _stub_orch_info(): + info = MagicMock() + tp = MagicMock() + tp.face_value = b"\x01\x00" # non-zero → payment required + info.ticket_params = tp + info.HasField = lambda field: field == "ticket_params" + info.SerializeToString = lambda: b"stub-orch-info-protobuf" + return info + + +@contextmanager +def _mock_signer(*, payment_body=b'{"payment":"TICKETS","segCreds":"SEG"}'): + """Mock the signer HTTP call and capture the request sent to it.""" + captured: list[Request] = [] + + class _MockResponse: + def __init__(self, body: bytes): + self._body = body + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + return _MockResponse(payment_body) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + yield captured + + +def _signer_body(captured: list[Request]) -> dict: + signer_reqs = [ + r for r in captured if "generate-live-payment" in r.full_url + ] + assert len(signer_reqs) == 1, f"expected 1 signer call, got {len(signer_reqs)}" + return json.loads(signer_reqs[0].data.decode("utf-8")) + + +# --------------------------------------------------------------------------- +# _extract_model_id +# --------------------------------------------------------------------------- + + +def test_extract_model_id_from_payload_model_id_key(): + assert _extract_model_id({"model_id": "google/nano-banana"}) == "google/nano-banana" + + +def test_extract_model_id_from_payload_model_key(): + assert _extract_model_id({"prompt": "a cat", "model": "flux-dev"}) == "flux-dev" + + +def test_extract_model_id_prefers_payload_over_parameters(): + assert ( + _extract_model_id({"model_id": "from-payload"}, {"model_id": "from-params"}) + == "from-payload" + ) + + +def test_extract_model_id_falls_back_to_parameters(): + assert _extract_model_id({"prompt": "x"}, {"model": "from-params"}) == "from-params" + + +def test_extract_model_id_empty_when_absent(): + assert _extract_model_id({"prompt": "no model here"}) == "" + assert _extract_model_id(None, None) == "" + assert _extract_model_id({"model_id": " "}) == "" # whitespace-only ignored + + +# --------------------------------------------------------------------------- +# _create_byoc_payment wire contract +# --------------------------------------------------------------------------- + + +def test_payment_body_carries_real_capability_and_model_id(): + """A representative BYOC capability sends the REAL capability + model_id.""" + with _mock_signer() as captured: + result = _create_byoc_payment( + orch_origin="https://orch.test:8936", + capability="nano-banana", + livepeer_hdr="", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + model_id="google/nano-banana", + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert body["model_id"] == "google/nano-banana" + # type stays lv2v so fee/pixel routing is unchanged + assert body["type"] == "lv2v" + assert body["orchestrator"] # base64 orch info present + # payment headers still returned as before + assert result["Livepeer-Payment"] == "TICKETS" + assert result["Livepeer-Segment"] == "SEG" + + +def test_payment_body_omits_model_id_when_absent(): + """Backward-compatible: no model id → field omitted (signer falls back).""" + with _mock_signer() as captured: + _create_byoc_payment( + orch_origin="https://orch.test:8936", + capability="nano-banana", + livepeer_hdr="", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + # model_id omitted → defaults to "" + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert "model_id" not in body, "empty model_id must be omitted for byte-identical body" + assert body["type"] == "lv2v" + + +# --------------------------------------------------------------------------- +# End-to-end through submit_byoc_job +# --------------------------------------------------------------------------- + + +def test_submit_byoc_job_threads_model_id_from_payload(): + """submit_byoc_job extracts the model id from the request payload and + forwards it (and the capability) to the signer payment request.""" + from livepeer_gateway import byoc as byoc_mod + + captured: list[Request] = [] + + class _MockResponse: + def __init__(self, body: bytes, status: int = 200): + self._body = body + self.status = status + self.headers = {} + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + url = req.full_url if hasattr(req, "full_url") else req.get_full_url() + if "generate-live-payment" in url: + return _MockResponse(b'{"payment":"TICKETS","segCreds":"SEG"}') + if "sign-byoc-job" in url: + return _MockResponse(b'{"sender":"0xabc","signature":"0xsig"}') + # orchestrator /process/request/ + return _MockResponse(b'{"images":[{"url":"https://img"}]}', 200) + + req = ByocJobRequest( + capability="nano-banana", + payload={"prompt": "a dragon", "model_id": "google/nano-banana"}, + ) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + byoc_mod.submit_byoc_job( + req, + orch_url="https://orch.test:8936", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert body["model_id"] == "google/nano-banana" + assert body["type"] == "lv2v" + + +# --------------------------------------------------------------------------- +# End-to-end through submit_training_job +# --------------------------------------------------------------------------- + + +def test_submit_training_job_uses_top_level_model_id_over_params(): + """submit_training_job must attribute payment to the explicit top-level + training ``model_id`` (which also builds the orchestrator body), not a + ``model_id`` nested in ``params`` that would otherwise win the payload + merge and send the wrong model to the signer.""" + from livepeer_gateway import byoc as byoc_mod + from livepeer_gateway.byoc import ByocTrainingRequest + + captured: list[Request] = [] + + class _MockResponse: + def __init__(self, body: bytes, status: int = 200): + self._body = body + self.status = status + self.headers = {} + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + url = req.full_url if hasattr(req, "full_url") else req.get_full_url() + if "generate-live-payment" in url: + return _MockResponse(b'{"payment":"TICKETS","segCreds":"SEG"}') + if "sign-byoc-job" in url: + return _MockResponse(b'{"sender":"0xabc","signature":"0xsig"}') + # orchestrator /process/train/ + return _MockResponse(b'{"job_id":"train-123","status":"submitted"}', 200) + + # Top-level training model differs from a stray params["model_id"]. + req = ByocTrainingRequest( + capability="flux-lora-trainer", + model_id="fal-ai/flux-lora-fast-training", + params={"model_id": "wrong/params-model", "trigger_word": "TOK"}, + ) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + byoc_mod.submit_training_job( + req, + orch_url="https://orch.test:8936", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + body = _signer_body(captured) + assert body["capability"] == "flux-lora-trainer" + # The explicit training model_id wins — NOT the nested params value. + assert body["model_id"] == "fal-ai/flux-lora-fast-training" + assert body["type"] == "lv2v"