From 513c86028784f60063c3b8328c5098b3cdb746d3 Mon Sep 17 00:00:00 2001 From: GraysonWu Date: Fri, 23 Oct 2026 16:09:56 -0700 Subject: [PATCH] Avoid encoding context for cached tabular predictions Signed-off-by: GraysonWu --- .../src/sdm_client/adapters/tabular.py | 13 ++++++-- .../tests/test_envelope_and_result.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/packages/structured-data-models-client/src/sdm_client/adapters/tabular.py b/packages/structured-data-models-client/src/sdm_client/adapters/tabular.py index ce1e2bf1c..bd906529a 100644 --- a/packages/structured-data-models-client/src/sdm_client/adapters/tabular.py +++ b/packages/structured-data-models-client/src/sdm_client/adapters/tabular.py @@ -110,7 +110,9 @@ def session_prediction_plan( ) -> Iterator[dict[str, Any]]: declared = _declared_by(pinned) try: - envelope = self._envelope(request, outputs, declared) + envelope = self._envelope( + request, outputs, declared, include_context=False + ) except Exception as error: # Writing the rows to the pinned schema is what makes the call # legible to the endpoint, and it is also the only thing that can @@ -167,6 +169,8 @@ def _envelope( request: Request, outputs: tuple[str, ...], pinned_dtypes: Mapping[str, str] | None = None, + *, + include_context: bool = True, ) -> dict[str, Any]: if request.payloads is not None: raise UnsupportedRequestError( @@ -266,11 +270,16 @@ def _envelope( }, }, 'schema': {'instance_table': instance_schema}, - 'context': {'instance_table': encode_table(context, dtypes)}, 'predict': {'instance_table': encode_table(predict, dtypes)}, 'output': {'fields': list(outputs)}, 'metadata': {'request_id': _request_id()}, } + # Cached scoring keeps the schema checks above, but the context is + # already on the server. Do not encode rows that will not be sent. + if include_context: + envelope['context'] = { + 'instance_table': encode_table(context, dtypes) + } if request.inference: envelope['inference'] = json_ready(request.inference) return envelope diff --git a/packages/structured-data-models-client/tests/test_envelope_and_result.py b/packages/structured-data-models-client/tests/test_envelope_and_result.py index 18b6ea232..8932d3874 100644 --- a/packages/structured-data-models-client/tests/test_envelope_and_result.py +++ b/packages/structured-data-models-client/tests/test_envelope_and_result.py @@ -141,6 +141,38 @@ def test_a_session_pins_the_context_and_scoring_sends_only_the_rows(): assert 'context' not in scoring +@pytest.mark.parametrize('task', ['classification', 'regression']) +@pytest.mark.parametrize('with_pinned_schema', [False, True]) +def test_session_scoring_does_not_encode_the_context( + monkeypatch, task, with_pinned_schema +): + from sdm_client.adapters import tabular + + adapter = TabularAdapter() + request = request_for(task, inference={'num_context_samples': 500}) + outputs = ('prediction',) + full = next(adapter.prediction_plan(request, outputs)) + pinned = adapter.session_plan(request, outputs) + expected = {key: full[key] for key in ('predict', 'output', 'inference')} + encode = tabular.encode_table + encoded_rows = [] + + def encode_scoring_rows(frame, dtypes): + assert frame is not request.context.instance_table + encoded_rows.append(len(frame)) + return encode(frame, dtypes) + + monkeypatch.setattr(tabular, 'encode_table', encode_scoring_rows) + actual = next( + adapter.session_prediction_plan( + request, outputs, pinned if with_pinned_schema else None + ) + ) + assert encoded_rows == [len(request.predict.instance_table)] + assert actual.pop('metadata')['request_id'].startswith('sdm-') + assert actual == expected + + def test_every_request_carries_an_id_that_correlates_it_with_server_logs(): assert envelope()['metadata']['request_id'].startswith('sdm-')