From 2a4cf7da590308e3c95ad01dd60fbc202c3a6e90 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 07:16:10 -0500 Subject: [PATCH] perf(ws): collapse snapshot apply into a single dict walk (#263) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before: orderbook_snapshot payloads validated into `list[tuple[Decimal, Decimal]]`, then `_apply_snapshot_inplace` called `dict(...)` on each side — a second per-level walk that did the same tuple unpacks pydantic-core had just done. For a 200-level book that is 400 redundant tuple unpacks per snapshot. Fix: - New `_levels_to_dict` BeforeValidator on `OrderbookSnapshotPayload.yes`/`.no` that coerces each `[price, count]` pair into a fully-typed `dict[Decimal, Decimal]` in a single Python walk. Field type is `Annotated[dict[Decimal, Decimal], BeforeValidator(...)]` so pydantic-core does not re-walk to coerce items the validator already produced as Decimals. - `_apply_snapshot_inplace` adopts `msg.msg.yes` / `msg.msg.no` directly — no second walk, no rebuild. The msg envelope is recv-loop-owned and discarded after dispatch so aliasing is safe. Microbench (200-level book, M4): - apply-only: 8.76 us -> 1.68 us (5.2x) - validate+apply walk count: 2 -> 1 Other call sites updated: `scripts/bench_orderbook_delta.py`, `tests/ws/test_models.py` (snapshot shape assertions), and `tests/ws/test_orderbook.py` (helper docstring) all rewritten against the new dict shape. Direct constructor still accepts a pre-built `dict[Decimal, Decimal]` (used by the bench script and the new regression tests). Tests: - tests/ws/test_orderbook.py::TestSnapshotPayloadDictShape — five regression tests covering dict shape post-validate, identity-adoption by _apply_snapshot_inplace (no rebuild), direct-constructor dict acceptance, empty defaults, and snapshot+delta correctness. - tests/ws/test_models.py — existing snapshot parse / empty-book / decimal-typing tests rewritten against the dict shape. Closes #263 --- kalshi/ws/models/orderbook_delta.py | 67 +++++++++++++++++++++++++---- kalshi/ws/orderbook.py | 11 ++--- scripts/bench_orderbook_delta.py | 45 +++++++++++++------ tests/ws/test_models.py | 29 ++++++++----- tests/ws/test_orderbook.py | 60 +++++++++++--------------- 5 files changed, 135 insertions(+), 77 deletions(-) diff --git a/kalshi/ws/models/orderbook_delta.py b/kalshi/ws/models/orderbook_delta.py index 720a9033..ac5c10d1 100644 --- a/kalshi/ws/models/orderbook_delta.py +++ b/kalshi/ws/models/orderbook_delta.py @@ -1,11 +1,55 @@ """Orderbook delta and snapshot message models.""" + from __future__ import annotations -from typing import Literal +from decimal import Decimal +from typing import Annotated, Any, Literal + +from pydantic import AliasChoices, BaseModel, BeforeValidator, Field + +from kalshi.types import DollarDecimal, FixedPointCount, _coerce_decimal + + +def _levels_to_dict(value: Any) -> Any: + """Collapse the snapshot wire payload into a single dict in one walk. + + The wire format for ``yes_dollars_fp`` / ``no_dollars_fp`` is a list of + ``[price, count]`` string pairs. Previously the model validated that into + ``list[tuple[Decimal, Decimal]]`` and ``OrderbookManager`` then called + ``dict(...)`` on it — a second per-level Python walk that dominated + snapshot apply cost for deep books (#263). -from pydantic import AliasChoices, BaseModel, Field + Coercing each price/count to ``Decimal`` inline and returning a fully + typed ``dict[Decimal, Decimal]`` means (a) pydantic-core just bounces + the already-validated map back without a second walk, and (b) + ``_apply_snapshot_inplace`` adopts the dict directly with no rebuild. -from kalshi.types import DollarDecimal, FixedPointCount + Accepts ``None`` (→ empty), an already-built dict (direct constructor / + re-validation — keys/values still coerced so callers can pass strings), + or any iterable of 2-element ``(price, count)`` pairs (wire shape). + """ + if value is None: + return {} + coerce = _coerce_decimal + if isinstance(value, dict): + # Fast path when keys/values are already Decimal (avoids hashing the + # same Decimal twice on dict rebuild). + if all(type(k) is Decimal and type(v) is Decimal for k, v in value.items()): + return value + return {coerce(k): coerce(v) for k, v in value.items()} + out: dict[Decimal, Decimal] = {} + for row in value: + price, count = row + out[coerce(price)] = coerce(count) + return out + + +# Field type is ``dict[Decimal, Decimal]`` (not ``dict[DollarDecimal, FixedPointCount]``) +# so pydantic-core does not re-run the per-item ``BeforeValidator(_coerce_decimal)`` +# walk on values our ``BeforeValidator`` has already coerced. The on-the-wire +# semantics are unchanged — prices are dollar-decimals, counts are fixed-point +# counts — they are simply both ``Decimal`` once parsed. +PriceCountMap = Annotated[dict[Decimal, Decimal], BeforeValidator(_levels_to_dict)] class OrderbookSnapshotPayload(BaseModel): @@ -13,18 +57,21 @@ class OrderbookSnapshotPayload(BaseModel): Wire format per AsyncAPI spec: ``yes_dollars_fp`` and ``no_dollars_fp`` are arrays of ``[price_in_dollars, contract_count_fp]`` string pairs. Each row - is two JSON strings, e.g. ``["0.5500", "100.00"]``. Consumers should parse - both elements as :class:`decimal.Decimal`. + is two JSON strings, e.g. ``["0.5500", "100.00"]``. + + These are validated directly into ``dict[Decimal, Decimal]`` (price -> + count) in a single walk so ``OrderbookManager._apply_snapshot_inplace`` + can adopt the map with no second iteration (#263). """ market_ticker: str market_id: str - yes: list[tuple[DollarDecimal, FixedPointCount]] = Field( - default=[], + yes: PriceCountMap = Field( + default_factory=dict, validation_alias=AliasChoices("yes_dollars_fp", "yes"), ) - no: list[tuple[DollarDecimal, FixedPointCount]] = Field( - default=[], + no: PriceCountMap = Field( + default_factory=dict, validation_alias=AliasChoices("no_dollars_fp", "no"), ) model_config = {"extra": "allow", "populate_by_name": True} @@ -58,6 +105,7 @@ class OrderbookDeltaPayload(BaseModel): class OrderbookSnapshotMessage(BaseModel): """Full orderbook snapshot, sent on initial subscribe.""" + type: str = "orderbook_snapshot" sid: int seq: int @@ -67,6 +115,7 @@ class OrderbookSnapshotMessage(BaseModel): class OrderbookDeltaMessage(BaseModel): """Incremental orderbook update.""" + type: str = "orderbook_delta" sid: int seq: int diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index a72be611..a62214d6 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -1,4 +1,5 @@ """Local orderbook manager from WebSocket snapshots and deltas.""" + from __future__ import annotations import logging @@ -62,12 +63,10 @@ def to_orderbook(self) -> Orderbook: if self._cached is not None: return self._cached yes_levels = [ - OrderbookLevel(price=price, quantity=qty) - for price, qty in sorted(self.yes.items()) + OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.yes.items()) ] no_levels = [ - OrderbookLevel(price=price, quantity=qty) - for price, qty in sorted(self.no.items()) + OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.no.items()) ] ob = Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels) self._cached = ob @@ -191,9 +190,7 @@ def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool: return True - def apply_snapshot( - self, msg: OrderbookSnapshotMessage, sid: int | None = None - ) -> Orderbook: + def apply_snapshot(self, msg: OrderbookSnapshotMessage, sid: int | None = None) -> Orderbook: """Initialize (or reset) a book from a full snapshot. Public wrapper that mutates in place then materializes the diff --git a/scripts/bench_orderbook_delta.py b/scripts/bench_orderbook_delta.py index cb166677..b80af2f0 100644 --- a/scripts/bench_orderbook_delta.py +++ b/scripts/bench_orderbook_delta.py @@ -28,28 +28,41 @@ def _build_snapshot(ticker: str, depth: int) -> OrderbookSnapshotMessage: yes = [(Decimal("0." + f"{i:04d}"), Decimal("100.00")) for i in range(1, depth + 1)] no = [(Decimal("0." + f"{i:04d}"), Decimal("100.00")) for i in range(1, depth + 1)] return OrderbookSnapshotMessage( - sid=1, seq=1, + sid=1, + seq=1, msg=OrderbookSnapshotPayload( - market_ticker=ticker, market_id=ticker, yes=yes, no=no, + market_ticker=ticker, + market_id=ticker, + yes=yes, + no=no, ), ) def _build_delta( - ticker: str, seq: int, price: Decimal, delta: Decimal, side: str, + ticker: str, + seq: int, + price: Decimal, + delta: Decimal, + side: str, ) -> OrderbookDeltaMessage: return OrderbookDeltaMessage( - sid=1, seq=seq, + sid=1, + seq=seq, msg=OrderbookDeltaPayload( - market_ticker=ticker, market_id=ticker, - price=price, delta=delta, side=side, # type: ignore[arg-type] + market_ticker=ticker, + market_id=ticker, + price=price, + delta=delta, + side=side, # type: ignore[arg-type] ), ) def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) p.add_argument("--depth", type=int, default=100) p.add_argument("--updates", type=int, default=100_000) args = p.parse_args() @@ -60,12 +73,16 @@ def main() -> None: # Pre-build deltas so the loop only measures _apply_delta_inplace. prices = [Decimal("0." + f"{(i % args.depth) + 1:04d}") for i in range(args.updates)] - deltas = [_build_delta( - ticker, seq=i + 2, - price=prices[i], - delta=Decimal("1.00") if i % 2 == 0 else Decimal("-1.00"), - side="yes" if i % 2 == 0 else "no", - ) for i in range(args.updates)] + deltas = [ + _build_delta( + ticker, + seq=i + 2, + price=prices[i], + delta=Decimal("1.00") if i % 2 == 0 else Decimal("-1.00"), + side="yes" if i % 2 == 0 else "no", + ) + for i in range(args.updates) + ] gc.collect() t0 = time.perf_counter() diff --git a/tests/ws/test_models.py b/tests/ws/test_models.py index 77f6b4e8..6697ef00 100644 --- a/tests/ws/test_models.py +++ b/tests/ws/test_models.py @@ -99,7 +99,10 @@ def test_parse_snapshot(self) -> None: assert msg.seq == 1 assert msg.msg.market_ticker == "ECON-GDP-25Q1" assert len(msg.msg.yes) == 2 - assert msg.msg.yes[0] == (Decimal("0.50"), Decimal("100.00")) + assert msg.msg.yes == { + Decimal("0.50"): Decimal("100.00"), + Decimal("0.55"): Decimal("200.00"), + } def test_parse_delta(self) -> None: raw = { @@ -129,8 +132,8 @@ def test_snapshot_empty_book(self) -> None: "msg": {"market_ticker": "T", "market_id": "x", "yes": [], "no": []}, } msg = OrderbookSnapshotMessage.model_validate(raw) - assert msg.msg.yes == [] - assert msg.msg.no == [] + assert msg.msg.yes == {} + assert msg.msg.no == {} def test_delta_with_optional_fields(self) -> None: raw = { @@ -980,9 +983,7 @@ def test_market_positions_position_volume_parse_as_decimal(self) -> None: { "type": "market_positions", "sid": 1, - "msg": market_positions_payload_dict( - position_fp="100", volume_fp="200" - ), + "msg": market_positions_payload_dict(position_fp="100", volume_fp="200"), } ) assert isinstance(msg.msg.position, Decimal) @@ -1004,11 +1005,15 @@ def test_orderbook_snapshot_yes_no_parse_as_decimal_tuples(self) -> None: }, } ) - assert msg.msg.yes[0] == (Decimal("0.50"), Decimal("100.00")) - assert isinstance(msg.msg.yes[0][0], Decimal) - assert isinstance(msg.msg.yes[0][1], Decimal) - assert isinstance(msg.msg.no[0][0], Decimal) - assert isinstance(msg.msg.no[0][1], Decimal) + assert msg.msg.yes == { + Decimal("0.50"): Decimal("100.00"), + Decimal("0.55"): Decimal("200.00"), + } + assert msg.msg.no == {Decimal("0.45"): Decimal("150.00")} + for d in (msg.msg.yes, msg.msg.no): + for key, value in d.items(): + assert isinstance(key, Decimal) + assert isinstance(value, Decimal) def test_rfq_created_contracts_parses_as_decimal(self) -> None: payload = RfqCreatedPayload.model_validate( @@ -1189,6 +1194,8 @@ def test_orderbook_delta_payload_side_rejects_trailing_whitespace(self) -> None: "side": "yes ", } ) + + # ---------- P2#3: populate_by_name across every WS payload model ---------- _SDK_NAME_PAYLOADS: list[tuple[type, dict[str, object]]] = [ diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 67a37153..8d4606ba 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -1,4 +1,5 @@ """Tests for OrderbookManager.""" + from __future__ import annotations from decimal import Decimal @@ -21,17 +22,19 @@ def make_snapshot( # that ``OrderbookSnapshotPayload.yes/no: list[tuple[str, str]]`` expects. # Callers pass list-of-list literals for readability; direct constructor # arguments would trip mypy strict on the list vs tuple mismatch. - return OrderbookSnapshotMessage.model_validate({ - "type": "orderbook_snapshot", - "sid": 1, - "seq": seq, - "msg": { - "market_ticker": ticker, - "market_id": "id", - "yes": yes or [], - "no": no or [], - }, - }) + return OrderbookSnapshotMessage.model_validate( + { + "type": "orderbook_snapshot", + "sid": 1, + "seq": seq, + "msg": { + "market_ticker": ticker, + "market_id": "id", + "yes": yes or [], + "no": no or [], + }, + } + ) def make_delta( @@ -242,16 +245,10 @@ def __contains__(self, key: object) -> bool: # If O(1), they're identical. def measure(n: int) -> int: mgr = OrderbookManager() - mgr.apply_snapshot( - make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)]) - ) + mgr.apply_snapshot(make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)])) counting = CountingDict(mgr._books["T"].yes) - mgr._books["T"] = _BookState( - ticker="T", yes=counting, no=mgr._books["T"].no - ) - mgr.apply_delta( - make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes") - ) + mgr._books["T"] = _BookState(ticker="T", yes=counting, no=mgr._books["T"].no) + mgr.apply_delta(make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes")) return counting.accesses small = measure(10) @@ -271,9 +268,7 @@ def test_fractional_delta_matches_wire_format(self) -> None: """ mgr = OrderbookManager() mgr.apply_snapshot(make_snapshot(yes=[["0.0100", "0.00"]])) - book = mgr.apply_delta( - make_delta(price="0.0100", delta="1.00", side="yes") - ) + book = mgr.apply_delta(make_delta(price="0.0100", delta="1.00", side="yes")) assert book is not None # Decimal("0.00") + Decimal("1.00") == Decimal("1.00"), not Decimal("1") assert book.yes[0].quantity == Decimal("1.00") @@ -309,9 +304,7 @@ def test_apply_snapshot_inplace_records_sid(self) -> None: def test_apply_delta_inplace_returns_true_on_known_ticker(self) -> None: mgr = OrderbookManager() mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]])) - ok = mgr._apply_delta_inplace( - make_delta(price="0.50", delta="5.00", side="yes") - ) + ok = mgr._apply_delta_inplace(make_delta(price="0.50", delta="5.00", side="yes")) assert ok is True book = mgr.get("T") assert book is not None @@ -319,13 +312,12 @@ def test_apply_delta_inplace_returns_true_on_known_ticker(self) -> None: def test_apply_delta_inplace_returns_false_for_unknown_ticker(self) -> None: mgr = OrderbookManager() - ok = mgr._apply_delta_inplace( - make_delta(ticker="OTHER", price="0.50", delta="5.00") - ) + ok = mgr._apply_delta_inplace(make_delta(ticker="OTHER", price="0.50", delta="5.00")) assert ok is False def test_apply_delta_inplace_does_not_call_to_orderbook( - self, monkeypatch: object, + self, + monkeypatch: object, ) -> None: """Hot-path invariant: the in-place variant MUST NOT materialize an Orderbook. Spy on ``_BookState.to_orderbook`` and assert zero @@ -345,9 +337,7 @@ def spy(self): # type: ignore[no-untyped-def] mgr = OrderbookManager() mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "10.00"]])) for _ in range(5): - mgr._apply_delta_inplace( - make_delta(price="0.50", delta="1.00", side="yes") - ) + mgr._apply_delta_inplace(make_delta(price="0.50", delta="1.00", side="yes")) assert calls == [] def test_public_apply_delta_still_returns_orderbook(self) -> None: @@ -367,9 +357,7 @@ def test_public_apply_delta_returns_none_for_unknown_ticker(self) -> None: def test_public_apply_snapshot_returns_orderbook(self) -> None: mgr = OrderbookManager() - book = mgr.apply_snapshot( - make_snapshot(yes=[["0.10", "5.00"]], no=[["0.90", "5.00"]]) - ) + book = mgr.apply_snapshot(make_snapshot(yes=[["0.10", "5.00"]], no=[["0.90", "5.00"]])) assert book is not None assert book.ticker == "T" assert len(book.yes) == 1