Skip to content

Commit 87c751f

Browse files
committed
Enable gradient explanations for kumo-tabular and kumo-relational
Allow the explanation output for kumo-tabular (one scored row), return it on result.explanation, and drop the old 1,000-row context rule for relational explanations, which the SDM NIM does not have. Signed-off-by: Jana Gagacheva <jgagacheva@nvidia.com>
1 parent 5aa6620 commit 87c751f

13 files changed

Lines changed: 265 additions & 88 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,18 @@ these packages. All five distributions share one version and one entry.
77

88
## Unreleased
99

10+
### Added
11+
12+
- `kumo-tabular` accepts `explanation` in `outputs`, over one scored row at a
13+
time, and returns it on `result.explanation`. Asking for one over more rows
14+
is refused before sending.
15+
1016
### Changed
1117

18+
- A `kumo-relational` explanation is no longer refused above the 1,000-row
19+
context budget. That limit belonged to an earlier deployment; the Kumo
20+
Relational SDM NIM explains under either budget.
21+
1222
- The Databricks example notebook is now a get-started notebook,
1323
`examples/databricks/kumo_relational_quickstart.ipynb`, in place of
1424
`adventureworks_relational.ipynb`. It deploys the model to a serving

‎docs/reference/limits.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ Both endpoints enforce sizes, and most of these are the relational one's.
2424
| context rows | tabular | 10,000 | sampled at `materialize()`, refused before sending on a direct `Request` |
2525
| rows scored per call | relational | 1,000, or 200 for link prediction | at `materialize()` |
2626
| an explanation | relational | exactly one entity | before sending |
27-
| an explanation | relational | only within the 1,000-row context budget | before sending |
27+
| an explanation | tabular | exactly one row | before sending |
2828
| hops | relational | 6 | at `materialize()` |
2929
| tables | relational | 15 | against the graph the request carries, not the schema it was built from |
3030

‎packages/structured-data-models-client/src/sdm_client/adapters/relational.py‎

Lines changed: 1 addition & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -340,8 +340,7 @@ def _preflight(
340340
"""Refuse what the endpoint enforces, before spending the round trip.
341341
342342
An explanation is the narrow case: the model explains one entity at a
343-
time, and only in the mode the smaller context budget selects. Sending
344-
more gets a 500 that names neither limit, so both are checked here.
343+
time, and the endpoint refuses more, so that limit is checked here.
345344
346345
`sends` names the halves this call puts on the wire. A session pins its
347346
context at `fit()` and a scoring call carries only rows, so judging the
@@ -373,19 +372,6 @@ def _preflight(
373372
details={'model': self.spec.model_id, 'scored': scored},
374373
)
375374

376-
cap = limits['context_rows_fast']
377-
if 'context' in sends and _asks_for_the_larger_context(payloads, cap):
378-
raise UnsupportedRequestError(
379-
f'an explanation is only produced in the mode the smaller '
380-
f'context budget selects, and this request was gathered with '
381-
f'num_context_samples above {cap:,}. Re-materialize with '
382-
f'num_context_samples={cap} or fewer',
383-
details={
384-
'model': self.spec.model_id,
385-
'context_rows_fast': cap,
386-
},
387-
)
388-
389375
def _check_class_count(
390376
self, payloads: tuple[dict[str, Any], ...], limits: Mapping[str, int]
391377
) -> None:
@@ -607,41 +593,6 @@ def _columns_over(payload: Mapping[str, Any], limit: int) -> list[str]:
607593
]
608594

609595

610-
def _asks_for_the_larger_context(
611-
payloads: tuple[dict[str, Any], ...], cap: int
612-
) -> bool:
613-
r"""Whether this request selected the mode that will not explain.
614-
615-
The endpoint gates explaining on ``inference.run_mode``, which the client
616-
derives from the *requested* ``num_context_samples``. Counting the gathered
617-
rows instead misses every graph that supplies less context than was asked
618-
for: 120 rows gathered under a request for 1,500 passed this guard and the
619-
endpoint's refusal escaped, naming a knob the caller never set.
620-
621-
A payload the engine did not gather carries no ``run_mode``, and for that
622-
one the rows are all there is to read.
623-
"""
624-
modes = {
625-
str(payload['inference']['run_mode'])
626-
for payload in payloads
627-
if isinstance(payload.get('inference'), Mapping)
628-
and payload['inference'].get('run_mode') is not None
629-
}
630-
if modes:
631-
return modes != {'fast'}
632-
return (
633-
sum(
634-
len(
635-
payload.get('context', {})
636-
.get('instance_table', {})
637-
.get('rows', [])
638-
)
639-
for payload in payloads
640-
)
641-
> cap
642-
)
643-
644-
645596
_WIRE_TARGET = 'TARGET'
646597
_A_PREDICTED_COLUMN = re.compile(
647598
r'^\s*PREDICT\s+([A-Za-z_]\w*\.[A-Za-z_]\w*)\s', re.IGNORECASE

‎packages/structured-data-models-client/src/sdm_client/adapters/tabular.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,9 @@ def prediction_plan(
9090
have to carry the same context, so splitting would multiply the upload
9191
without raising the cap.
9292
"""
93-
yield self._envelope(request, outputs)
93+
envelope = self._envelope(request, outputs)
94+
self._check_explanation_rows(request, outputs)
95+
yield envelope
9496

9597
def session_plan(
9698
self, request: Request, outputs: tuple[str, ...]
@@ -119,6 +121,7 @@ def session_prediction_plan(
119121
raise _a_row_the_pinned_schema_cannot_hold(
120122
request, declared, error
121123
) from error
124+
self._check_explanation_rows(request, outputs)
122125
yield {
123126
key: envelope[key] for key in _SESSION_SCORING if key in envelope
124127
}
@@ -140,6 +143,14 @@ def normalize(
140143
f'match the contract: {error}',
141144
code='INVALID_RESPONSE',
142145
) from error
146+
explanation = next(
147+
(
148+
item.explanation
149+
for item in parsed.predictions
150+
if item.explanation
151+
),
152+
None,
153+
)
143154
return PredictionResult(
144155
model=parsed.model,
145156
request_id=parsed.id,
@@ -149,6 +160,7 @@ def normalize(
149160
),
150161
metadata=parsed.metadata,
151162
outputs=outputs,
163+
explanation=explanation,
152164
)
153165

154166
def explain_error(
@@ -162,6 +174,27 @@ def close(self) -> None:
162174
r"""Nothing to release: this adapter sends through the client's own
163175
transport and opens nothing of its own."""
164176

177+
def _check_explanation_rows(
178+
self, request: Request, outputs: tuple[str, ...]
179+
) -> None:
180+
r"""Refuse an explanation over more rows than the endpoint explains.
181+
182+
The endpoint explains one row at a time and refuses more with
183+
``UNSUPPORTED_OUTPUT``, so the limit is named here before anything is
184+
sent. Only the rows being scored count: a session's pinned context is
185+
never judged by it.
186+
"""
187+
if 'explanation' not in outputs:
188+
return
189+
allowed = self.spec.runtime.limits['explanation_predict_rows']
190+
scored = len(request.predict.instance_table)
191+
if scored != allowed:
192+
raise UnsupportedRequestError(
193+
f'an explanation covers exactly {allowed} row, and this '
194+
f'request scores {scored}; ask for one at a time',
195+
details={'model': self.spec.model_id, 'scored': scored},
196+
)
197+
165198
def _envelope(
166199
self,
167200
request: Request,

‎packages/structured-data-models-client/src/sdm_client/models.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
# large context to be told the limit at the far end.
7878
TABULAR_LIMITS = {
7979
'context_rows_max': 10_000,
80+
'explanation_predict_rows': 1,
8081
}
8182

8283
KUMO_TABULAR = ModelSpec(
@@ -86,7 +87,9 @@
8687
# The endpoint takes the coarse kind only and decides binary versus
8788
# multiclass itself, so this client never narrows it here.
8889
tasks=frozenset({'classification', 'regression'}),
89-
outputs=frozenset({'prediction', 'probabilities', 'quantiles'}),
90+
outputs=frozenset(
91+
{'prediction', 'probabilities', 'quantiles', 'explanation'}
92+
),
9093
limits=TABULAR_LIMITS,
9194
),
9295
)
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
2+
# All rights reserved.
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
"""Fixtures shared across the client's test modules."""
6+
7+
from __future__ import annotations
8+
9+
from collections.abc import Callable
10+
11+
import pandas as pd
12+
import pytest
13+
14+
from sdm_client import Request
15+
16+
17+
@pytest.fixture
18+
def tabular_request() -> Callable[..., Request]:
19+
"""Build a single-table request: 20 labelled context rows, ``rows`` to score.
20+
21+
A classification target alternates 0/1; a regression target is a float.
22+
"""
23+
24+
def build(rows: int = 1, task: str = 'classification') -> Request:
25+
labels = [
26+
i % 2 if task == 'classification' else float(i) for i in range(20)
27+
]
28+
context = pd.DataFrame(
29+
{
30+
'row_id': range(20),
31+
'age': [20.0 + i for i in range(20)],
32+
'churn': labels,
33+
}
34+
)
35+
predict = pd.DataFrame(
36+
{
37+
'row_id': range(90, 90 + rows),
38+
'age': [33.0 + i for i in range(rows)],
39+
}
40+
)
41+
return Request(
42+
context=context, predict=predict, target='churn', task=task
43+
)
44+
45+
return build

‎packages/structured-data-models-client/tests/test_limits_and_absent_outputs.py‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -324,24 +324,29 @@ def explanation_plan(request):
324324
)
325325

326326

327-
def test_an_explanation_is_refused_on_the_mode_the_endpoint_refuses_it_on():
328-
"""A graph that supplies 120 rows for a request of 1,500 still selects the
329-
larger-context mode, and that is what the endpoint gates explaining on."""
330-
with pytest.raises(StructuredDataError) as raised:
331-
explanation_plan(gathered(context_rows=120, run_mode='best'))
332-
assert 'num_context_samples' in str(raised.value)
327+
def test_explains_larger_context():
328+
"""An explanation is sent under the larger-context run mode.
329+
330+
Assumptions:
331+
- The Kumo Relational SDM NIM explains under either context budget.
332+
- The payload scores one entity.
333+
"""
334+
assert explanation_plan(gathered(context_rows=120, run_mode='best'))
333335

334336

335337
def test_a_small_context_asked_for_smally_still_explains():
336338
assert explanation_plan(gathered(context_rows=120, run_mode='fast'))
337339

338340

339-
def test_a_payload_naming_no_mode_is_judged_on_the_rows_it_carries():
340-
"""Nothing the engine gathers arrives without a run_mode, but a hand-built
341-
payload can, and the rows are then all there is to read."""
341+
def test_explains_without_mode():
342+
"""An explanation is sent whatever the context size when no run mode is
343+
named.
344+
345+
Assumptions:
346+
- The payload scores one entity and names no run_mode.
347+
"""
342348
assert explanation_plan(gathered(context_rows=120, run_mode=None))
343-
with pytest.raises(StructuredDataError):
344-
explanation_plan(gathered(context_rows=1_001, run_mode=None))
349+
assert explanation_plan(gathered(context_rows=1_001, run_mode=None))
345350

346351

347352
# --- a context that teaches nothing, on both routes -------------------------

‎packages/structured-data-models-client/tests/test_limits_and_sessions.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -83,18 +83,20 @@ def test_an_explanation_covering_more_than_one_entity_is_refused():
8383
assert 'one at a time' in str(error.value)
8484

8585

86-
def test_an_explanation_over_the_smaller_context_budget_is_refused():
87-
"""Explaining forces the smaller budget, which caps context at 1,000."""
86+
def test_explains_large_context():
87+
"""An explanation over a 5,000-row context is sent.
88+
89+
Assumptions:
90+
- The payload scores one entity.
91+
- The SDM NIM has no 1,000-row explanation limit.
92+
"""
8893
request = _Gathered(payload(predict_rows=1, context_rows=5_000))
89-
with pytest.raises(UnsupportedRequestError) as error:
90-
list(
91-
RelationalAdapter().prediction_plan(
92-
request, ('prediction', 'explanation')
93-
)
94+
plans = list(
95+
RelationalAdapter().prediction_plan(
96+
request, ('prediction', 'explanation')
9497
)
95-
message = str(error.value)
96-
assert 'num_context_samples' in message
97-
assert '1,000' in message
98+
)
99+
assert plans[0]['output']['fields'] == ['prediction', 'explanation']
98100

99101

100102
def test_an_explanation_within_both_limits_is_allowed():

‎packages/structured-data-models-client/tests/test_request_and_dispatch.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,14 +177,20 @@ def test_a_task_kind_the_model_does_not_serve_is_refused():
177177
check_request_against_model(request, KUMO_RELATIONAL, ('prediction',))
178178

179179

180-
def test_an_output_field_the_model_cannot_produce_is_refused_by_name():
180+
def test_unserved_output_named():
181+
"""An output the model cannot produce is refused, naming it and the
182+
outputs the model does serve.
183+
184+
Assumptions:
185+
- kumo-tabular does not serve rankings.
186+
"""
181187
request = simple_request()
182188
with pytest.raises(UnsupportedRequestError) as error:
183189
check_request_against_model(
184-
request, KUMO_TABULAR, ('prediction', 'explanation')
190+
request, KUMO_TABULAR, ('prediction', 'rankings')
185191
)
186192
message = str(error.value)
187-
assert 'explanation' in message
193+
assert 'rankings' in message
188194
assert 'probabilities' in message
189195

190196

0 commit comments

Comments
 (0)