From 7c730ce070bb2ef67867cd5aa3461d266ccc6212 Mon Sep 17 00:00:00 2001 From: DG-git-dev Date: Wed, 23 Sep 2026 17:44:59 +0800 Subject: [PATCH 1/3] Fix order-model serving wire format and request bounds --- README.md | 5 ++ market_simulation/rollout/model_client.py | 15 +++-- .../rollout/order_model_serving.py | 28 ++++++--- market_simulation/rollout/wire_format.py | 39 ++++++++++++ ray_serving.yaml | 2 +- tests/test_wire_format.py | 62 +++++++++++++++++++ 6 files changed, 135 insertions(+), 16 deletions(-) create mode 100644 market_simulation/rollout/wire_format.py create mode 100644 tests/test_wire_format.py diff --git a/README.md b/README.md index fad3712..4419206 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,11 @@ To start the order model Ray server: bash scripts/start-order-model.sh ``` +The order-model HTTP endpoint listens on localhost by default. Expose it to +remote clients only behind network access controls and authentication. The +client and server must use the same version: requests and responses use a +fixed-size little-endian int32 binary format. + > **Prerequisites:** > - The Ray server must be running and accessible at the configured IP and port > - Sufficient computational resources are required to run the model diff --git a/market_simulation/rollout/model_client.py b/market_simulation/rollout/model_client.py index 1a71e51..9a08549 100644 --- a/market_simulation/rollout/model_client.py +++ b/market_simulation/rollout/model_client.py @@ -1,24 +1,27 @@ from __future__ import annotations import logging -import pickle from typing import TYPE_CHECKING +import numpy as np import ray import requests +from market_simulation.conf import C +from market_simulation.rollout.wire_format import decode_int32, encode_int32 + if TYPE_CHECKING: - import numpy as np import numpy.typing as npt @ray.remote def send_query(url: str, state: npt.NDArray[np.int32]) -> npt.NDArray[np.int32]: """Send query to serving model.""" - pdata = pickle.dumps(state) - response = requests.post(url, data=pdata) - output = pickle.loads(response.content) - return output + expected_elements = C.order_model.seq_len * C.order_model.token_dim + pdata = encode_int32(state, expected_elements) + response = requests.post(url, data=pdata, headers={"Content-Type": "application/octet-stream"}) + response.raise_for_status() + return decode_int32(response.content, 1) class ModelClient: diff --git a/market_simulation/rollout/order_model_serving.py b/market_simulation/rollout/order_model_serving.py index f12a723..6d48a20 100644 --- a/market_simulation/rollout/order_model_serving.py +++ b/market_simulation/rollout/order_model_serving.py @@ -1,15 +1,16 @@ from __future__ import annotations import logging -import pickle import numpy as np import numpy.typing as npt import torch from ray import serve +from starlette.responses import Response from market_simulation.conf import C from market_simulation.models.order_model import OrderModel +from market_simulation.rollout.wire_format import PayloadTooLarge, decode_int32, encode_int32, read_fixed_body @serve.deployment( @@ -35,7 +36,7 @@ def _load_model(self) -> OrderModel: return order_model @serve.batch(max_batch_size=C.model_serving.max_batch_size) # type: ignore - async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[npt.NDArray[np.int32]]: + async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[bytes]: """Batch inference.""" batch_size = len(requests) input_tensor = torch.from_numpy(np.asarray(requests)).cuda() @@ -45,18 +46,27 @@ async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[n output_tensor: np.ndarray = self.model.sample(input_tensor, self.temperature).int().cpu().reshape((batch_size, -1)).numpy() logging.info(f"output shape: {output_tensor.shape}") - results: list[npt.NDArray[np.int32]] = [] + results: list[bytes] = [] for i in range(batch_size): output = output_tensor[i] - arr = np.array([output], dtype=np.int32) - results.append(pickle.dumps(arr)) # type: ignore + results.append(encode_int32(output, 1)) return results - async def __call__(self, request) -> list[npt.NDArray[np.int32]]: # noqa: ANN001 + async def __call__(self, request) -> Response: # noqa: ANN001 """Handle request.""" - request_bytes = await request.body() - arr = pickle.loads(request_bytes) - return await self.batch_inference(arr) # type: ignore + expected_elements = C.order_model.seq_len * C.order_model.token_dim + expected_bytes = expected_elements * np.dtype(np.int32).itemsize + if request.headers.get("content-type") != "application/octet-stream": + return Response("Expected application/octet-stream", status_code=415) + try: + request_bytes = await read_fixed_body(request.stream(), expected_bytes) + arr = decode_int32(request_bytes, expected_elements) + except PayloadTooLarge as error: + return Response(str(error), status_code=413) + except ValueError as error: + return Response(str(error), status_code=400) + result = await self.batch_inference(arr) # type: ignore + return Response(result, media_type="application/octet-stream") order_model_app = OrderModelServing.bind() # type: ignore diff --git a/market_simulation/rollout/wire_format.py b/market_simulation/rollout/wire_format.py new file mode 100644 index 0000000..b06df72 --- /dev/null +++ b/market_simulation/rollout/wire_format.py @@ -0,0 +1,39 @@ +"""Fixed-size int32 wire format for order-model requests and responses.""" + +from collections.abc import AsyncIterable + +import numpy as np +import numpy.typing as npt + +WIRE_DTYPE = np.dtype(" bytes: + """Read a request body without buffering more than the expected size.""" + body = bytearray() + async for chunk in chunks: + if len(body) + len(chunk) > expected_size: + raise PayloadTooLarge("Order-model request is too large") + body.extend(chunk) + if len(body) != expected_size: + raise ValueError("Invalid int32 payload length") + return bytes(body) + + +def encode_int32(values: npt.NDArray[np.int32], expected_elements: int) -> bytes: + """Encode exactly the expected number of int32 values in little-endian order.""" + array = np.asarray(values) + if array.dtype != np.dtype(np.int32) or array.size != expected_elements: + raise ValueError("Expected a fixed-size int32 array") + return array.astype(WIRE_DTYPE, copy=False).tobytes(order="C") + + +def decode_int32(payload: bytes, expected_elements: int) -> npt.NDArray[np.int32]: + """Decode a fixed-size int32 payload without object deserialization.""" + if len(payload) != expected_elements * WIRE_DTYPE.itemsize: + raise ValueError("Invalid int32 payload length") + return np.frombuffer(payload, dtype=WIRE_DTYPE).astype(np.int32, copy=True) diff --git a/ray_serving.yaml b/ray_serving.yaml index 620d111..34ea77f 100644 --- a/ray_serving.yaml +++ b/ray_serving.yaml @@ -1,7 +1,7 @@ proxy_location: EveryNode http_options: - host: 0.0.0.0 + host: 127.0.0.1 port: 8000 diff --git a/tests/test_wire_format.py b/tests/test_wire_format.py new file mode 100644 index 0000000..3b9e57e --- /dev/null +++ b/tests/test_wire_format.py @@ -0,0 +1,62 @@ +"""Checks for the order-model network format and bounded request reader.""" + +import asyncio +import pickle +import unittest + +import numpy as np + +from market_simulation.rollout.wire_format import PayloadTooLarge, decode_int32, encode_int32, read_fixed_body + + +async def chunks(*parts: bytes): + """Yield a request body in arbitrary network chunks.""" + for part in parts: + yield part + + +def fail_if_unpickled() -> None: + """Make an unsafe deserialization observable to the test.""" + raise AssertionError("Untrusted pickle was executed") + + +class PickleTrap: + """A payload that must remain inert on the wire.""" + + def __reduce__(self): + """Return a callable that fails if deserialized.""" + return fail_if_unpickled, () + + +class WireFormatTests(unittest.TestCase): + """Exercise fixed-size int32 transport at the network boundary.""" + + def test_round_trip(self) -> None: + values = np.array([1, -2, 123456], dtype=np.int32) + payload = encode_int32(values, 3) + self.assertEqual(payload, b"\x01\x00\x00\x00\xfe\xff\xff\xff\x40\xe2\x01\x00") + np.testing.assert_array_equal(decode_int32(payload, 3), values) + + def test_rejects_wrong_type_and_size(self) -> None: + with self.assertRaises(ValueError): + encode_int32(np.array([1], dtype=object), 1) + with self.assertRaises(ValueError): + encode_int32(np.array([1], dtype=np.int32), 2) + with self.assertRaises(ValueError): + decode_int32(b"\x00" * 3, 1) + + def test_stream_is_bounded_and_exact(self) -> None: + self.assertEqual(asyncio.run(read_fixed_body(chunks(b"ab", b"cd"), 4)), b"abcd") + with self.assertRaises(PayloadTooLarge): + asyncio.run(read_fixed_body(chunks(b"abcd", b"e" * 1000000), 4)) + with self.assertRaises(ValueError): + asyncio.run(read_fixed_body(chunks(b"abc"), 4)) + + def test_pickle_payload_is_never_deserialized(self) -> None: + payload = pickle.dumps(PickleTrap()).ljust(64, b"\x00") + decoded = decode_int32(payload, 16) + self.assertEqual(decoded.size, 16) + + +if __name__ == "__main__": + unittest.main() From b9590de6efde7e1890b49a65ccd25d388a31ddf4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:25:20 +0000 Subject: [PATCH 2/3] Fix Ruff lint failures in wire-format changes Co-authored-by: DG-git-dev <78746494+DG-git-dev@users.noreply.github.com> --- market_simulation/rollout/model_client.py | 2 +- .../rollout/order_model_serving.py | 4 +- market_simulation/rollout/wire_format.py | 4 +- tests/__init__.py | 1 + tests/test_wire_format.py | 41 +++++++++++-------- 5 files changed, 31 insertions(+), 21 deletions(-) create mode 100644 tests/__init__.py diff --git a/market_simulation/rollout/model_client.py b/market_simulation/rollout/model_client.py index 9a08549..70a3480 100644 --- a/market_simulation/rollout/model_client.py +++ b/market_simulation/rollout/model_client.py @@ -3,7 +3,6 @@ import logging from typing import TYPE_CHECKING -import numpy as np import ray import requests @@ -11,6 +10,7 @@ from market_simulation.rollout.wire_format import decode_int32, encode_int32 if TYPE_CHECKING: + import numpy as np import numpy.typing as npt diff --git a/market_simulation/rollout/order_model_serving.py b/market_simulation/rollout/order_model_serving.py index 6d48a20..36470e0 100644 --- a/market_simulation/rollout/order_model_serving.py +++ b/market_simulation/rollout/order_model_serving.py @@ -10,7 +10,7 @@ from market_simulation.conf import C from market_simulation.models.order_model import OrderModel -from market_simulation.rollout.wire_format import PayloadTooLarge, decode_int32, encode_int32, read_fixed_body +from market_simulation.rollout.wire_format import PayloadTooLargeError, decode_int32, encode_int32, read_fixed_body @serve.deployment( @@ -61,7 +61,7 @@ async def __call__(self, request) -> Response: # noqa: ANN001 try: request_bytes = await read_fixed_body(request.stream(), expected_bytes) arr = decode_int32(request_bytes, expected_elements) - except PayloadTooLarge as error: + except PayloadTooLargeError as error: return Response(str(error), status_code=413) except ValueError as error: return Response(str(error), status_code=400) diff --git a/market_simulation/rollout/wire_format.py b/market_simulation/rollout/wire_format.py index b06df72..d2c79ac 100644 --- a/market_simulation/rollout/wire_format.py +++ b/market_simulation/rollout/wire_format.py @@ -8,7 +8,7 @@ WIRE_DTYPE = np.dtype(" b body = bytearray() async for chunk in chunks: if len(body) + len(chunk) > expected_size: - raise PayloadTooLarge("Order-model request is too large") + raise PayloadTooLargeError("Order-model request is too large") body.extend(chunk) if len(body) != expected_size: raise ValueError("Invalid int32 payload length") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..7744664 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for MarS.""" diff --git a/tests/test_wire_format.py b/tests/test_wire_format.py index 3b9e57e..0ae87eb 100644 --- a/tests/test_wire_format.py +++ b/tests/test_wire_format.py @@ -3,13 +3,14 @@ import asyncio import pickle import unittest +from collections.abc import AsyncIterator, Callable import numpy as np -from market_simulation.rollout.wire_format import PayloadTooLarge, decode_int32, encode_int32, read_fixed_body +from market_simulation.rollout.wire_format import PayloadTooLargeError, decode_int32, encode_int32, read_fixed_body -async def chunks(*parts: bytes): +async def chunks(*parts: bytes) -> AsyncIterator[bytes]: """Yield a request body in arbitrary network chunks.""" for part in parts: yield part @@ -23,7 +24,7 @@ def fail_if_unpickled() -> None: class PickleTrap: """A payload that must remain inert on the wire.""" - def __reduce__(self): + def __reduce__(self) -> tuple[Callable[[], None], tuple[()]]: """Return a callable that fails if deserialized.""" return fail_if_unpickled, () @@ -32,30 +33,38 @@ class WireFormatTests(unittest.TestCase): """Exercise fixed-size int32 transport at the network boundary.""" def test_round_trip(self) -> None: + """Encode and decode fixed-width little-endian int32 payloads.""" values = np.array([1, -2, 123456], dtype=np.int32) payload = encode_int32(values, 3) - self.assertEqual(payload, b"\x01\x00\x00\x00\xfe\xff\xff\xff\x40\xe2\x01\x00") + assert payload == b"\x01\x00\x00\x00\xfe\xff\xff\xff\x40\xe2\x01\x00" np.testing.assert_array_equal(decode_int32(payload, 3), values) def test_rejects_wrong_type_and_size(self) -> None: - with self.assertRaises(ValueError): - encode_int32(np.array([1], dtype=object), 1) - with self.assertRaises(ValueError): - encode_int32(np.array([1], dtype=np.int32), 2) - with self.assertRaises(ValueError): - decode_int32(b"\x00" * 3, 1) + """Reject payloads with an unexpected dtype, element count, or byte size.""" + expect_raises(ValueError, encode_int32, np.array([1], dtype=object), 1) + expect_raises(ValueError, encode_int32, np.array([1], dtype=np.int32), 2) + expect_raises(ValueError, decode_int32, b"\x00" * 3, 1) def test_stream_is_bounded_and_exact(self) -> None: - self.assertEqual(asyncio.run(read_fixed_body(chunks(b"ab", b"cd"), 4)), b"abcd") - with self.assertRaises(PayloadTooLarge): - asyncio.run(read_fixed_body(chunks(b"abcd", b"e" * 1000000), 4)) - with self.assertRaises(ValueError): - asyncio.run(read_fixed_body(chunks(b"abc"), 4)) + """Read exactly the expected stream length and reject over/underflow.""" + assert asyncio.run(read_fixed_body(chunks(b"ab", b"cd"), 4)) == b"abcd" + expect_raises(PayloadTooLargeError, asyncio.run, read_fixed_body(chunks(b"abcd", b"e" * 1000000), 4)) + expect_raises(ValueError, asyncio.run, read_fixed_body(chunks(b"abc"), 4)) def test_pickle_payload_is_never_deserialized(self) -> None: + """Treat pickle bytes as inert int32 wire data.""" payload = pickle.dumps(PickleTrap()).ljust(64, b"\x00") decoded = decode_int32(payload, 16) - self.assertEqual(decoded.size, 16) + assert decoded.size == 16 + + +def expect_raises(exception_type: type[Exception], function: Callable[..., object], *args: object) -> None: + """Assert that a callable raises the expected exception type.""" + try: + function(*args) + except exception_type: + return + raise AssertionError(f"Expected {exception_type.__name__}") if __name__ == "__main__": From 45201b3ebd8b41709ecd6c5a991a8ea55016213c Mon Sep 17 00:00:00 2001 From: DG-git-dev Date: Wed, 23 Sep 2026 18:33:00 +0800 Subject: [PATCH 3/3] Pin Ruff action to repository lint version --- .github/workflows/linting.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml index 38e55b8..d718cb4 100644 --- a/.github/workflows/linting.yml +++ b/.github/workflows/linting.yml @@ -38,6 +38,8 @@ jobs: - name: Run Ruff uses: chartboost/ruff-action@v1 + with: + version: "0.11.2" - name: Run ShellCheck uses: ludeeus/action-shellcheck@master