structured-data-models-client 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
sdm_client/client.py ADDED
@@ -0,0 +1,1416 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
2
+ # All rights reserved.
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ r"""One connection to one NIM, serving every model it can reach.
6
+
7
+ The model is an argument:
8
+
9
+ >>> with SDMClient(url='http://localhost:8000') as client:
10
+ ... result = client.predict(request, model=KumoTabular())
11
+
12
+ Each client owns its own transport, so several can target different endpoints
13
+ or tenants at once, including from several threads: a prediction always goes to
14
+ the endpoint and credential of the client that started it.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import dataclasses
20
+ import warnings
21
+ from collections.abc import Callable, Iterable, Mapping, Sequence
22
+ from types import TracebackType
23
+ from typing import Any
24
+
25
+ from sdm_client.base import (
26
+ ModelAdapter,
27
+ ModelCapabilities,
28
+ ModelSpec,
29
+ check_request_against_model,
30
+ describe,
31
+ )
32
+ from sdm_client.core.transport import Transport, WireTransport, redact_url
33
+ from sdm_client.core.warnings import SDMWarning, caller_stacklevel
34
+ from sdm_client.decompose import RANKING, Decomposition, about_the_part
35
+ from sdm_client.errors import (
36
+ StructuredDataError,
37
+ UnknownModelError,
38
+ UnsupportedRequestError,
39
+ )
40
+ from sdm_client.models import MODELS, Model, spec_for
41
+ from sdm_client.request import (
42
+ Question,
43
+ Request,
44
+ _envelope_rows,
45
+ warn_of_rows_scored_from_their_own_answers,
46
+ )
47
+ from sdm_client.result import PredictionResult
48
+ from sdm_client.schema import _as_key_tuple
49
+ from sdm_client.wire.keys import (
50
+ _named_by,
51
+ keys_in,
52
+ scored_rows,
53
+ without_the_keys,
54
+ )
55
+
56
+ # A NIM that does not serve the session routes at all, as opposed to one that
57
+ # has merely forgotten this session.
58
+ _SESSIONS_UNSUPPORTED = frozenset({404, 405, 501})
59
+
60
+
61
+ def _adapter_for(spec: ModelSpec) -> ModelAdapter:
62
+ r"""The adapter that shapes requests for ``spec``.
63
+
64
+ Imported here rather than at module scope so that importing the client does
65
+ not pull in an adapter that reaches for an optional driver.
66
+ """
67
+ if spec.model_id == 'kumo-tabular':
68
+ from sdm_client.adapters.tabular import TabularAdapter
69
+
70
+ return TabularAdapter()
71
+ if spec.model_id == 'kumo-relational':
72
+ from sdm_client.adapters.relational import RelationalAdapter
73
+
74
+ return RelationalAdapter()
75
+ raise UnknownModelError(spec.model_id, list(MODELS))
76
+
77
+
78
+ class FittedModel:
79
+ r"""A context pinned on the endpoint, ready to score against.
80
+
81
+ Returned by :meth:`SDMClient.fit`. The contract's
82
+ ``POST /v1/sessions`` uploads a context and runs no inference; scoring
83
+ against it is a second call. That is what ``fit`` and ``predict`` are here.
84
+
85
+ Close it when done, or use it as a context manager. A session holds memory
86
+ on the endpoint until it is released or expires.
87
+
88
+ ``pinned`` is what the session pinned, minus the context itself. Reading an
89
+ answer needs the task and the schema, where the key an id stands for is
90
+ named. The rows have been uploaded and live on the endpoint, so
91
+ keeping a second encoded copy for the lifetime of the handle would hold a
92
+ whole dataset for nothing.
93
+
94
+ It is not constructed directly. What it is built from is how the client
95
+ talks to a model and what the session pinned, none of which is public.
96
+ """
97
+
98
+ _client: SDMClient
99
+ _adapter: ModelAdapter
100
+ _session_id: str
101
+ _outputs: tuple[str, ...]
102
+ _named_outputs: tuple[str, ...] | None
103
+ _pinned_keys: frozenset[Any]
104
+ _pinned_over_a_window: bool
105
+ _held_back: int
106
+ _pinned: dict[str, Any]
107
+ _question: Question | None
108
+ _closed: bool
109
+
110
+ def __init__(self) -> None:
111
+ raise TypeError(
112
+ 'a FittedModel is returned by SDMClient.fit(), not constructed '
113
+ 'directly'
114
+ )
115
+
116
+ @classmethod
117
+ def _open(
118
+ cls,
119
+ client: SDMClient,
120
+ adapter: ModelAdapter,
121
+ session_id: str,
122
+ outputs: tuple[str, ...],
123
+ pinned: dict[str, Any] | None = None,
124
+ named_outputs: Sequence[str] | None = None,
125
+ pinned_keys: frozenset[Any] = frozenset(),
126
+ pinned_over_a_window: bool = False,
127
+ held_back: int = 0,
128
+ question: Question | None = None,
129
+ ) -> FittedModel:
130
+ r"""The handle for a session ``SDMClient.fit`` has just created."""
131
+ fitted = cls.__new__(cls)
132
+ fitted._client = client
133
+ fitted._adapter = adapter
134
+ fitted._session_id = session_id
135
+ fitted._outputs = outputs
136
+ fitted._named_outputs = (
137
+ None if named_outputs is None else tuple(named_outputs)
138
+ )
139
+ fitted._pinned_keys = pinned_keys
140
+ fitted._pinned_over_a_window = pinned_over_a_window
141
+ fitted._held_back = held_back
142
+ fitted._pinned = pinned or {}
143
+ fitted._question = question
144
+ fitted._closed = False
145
+ return fitted
146
+
147
+ @property
148
+ def session_id(self) -> str:
149
+ return self._session_id
150
+
151
+ @property
152
+ def model(self) -> str:
153
+ return self._adapter.spec.model_id
154
+
155
+ def predict(self, request: Request) -> PredictionResult:
156
+ r"""Score ``request``'s rows against the pinned context.
157
+
158
+ Only the rows travel: the context, schema and task were pinned when
159
+ this handle was created. Reading the answer needs both halves because
160
+ the key an id stands for is named in the pinned schema. The
161
+ adapter is handed the whole request while the wire gets only the half
162
+ the contract allows on a scoring call.
163
+ """
164
+ _require_a_request(request)
165
+ if self._closed:
166
+ raise StructuredDataError(
167
+ 'this fitted model is closed; fit again to score more rows',
168
+ code='INVALID_CONFIGURATION',
169
+ )
170
+ parts = getattr(request, 'parts', None)
171
+ if parts is not None:
172
+ raise StructuredDataError(
173
+ f'this request is asked as {len(parts)} questions and a session '
174
+ f'pins one context, so scoring it here would answer the pinned '
175
+ f'question that many times. Ask it with predict()',
176
+ code='UNSUPPORTED_REQUEST',
177
+ )
178
+ asked = getattr(request, '_question', None)
179
+ if self._question is not None and asked is not None:
180
+ _refuse_a_question_the_session_did_not_pin(self._question, asked)
181
+ if request.nothing_to_score:
182
+ return _empty_answer(self.model, self._outputs, request)
183
+ _warn_scoring_what_was_pinned(
184
+ request,
185
+ self._pinned_keys,
186
+ self._pinned_over_a_window,
187
+ _pinned_dtypes(self._pinned),
188
+ )
189
+ parts = _numbered_across_calls(
190
+ self._adapter.session_prediction_plan(
191
+ request, self._outputs, self._pinned
192
+ ),
193
+ lambda payload: self._adapter.normalize(
194
+ self._client._transport.session_predict(
195
+ self._session_id, payload
196
+ ),
197
+ self._outputs,
198
+ {**self._pinned, **payload},
199
+ ),
200
+ )
201
+ if not parts:
202
+ raise StructuredDataError(
203
+ 'the adapter produced no scoring call for this request',
204
+ code='INVALID_REQUEST',
205
+ )
206
+ return _report_absent(
207
+ PredictionResult.merge(parts), self._named_outputs, self.model
208
+ )
209
+
210
+ def close(self) -> None:
211
+ r"""Release the pinned context. Safe to call twice.
212
+
213
+ The handle is marked closed only once the endpoint has released the
214
+ session. Marking it first made a failed release permanent: the next
215
+ call returned early and sent nothing, so the one thing the endpoint
216
+ asks for when it answers `429 Session maintenance is already pending;
217
+ retry the request later` was the one thing the handle could not do.
218
+ """
219
+ if self._closed:
220
+ return
221
+ self._client._transport.delete_session(self._session_id)
222
+ self._closed = True
223
+ pinned = getattr(self._client, '_pinned', None)
224
+ if pinned is not None:
225
+ pinned.discard(self)
226
+
227
+ def __enter__(self) -> FittedModel:
228
+ return self
229
+
230
+ def __exit__(self, *exc: object) -> None:
231
+ r"""Release the context without masking an active block exception.
232
+
233
+ A failure to release is a real failure and is raised, which is what a
234
+ caller wrapping the block in `try` relies on.
235
+
236
+ The exception is a block that is already raising, where nothing is
237
+ said at all. Anything leaving `__exit__` replaces the exception the
238
+ block was raising, and under `-W error` that includes a warning, so
239
+ the one place this must not speak is the one where speaking would
240
+ lose the failure the caller actually needs to see. There the session
241
+ is left to the endpoint's own expiry, and `close()` can be called
242
+ again by a caller who wants to know.
243
+ """
244
+ try:
245
+ self.close()
246
+ except Exception:
247
+ if not (exc and exc[0] is not None):
248
+ raise
249
+
250
+ def __repr__(self) -> str:
251
+ state = 'closed' if self._closed else 'open'
252
+ return f'FittedModel({self.model!r}, {state})'
253
+
254
+
255
+ class SDMClient:
256
+ r"""A connection to one Universal TFM NIM.
257
+
258
+ Use it as a context manager, or call :meth:`close`.
259
+ """
260
+
261
+ def __init__(
262
+ self,
263
+ url: str,
264
+ api_key: str | None = None,
265
+ *,
266
+ verify_ssl: bool = True,
267
+ timeout: float = 60.0,
268
+ max_retries: int = 3,
269
+ ) -> None:
270
+ r"""Open a client against one NIM.
271
+
272
+ No request is made here: the endpoint is first contacted by
273
+ :meth:`health_ready`, :meth:`models`, or a prediction.
274
+
275
+ Args:
276
+ url: Base URL of the NIM, e.g. ``'http://localhost:8000'``.
277
+ api_key: Sent as ``X-API-Key`` when an authenticating gateway
278
+ fronts the NIM. Refused over plaintext ``http://`` to a
279
+ non-local host.
280
+ verify_ssl: Whether to verify the server's TLS certificate.
281
+ timeout: Seconds to wait for each individual attempt, not for the
282
+ call as a whole.
283
+ max_retries: Transport-level retries of a transient failure (429,
284
+ 502, 503, 504, or a dropped connection) with exponential
285
+ backoff. A 500 is not retried: the payload and the model are
286
+ the same on every attempt. ``POST /v1/sessions`` is excluded:
287
+ a re-sent create would orphan a pinned context on the
288
+ endpoint.
289
+ """
290
+ self._start(
291
+ Transport(
292
+ url,
293
+ api_key,
294
+ verify_ssl=verify_ssl,
295
+ timeout=timeout,
296
+ max_retries=max_retries,
297
+ )
298
+ )
299
+
300
+ def _start(self, transport: WireTransport) -> None:
301
+ r"""Set everything a client holds, given what carries its requests.
302
+
303
+ The one place a client's state is set. :meth:`for_databricks_serving`
304
+ builds its client without ``__init__``, so state set only there was
305
+ missing from a serving client until something reached for it. Add
306
+ new state here, not to ``__init__``.
307
+ """
308
+ self._transport: WireTransport = transport
309
+ self._adapters: dict[str, ModelAdapter] = {}
310
+ self._pinned: set[FittedModel] = set()
311
+
312
+ @classmethod
313
+ def for_databricks_serving(
314
+ cls,
315
+ endpoint: str,
316
+ *,
317
+ workspace_client: Any | None = None,
318
+ timeout: float | None = None,
319
+ ) -> SDMClient:
320
+ r"""Open a client against a Databricks Model Serving endpoint.
321
+
322
+ The model is served by name inside a workspace, so there is no URL and
323
+ no key: calls go through the Databricks SDK with the workspace's own
324
+ identity, which in a notebook is the notebook's.
325
+
326
+ >>> with SDMClient.for_databricks_serving('kumo-relational') as client:
327
+ ... result = client.predict(request, model=KumoRelational())
328
+
329
+ A serving endpoint answers one request at a time, so :meth:`predict`
330
+ works and :meth:`fit` is refused: it has no sessions to pin a context
331
+ in.
332
+
333
+ Args:
334
+ endpoint: The serving endpoint's name, not its URL. A URL is
335
+ refused without being repeated, since one can carry a token.
336
+ workspace_client: A ``databricks.sdk.WorkspaceClient`` to call
337
+ through, for a workspace or credential other than the ambient
338
+ one. A ``WorkspaceClient`` itself is called through a copy
339
+ that names ``structured-data-models-client`` in the
340
+ User-Agent, so it is not changed; a subclass or test double
341
+ is used as given, without the name.
342
+ timeout: Seconds one attempt may take, and the longest the SDK
343
+ keeps retrying. When omitted, the ambient workspace client
344
+ gets 300, far above a NIM's default because a serving
345
+ endpoint that has scaled to zero takes minutes to start, and
346
+ a passed ``workspace_client`` keeps its own ``Config``
347
+ timeouts. Given with a ``WorkspaceClient``, it is set on the
348
+ copy, never on the caller's client; with a subclass or test
349
+ double, which cannot be copied, it is refused.
350
+
351
+ Raises:
352
+ MissingExtraError: If ``databricks-sdk`` is not installed; install
353
+ ``structured-data-models-client[databricks-serving]``.
354
+ StructuredDataError: With code ``INVALID_CONFIGURATION``, if
355
+ ``endpoint`` is not a name, ``timeout`` is not a positive
356
+ number of seconds or is passed with a ``workspace_client``
357
+ that cannot be copied, or the ambient configuration cannot
358
+ authenticate.
359
+ """
360
+ from sdm_client.core.databricks_serving import (
361
+ DatabricksServingTransport,
362
+ )
363
+
364
+ client = cls.__new__(cls)
365
+ client._start(
366
+ DatabricksServingTransport(
367
+ endpoint, workspace_client, timeout=timeout
368
+ )
369
+ )
370
+ return client
371
+
372
+ @property
373
+ def url(self) -> str:
374
+ r"""Where this client sends requests.
375
+
376
+ A NIM's base URL, or ``'databricks-serving:<name>'`` for a client
377
+ opened with :meth:`for_databricks_serving`.
378
+ """
379
+ return self._transport.url
380
+
381
+ def models(self) -> list[str]:
382
+ r"""The model ids **the endpoint** advertises.
383
+
384
+ Read from ``/v1/models``, not from this client's own list. The previous
385
+ generation reported its own registry here, which meant it claimed a
386
+ model the endpoint did not serve.
387
+ """
388
+ return self._transport.list_models()
389
+
390
+ def capabilities(self, model: str | Model) -> ModelCapabilities:
391
+ r"""What ``model`` supports, as far as this client knows.
392
+
393
+ Describes the client side and is tagged ``source='sdk'``. It needs no
394
+ live connection, and it is not the endpoint's own answer. Use
395
+ :meth:`models` for that.
396
+ """
397
+ return describe(self._spec(model))
398
+
399
+ def health_ready(self) -> bool:
400
+ r"""Whether the NIM answers ``GET /v1/health/ready`` with 200."""
401
+ return self._transport.health_ready()
402
+
403
+ def predict(
404
+ self,
405
+ request: Request,
406
+ *,
407
+ model: str | Model,
408
+ outputs: Sequence[str] | None = None,
409
+ ) -> PredictionResult:
410
+ r"""Answer ``request`` with ``model``.
411
+
412
+ Args:
413
+ request: What to ask. Built directly, or by a predictive query.
414
+ model: A model handle such as ``KumoTabular()``, or a model id.
415
+ outputs: The output fields to ask for. Defaults to what the model
416
+ produces for this task kind.
417
+
418
+ Returns:
419
+ Always a :class:`~sdm_client.result.PredictionResult`, never a
420
+ frame-or-explanation union.
421
+
422
+ Raises:
423
+ UnknownModelError: If ``model`` is not one the contract carries.
424
+ UnsupportedRequestError: If the model cannot answer this request.
425
+ Raised before anything is sent.
426
+ NimRequestError: If the endpoint answered with an error.
427
+ """
428
+ _require_a_request(request)
429
+ _require_output_names(outputs)
430
+ if request.parts is not None:
431
+ return self._answer_in_parts(request, model, outputs)
432
+ adapter, fields = self._prepare(request, model, outputs)
433
+ if request.nothing_to_score:
434
+ return _empty_answer(adapter.spec.model_id, fields, request)
435
+ # Asked here, where both halves of the request travel together, which
436
+ # is the only call at which the pairing is a leak.
437
+ warn_of_rows_scored_from_their_own_answers(request)
438
+ # Every call's rows are kept. A request past the endpoint's per-call
439
+ # cap is split by the adapter, and each call answers about its own
440
+ # rows, so returning only the last would drop the earlier batches.
441
+ parts = _numbered_across_calls(
442
+ adapter.prediction_plan(request, fields),
443
+ lambda payload: adapter.normalize(
444
+ self._send(adapter, payload), fields, payload
445
+ ),
446
+ )
447
+ if not parts:
448
+ raise StructuredDataError(
449
+ 'the adapter produced no call for this request',
450
+ code='INVALID_REQUEST',
451
+ )
452
+ return _report_absent(
453
+ PredictionResult.merge(parts), outputs, adapter.spec.model_id
454
+ )
455
+
456
+ def _answer_in_parts(
457
+ self,
458
+ request: Request,
459
+ model: str | Model,
460
+ outputs: Sequence[str] | None,
461
+ ) -> PredictionResult:
462
+ r"""Ask the questions this one is asked as, and give back one answer.
463
+
464
+ A ranking and a forecast have no task on the endpoint. Each is asked
465
+ as questions it does serve, and what comes back has to be put together
466
+ as the answer that was actually asked for: an order over candidates
467
+ or a value per step, rather than handed back as the several answers
468
+ it took.
469
+ """
470
+ rewritten = request._decomposition
471
+ assert rewritten is not None
472
+ ranking = rewritten.kind == RANKING
473
+ unscored: frozenset[int] = frozenset()
474
+ if ranking:
475
+ _refuse_outputs_a_ranking_cannot_give(outputs)
476
+ unscored = self._candidates_no_context_teaches(
477
+ request, model, rewritten
478
+ )
479
+ # A ranking needs how sure each answer was, which is what orders the
480
+ # candidates, and the fields it gives back are the merge's own. A step
481
+ # of a forecast is an ordinary regression, so the caller's fields are
482
+ # what each part is asked for.
483
+ wanted = _RANKING_PARTS if ranking else outputs
484
+ answers: list[PredictionResult] = []
485
+ answered: list[int] = []
486
+ for position, (part, (_, about)) in enumerate(
487
+ zip(request.parts or (), rewritten.parts, strict=True)
488
+ ):
489
+ if position in unscored:
490
+ continue
491
+ answered.append(position)
492
+ try:
493
+ answers.append(self.predict(part, model=model, outputs=wanted))
494
+ except Exception as error:
495
+ answered.pop()
496
+ _carry_the_steps_already_paid_for(
497
+ error, answers, rewritten, answered, ranking
498
+ )
499
+ if not isinstance(error, StructuredDataError):
500
+ raise
501
+ raise about_the_part(error, rewritten, about) from error
502
+ if ranking:
503
+ parts = rewritten.parts
504
+ scored = [part for i, part in enumerate(parts) if i not in unscored]
505
+ return _ranked(
506
+ answers,
507
+ dataclasses.replace(rewritten, parts=tuple(scored)),
508
+ outputs,
509
+ unscored=[parts[i][1] for i in sorted(unscored)],
510
+ )
511
+ merged = _by_step(answers, rewritten)
512
+ _warn_of_steps_that_carry_no_explanation(
513
+ answers, rewritten, outputs, spec_for(model).model_id
514
+ )
515
+ return _report_absent(merged, outputs, spec_for(model).model_id)
516
+
517
+ def _candidates_no_context_teaches(
518
+ self,
519
+ request: Request,
520
+ model: str | Model,
521
+ rewritten: Decomposition,
522
+ ) -> frozenset[int]:
523
+ r"""Which of a ranking's parts ask about a one-class candidate.
524
+
525
+ The endpoint refuses a classification whose context holds one class,
526
+ so such a candidate cannot be scored. Checking every part before any
527
+ call sends none for those, names them all at once, and refuses a
528
+ ranking with nothing left before it spends a call.
529
+
530
+ A candidate is scored or not for the whole ranking, because each
531
+ candidate is one call: a part is never batched, and a ranking over
532
+ more entities than one call scores is refused when it is gathered.
533
+ Were a part ever split, the engine would sample a context per batch,
534
+ and one cold batch would drop the candidate from every row rather
535
+ than from its own, since ``unscored_candidates`` names a candidate
536
+ for the whole result rather than per row.
537
+ """
538
+ positions = []
539
+ for position, (part, (_, about)) in enumerate(
540
+ zip(request.parts or (), rewritten.parts, strict=True)
541
+ ):
542
+ try:
543
+ adapter, fields = self._prepare(part, model, _RANKING_PARTS)
544
+ for _ in adapter.prediction_plan(part, fields):
545
+ pass
546
+ except StructuredDataError as error:
547
+ details = getattr(error, 'details', None) or {}
548
+ if details.get('classes_in_context') != 1:
549
+ raise about_the_part(error, rewritten, about) from error
550
+ positions.append(position)
551
+ if not positions:
552
+ return frozenset()
553
+ cold = [rewritten.parts[position][1] for position in positions]
554
+ many = len(cold) > 1
555
+ named = ', '.join(repr(candidate) for candidate in cold[:-1])
556
+ if len(cold) == len(rewritten.parts):
557
+ which = (
558
+ f'the contexts gathered for candidates {named} and '
559
+ f'{cold[-1]!r} each hold'
560
+ if many
561
+ else f'the context gathered for candidate {cold[0]!r} holds'
562
+ )
563
+ raise UnsupportedRequestError(
564
+ f'{which} one answer only, and a classification needs at '
565
+ f'least two classes to learn from, so there is nothing to '
566
+ f'rank and nothing was sent. {_SCORE_THE_UNSCORED}',
567
+ details={
568
+ 'model': spec_for(model).model_id,
569
+ 'candidates': cold,
570
+ 'classes_in_context': 1,
571
+ },
572
+ )
573
+ which = (
574
+ f'candidates {named} and {cold[-1]!r} are'
575
+ if many
576
+ else f'candidate {cold[0]!r} is'
577
+ )
578
+ them = 'them' if many else 'it'
579
+ warnings.warn(
580
+ f'{which} left out of this ranking, unscored: the context '
581
+ f'sampled for {"each" if many else "it"} holds one answer only, '
582
+ f'nearly always because no entity in it linked to the candidate '
583
+ f'in the window, and a classification needs two to learn from. '
584
+ f'No call was sent for {them}, so each row ranks only the rest, '
585
+ f"and the result's unscored_candidates names {them}. "
586
+ f'{_SCORE_THE_UNSCORED}',
587
+ category=SDMWarning,
588
+ stacklevel=caller_stacklevel(),
589
+ )
590
+ return frozenset(positions)
591
+
592
+ def _send(self, adapter: Any, payload: dict[str, Any]) -> dict[str, Any]:
593
+ r"""Post one payload, letting the adapter explain a refusal."""
594
+ try:
595
+ return self._transport.predict(payload)
596
+ except StructuredDataError as error:
597
+ explained = adapter.explain_error(error, payload)
598
+ if explained is error:
599
+ # An adapter that had nothing to add returns the same object,
600
+ # and `raise error from error` makes it its own cause, which
601
+ # is a cycle for anything that walks the chain.
602
+ raise
603
+ raise explained from error
604
+
605
+ def fit(
606
+ self,
607
+ request: Request,
608
+ *,
609
+ model: str | Model,
610
+ outputs: Sequence[str] | None = None,
611
+ holdout: Iterable[Any] | None = None,
612
+ ) -> FittedModel:
613
+ r"""Pin ``request``'s context on the endpoint and return a handle.
614
+
615
+ The pinned half is model, task, schema and context; later calls send
616
+ only the rows to score. Use this when asking several questions of the
617
+ same context, which is what makes uploading it once worthwhile.
618
+
619
+ `holdout` names every row the session will go on to score. The context
620
+ is chosen once, here, and `request` only excludes the rows *it* names,
621
+ so every later call scores rows that are sitting in the pinned context
622
+ with their answers. Unlike the one-shot path, there is nothing the
623
+ caller can do about it afterwards, because the context cannot be
624
+ changed once it is pinned. Naming the population up front is what lets
625
+ it be excluded while there is still a context to exclude it from.
626
+
627
+ A row is held back at every moment it appears. For a target that is a
628
+ window, the same entity at an *earlier* anchor is legitimate context
629
+ rather than a leak, so this gives up examples it did not have to. It
630
+ gives them up because a session cannot know which moments it will be
631
+ asked about: holding the entity back everywhere is the only choice
632
+ that cannot leak, and a leak here has no remedy afterwards. Score at
633
+ known moments with `predict()` if those examples are worth more than
634
+ the certainty.
635
+
636
+ Args:
637
+ request: The request whose context is pinned.
638
+ model: Which model answers.
639
+ outputs: The fields to ask for.
640
+ holdout: Keys of every row this session may score. They are taken
641
+ out of the pinned context before it is sent.
642
+
643
+ Raises:
644
+ StructuredDataError: If the endpoint does not serve sessions,
645
+ with code ``'SESSIONS_UNSUPPORTED'``.
646
+ """
647
+ if not getattr(self._transport, 'serves_sessions', True):
648
+ raise StructuredDataError(
649
+ f'{self._transport.url} cannot pin a context: it answers one '
650
+ f'request at a time. Use predict(), which sends the context '
651
+ f'with every call',
652
+ code='SESSIONS_UNSUPPORTED',
653
+ )
654
+ _require_a_request(request)
655
+ _require_output_names(outputs)
656
+ _refuse_to_pin_a_rewritten_question(request)
657
+ adapter, fields = self._prepare(request, model, outputs)
658
+ payload = adapter.session_plan(request, fields)
659
+ held = _held_back(payload, holdout)
660
+ try:
661
+ try:
662
+ created = self._transport.create_session(payload)
663
+ except StructuredDataError as error:
664
+ explained = adapter.explain_error(error, payload)
665
+ if explained is error:
666
+ raise
667
+ raise explained from error
668
+ except Exception as error:
669
+ status = getattr(error, 'status_code', None)
670
+ if status in _SESSIONS_UNSUPPORTED:
671
+ raise StructuredDataError(
672
+ f'this endpoint does not serve sessions (HTTP {status}); '
673
+ f'use predict() instead, which sends the context with '
674
+ f'every call',
675
+ code='SESSIONS_UNSUPPORTED',
676
+ ) from error
677
+ raise
678
+ session_id = created.get('session_id') or created.get('id')
679
+ if not session_id:
680
+ raise StructuredDataError(
681
+ 'the endpoint created a session but named no session id',
682
+ code='INVALID_RESPONSE',
683
+ details={'response_keys': sorted(created)},
684
+ )
685
+ fitted = FittedModel._open(
686
+ self,
687
+ adapter,
688
+ str(session_id),
689
+ fields,
690
+ {key: value for key, value in payload.items() if key != 'context'},
691
+ named_outputs=outputs,
692
+ pinned_keys=keys_in(
693
+ payload, 'context', with_anchor=request.target_moves_with_time
694
+ ),
695
+ held_back=held,
696
+ pinned_over_a_window=request.target_moves_with_time,
697
+ question=request._question,
698
+ )
699
+ self._pinned.add(fitted)
700
+ return fitted
701
+
702
+ def _prepare(
703
+ self,
704
+ request: Request,
705
+ model: str | Model,
706
+ outputs: Sequence[str] | None,
707
+ ) -> tuple[ModelAdapter, tuple[str, ...]]:
708
+ r"""Resolve the adapter and output fields, and refuse early.
709
+
710
+ Every check that can be made without the network is made here, so a
711
+ request the model cannot answer fails with a message naming the model
712
+ rather than as an opaque 422.
713
+ """
714
+ _require_a_request(request)
715
+ _require_output_names(outputs)
716
+ spec = self._spec(model)
717
+ adapter = self._adapters.get(spec.model_id)
718
+ if adapter is None:
719
+ adapter = _adapter_for(spec)
720
+ self._adapters[spec.model_id] = adapter
721
+
722
+ fields = (
723
+ tuple(outputs)
724
+ if outputs is not None
725
+ else adapter.default_output(request)
726
+ )
727
+ if not fields:
728
+ raise StructuredDataError(
729
+ 'outputs must name at least one field',
730
+ code='INVALID_REQUEST',
731
+ )
732
+ if isinstance(outputs, (str, bytes)):
733
+ raise StructuredDataError(
734
+ f'outputs must be a sequence of field names, not a single '
735
+ f'string; pass [{outputs!r}]',
736
+ code='INVALID_REQUEST',
737
+ )
738
+ check_request_against_model(request, spec, fields)
739
+ return adapter, fields
740
+
741
+ @staticmethod
742
+ def _spec(model: str | Model) -> ModelSpec:
743
+ return spec_for(model)
744
+
745
+ def close(self) -> None:
746
+ r"""Release pooled connections and retire this client.
747
+
748
+ Contexts this client pinned are released first, because the transport
749
+ that would release them is about to close and `FittedModel.close()`
750
+ is the only call that can. Adapters follow, then the transport itself
751
+ even if one of them fails. Its pool is the one that cannot be
752
+ rebuilt on demand.
753
+ """
754
+ try:
755
+ error: BaseException | None = None
756
+ for fitted in list(self._pinned):
757
+ try:
758
+ fitted.close()
759
+ except BaseException as session_error: # noqa: BLE001
760
+ error = error or session_error
761
+ self._pinned.clear()
762
+ for adapter in self._adapters.values():
763
+ try:
764
+ adapter.close()
765
+ except BaseException as adapter_error: # noqa: BLE001
766
+ error = error or adapter_error
767
+ if error is not None:
768
+ raise error
769
+ finally:
770
+ self._transport.close()
771
+
772
+ def __enter__(self) -> SDMClient:
773
+ return self
774
+
775
+ def __exit__(
776
+ self,
777
+ exc_type: type[BaseException] | None,
778
+ exc: BaseException | None,
779
+ tb: TracebackType | None,
780
+ ) -> None:
781
+ self.close()
782
+
783
+ def __repr__(self) -> str:
784
+ return f'SDMClient(url={redact_url(self._transport.url)!r})'
785
+
786
+
787
+ def _numbered_across_calls(
788
+ plan: Iterable[dict[str, Any]],
789
+ answer: Callable[[dict[str, Any]], PredictionResult],
790
+ ) -> list[PredictionResult]:
791
+ r"""Answer each call, numbering its rows from where the last call's ended.
792
+
793
+ The endpoint numbers every call's rows from zero, so a request sent in
794
+ several calls would otherwise give each ``row_index`` once per call, and
795
+ ``to_pandas()``, which orders by it, would interleave the batches.
796
+
797
+ Every call is counted before the first is sent, because a call whose rows
798
+ cannot be counted leaves the ones after it unnumberable, and finding that
799
+ out between two sends spends calls on an answer that cannot be returned.
800
+
801
+ A call that fails ends the request, and the answers already paid for are
802
+ carried on the error as ``details['answered']``, with the rows they cover
803
+ as ``details['rows_answered']``: the caller can keep them and ask again
804
+ for the rest rather than resend the whole population.
805
+ """
806
+ calls = list(plan)
807
+ counted = [_envelope_rows(payload, 'predict') for payload in calls]
808
+ if any(rows is None for rows in counted[:-1]):
809
+ raise StructuredDataError(
810
+ 'a request sent in several calls has a call whose rows could not '
811
+ 'be counted, so the rows of the calls after it cannot be numbered',
812
+ code='INTERNAL_ERROR',
813
+ )
814
+ answers: list[PredictionResult] = []
815
+ start = 0
816
+ for payload, rows in zip(calls, counted, strict=True):
817
+ try:
818
+ answers.append(_numbered_from(answer(payload), start))
819
+ except Exception as error:
820
+ _carry_the_answers_already_paid_for(error, answers, start)
821
+ raise
822
+ start += rows or 0
823
+ return answers
824
+
825
+
826
+ def _carry_the_steps_already_paid_for(
827
+ error: BaseException,
828
+ answers: list[PredictionResult],
829
+ rewritten: Decomposition,
830
+ answered: Sequence[int],
831
+ ranking: bool,
832
+ ) -> None:
833
+ r"""Record on ``error`` the parts answered before the one that failed.
834
+
835
+ A forecast's steps are kept as the answer they would have been: numbered
836
+ by step, since a step's rows are numbered from its own zero and only the
837
+ step tells two of them apart.
838
+
839
+ A ranking's are counted and not kept. Its parts are not an answer until
840
+ they are all in: what a caller asked for is an order over the candidates,
841
+ and an order over some of them is a different claim, not a partial one.
842
+ """
843
+ if not answers:
844
+ return
845
+ held = getattr(error, 'details', None)
846
+ error.details = { # type: ignore[attr-defined]
847
+ **(held if isinstance(held, dict) else {}),
848
+ 'parts_answered': len(answers),
849
+ **(
850
+ {}
851
+ if ranking
852
+ else {
853
+ 'answered': _by_step(
854
+ answers,
855
+ dataclasses.replace(
856
+ rewritten,
857
+ parts=tuple(rewritten.parts[at] for at in answered),
858
+ ),
859
+ )
860
+ }
861
+ ),
862
+ }
863
+
864
+
865
+ def _carry_the_answers_already_paid_for(
866
+ error: BaseException, answers: list[PredictionResult], rows: int
867
+ ) -> None:
868
+ r"""Record on ``error`` the calls that were answered before it.
869
+
870
+ A request sent in fifty calls that fails on the thirty-seventh has
871
+ thirty-six answers the endpoint was already paid for. Dropping them makes
872
+ the caller buy them again; the error keeps its own type and code, so
873
+ nothing branching on either has to know they are there.
874
+
875
+ The same holds of a question asked in parts, where what was answered is
876
+ whole parts rather than batches of one. ``about_the_part`` copies the
877
+ error's ``__dict__``, so what is recorded here survives being said of the
878
+ part it happened in.
879
+ """
880
+ if not answers:
881
+ return
882
+ # Not every failure is this client's own: the Databricks transport lets
883
+ # its OAuth refresh raise, and those calls were paid for too.
884
+ held = getattr(error, 'details', None)
885
+ error.details = { # type: ignore[attr-defined]
886
+ **(held if isinstance(held, dict) else {}),
887
+ 'answered': PredictionResult.merge(answers),
888
+ 'calls_answered': len(answers),
889
+ 'rows_answered': rows,
890
+ }
891
+
892
+
893
+ def _numbered_from(result: PredictionResult, start: int) -> PredictionResult:
894
+ r"""``result`` with each row numbered ``start`` places further on."""
895
+ if start:
896
+ result.items = tuple(
897
+ dataclasses.replace(item, row_index=item.row_index + start)
898
+ if item.row_index is not None
899
+ else item
900
+ for item in result.items
901
+ )
902
+ return result
903
+
904
+
905
+ def _report_absent(
906
+ result: PredictionResult, requested: Sequence[str] | None, model: str
907
+ ) -> PredictionResult:
908
+ r"""Say when the answer does not carry a field the caller named.
909
+
910
+ Only for fields the caller named. The defaults are chosen per task kind and
911
+ a model is entitled not to fill one of them; a caller who named a field is
912
+ waiting for a column that is not going to appear.
913
+
914
+ Not raised, because the answer is a real answer to the rest of the request
915
+ and the endpoint accepted it. A caller who wants this to be fatal can turn
916
+ warnings into errors; a caller who does not now at least knows.
917
+
918
+ It reports what the answer did not contain and does not say why. The reason
919
+ varies and is not knowable from here: `kumo-relational` returns
920
+ `prediction` for a regression asked on its own and omits it when quantiles
921
+ are asked for alongside, so "the model does not produce it for this task
922
+ kind" would be false in exactly the case that prompted this.
923
+ """
924
+ if not requested:
925
+ return result
926
+ absent = result.absent(requested)
927
+ if absent:
928
+ warnings.warn(
929
+ f'{model} answered without {list(absent)}, which this request '
930
+ f'asked for, so no column of to_pandas() will carry it. The '
931
+ f'answer to the rest of the request is unaffected.',
932
+ category=SDMWarning,
933
+ stacklevel=caller_stacklevel(),
934
+ )
935
+ return result
936
+
937
+
938
+ def _names_its_rows(request: Request) -> bool:
939
+ r"""Whether answers to this request will say which row each one is about.
940
+
941
+ An ``id`` is the entity a row names, so it is there when the rows being
942
+ scored carry one. On the route that builds its own frames that needs both
943
+ a key in the schema and every one of its columns in the half being sent:
944
+ asking only the schema gave an empty answer an ``id`` where a full one had
945
+ none. On the gathered route the entity is what was asked for, and the
946
+ envelope spells it as a transport reference rather than as the schema's
947
+ key. Asking the same question there answered no about requests whose
948
+ every answer names its entity.
949
+ """
950
+ if request.payloads:
951
+ # A gathered request is about entities by construction: the envelope
952
+ # carries an entity reference for every row it scores, and the adapter
953
+ # names each answer with the entity it asked about.
954
+ return True
955
+ table = request.entity_table
956
+ if table is None:
957
+ return False
958
+ try:
959
+ key = _as_key_tuple(request.schema[table].primary_key)
960
+ if not key:
961
+ return False
962
+ scored = request.predict.instance_table.columns
963
+ return all(name in scored for name in key)
964
+ except Exception: # noqa: BLE001 - a schema too partial to ask is a no
965
+ return False
966
+
967
+
968
+ def _empty_answer(
969
+ model: str, outputs: Sequence[str], request: Request | None = None
970
+ ) -> PredictionResult:
971
+ r"""The answer to a request with no rows in it, without asking.
972
+
973
+ `Request` allows a zero-row `predict` on purpose: the contract sets no
974
+ minimum and an empty batch inside a loop is a reasonable no-op. Sent, it
975
+ was not a no-op. The endpoint cannot encode zero rows and answered
976
+ `500 Unexpected server error`, so the one shape documented as harmless was
977
+ the one that failed, and failed opaquely.
978
+
979
+ Nothing is asked, because there is nothing to ask about. The result carries
980
+ the columns the request would have produced, which `to_pandas()` already
981
+ knows how to build with no rows in them.
982
+ """
983
+ built = PredictionResult(
984
+ model=model, request_id='', items=(), outputs=tuple(outputs)
985
+ )
986
+ built._names_its_rows = request is None or _names_its_rows(request)
987
+ return built
988
+
989
+
990
+ def _held_back(payload: dict[str, Any], holdout: Iterable[Any] | None) -> int:
991
+ r"""Take the rows a session will score out of the context it pins.
992
+
993
+ Says how many went, because a caller who names a population expects it
994
+ gone and a silent nothing would look the same as a silent everything.
995
+ """
996
+ if holdout is None:
997
+ return 0
998
+ if isinstance(holdout, (str, bytes)):
999
+ raise StructuredDataError(
1000
+ f'holdout names the rows this session may score, and {holdout!r} '
1001
+ f'is one string rather than a list of keys; pass [{holdout!r}] if '
1002
+ f'that is the key you mean',
1003
+ code='INVALID_REQUEST',
1004
+ )
1005
+ keys = frozenset(
1006
+ _named_by(key if isinstance(key, (tuple, list)) else [key])
1007
+ for key in holdout
1008
+ )
1009
+ if not keys:
1010
+ return 0
1011
+ removed = without_the_keys(payload, keys)
1012
+ if not removed:
1013
+ warnings.warn(
1014
+ f'none of the {len(keys)} keys named as holdout matched a row of '
1015
+ f'the context about to be pinned, so nothing was held back. The '
1016
+ f'keys are compared as the schema spells them, so a key given as '
1017
+ f'a different type matches nothing. Check they are the same '
1018
+ f'kind of value as the entity column holds.',
1019
+ category=SDMWarning,
1020
+ stacklevel=caller_stacklevel(),
1021
+ )
1022
+ if removed:
1023
+ warnings.warn(
1024
+ f'{removed} of the rows named as holdout were in the context '
1025
+ f'about to be pinned, carrying the answers this session will be '
1026
+ f'asked for. They have been taken out of it, so the context is '
1027
+ f'that many rows smaller than the request gathered.',
1028
+ category=SDMWarning,
1029
+ stacklevel=caller_stacklevel(),
1030
+ )
1031
+ return removed
1032
+
1033
+
1034
+ def _pinned_dtypes(pinned: Mapping[str, Any]) -> dict[str, str] | None:
1035
+ r"""The dtypes a session pinned, read back off the envelope it sent.
1036
+
1037
+ The pinned keys were written with these, so the scored ones have to be
1038
+ read with them or the same entity has two names.
1039
+ """
1040
+ columns = ((pinned.get('schema') or {}).get('instance_table') or {}).get(
1041
+ 'columns'
1042
+ )
1043
+ if not isinstance(columns, Mapping):
1044
+ return None
1045
+ return {
1046
+ name: str(described['dtype'])
1047
+ for name, described in columns.items()
1048
+ if isinstance(described, Mapping) and described.get('dtype')
1049
+ }
1050
+
1051
+
1052
+ def _warn_scoring_what_was_pinned(
1053
+ request: Request,
1054
+ pinned: frozenset[Any],
1055
+ over_a_window: bool = False,
1056
+ dtypes: Mapping[str, str] | None = None,
1057
+ ) -> None:
1058
+ r"""Say when a call scores rows the pinned context already answers.
1059
+
1060
+ A session pins the context gathered by the `fit()` request, which excluded
1061
+ the rows *that* request named. Every later call names different rows, and
1062
+ nothing excluded those, so they sit in the pinned context with their target
1063
+ values and the model is asked to predict what it has been shown. Measured
1064
+ on a target that is pure noise, a session scored 0.595 where the same rows
1065
+ scored one-shot came back 0.455, against a chance rate of 0.5.
1066
+
1067
+ It cannot be fixed here: the context lives on the endpoint and the rows
1068
+ cannot be taken out of it. Nor can the caller fix it after the fact. The
1069
+ equivalent of the one-shot remedy is to blank the targets before `fit()`,
1070
+ which means knowing every row the session will ever score, and a caller who
1071
+ knows that has no reason to open a session. So it is reported, with the two
1072
+ things that do work.
1073
+
1074
+ An entity asked about over a window is only given away at the moment it
1075
+ was pinned at, so there the moment is part of what is compared. Comparing
1076
+ the entity alone called every later call about the same entity a leak,
1077
+ which under warnings-as-errors is a working session that stops working.
1078
+ """
1079
+ if not pinned:
1080
+ return
1081
+ # Rows, not the keys among them: two rows asking about one entity are two
1082
+ # rows being scored, and a set had already lost one of them before either
1083
+ # number was taken.
1084
+ scored = scored_rows(request, with_anchor=over_a_window, dtypes=dtypes)
1085
+ shown = sum(1 for named in scored if named in pinned)
1086
+ if not shown:
1087
+ return
1088
+ warnings.warn(
1089
+ f'{shown} of the {len(scored)} rows being scored are in this '
1090
+ f"session's pinned context, with their target values, so the model "
1091
+ f'is being asked about rows it has already been shown the answers '
1092
+ f'for. The context was chosen when fit() ran and cannot be changed '
1093
+ f'now. Either clear the target on every row this session may score '
1094
+ f'before calling fit(), or use predict(), which excludes the rows it '
1095
+ f'scores from the context it sends.',
1096
+ category=SDMWarning,
1097
+ stacklevel=caller_stacklevel(),
1098
+ )
1099
+
1100
+
1101
+ #: What a ranking's parts are asked for: how sure each answer was is what
1102
+ #: orders the candidates, not which side of the line it fell.
1103
+ _RANKING_PARTS = ('prediction', 'probabilities')
1104
+ _RANKING_OUTPUTS = ('prediction', 'rankings')
1105
+ #: What can still rank a candidate whose sampled context holds one answer.
1106
+ _SCORE_THE_UNSCORED = (
1107
+ 'Gather a wider context with num_context_samples if other entities in '
1108
+ 'the data answer differently, or ask over a table of (entity, '
1109
+ 'candidate) pairs, which scores every candidate from its features.'
1110
+ )
1111
+
1112
+
1113
+ def _ranked(
1114
+ answers: list[PredictionResult],
1115
+ rewritten: Decomposition,
1116
+ outputs: Sequence[str] | None,
1117
+ unscored: Sequence[Any] = (),
1118
+ ) -> PredictionResult:
1119
+ r"""The candidates each entity reached, likeliest first.
1120
+
1121
+ Ordered by how sure the answer to that candidate's question was, which is
1122
+ what a ranking is: the questions were asked separately and only their
1123
+ order matters between them.
1124
+
1125
+ Each entity keeps the row it was asked about on, so the answer joins back
1126
+ onto the frame the caller scored. The candidates are the values the caller
1127
+ named, not what they looked like once written into a query.
1128
+
1129
+ ``rewritten`` holds only the parts that were asked. The candidates that
1130
+ were not are carried beside the ranking, never in it, since any score
1131
+ given to them would be invented.
1132
+ """
1133
+
1134
+ from sdm_api.generated.tfm_api import PredictionItem
1135
+
1136
+ # Keyed by the row, not by the entity on it: the same entity may be asked
1137
+ # about more than once, at different moments, and folding those together
1138
+ # would answer with one row where two were asked and could put the same
1139
+ # candidate in it twice.
1140
+ scores: dict[Any, list[tuple[float, Any]]] = {}
1141
+ asked_about: dict[Any, tuple[int | None, Any]] = {}
1142
+ for answer, (_, candidate) in zip(answers, rewritten.parts, strict=True):
1143
+ for item in answer.items:
1144
+ where = (item.row_index, item.id)
1145
+ scores.setdefault(where, []).append(
1146
+ (_how_sure(item, candidate), candidate)
1147
+ )
1148
+ asked_about.setdefault(where, (item.row_index, item.id))
1149
+
1150
+ top = rewritten.top
1151
+ wanted = tuple(outputs or _RANKING_OUTPUTS)
1152
+ items = []
1153
+ for position, (where, scored) in enumerate(scores.items()):
1154
+ row, entity = asked_about[where]
1155
+ # Sort only on the score: candidates are of the caller's type and two
1156
+ # of them need not be comparable to each other.
1157
+ scored.sort(key=lambda pair: pair[0], reverse=True)
1158
+ chosen = scored[:top] if top else scored
1159
+ items.append(
1160
+ PredictionItem(
1161
+ row_index=position if row is None else row,
1162
+ id=entity,
1163
+ prediction=(
1164
+ [candidate for _, candidate in chosen]
1165
+ if 'prediction' in wanted
1166
+ else None
1167
+ ),
1168
+ rankings=(
1169
+ tuple(
1170
+ {'candidate': candidate, 'score': score}
1171
+ for score, candidate in chosen
1172
+ )
1173
+ if 'rankings' in wanted
1174
+ else None
1175
+ ),
1176
+ )
1177
+ )
1178
+ return PredictionResult(
1179
+ model=answers[0].model if answers else '',
1180
+ request_id=answers[0].request_id if answers else '',
1181
+ items=items,
1182
+ outputs=wanted,
1183
+ unscored_candidates=unscored,
1184
+ )
1185
+
1186
+
1187
+ def _how_sure(item: Any, candidate: Any) -> float:
1188
+ r"""How sure the endpoint was that this entity reaches this candidate.
1189
+
1190
+ A ranking is nothing but the order of these, so a missing one cannot be
1191
+ read as zero: that would sort a candidate the endpoint was confident about
1192
+ to the bottom and give back an order that looks like an answer.
1193
+ """
1194
+ given = item.probabilities or {}
1195
+ for spelling in ('true', 'True', 'TRUE', '1'):
1196
+ if spelling in given:
1197
+ return float(given[spelling])
1198
+ raise StructuredDataError(
1199
+ f'the endpoint answered the question about candidate {candidate!r} '
1200
+ f'without saying how sure it was, and a ranking is only the order of '
1201
+ f'those. It gave back {sorted(given) or "no probabilities"}, and this '
1202
+ f'client knows a positive class spelled "true", "True", "TRUE" or '
1203
+ f'"1".',
1204
+ code='INTERNAL_ERROR',
1205
+ )
1206
+
1207
+
1208
+ def _top_k_of(asked: str) -> int | None:
1209
+ r"""How many a ranking asked for, as the query says."""
1210
+ import re
1211
+
1212
+ found = re.search(r'\bRANK\s+TOP\s+(\d+)', asked, flags=re.IGNORECASE)
1213
+ return int(found.group(1)) if found else None
1214
+
1215
+
1216
+ def _warn_of_steps_that_carry_no_explanation(
1217
+ answers: list[PredictionResult],
1218
+ rewritten: Decomposition,
1219
+ requested: Sequence[str] | None,
1220
+ model: str,
1221
+ ) -> None:
1222
+ """Say which steps of a forecast came back without an explanation.
1223
+
1224
+ Each step is explained on its own, so some can carry one and others not.
1225
+ The merged answer holds whatever came back, and a check that asks only
1226
+ whether the whole result carries an explanation sees the ones that did and
1227
+ says nothing about the ones that did not.
1228
+ """
1229
+ if not requested or 'explanation' not in requested:
1230
+ return
1231
+ missing = [
1232
+ step
1233
+ for (_, step), answer in zip(rewritten.parts, answers, strict=True)
1234
+ if answer.explanation is None
1235
+ ]
1236
+ if not missing or len(missing) == len(answers):
1237
+ return
1238
+ warnings.warn(
1239
+ f'{model} explained {len(answers) - len(missing)} of '
1240
+ f'{len(answers)} forecast steps, and step '
1241
+ f'{", ".join(str(step) for step in missing)} carries none, so '
1242
+ f'the explanation this request asked for is not there for every step.',
1243
+ category=SDMWarning,
1244
+ stacklevel=caller_stacklevel(),
1245
+ )
1246
+
1247
+
1248
+ def _by_step(
1249
+ answers: list[PredictionResult], rewritten: Decomposition
1250
+ ) -> PredictionResult:
1251
+ r"""What each step holds, in the order the steps come.
1252
+
1253
+ `forecast_step` is the field the contract already carries for this, so a
1254
+ reader of the frame gets the series it asked for without knowing that it
1255
+ arrived as several answers. It is a column of every such answer, so an
1256
+ empty one carries it too. An entity keeps its own row across the steps:
1257
+ the pair of row and step is what names a point in the series, which is how
1258
+ the contract reads it.
1259
+
1260
+ Each step is explained on its own, so the explanations are kept under the
1261
+ step they belong to. Dropping them left the result saying it carried an
1262
+ explanation while holding none.
1263
+ """
1264
+ items = [
1265
+ dataclasses.replace(item, forecast_step=step)
1266
+ for (_, step), answer in zip(rewritten.parts, answers, strict=True)
1267
+ for item in answer.items
1268
+ ]
1269
+ by_step = {
1270
+ step: answer.explanation
1271
+ for (_, step), answer in zip(rewritten.parts, answers, strict=True)
1272
+ if answer.explanation is not None
1273
+ }
1274
+ built = PredictionResult(
1275
+ model=answers[0].model if answers else '',
1276
+ request_id=answers[0].request_id if answers else '',
1277
+ items=items,
1278
+ outputs=tuple(answers[0].outputs) if answers else (),
1279
+ explanation={'by_forecast_step': by_step} if by_step else None,
1280
+ )
1281
+ # `forecast_step` is put on the items above, so with no items there is
1282
+ # nothing to read it from and an empty forecast lost the column that says
1283
+ # which step an answer is about, just as `id` loses the column that
1284
+ # says which row.
1285
+ built._names_its_step = True
1286
+ built._names_its_rows = all(part._names_its_rows for part in answers)
1287
+ return built
1288
+
1289
+
1290
+ _WHAT_DIFFERS = (
1291
+ ('query', 'the query'),
1292
+ ('kind', 'the task kind'),
1293
+ ('target', 'the target column'),
1294
+ ('dtype', "the target column's dtype"),
1295
+ ('entity_table', 'the entity table'),
1296
+ )
1297
+
1298
+
1299
+ def _refuse_a_question_the_session_did_not_pin(
1300
+ pinned: Question, asked: Question
1301
+ ) -> None:
1302
+ r"""Require a scoring call to use the pinned question.
1303
+
1304
+ Said field by field: the rest of the request is identical, so printing both
1305
+ whole does not show which one differs.
1306
+ """
1307
+ if pinned == asked:
1308
+ return
1309
+ for field_name, described in _WHAT_DIFFERS:
1310
+ mine = getattr(pinned, field_name)
1311
+ theirs = getattr(asked, field_name)
1312
+ if mine != theirs:
1313
+ said = f'{described} pinned is {mine!r} and this asks {theirs!r}'
1314
+ break
1315
+ else:
1316
+ said = (
1317
+ 'the context pinned is not the one this request carries, and the '
1318
+ 'context is what the answer is learnt from'
1319
+ )
1320
+ raise StructuredDataError(
1321
+ f'this session pinned a different question: {said}. Only the rows '
1322
+ f'travel on a scoring call, so the pinned question is the one that '
1323
+ f'would be answered. Fit again to ask something else',
1324
+ code='INVALID_REQUEST',
1325
+ )
1326
+
1327
+
1328
+ def _require_a_request(request: Any) -> None:
1329
+ """Refuse anything that is not a request, before reading a field off it.
1330
+
1331
+ `errors.py` promises that `except StructuredDataError` catches everything
1332
+ this client reports. `_prepare` held the right message and ran after the
1333
+ first attribute was dereferenced, so the wrong type escaped as an
1334
+ `AttributeError` naming a field the caller never heard of.
1335
+
1336
+ Asked at the entry rather than in `_prepare`, because a rewritten question
1337
+ is answered in parts and never reaches it.
1338
+ """
1339
+ if not isinstance(request, Request):
1340
+ raise StructuredDataError(
1341
+ f'request must be a Request, got {type(request).__name__}; build '
1342
+ f'one with Request(context=..., predict=...) or '
1343
+ f'PredictiveQuery(...).materialize(...)',
1344
+ code='INVALID_REQUEST',
1345
+ )
1346
+
1347
+
1348
+ def _require_output_names(outputs: Any) -> None:
1349
+ """Refuse an `outputs` that is not a sequence of field names."""
1350
+ if outputs is None:
1351
+ return
1352
+ if isinstance(outputs, str):
1353
+ raise StructuredDataError(
1354
+ f'outputs must be a sequence of field names, not a single '
1355
+ f'string; pass [{outputs!r}]',
1356
+ code='INVALID_REQUEST',
1357
+ )
1358
+ if not isinstance(outputs, Sequence) or not all(
1359
+ isinstance(name, str) for name in outputs
1360
+ ):
1361
+ raise StructuredDataError(
1362
+ f'outputs must be a sequence of field names, got '
1363
+ f'{type(outputs).__name__}',
1364
+ code='INVALID_REQUEST',
1365
+ )
1366
+
1367
+
1368
+ def _refuse_to_pin_a_rewritten_question(request: Request) -> None:
1369
+ r"""A session pins one context, and a rewritten question has several.
1370
+
1371
+ A ranking's parts each carry their own labels for whether an entity reached
1372
+ *this* candidate. A forecast's parts each cover their own window. One
1373
+ pinned context can only be one of them, so pinning would answer the first
1374
+ part and give back something shaped like the whole question. Refusing says
1375
+ what to do instead.
1376
+ """
1377
+ rewritten = request._decomposition
1378
+ if rewritten is None:
1379
+ return
1380
+ raise StructuredDataError(
1381
+ f'a {rewritten.kind} is asked as {len(rewritten.parts)} questions, '
1382
+ f'each needing the context its own labels come from, and a session '
1383
+ f'pins one context. Ask it with predict(), which sends each and puts '
1384
+ f'the answers back together.',
1385
+ code='INVALID_REQUEST',
1386
+ )
1387
+
1388
+
1389
+ def _refuse_outputs_a_ranking_cannot_give(
1390
+ outputs: Sequence[str] | None,
1391
+ ) -> None:
1392
+ r"""A ranking gives back an order, and the fields that describe one.
1393
+
1394
+ The parts it was asked as are classifications, so their fields are the
1395
+ endpoint's; the merged answer's fields are this client's, and asking it
1396
+ for one it does not build would otherwise be answered with a result that
1397
+ names the field and does not carry it.
1398
+
1399
+ Naming none is refused the way every other request refuses it. A ranking
1400
+ was the one shape that read an empty list as *give me the usual*, so the
1401
+ one caller who had said what they wanted was the one who did not get it.
1402
+ """
1403
+ if outputs is not None and not outputs:
1404
+ raise StructuredDataError(
1405
+ 'outputs must name at least one field', code='INVALID_REQUEST'
1406
+ )
1407
+ unknown = [name for name in outputs or () if name not in _RANKING_OUTPUTS]
1408
+ if unknown:
1409
+ raise StructuredDataError(
1410
+ f'a ranking gives back {", ".join(_RANKING_OUTPUTS)}, and '
1411
+ f'{", ".join(repr(name) for name in unknown)} '
1412
+ f'{"are" if len(unknown) > 1 else "is"} not among them. The '
1413
+ f'per-candidate questions it is asked as are classifications; ask '
1414
+ f'one of those directly for the fields a classification gives.',
1415
+ code='INVALID_REQUEST',
1416
+ )