# Databricks notebook source
# MAGIC %md
# MAGIC # AI Search endpoint load test
# MAGIC
# MAGIC Generated on 2026-01-01.
# MAGIC
# MAGIC Sends queries to one index at a fixed rate in five stages: 25%, 50%, 75% and 100% of your Target QPS, then
# MAGIC 120% to show headroom. Each stage sends 5% more than its target so brief hiccups don't cause a shortfall;
# MAGIC results are judged on the target. It reports throughput and latency per stage, then a verdict.
# MAGIC
# MAGIC **These are real queries.** If the endpoint serves production traffic, run this when a few minutes of extra
# MAGIC load, up to 120% of Target QPS, are acceptable.
# MAGIC
# MAGIC Queries are sampled from the index's source table, and no query repeats within a run. Each run uses a new random
# MAGIC seed:
# MAGIC - For an index with self-managed embeddings, it sends vectors sampled from that table: first as they are, then
# MAGIC   each with a small random offset. If the table can't be read, it sends random vectors instead.
# MAGIC - For an index that computes its embeddings with a model serving endpoint, it sends text sampled from that
# MAGIC   table: first as it is, then with a number added. That endpoint embeds every query, so it gets this load too.
# MAGIC
# MAGIC Set **Query type** to HYBRID to test hybrid search (vector and keyword) instead of ANN. Hybrid queries send
# MAGIC the first 16 words of each sampled text. For an index that computes its embeddings, they send only that
# MAGIC text. For an index with self-managed embeddings, they send it beside the vector from the same row; the text
# MAGIC comes from the source table column named in **Text column**, which must be a string column the index syncs.
# MAGIC
# MAGIC It takes about three to five minutes and needs Databricks Runtime 13.3 LTS or above and permission to query the
# MAGIC index. An index that computes its embeddings also needs permission to read its source table, and may need
# MAGIC permission to query its embedding endpoint. Pick the index in the **Index** widget, set **Target QPS** if
# MAGIC needed, then click **Run all**.
# MAGIC
# MAGIC Set **Show diagnostics** to true to print more for each stage:
# MAGIC - Failures by kind, with a few sample error replies, plus connection, CPU and network figures.
# MAGIC - For an index that computes its embeddings, how server time splits between embedding and search.
# MAGIC
# MAGIC To check whether the embedding endpoint limits QPS, turn it on and look at the stages that fell short. If
# MAGIC embedding time is most of server time, or the sample error replies name a Model Serving endpoint (such as
# MAGIC "Request limit exceeded for Model Serving endpoint"), the embedding endpoint is the limit, not the AI Search
# MAGIC endpoint.
# MAGIC
# MAGIC The code cells after **Console run record** run inside a `try` so a run started from the console can record
# MAGIC its errors; if you edit them, keep their indentation.

# COMMAND ----------

# DBTITLE 1,Imports
import asyncio
import base64
import gc
import itertools
import json
import math
import multiprocessing
import os
import re
import secrets
import ssl
import textwrap
import threading
import time
import urllib.parse
import urllib.request
import warnings
from array import array
from collections import Counter

import numpy as np
import requests
from pyspark.sql import functions as F
from pyspark.sql.types import StringType
from requests.certs import where as requests_ca_bundle

# COMMAND ----------

# DBTITLE 1,Settings
# Set the values in the widgets above. The Advanced values at the end of this cell are optional.
# The names are base64-encoded so any characters in them are safe to embed in this code.
CONFIG = json.loads(base64.b64decode("eyJlbmRwb2ludF9uYW1lIjoiZSIsImluZGV4X25hbWVzIjpbImkiXX0=").decode("utf-8"))
ENDPOINT_NAME = CONFIG["endpoint_name"]
CONFIGURED_TARGET_QPS = 10

dbutils.widgets.dropdown("index", CONFIG["index_names"][0], CONFIG["index_names"], "Index")
dbutils.widgets.text("target_qps", "10", "Target QPS")
dbutils.widgets.text("num_results", "10", "Results per query")
dbutils.widgets.dropdown("query_type", "ANN", ["ANN", "HYBRID"], "Query type")
# Only for a hybrid test of an index with self-managed embeddings: the source table column its query text comes from.
dbutils.widgets.text("text_column", "", "Text column")
# Used when all three are filled in. Recommended, and usually required for an index that computes its embeddings.
dbutils.widgets.text("sp_client_id", "", "Service principal client ID (recommended)")
dbutils.widgets.text("sp_secret_scope", "", "Service principal secret scope (recommended)")
dbutils.widgets.text("sp_secret_key", "", "Secret key holding the client secret (recommended)")
dbutils.widgets.dropdown("debug", "false", ["false", "true"], "Show diagnostics")
dbutils.widgets.text("stage_fractions", "0.25,0.5,0.75,1.0,1.2", "Stages as shares of Target QPS (advanced)")
dbutils.widgets.text("stage_seconds", "24", "Seconds per stage (advanced)")
# Set only when the console runs this notebook as a job; it names the run in the console's run record.
dbutils.widgets.text("run_key", "", "Console run key (leave empty)")

def _int_widget(name, label, low, high):
    text = dbutils.widgets.get(name).strip()
    try:
        value = int(text)
    except ValueError:
        raise ValueError("%s must be a whole number, got '%s'." % (label, text)) from None
    if value < low:
        raise ValueError("%s must be at least %d." % (label, low))
    if value > high:
        print("%s is capped at %d." % (label, high))
    return min(value, high)

# Helpers the Settings cell needs: parsing for the advanced widgets, and redaction for any text the notebook prints or
# records. A widget value is only ever read as data and checked here; an invalid one falls back to its default with a
# note, so a console run never fails on a stage setting or query type.
import json
import re

MAX_STAGES = 10
MAX_STAGE_FRACTION = 5
MIN_STAGE_SECONDS = 5
MAX_STAGE_SECONDS = 600
STAGE_FRACTIONS_LABEL = "Stages as shares of Target QPS"
STAGE_SECONDS_LABEL = "Seconds per stage"
# The query types this test can send; the first is the default.
QUERY_TYPES = ("ANN", "HYBRID")
QUERY_TYPE_LABEL = "Query type"


def parse_query_type(text):
    """(query type, note): one of QUERY_TYPES, in any case; otherwise the default and a note saying why."""
    value = str(text).strip().upper()
    if value in QUERY_TYPES:
        return value, None
    return QUERY_TYPES[0], "The %s widget must be one of %s, got '%s'; using %s." % (
        QUERY_TYPE_LABEL, ", ".join(QUERY_TYPES), str(text).strip(), QUERY_TYPES[0])


def parse_stage_fractions(text, default):
    """(fractions, note): 1 to MAX_STAGES comma-separated shares in (0, MAX_STAGE_FRACTION], strictly ascending and
    including 1.0, the stage the verdict judges; otherwise default and a note saying why."""
    text = str(text).strip()
    try:
        values = [float(part) for part in text.split(",")]
    except ValueError:
        values = []
    if (1 <= len(values) <= MAX_STAGES and all(0 < value <= MAX_STAGE_FRACTION for value in values)
            and 1.0 in values and all(low < high for low, high in zip(values, values[1:]))):
        return values, None
    return default, ("The %s widget must be 1 to %d comma-separated numbers above 0 and at most %d, in increasing "
                     "order and including 1.0, got '%s'; using %s." % (
                         STAGE_FRACTIONS_LABEL, MAX_STAGES, MAX_STAGE_FRACTION, text,
                         ",".join(str(value) for value in default)))


def parse_stage_seconds(text, default):
    """(seconds, note): a whole number, held to MIN_STAGE_SECONDS..MAX_STAGE_SECONDS with a note; text that isn't a
    whole number gives default and a note."""
    text = str(text).strip()
    try:
        value = int(text)
    except ValueError:
        return default, "The %s widget must be a whole number, got '%s'; using %d." % (STAGE_SECONDS_LABEL, text, default)
    bounded = min(max(value, MIN_STAGE_SECONDS), MAX_STAGE_SECONDS)
    if bounded == value:
        return value, None
    return bounded, "The %s widget must be from %d to %d, got %d; using %d." % (
        STAGE_SECONDS_LABEL, MIN_STAGE_SECONDS, MAX_STAGE_SECONDS, value, bounded)


# Only this much of an error reply is searched for what it says or printed, since its start says what it is.
REPLY_PREFIX_BYTES = 4096
# A token-like word must not follow a letter or digit; unlike \b this also matches after "_", as in access_token.
NOT_AFTER_ALPHANUMERIC = r"(?<![A-Za-z0-9])"
# Authorization schemes whose next word is the credential.
AUTH_SCHEMES = ("bearer", "basic", "token", "digest")
# The only words kept after a scheme or an Authorization header, since error text uses them there ("Basic auth is not
# supported"); anything else that follows is taken for a credential, whatever its length or case.
NOT_CREDENTIALS = ("auth", "authentication", "token", "header", "credentials", "required", "failed", "invalid",
                   "missing", "expired", "is", "not", "supported")
# A value runs to the next space, quote, comma or semicolon.
VALUE = r'([^ "\',;]+)'
# Between a scheme or header and its value: a space, or one or two of ":" and "=" (as in ":=" or "==") with optional
# spaces, then an optional quote, after a "[" when the value is a quoted array item (refresh_token=["abc"]).
SEPARATOR = r'((?: ?[:=](?: ?[:=])? ?| )(?:\[(?=[\'"]))?[\'"]?)'
# One of NOT_CREDENTIALS as a whole word, matched case-sensitively like the checks below; a ":" or "=" may follow it
# ("token is: abc").
NOT_CREDENTIAL_WORD = r'(?-i:%s)(?![^ "\',;:=])' % "|".join(sorted(NOT_CREDENTIALS, key=len, reverse=True))
# An Authorization header, its first word, and the word after it if any.
HEADER_VALUE = (NOT_AFTER_ALPHANUMERIC + r'(authorization[\'"]? ?[:=] ?[\'"]?)' + VALUE + "(?:" + SEPARATOR + VALUE
                + ")?")
# "token" is also an ordinary word in error text ("the OAuth token is bound to the caller's client certificate"), so
# after it only a value shaped like a credential is taken for one: quoted, or with a digit or "=", or 16 characters,
# which every Databricks token and secret has while English words almost never do. "token=" and "token:" are
# key=value secrets below, which redact any value. A "token" ending a name (refresh_token, x-databricks-token) is a
# header or field, never prose, so it keeps the scheme rule.
PROSE_SCHEMES = ("token",)
# Error codes ending in a key name, which name no key: "INVALID_TOKEN: The token expired." is printed, then redacted
# again with the rest of an error message, so its text must stay. The platform's codes, and OAuth's invalid_token
# (RFC 6750). Only as "<code>: " and before a value not shaped like a credential (no quote or digit or "=", under 16
# characters); any other name (API_TOKEN: x) or form (INVALID_TOKEN=x) is a key.
KEY_NAMED_ERROR_CODES = ("INVALID_TOKEN", "EXPIRED_TOKEN", "invalid_token")
NOT_ERROR_CODE_KEY = "(?!%s)" % "|".join(
    r"(?-i:(?<=(?<![\w-])%s_)%s): (?![\'\"])(?![^ \"\',;]*[0-9=])(?![^ \"\',;]{16})" % tuple(code.rsplit("_", 1))
    for code in KEY_NAMED_ERROR_CODES)
NAME_SCHEMES = (r"(?<=[_-])" + NOT_ERROR_CODE_KEY + "token",)
CREDENTIAL_SHAPED_VALUE = r'((?<=[\'"])[^ "\',;]+|(?=[^ "\',;]*[0-9=])[^ "\',;]+|[^ "\',;]{16,})'


def _scheme_value(schemes, value):
    """A scheme, any run of allowed words or repeated scheme words, then the credential: the first word that isn't
    allowed and matches value."""
    return (NOT_AFTER_ALPHANUMERIC + "(" + "|".join(schemes) + ")" + SEPARATOR
            # An allowed word that is also a scheme ("token") is matched only as allowed, so a run of them can't
            # backtrack exponentially.
            + "((?:(?:" + NOT_CREDENTIAL_WORD + "|(?!" + NOT_CREDENTIAL_WORD + ")(?:" + "|".join(AUTH_SCHEMES)
            + r')(?![^ "\',;:=]))[:=]*'
            # Any ":" or "=" right before the value is kept with the allowed words, all of it, so the value is never a
            # separator ("token = abc" would print abc) and a credential starting with them ("bearer '=abc'") is
            # still redacted after them.
            + r' [\'"]?)*[:=]*)(?![:=]|' + NOT_CREDENTIAL_WORD + ")" + value)


SCHEME_VALUE = _scheme_value([scheme for scheme in AUTH_SCHEMES if scheme not in PROSE_SCHEMES] + list(NAME_SCHEMES),
                             VALUE)
PROSE_SCHEME_VALUE = _scheme_value(PROSE_SCHEMES, CREDENTIAL_SHAPED_VALUE)


def _redact_scheme_value(match):
    """The credential after a scheme redacted, keeping the allowed words before it, as in "Bearer token <credential>"."""
    return match.group(1) + match.group(2) + match.group(3) + "[redacted]"


def _redact_header_value(match):
    """The header's value redacted. After a known scheme or an allowed word it stays readable and the next word is
    redacted unless allowed; after any other first word that word and the next are redacted, since the first may be
    an unknown scheme and the next its credential."""
    prefix, first, separator, second = match.groups()
    readable = first in NOT_CREDENTIALS or first.lower() in AUTH_SCHEMES
    shown = first if readable else "[redacted]"
    if second is None:
        return prefix + shown
    return prefix + shown + separator + (second if second in NOT_CREDENTIALS else "[redacted]")


# Anything that looks like a credential: the value after an Authorization scheme or header, Databricks tokens and
# secrets (the prefixes and alphabet of common/logging/redactor/filter/src/rules/gitleaks.toml; dkea and dose at their
# full 32 characters, so words like "dosage" stay), JWTs, and key=value secrets, whose values (passwords especially)
# can be short.
SECRET_PATTERNS = [(re.compile(pattern, re.IGNORECASE), replacement) for pattern, replacement in (
    (HEADER_VALUE, _redact_header_value),
    (SCHEME_VALUE, _redact_scheme_value),
    (PROSE_SCHEME_VALUE, _redact_scheme_value),
    (NOT_AFTER_ALPHANUMERIC + r"(dapi|dsapi)[a-h0-9]{8,}(-\d)?", "[redacted]"),
    (NOT_AFTER_ALPHANUMERIC + r"(dkea|dose)[a-h0-9]{32}", "[redacted]"),
    (r"eyJ[\w-]+\.[\w-]+\.[\w-]*", "[redacted]"),
    # A value a rule above already redacted isn't redacted again into "[[", while an array (token=["abc"]) still is.
    # Allowed words may come before an "=" or ":=" ("token is = abc"), but not before a lone ":", which ends prose too
    # ("the token expired: sign in again").
    (NOT_AFTER_ALPHANUMERIC + NOT_ERROR_CODE_KEY + r'(token|secret|password|api[_-]?key)((?:(?: '
     + NOT_CREDENTIAL_WORD + r')+(?= ?:?=))?'
     r'\W{0,3}[:=](?!\W{0,3}redacted\])\W{0,3})[^\s",;&]+', r"\1\2[redacted]"))]


def _printable(data, limit):
    """Up to limit printable characters of data (bytes), with anything token-like redacted. Every character that isn't
    printable ASCII (tabs, NUL, a no-break space, any other whitespace) becomes one space first, so none can split a
    scheme from its credential and hide it from the patterns."""
    text = "".join(char if " " <= char <= "~" else " " for char in data[:REPLY_PREFIX_BYTES].decode("utf-8", "replace"))
    text = re.sub(" +", " ", text)
    for pattern, replacement in SECRET_PATTERNS:
        text = pattern.sub(replacement, text)
    return text.strip()[:limit]


# The fields of a JSON error reply that say what went wrong, most specific first; the rest (request IDs, details) would
# only push the explanation past the limit.
REPLY_TEXT_FIELDS = ("message", "error_description", "error")
# A reply is parsed up to this size, beyond the printed REPLY_PREFIX_BYTES, since a long details list can push a JSON
# reply past it and leave its message unparsed.
REPLY_JSON_MAX_BYTES = 65536
# Only an error code of this shape is shown, as _reply_error_code does: anything else may be a value from the request.
REPLY_ERROR_CODE = re.compile(r"[A-Z_]{1,64}")
# An OAuth error reply has no error_code but names its code in "error" (invalid_client), in this shape.
REPLY_OAUTH_ERROR = re.compile(r"[a-z_]{1,64}")


def _reply_text(data, limit, with_error_code=True):
    """What an error reply (bytes) says, for an error message, redacted like _printable:
    - from a JSON object, "<code>: <text>": its error_code (unless with_error_code is False) or else an OAuth "error",
      and the first non-empty of REPLY_TEXT_FIELDS;
    - from a JSON reply with neither, the whole reply as it is, or nothing when its error_code is left out;
    - from other text, the whole reply ended as a sentence, so text after it reads apart from it.
    At most limit characters, ending in "..." when cut."""
    try:
        body = json.loads(bytes(data[:REPLY_JSON_MAX_BYTES]).decode("utf-8"))
        is_json = True
    # A deeply nested reply exceeds the parser's recursion limit.
    except (ValueError, RecursionError):
        body, is_json = None, False
    parts = [data]
    if isinstance(body, dict):
        oauth_error = body.get("error") if "error_code" not in body else None
        oauth_error = oauth_error if isinstance(oauth_error, str) and REPLY_OAUTH_ERROR.fullmatch(oauth_error) else None
        text = next((body[field] for field in REPLY_TEXT_FIELDS
                     if isinstance(body.get(field), str) and body[field].strip() and body[field] != oauth_error), None)
        code = body.get("error_code") if with_error_code else None
        code = (code if isinstance(code, str) and REPLY_ERROR_CODE.fullmatch(code) else None) or oauth_error
        if text or code:
            parts, is_json = [part.encode("utf-8") for part in (code, text) if part], False
        elif not with_error_code and isinstance(body.get("error_code"), str) and REPLY_ERROR_CODE.fullmatch(
                body["error_code"]):
            # The caller already shows the code, which is all the reply says.
            return ""
    # The code and the text are redacted apart, so a code ending in a key name (INVALID_TOKEN) can't redact the text's
    # first word as its value.
    text = ": ".join(part for part in (_printable(part, REPLY_PREFIX_BYTES) for part in parts) if part)
    # A JSON reply shown whole already ends in its closing bracket or quote, which a full stop would only clutter.
    if text and not is_json and text[-1] not in ".!?" and len(text) < limit:
        text += "."
    if len(text) > limit:
        return text[:limit - 3].rstrip() + "..." if limit >= 3 else text[:limit]
    return text

INDEX_NAME = dbutils.widgets.get("index")
# Not capped at the endpoint's Target QPS: testing above it shows where it starts turning load away, and a warning
# below says so once the endpoint's current Target QPS is read.
TARGET_QPS = _int_widget("target_qps", "Target QPS", 1, 1000000)
NUM_RESULTS = _int_widget("num_results", "Results per query", 1, 1000)
QUERY_TYPE, _query_type_note = parse_query_type(dbutils.widgets.get("query_type"))
if _query_type_note:
    print(_query_type_note)
HYBRID_QUERIES = QUERY_TYPE == "HYBRID"
TEXT_COLUMN = dbutils.widgets.get("text_column").strip() if HYBRID_QUERIES else ""
SP_CLIENT_ID = dbutils.widgets.get("sp_client_id").strip()
SP_SECRET_SCOPE = dbutils.widgets.get("sp_secret_scope").strip()
SP_SECRET_KEY = dbutils.widgets.get("sp_secret_key").strip()
SP_CONFIGURED = bool(SP_CLIENT_ID and SP_SECRET_SCOPE and SP_SECRET_KEY)
DEBUG = dbutils.widgets.get("debug") == "true"
RUN_KEY = dbutils.widgets.get("run_key").strip()
# A console run can't run copies of itself or show diagnostics, so its advice differs.
CONSOLE_RUN = bool(RUN_KEY)

# Advanced (optional): a fixed seed to repeat a run's queries (None picks a new one). The advanced widgets set each
# stage as a share of Target QPS (keep 1.0, which the verdict judges) and how long each stage runs; these defaults are
# used when a widget isn't valid.
DEFAULT_STAGE_FRACTIONS = [0.25, 0.5, 0.75, 1.0, 1.2]
DEFAULT_STAGE_SECONDS = 24
QUERY_SEED = None
STAGE_FRACTIONS, _stage_note = parse_stage_fractions(dbutils.widgets.get("stage_fractions"), DEFAULT_STAGE_FRACTIONS)
if _stage_note:
    print(_stage_note)
STAGE_SECONDS, _stage_note = parse_stage_seconds(dbutils.widgets.get("stage_seconds"), DEFAULT_STAGE_SECONDS)
if _stage_note:
    print(_stage_note)

FULL_STAGE = STAGE_FRACTIONS.index(1.0)
STAGE_QPS = [max(1, int(TARGET_QPS * fraction)) for fraction in STAGE_FRACTIONS]
SEND_RATE_FACTOR = 1.05
SEED = QUERY_SEED if QUERY_SEED is not None else secrets.randbits(32)
PRINT_WIDTH = 100
# Error replies printed in messages are cut to this many characters.
ERROR_TEXT_CHARS = 300


def print_wrapped(text, indent=""):
    """Prints prose wrapped to PRINT_WIDTH, with continuation lines indented two more spaces."""
    print(textwrap.fill(text, width=PRINT_WIDTH, initial_indent=indent, subsequent_indent=indent + "  ",
                        break_long_words=False, break_on_hyphens=False))

def _capitalized(text):
    return text[:1].upper() + text[1:]

print_wrapped("Load testing %s on endpoint %s%s: stages of %d s at %s QPS (query seed %d)." % (
    INDEX_NAME, ENDPOINT_NAME, " with hybrid queries" if HYBRID_QUERIES else "", STAGE_SECONDS,
    ", ".join("{:,}".format(qps) for qps in STAGE_QPS), SEED))
# Default endpoint cap on hybrid results; endpoints can change it, so only warn: the test query's 400 is authoritative.
HYBRID_MAX_RESULTS = 200
if HYBRID_QUERIES and NUM_RESULTS > HYBRID_MAX_RESULTS:
    print_wrapped("WARNING: Endpoints allow at most {} results per hybrid query unless set otherwise, and this test "
                  "asks for {:,}.".format(HYBRID_MAX_RESULTS, NUM_RESULTS))

# COMMAND ----------

# DBTITLE 1,Console run record
# The run record the console reads when it runs this notebook as a job (see "Running from the console" in
# loadTestNotebookChecks.md). Each run rewrites only its own entry of one shared per-user file.
import datetime
import glob
import json
import math
import os
import secrets
import time

PROGRESS_FILE_NAME = ".ai_search_load_tests.json"
PROGRESS_FILE_VERSION = 1
# Beyond this many runs, the oldest droppable entries are removed; recent and live runs are kept.
MAX_RECORDED_RUNS = 20
# An entry still in progress after this long without an update belongs to a run that was killed, so it can be dropped.
STALE_RUN_SECONDS = 24 * 60 * 60
# A finished entry is kept at least this long after its last update, so an automatic retry still finds it.
FINISHED_GRACE_SECONDS = 15 * 60
# Temporary files older than this were left by a write that was interrupted, so they are removed.
STALE_TEMPORARY_FILE_SECONDS = 60 * 60
# Statuses after which nothing else is recorded, so a later exception can't turn a finished run into a failed one.
FINAL_STATUSES = ("succeeded", "failed", "aborted_retry", "cancelled")
RETRY_SKIPPED_ERROR = ("Automatic retry skipped so the endpoint isn't loaded twice. Start a new run from the "
                       "console.")
RUN_ERROR_CHARS = 500
RUN_CANCELLED_ERROR = "Run was cancelled."
RUN_TIMED_OUT_ERROR = "Run timed out."
# A job cancel reaches the notebook as one of these. NotebookExit's message carries cluster details, so it is never
# stored; it is matched by name because it can't be imported outside Databricks.
CANCEL_EXCEPTION_NAMES = ("KeyboardInterrupt", "NotebookExit", "RunCancelled")
# How often a console run asks the Jobs API whether it was cancelled, since the platform's own cancel isn't guaranteed
# to reach the notebook while it sends load.
CANCEL_CHECK_SECONDS = 5.0
# Where the notebook context keeps a run id, tried in order: the parent run first, since the console submits and cancels
# that one, then this task's own run.
RUN_ID_ATTRIBUTES = ("multitaskParentRunId", "rootRunId", "currentRunId", "runId")
# The same in the context's toJson, which serverless doesn't offer; kept for other compute.
RUN_ID_PATHS = (("tags", "multitaskParentRunId"), ("rootRunId", "id"), ("currentRunId", "id"), ("tags", "jobRunId"),
                ("tags", "runId"))
# A run record that exists but can't be read is read this many times, waiting longer each time, before giving up.
READ_ATTEMPTS = 3
READ_BACKOFF_SECONDS = 0.5

# Why requests or a run failed, as the record carries it (error_cause, preflight.cause, each stage's errors_by_cause);
# the console maps the codes to its own text, and the notebook prints these next steps for a failed test query. "other"
# is a stage's failures with none of the named causes; "unknown" a run that failed for a reason not classified.
CAUSE_NEXT_STEPS = {
    "throttled": ("the endpoint turned requests away, most likely under load; it is provisioned for its own Target QPS "
                  "of 10-result queries, so lower this test's Target QPS or Results per query, or raise the endpoint's "
                  "Target QPS"),
    "request_limited": ("Databricks' request rate limit turned these requests away before they reached the endpoint; "
                        "contact Databricks support about this workspace's limit for AI Search"),
    "embedding_rate_limited": ("the model serving endpoint that embeds each query turned requests away; increase its "
                               "throughput and check its rate limits, or compute query embeddings yourself and send "
                               "them as query_vector, alongside query_text for hybrid search "
                               "(https://docs.databricks.com/api/ai-search/v1/query-index)"),
    "capacity_limited": ("the endpoint or index reached its limit on stored data, which isn't about the request rate; "
                         "delete some data, or contact Databricks support"),
    "timeout": ("queries ran past the time limit, which expensive queries (many results per query, filters, hybrid "
                "search) or an overloaded endpoint cause; return fewer results per query or lower Target QPS"),
    "server_error": ("the endpoint, or something in front of it, returned 5xx errors; rerun, and if it repeats, contact "
                     "Databricks support"),
    "auth": ("the identity this test signed in as can't query this index; ask for permission to query it (and its "
             "embedding endpoint, if the index computes its embeddings), or use a service principal that has it, then "
             "rerun"),
    "service_principal_required": ("use a service principal whose OAuth secret has the vector-search and "
                                   "model-serving scopes (or all-apis), then rerun: in this notebook, fill in the "
                                   "three service principal widgets; in the console, add it under Advanced options > "
                                   "Service principal"),
    "service_principal_sign_in": ("check the service principal's client ID and the secret scope and key that hold "
                                  "its OAuth secret, then rerun"),
    "service_principal_scope": ("create an OAuth secret for the service principal with the vector-search and "
                                "model-serving scopes (or all-apis); if it has them, give the service principal Can "
                                "Query on the embedding endpoint; then rerun"),
    "not_found": "the index or endpoint wasn't found; check the index and endpoint names, then rerun",
    "bad_request": "the endpoint rejected the query itself; check Results per query and the index's settings, then rerun",
    "unreachable": ("this cluster couldn't connect to the endpoint; check that it can reach the endpoint, with no "
                    "proxy or firewall in the way, then rerun"),
    "no_reply": ("connections closed or failed to open before any reply; check the network between this cluster and "
                 "the endpoint, then rerun"),
    # The console's job timeout headroom (LOAD_TEST_TIMEOUT_HEADROOM_SECONDS in loadTestRunRecord.ts).
    "run_timed_out": ("the job ran past its time limit, which allows the stages' own time plus 15 minutes for starting "
                      "up and preparing queries; rerun, and if it repeats, contact Databricks support"),
    "workspace_rate_limited": ("the workspace's API rate limit was hit before any load was sent; wait a minute, then "
                               "rerun"),
    "unsupported_index": ("this test samples query text from the index's source table, and this index has none; test "
                          "an index that has a source table"),
    "source_table_access": ("this test samples query text from the index's source table; get permission to read it, "
                            "then rerun"),
    "no_query_text": ("this test samples query text from the index's source table, and its text column is empty; add "
                      "text to it, then rerun"),
    "index_not_ready": "the index has no query URL yet, so it may still be provisioning; rerun once it is online",
    "hybrid_too_many_results": ("the endpoint caps Results per query for hybrid search; lower Results per query to the "
                                "maximum the error names, then rerun"),
    "hybrid_needs_text_field": ("hybrid search needs at least one text column synced to the index, and this index has "
                                "none; sync a string column, or run an ANN load test"),
    "hybrid_text_column": ("the chosen text column isn't a string column that the index syncs; hybrid tests of "
                           "self-managed indexes take query text from such a column, so pick one, then rerun"),
    "other": "rerun, and if it repeats, contact Databricks support (Show diagnostics prints sample replies)",
    "unknown": "see the error; rerun, and if it repeats, contact Databricks support",
}
# Next steps narrower than their cause's, by the record's error_detail; keep in sync with getLoadTestCauseAction.
DETAIL_NEXT_STEPS = {
    # The sign-in's 401 invalid_client (WorkspaceOAuthClientAuthenticator.scala, OAuthClientAuthenticator.scala).
    "invalid_client": ("check that the client ID is the service principal's application ID (a UUID), not its numeric "
                       "ID or display name; that the secret is a current Databricks OAuth secret created for that "
                       "same service principal, not a Microsoft Entra ID secret, and not expired or disabled; "
                       "that the service principal is added to this workspace; and that the secret scope and key you "
                       "entered hold that secret; then rerun"),
}
# How the notebook names each cause when it prints failures.
CAUSE_LABELS = {"throttled": "429s from the endpoint",
                "request_limited": "429s from the request rate limit",
                "embedding_rate_limited": "429s from the embedding endpoint",
                "capacity_limited": "429s at the storage capacity limit", "timeout": "timeouts",
                "server_error": "5xx errors", "auth": "401 or 403 replies",
                "service_principal_required": "sign-ins that need a service principal",
                "service_principal_sign_in": "service principal sign-ins that failed",
                "service_principal_scope": "OAuth secrets missing a scope", "not_found": "404 replies",
                "bad_request": "other 4xx replies", "unreachable": "connection failures",
                "no_reply": "requests with no reply", "run_timed_out": "job timeouts",
                "workspace_rate_limited": "workspace API rate limits",
                "unsupported_index": "an index this test can't query",
                "source_table_access": "source table not readable", "no_query_text": "no query text",
                "index_not_ready": "index not ready",
                "hybrid_too_many_results": "too many results for hybrid search",
                "hybrid_needs_text_field": "no text column for hybrid search",
                "hybrid_text_column": "text column not usable for hybrid search",
                "other": "other failures", "unknown": "other failures"}
# Conditions the console shows above a run's results. The console mirrors these codes. client_connections_saturated is
# no longer written, since the verdict says it, but stays so older records still parse.
RUN_WARNINGS = ("queries_via_workspace", "above_endpoint_target_qps", "client_connections_saturated",
                "endpoint_unreadable", "workspace_route_fallback")
# The run warnings about the load sent, and the phases a run fails in before it sends any (ProgressFile.at). The route
# and endpoint_unreadable warnings stay: they are about the index and endpoint, and the route can explain a test query
# that was refused.
LOAD_WARNINGS = ("above_endpoint_target_qps", "client_connections_saturated")
PHASES_BEFORE_LOAD = ("setup", "sign_in", "lookup", "preflight")
# A test query failing with one of these would fail every query, so the run stops before its first stage.
FATAL_CAUSES = ("auth", "service_principal_required", "service_principal_scope", "not_found", "bad_request",
                "hybrid_too_many_results", "hybrid_needs_text_field")
# The order of a stage's errors_by_cause; each failed request is in exactly one. "throttled" is the endpoint's own
# 429s: overload replies and any 429 whose reply names no limit.
STAGE_CAUSES = ("throttled", "request_limited", "embedding_rate_limited", "capacity_limited", "timeout", "server_error",
                "auth", "not_found", "bad_request", "no_reply", "other")
# Lowercased reply text naming what refused a request; embedding-call errors name the "Model Serving endpoint".
EMBEDDING_REPLY_MARKER = b"model serving endpoint"
# Databricks' request rate limit: "Too many requests. Please wait a moment and try again."
REQUEST_LIMIT_REPLY_MARKER = b"too many requests. please wait a moment"
# "Request timed out..." (429 or 504) or "Exceeded request timeout..." (429): an expensive query or overload.
TIMEOUT_REPLY_MARKERS = (b"request timed out", b"exceeded request timeout")
# Overload 429s: "rejected due to heavy load", or a bare "tooManyRequests" or "tooManyConcurrentRequests".
OVERLOAD_REPLY_MARKERS = (b"rejected due to heavy load", b"toomanyrequests", b"toomanyconcurrentrequests")
# "Endpoint (or Index) capacity has reached capacity limit" is about the data stored, not the request rate.
CAPACITY_LIMIT_REPLY_MARKER = b"capacity has reached capacity limit"
# Lowercased 400 text for a hybrid query the index or endpoint can't serve (BrickindexCommonMessages.scala
# numResultsExceedsHybridSearchLimit and missingTextFieldsForHybridSearch), by cause.
HYBRID_REPLY_MARKERS = (("hybrid_too_many_results", b"exceeds the maximum allowed for hybrid search"),
                        ("hybrid_needs_text_field", b"hybrid search requires at least one text column in the index"))
# The cause each reply_counter counter belongs to, so a reply has one cause wherever it is classified.
COUNTER_CAUSES = {"request_limited": "request_limited", "embedding_rate_limited": "embedding_rate_limited",
                  "embedding_errors": "server_error", "timed_out": "timeout", "overloaded": "throttled",
                  "capacity_limited": "capacity_limited"}
PREFLIGHT_ERROR_CHARS = 300
# Why the ramp stopped early, as skipped_stages records it: almost every request of a stage failed, or enough of them
# timed out or failed with 5xx errors that the endpoint was in distress (_stop_reason in loadTestStageCounts.py).
SKIP_REASONS = ("almost_all_failed", "endpoint_distress")


class LoadTestError(RuntimeError):
    """A failure whose cause (a CAUSE_NEXT_STEPS code) is known, so the record carries it as error_cause. With
    next_step, the message only says what failed, and the notebook prints the cause's next step when it stops."""

    def __init__(self, message, cause, next_step=False, detail=None):
        super().__init__(message)
        self.cause = cause
        self.next_step = next_step
        # A DETAIL_NEXT_STEPS code, recorded as error_detail.
        self.detail = detail


def next_step_for(exc):
    """The next step the notebook prints for a LoadTestError with next_step."""
    return DETAIL_NEXT_STEPS.get(getattr(exc, "detail", None)) or CAUSE_NEXT_STEPS[exc.cause]


def reply_counter(status, body_start=b"", embedding=False):
    """The ramp's counter (a COUNTER_CAUSES key) for one error reply, from its HTTP status and the start of its body, or
    None when the body names nothing; embedding says whether the index computes its embeddings. A reply naming the
    embedding endpoint is that endpoint's even when it says it timed out, so the verdict can blame it."""
    text = bytes(body_start).lower()
    if status == 429 and REQUEST_LIMIT_REPLY_MARKER in text:
        return "request_limited"
    if embedding and (status == 429 or status >= 500) and EMBEDDING_REPLY_MARKER in text:
        return "embedding_rate_limited" if status == 429 else "embedding_errors"
    if status in (429, 504) and any(marker in text for marker in TIMEOUT_REPLY_MARKERS):
        return "timed_out"
    if status == 429 and any(marker in text for marker in OVERLOAD_REPLY_MARKERS):
        return "overloaded"
    if status == 429 and CAPACITY_LIMIT_REPLY_MARKER in text:
        return "capacity_limited"
    return None


# Lowercased 401/403 text the gateway gives when the embedding endpoint won't take the test's token
# (BrickindexCommonMessages.scala); gateways older than the first two give only the generic ones. The first avoids the
# apostrophe in "caller's", which a JSON reply may escape.
CERTIFICATE_BOUND_REPLY_MARKER = b"client certificate, so ai search cannot forward it"
MODEL_SERVING_SCOPE_REPLY_MARKER = b"does not carry the model-serving scope"
EMBEDDING_DENIED_REPLY_MARKER = b"permission denied to invoke model serving endpoint"
# The gateway's 401 for an embedding call: the credential or its scope, not a missing Can Query.
EMBEDDING_UNAUTHENTICATED_REPLY_MARKER = b"authentication failed while invoking model serving endpoint"
# One wording for the warning before load ("usually ") and the test query's refusal (""): the notebook user's token is
# bound to this compute's certificate, so a gateway that passes it on to the embedding endpoint is refused.
NOTEBOOK_USER_CANNOT_EMBED = "the notebook user's sign-in %scan't be used to embed query text for this index"
NOTEBOOK_USER_TOKEN_REFUSED = (NOTEBOOK_USER_CANNOT_EMBED % "").capitalize() + "."


def _allows_model_serving(scope):
    return scope == "all-apis" or "model-serving" in str(scope).split()


def falls_back_to_workspace_route(cause, notebook_user, via_workspace, route_optimized, same_embedding_endpoint):
    """Whether a test query refused for the notebook user's certificate-bound token is retried through the workspace
    URL, which embeds through Model Serving's proxy without passing the token on. Not for a route-optimized embedding
    endpoint, which that proxy rejects (BrickindexCommonMessages.scala), nor when the workspace route would embed
    with another endpoint than the one this run read. An unreadable endpoint (None) is eligible; the retry decides."""
    return (cause == "service_principal_required" and notebook_user and not via_workspace
            and route_optimized is not True and same_embedding_endpoint)


def embedding_token_refusal(status, body_start, identity, notebook_user, via_workspace, route_optimized,
                            embedding_endpoint_name=None, sp_scope=None):
    """(cause, message) for a 401 or 403 saying the embedding endpoint refused the token of identity (the notebook
    user when notebook_user), else None. The notebook user's token is certificate-bound, so a service principal is
    the fix when the reply says so, or on the endpoint's own URL to an endpoint that isn't route-optimized, the one
    route that passes the token on unchanged. An endpoint that couldn't be read (route_optimized None) counts as not
    route-optimized, the usual case: the reply already shows the token was refused. sp_scope is the scope the service
    principal signed in with, None when unknown."""
    if status not in (401, 403):
        return None
    text = bytes(body_start).lower()
    # The one route that passes the caller's token on to the embedding endpoint.
    token_passed_on = not via_workspace and not route_optimized
    certificate_bound_route = notebook_user and token_passed_on
    if CERTIFICATE_BOUND_REPLY_MARKER in text:
        return ("service_principal_required", NOTEBOOK_USER_TOKEN_REFUSED) if notebook_user else None
    if MODEL_SERVING_SCOPE_REPLY_MARKER in text:
        if certificate_bound_route:
            return "service_principal_required", NOTEBOOK_USER_TOKEN_REFUSED
        if not notebook_user:
            return ("service_principal_scope",
                    "The service principal's OAuth secret needs the model-serving scope to send query text.")
        return None
    denied, unauthenticated = EMBEDDING_DENIED_REPLY_MARKER in text, EMBEDDING_UNAUTHENTICATED_REPLY_MARKER in text
    if not (denied or unauthenticated):
        return None
    if certificate_bound_route:
        return "service_principal_required", NOTEBOOK_USER_TOKEN_REFUSED
    # Not the notebook's _capitalized: this module runs on its own in its tests, with no notebook names.
    who, endpoint = identity[:1].upper() + identity[1:], embedding_endpoint_name or "of this index"
    # Older gateways send the generic text for a missing scope too, and a passed-on token without model-serving is
    # refused whatever the permissions, so Can Query alone wouldn't fix it. The other routes don't pass this token on
    # unchanged, so they keep the Can Query advice.
    if not notebook_user and token_passed_on and sp_scope is not None and not _allows_model_serving(sp_scope):
        return "service_principal_scope", (
            "The service principal's OAuth secret has the %s scope but not model-serving, which the embedding "
            "endpoint %s needs. Use a secret with the model-serving scope (or all-apis), and give %s Can Query on "
            "that endpoint if it doesn't have it." % (sp_scope, endpoint, identity))
    if unauthenticated:
        return "auth", "%s's token was rejected by the embedding endpoint %s%s." % (
            who, endpoint, "" if notebook_user else "; check that the service principal's OAuth secret has the "
                                                    "model-serving scope (or all-apis)")
    return "auth", "%s needs Can Query on the embedding endpoint %s." % (who, endpoint)


def count_reply(counts, status, body_start, embedding):
    """Adds one error reply to a stage's counters (a Counter), as failures_by_cause reads them."""
    counter = reply_counter(status, body_start, embedding)
    if counter is not None:
        counts[counter] += 1
        # 429 timeouts are among the 429s and 504 ones among the errors, so they are kept apart.
        if counter == "timed_out" and status == 429:
            counts["timed_out_429"] += 1


def reply_cause(status, body_start=b"", embedding=False):
    """The cause code of one failed reply; the same reply counted by a stage's counters gets the same cause from
    failures_by_cause."""
    counter = reply_counter(status, body_start, embedding)
    if counter is not None:
        return COUNTER_CAUSES[counter]
    if status in (401, 403):
        return "auth"
    if status == 404:
        return "not_found"
    if status == 408:
        return "timeout"
    if status == 429:
        return "throttled"
    if status >= 500:
        return "server_error"
    if status == 400:
        # Not counted per reply: every hybrid query gets the same 400, so the test query already stops the run.
        text = bytes(body_start).lower()
        for cause, marker in HYBRID_REPLY_MARKERS:
            if marker in text:
                return cause
    if 400 <= status < 500:
        return "bad_request"
    return "unknown"


def _http_errors(kinds, include):
    return sum(count for kind, count in kinds.items()
               if isinstance(kind, str) and kind.startswith("HTTP ") and kind[5:].isdigit() and include(int(kind[5:])))


def stage_timeouts(stage):
    """A stage's timed-out requests: replies saying so (counted as "timed_out"), 408s, which reply_cause also calls
    timeouts, and requests this notebook timed out. The verdict counts timeouts with this too."""
    kinds = stage.get("errors_by_kind") or {}
    return (stage.get("timed_out") or 0) + kinds.get("timeout", 0) + _http_errors(kinds, lambda code: code == 408)


def failures_by_cause(stage):
    """A stage's failed requests by cause, from the counts the ramp keeps (count_reply); each failure is in exactly one
    cause, the one reply_cause gives it."""
    kinds = stage.get("errors_by_kind") or {}
    timed_out, timed_out_429 = stage.get("timed_out") or 0, stage.get("timed_out_429") or 0
    request_limited, embedding_rate_limited, capacity_limited = (
        stage.get(key) or 0 for key in ("request_limited", "embedding_rate_limited", "capacity_limited"))
    causes = {
        "throttled": ((stage.get("rate_limited") or 0) - timed_out_429 - request_limited - embedding_rate_limited
                      - capacity_limited),
        "request_limited": request_limited,
        "embedding_rate_limited": embedding_rate_limited,
        "capacity_limited": capacity_limited,
        "timeout": stage_timeouts(stage),
        # 504s that said the query timed out are among the timeouts.
        "server_error": _http_errors(kinds, lambda code: code >= 500) - (timed_out - timed_out_429),
        "auth": _http_errors(kinds, lambda code: code in (401, 403)),
        "not_found": _http_errors(kinds, lambda code: code == 404),
        "bad_request": _http_errors(kinds, lambda code: 400 <= code < 500 and code not in (401, 403, 404, 408)),
        "no_reply": stage.get("no_reply") or 0,
    }
    causes = {key: max(0, int(count)) for key, count in causes.items()}
    failures = (stage.get("rate_limited") or 0) + (stage.get("errors") or 0) + (stage.get("no_reply") or 0)
    causes["other"] = max(0, int(failures) - sum(causes.values()))
    return {key: causes[key] for key in STAGE_CAUSES}


def main_cause(counts):
    """The cause with the most failures (the first in STAGE_CAUSES on a tie), or None when there are none."""
    cause = max(STAGE_CAUSES, key=lambda key: counts.get(key, 0))
    return cause if counts.get(cause, 0) else None


def preflight_record(status, attempts, http_status=None, cause=None, error=None):
    """The test query's outcome: "ok"; "degraded", when it kept failing in a way load can explain, so the ramp ran
    anyway; or "failed", with a fatal cause. error must already be redacted."""
    return {"status": status, "attempts": attempts, "http_status": http_status, "cause": cause,
            "error": None if error is None else str(error)[:PREFLIGHT_ERROR_CHARS]}


def skipped_stage_record(stage_index, target_qps, cause, reason="almost_all_failed"):
    """A stage the ramp didn't run because of how the one before it failed, mostly for cause: reason is one of
    SKIP_REASONS."""
    return {"stage_index": stage_index, "target_qps": json_number(target_qps), "skipped": True,
            "skip_reason": reason, "skip_cause": cause}


class RunCancelled(Exception):
    """The Jobs API says this run was cancelled; raised so the ramp stops its workers and the run is recorded."""


class RunTimedOut(Exception):
    """The Jobs API says this run hit its timeout; recorded as a failure, since nobody cancelled it."""


class RunRecordUnavailable(RuntimeError):
    """The console's run record can't be read or written, so a console run sends no load: it couldn't tell a retry."""


class SigningInUnconfirmed(RunRecordUnavailable):
    """signing_in couldn't be confirmed, so this attempt stops before its first request."""


class NewerRunRecord(ValueError):
    """The run record was written by a newer version of this notebook, so this one must not rewrite it."""


def progress_path(user):
    """The record's path in the user's home folder, or None when the user name can't be a single folder name."""
    if not isinstance(user, str) or not user or "/" in user or "\\" in user or user.startswith("."):
        return None
    return "/Workspace/Users/%s/%s" % (user, PROGRESS_FILE_NAME)


def utc_now():
    return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def json_number(value, digits=1):
    """A JSON-safe number: ints as they are, floats rounded, and None for anything missing or not finite."""
    if value is None or isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    try:
        value = float(value)
    except (TypeError, ValueError):
        return None
    return round(value, digits) if math.isfinite(value) else None


def stage_record(stage):
    """The fields of one _stage_result the console shows."""
    record = {key: json_number(stage.get(key)) for key in ("target_qps", "achieved_qps", "skipped_qps")}
    for key in ("sent", "rate_limited", "errors", "no_reply", "timed_out"):
        record[key] = json_number(stage.get(key))
    # embedding_ms is the endpoint's embedding time within server_ms, sampled from some replies; all None for an
    # index queried with vectors.
    for key in ("client_ms", "server_ms", "embedding_ms"):
        values = list(stage.get(key) or [])
        record[key] = [json_number(value) for value in (values + [None] * 4)[:4]]
    record["errors_by_cause"] = failures_by_cause(stage)
    return record


def _param_record(value):
    """A recommendation param as JSON: booleans and text as they are, numbers JSON-safe, anything else None."""
    if isinstance(value, bool) or isinstance(value, str):
        return value
    return json_number(value)


def _recommendation_record(recommendation):
    return {"code": str(recommendation["code"]), "text": str(recommendation["text"]),
            "params": {str(key): _param_record(value) for key, value in recommendation["params"].items()},
            "alternatives": [_recommendation_record(alternative)
                             for alternative in recommendation.get("alternatives") or []]}


def verdict_record(verdict):
    """The verdict as the console reads it. A verdict without findings or recommendations records them as None, which
    the console reads as a verdict to show by its sentences."""
    targets = verdict.get("targets") or {}
    findings, recommendations = verdict.get("findings"), verdict.get("recommendations")
    # One per finding, so the console can pair them; a verdict without them records None for each.
    codes = list(verdict.get("finding_codes") or [])
    codes = (codes + [None] * len(findings))[:len(findings)] if findings is not None else None
    return {"kind": verdict["kind"], "sentences": [str(sentence) for sentence in verdict["sentences"]],
            "targets": {key: json_number(targets.get(key), 2) for key in ("p95_ms", "p99_ms", "success_pct")},
            "findings": None if findings is None else [str(finding) for finding in findings],
            "finding_codes": None if codes is None else [None if code is None else str(code) for code in codes],
            "recommendations": None if recommendations is None else [
                _recommendation_record(recommendation) for recommendation in recommendations]}


def new_run(run_key, endpoint, index, target_qps, num_results, stage_seconds, stage_fractions, now, job_run_id=None,
            query_type="ANN", text_column=None):
    # job_run_id is a string so the console can pass it to the Jobs API as is; stopped_early is set when the run ends.
    return {"run_key": run_key, "endpoint": endpoint, "index": index, "target_qps": target_qps,
            "num_results": num_results, "stage_seconds": stage_seconds,
            "stage_fractions": [json_number(fraction, 4) for fraction in stage_fractions], "attempt": 0,
            "status": "starting", "started_at": now, "updated_at": now, "total_stages": len(stage_fractions),
            "completed_stages": 0, "current_stage_target_qps": None, "stages": [], "verdict": None, "error": None,
            "job_run_id": None if job_run_id is None else str(job_run_id), "stopped_early": False,
            # Added after the first version of the entry, so readers treat them as optional.
            "skipped_stages": [], "preflight": None, "failed_at": None, "error_cause": None, "signed_in_as": None,
            "sign_in_fallback_reason": None, "error_detail": None,
            # Set once the index and endpoint are read: "workspace" or "endpoint" (the URL queries go to), RUN_WARNING
            # codes, the endpoint's Target QPS, the name of the endpoint that embeds queries, and whether the index
            # computes its embeddings, since that name can be unknown when it does.
            "query_route": None, "warnings": [], "endpoint_target_qps": None, "embedding_endpoint": None,
            "computes_embeddings": None,
            # "ANN" or "HYBRID", and for a hybrid test of an index with self-managed embeddings the column its query
            # text comes from (a name, never its text).
            "query_type": query_type, "text_column": text_column or None}


def load_runs(path):
    """The file's runs by run key, {} only when there is no file. Raises OSError or ValueError if it can't be read or
    is in another format, including another version, so no caller ever rewrites a file it couldn't read."""
    try:
        with open(path) as progress_file:
            data = json.load(progress_file)
    except FileNotFoundError:
        return {}
    runs = data.get("runs") if isinstance(data, dict) else None
    version = data.get("version") if isinstance(data, dict) else None
    if isinstance(version, int) and not isinstance(version, bool) and version > PROGRESS_FILE_VERSION:
        raise NewerRunRecord("version %d" % version)
    if not isinstance(runs, dict) or version != PROGRESS_FILE_VERSION:
        raise ValueError("not a version %d run record" % PROGRESS_FILE_VERSION)
    return {key: run for key, run in runs.items() if isinstance(run, dict)}


def _started_at(run):
    started_at = run.get("started_at")
    return started_at if isinstance(started_at, str) else ""


def _parsed_time(text):
    try:
        return datetime.datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ")
    except (TypeError, ValueError):
        return None


def _droppable(entry, now):
    """Finished and past FINISHED_GRACE_SECONDS, so a later automatic retry still finds it; or in progress with no
    update for STALE_RUN_SECONDS, since a run killed outright never finishes. An unreadable updated_at is droppable."""
    updated_at, now = _parsed_time(entry.get("updated_at")), _parsed_time(now)
    if updated_at is None:
        return True
    age = (now - updated_at).total_seconds() if now is not None else float("inf")
    if entry.get("status") in FINAL_STATUSES:
        return age > FINISHED_GRACE_SECONDS
    return age > STALE_RUN_SECONDS


def merged_runs(runs, run):
    """runs with run's entry replaced. Beyond MAX_RECORDED_RUNS the oldest droppable runs go first, so a run the
    console may still be showing is kept."""
    runs = dict(runs)
    runs[run["run_key"]] = run
    droppable = sorted((key for key in runs if key != run["run_key"] and _droppable(runs[key], run.get("updated_at"))),
                       key=lambda key: _started_at(runs[key]))
    for key in droppable[:max(0, len(runs) - MAX_RECORDED_RUNS)]:
        del runs[key]
    return runs


def write_run(path, run, load=None, sleep=None, attempts=READ_ATTEMPTS):
    """Re-reads the file (up to attempts times) and replaces only this run's entry; the rename keeps a reader from
    seeing half a file. Raises RunRecordUnavailable, writing nothing, when an existing file can't be read, so other
    runs' entries survive."""
    content = json.dumps({"version": PROGRESS_FILE_VERSION,
                          "runs": merged_runs(_load_runs_patiently(path, load, sleep, attempts), run)},
                         allow_nan=False, separators=(",", ":"))
    # Random, so concurrent runs on different machines never share a temporary file.
    temporary_path = "%s.%s.tmp" % (path, secrets.token_hex(8))
    try:
        with open(temporary_path, "w") as temporary_file:
            temporary_file.write(content)
        os.replace(temporary_path, path)
    finally:
        if os.path.exists(temporary_path):
            os.remove(temporary_path)


def sweep_temporary_files(path, now=time.time):
    """Removes temporary files of interrupted writes; one still being written is younger than the cutoff."""
    for temporary_path in glob.glob(glob.escape(path) + ".*.tmp"):
        try:
            if now() - os.path.getmtime(temporary_path) > STALE_TEMPORARY_FILE_SECONDS:
                os.remove(temporary_path)
        except OSError:
            pass


def _as_json(run):
    return json.loads(json.dumps(run))


def _attempt(run):
    attempt = run.get("attempt")
    return attempt if isinstance(attempt, int) and not isinstance(attempt, bool) else 0


def _load_runs_patiently(path, load=None, sleep=None, attempts=READ_ATTEMPTS):
    # Looked up when called, so tests can stand in for the file system and the clock.
    load, sleep = load or load_runs, sleep or time.sleep
    for attempt in range(attempts):
        try:
            return load(path)
        except NewerRunRecord:
            # Deleting it would lose the newer notebook's runs, so the message says to regenerate this one instead.
            raise RunRecordUnavailable(
                "The console's run record %s in your home folder was written by a newer notebook version, so this "
                "run sends no load. Regenerate the notebook from the console." % PROGRESS_FILE_NAME) from None
        except (OSError, ValueError) as exc:
            failure = exc
            if attempt + 1 < attempts:
                sleep(READ_BACKOFF_SECONDS * 2 ** attempt)
    # Only a ValueError (bad JSON or another format) means the file is damaged and deleting it helps; the console
    # matches "Could not read the console's run record" to say so, so an OSError is worded apart.
    if isinstance(failure, ValueError):
        raise RunRecordUnavailable(
            "Could not read the console's run record %s in your home folder because it is damaged (%s), so this run "
            "sends no load: it can't tell whether it is an automatic retry. Delete the file, then start a new run "
            "from the console." % (PROGRESS_FILE_NAME, type(failure).__name__))
    raise RunRecordUnavailable(
        "Could not open the console's run record %s in your home folder (%s), so this run sends no load: it can't "
        "tell whether it is an automatic retry. Start a new run from the console." % (
            PROGRESS_FILE_NAME, type(failure).__name__))


def retry_entry(previous, now):
    """The entry for an automatic retry of a run that got past "starting": a finished run keeps its status and error,
    and one cut off mid-run becomes aborted_retry."""
    entry = dict(previous)
    entry.update(attempt=_attempt(previous) + 1, retry_skipped_at=now, updated_at=now)
    if previous.get("status") not in FINAL_STATUSES:
        entry.update(status="aborted_retry", error=RETRY_SKIPPED_ERROR)
    return entry


def _quick_write(path, run):
    write_run(path, run, attempts=1)


def confirmed_write(path, run_key, entry, load=None, write=None, sleep=None):
    """Writes entry and reads it back; True once the file holds it. Tried twice, since a concurrent writer can replace
    the file between the write and the read."""
    write = write or write_run
    for _ in range(2):
        try:
            write(path, entry)
            if _load_runs_patiently(path, load, sleep).get(run_key) == _as_json(entry):
                return True
        except Exception:
            # The same failures _write tolerates, so a final update never raises from here.
            pass
    return False


def begin_run(path, run_key, entry, now, load=None, write=None, sleep=None,
              sweep=sweep_temporary_files):
    """Starts this attempt's entry and says what to do: "run" sends the load; on an automatic retry "exit" (the run
    already succeeded) ends the job as it was and "fail" (it failed or was cut off) fails the job, both without
    sending any query. Raises RunRecordUnavailable when the record can't be read, or the new entry can't be confirmed,
    since without it the next retry couldn't be told apart."""
    write = write or write_run
    sweep(path)
    previous = _load_runs_patiently(path, load, sleep).get(run_key)
    if previous is not None and previous.get("status") != "starting":
        retry = retry_entry(previous, now)
        # An entry from before job_run_id existed takes this attempt's.
        if not retry.get("job_run_id"):
            retry["job_run_id"] = entry.get("job_run_id")
        entry = retry
        # Read back so a concurrent writer can't leave the console showing the earlier status; no load is sent
        # either way, so a write that still can't be confirmed is ignored.
        confirmed_write(path, run_key, entry, load, write, sleep)
        return entry, "exit" if previous.get("status") == "succeeded" else "fail"
    entry = dict(entry, status="starting", updated_at=now)
    if previous is not None:
        # That attempt sent nothing, but it is still counted.
        entry["attempt"] = _attempt(previous) + 1
    if confirmed_write(path, run_key, entry, load, write, sleep):
        return entry, "run"
    raise RunRecordUnavailable(
        "Could not save the console's run record %s in your home folder, so this run sends no load: an automatic "
        "retry couldn't tell that it already ran. Start a new run from the console." % PROGRESS_FILE_NAME)


def _run_id(value):
    if isinstance(value, str) and value.isdigit():
        value = int(value)
    return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None


def job_run_id(repl_context=None, context_json=None):
    """The run id to watch for a cancel: from the REPL context's attributes (dbruntime get_context()), else from the
    context's toJson, called only when needed since it can raise; None when neither has one."""
    for name in RUN_ID_ATTRIBUTES:
        try:
            value = _run_id(getattr(repl_context, name, None))
        except Exception:
            value = None
        if value:
            return value
    try:
        context = context_json() if context_json else None
    except Exception:
        return None
    for path in RUN_ID_PATHS:
        value = context
        for key in path:
            value = value.get(key) if isinstance(value, dict) else None
        if _run_id(value):
            return _run_id(value)
    return None


def run_end_reason(job_run):
    """Why a Jobs API runs/get reply shows the run ending: "cancelled", "timed_out", or None while it is live.

    Jobs API 2.2 (elastic-spark-common/api/messages/runs.proto, run_status.proto) gives termination_details.code
    USER_CANCELED (a user) or CANCELED (the platform), and result_state CANCELED or TIMEDOUT. A user's cancel is
    checked first, then a TIMEDOUT result, then any other cancel; loadTestRunApi.ts getJobOutcome uses the same order.
    Before either is set, only user_cancelled_or_timedout or the life cycle may show the run ending; that counts as a
    cancel unless the state message names a timeout, which is only trusted once the run is ending.
    """
    if not isinstance(job_run, dict):
        return None
    state = job_run.get("state") if isinstance(job_run.get("state"), dict) else {}
    status = job_run.get("status") if isinstance(job_run.get("status"), dict) else {}
    termination = status.get("termination_details") if isinstance(status.get("termination_details"), dict) else {}
    code, result_state = termination.get("code"), state.get("result_state")
    if code == "USER_CANCELED":
        return "cancelled"
    if result_state == "TIMEDOUT":
        return "timed_out"
    if code == "CANCELED" or result_state == "CANCELED":
        return "cancelled"
    if not (state.get("life_cycle_state") in ("TERMINATING", "TERMINATED")
            or state.get("user_cancelled_or_timedout") is True
            or status.get("state") in ("TERMINATING", "TERMINATED")):
        return None
    state_message = state.get("state_message") if isinstance(state.get("state_message"), str) else ""
    return "timed_out" if "timed out" in state_message.lower() or "timeout" in state_message.lower() else "cancelled"


def succeeded_fields(verdict, completed_stages, planned_stages):
    """The update that ends a run that finished its ramp; fewer stages than planned means it stopped early (see
    SKIP_REASONS)."""
    return {"status": "succeeded", "current_stage_target_qps": None, "verdict": verdict_record(verdict),
            "stopped_early": completed_stages < planned_stages}


class CancelWatch:
    """Raises RunCancelled or RunTimedOut from check() once the Jobs API shows the run ending; with no fetch it
    never does."""

    def __init__(self, fetch, interval=CANCEL_CHECK_SECONDS, clock=time.monotonic, printer=print, debug=False):
        self.fetch = fetch
        self.interval = interval
        self.clock = clock
        self.printer = printer
        self.debug = debug
        self.last_check = None
        self.warned = False

    def check(self, force=False):
        """Asks at most once per interval unless forced; an API error counts as not cancelled."""
        if self.fetch is None:
            return
        now = self.clock()
        if not force and self.last_check is not None and now - self.last_check < self.interval:
            return
        self.last_check = now
        try:
            job_run = self.fetch()
        except Exception as exc:
            if self.debug and not self.warned:
                self.warned = True
                self.printer("Could not check whether this run was cancelled (%s)." % type(exc).__name__)
            return
        reason = run_end_reason(job_run)
        if reason == "timed_out":
            raise RunTimedOut(RUN_TIMED_OUT_ERROR)
        if reason == "cancelled":
            raise RunCancelled(RUN_CANCELLED_ERROR)


class ProgressFile:
    """This run's entry, written on every update; with no path it only keeps the entry in memory."""

    def __init__(self, path, run, write=None, load=None, now=utc_now, printer=print, debug=False, quick_write=None):
        self.path = path
        self.run = run
        self.write = write
        # Updates mid-run read the file once, so a slow or failing read doesn't hold up the next stage; statuses that
        # must stick go through confirmed_write with the patient read.
        self.quick_write = quick_write or write or _quick_write
        self.load = load
        self.now = now
        self.printer = printer
        self.debug = debug
        self.warned = False
        # Where the run is, for failed_at; kept in memory only, so marking it costs no write.
        self.phase, self.stage_index = "setup", None

    def at(self, phase, stage_index=None):
        """Marks where the run is: "setup", "sign_in", "lookup" (reading the index and endpoint), "preflight",
        "stage" (with its index) or "results"."""
        self.phase, self.stage_index = phase, stage_index

    def note(self, **fields):
        """Sets fields without writing them; the next update or the final status carries them."""
        self.run.update(fields)

    def _write(self):
        started = time.monotonic()
        try:
            self.quick_write(self.path, self.run)
        except Exception as exc:
            # Only the type, as elsewhere in this notebook, because messages can carry paths or credentials.
            if not self.warned:
                self.warned = True
                self.printer("WARNING: Could not update the console's run record (%s); the load test continues."
                             % type(exc).__name__)
            return
        if self.debug:
            self.printer("Run record written in %.0f ms." % (1000.0 * (time.monotonic() - started)))

    def update(self, **fields):
        self.run.update(fields)
        self.run["updated_at"] = self.now()
        if self.path is None:
            return
        if self.run.get("status") not in FINAL_STATUSES:
            self._write()
            return
        # The console learns a run failed or was cancelled only from this entry, so the last status is read back and
        # written again if a failed read or a concurrent writer kept it out of the file.
        started = time.monotonic()
        if confirmed_write(self.path, self.run["run_key"], self.run, self.load, self.write):
            if self.debug:
                self.printer("Run record written in %.0f ms." % (1000.0 * (time.monotonic() - started)))
            return
        self.printer("WARNING: Could not save this run's final status (%s) to the console's run record, so the console "
                     "may show it as still running." % self.run["status"])

    def update_confirmed(self, **fields):
        """Like update, but read back and written again if a concurrent writer undid it; never raises. Used for the last
        status before the first query, so a stale write can't leave the entry at "starting" once load may be sent."""
        self.run.update(fields)
        self.run["updated_at"] = self.now()
        if self.path is not None:
            # A False is ignored: signing_in was already confirmed, and the next update writes the entry again.
            confirmed_write(self.path, self.run["run_key"], self.run, self.load, self.write)

    def confirm(self, **fields):
        """Like update, but raises RunRecordUnavailable unless the file is read back holding the entry. Used for the
        first status past "starting", so an automatic retry can never find "starting" after load was sent."""
        self.run.update(fields)
        self.run["updated_at"] = self.now()
        if self.path is not None and not confirmed_write(self.path, self.run["run_key"], self.run, self.load,
                                                         self.write):
            raise SigningInUnconfirmed(
                "Could not save the console's run record %s in your home folder, so this run sends no load: an "
                "automatic retry couldn't tell that it already started. Start a new run from the console."
                % PROGRESS_FILE_NAME)

    def fail(self, message, cause="unknown", detail=None):
        if self.run.get("status") not in FINAL_STATUSES:
            fields = {}
            if self.phase in PHASES_BEFORE_LOAD and not self.run.get("stages"):
                fields["warnings"] = [code for code in self.run.get("warnings") or [] if code not in LOAD_WARNINGS]
            self.update(status="failed", error=message[:RUN_ERROR_CHARS], error_cause=cause, error_detail=detail,
                        failed_at={"phase": self.phase, "stage_index": self.stage_index}, **fields)

    def record_exception(self, exc, redact):
        """Records what ended the run: a cancel, or any other exception, redacted. A run that already finished keeps
        its status, so the notebook's own exit after success is never taken for a cancel."""
        if self.run.get("status") in FINAL_STATUSES:
            return
        # Raised before this attempt's first request. The file holds "starting" if no write landed, so the retry may run,
        # or "signing_in" if one landed but its read-back failed, so the retry stops as aborted_retry; both are safe,
        # and recording "failed" here would make the retry refuse a run that sent nothing.
        if isinstance(exc, SigningInUnconfirmed):
            return
        if isinstance(exc, KeyboardInterrupt) or type(exc).__name__ in CANCEL_EXCEPTION_NAMES:
            self.update(status="cancelled", error=RUN_CANCELLED_ERROR)
        elif isinstance(exc, RunTimedOut):
            self.fail(RUN_TIMED_OUT_ERROR, "run_timed_out")
        elif isinstance(exc, LoadTestError) and exc.cause in CAUSE_NEXT_STEPS:
            # The message already says what failed, and error_cause what kind of failure it was.
            self.fail(redact(str(exc)), exc.cause, getattr(exc, "detail", None))
        else:
            self.fail(redact("%s: %s" % (type(exc).__name__, exc)))

# Without a run key (any run not started by the console) nothing is written, and progress only lives in memory.
JOB_RUN_ID = None
if RUN_KEY:
    try:
        from dbruntime.databricks_repl_context import get_context
        _repl_context = get_context()
    except Exception:
        _repl_context = None
    JOB_RUN_ID = job_run_id(_repl_context, lambda: json.loads(
        dbutils.notebook.entry_point.getDbutils().notebook().getContext().toJson()))
    if JOB_RUN_ID is None:
        print_wrapped("Could not find this job run's id, so a cancel from the console may only stop the load when "
                      "the platform ends this notebook.")
progress = ProgressFile(None, new_run(RUN_KEY, ENDPOINT_NAME, INDEX_NAME, TARGET_QPS, NUM_RESULTS, STAGE_SECONDS,
                                      STAGE_FRACTIONS, utc_now(), JOB_RUN_ID, QUERY_TYPE),
                        printer=print_wrapped, debug=DEBUG)

def _redacted(text, limit=RUN_ERROR_CHARS):
    """Text for the run record, redacted like printed error replies."""
    return _printable(str(text).encode("utf-8", "replace"), limit)

def _record_cell_error(exc):
    """Called by every later cell when it fails or is cancelled."""
    if getattr(exc, "next_step", False) and exc.cause in CAUSE_NEXT_STEPS:
        # The error says only what failed; the console shows the same next step for the recorded cause.
        print_wrapped("Next step: %s." % next_step_for(exc))
    progress.record_exception(exc, _redacted)

if RUN_KEY:
    # A console run sends no load without a working record: serverless reruns a failed job task from the start, and
    # only the record can tell that rerun apart.
    try:
        progress.path = progress_path(spark.sql("select current_user()").first()[0])
    except Exception as exc:
        raise RunRecordUnavailable("Could not find who is running this notebook for the console's run record (%s)."
                                   % type(exc).__name__) from None
    if progress.path is None:
        raise RunRecordUnavailable("This user's home folder can't hold the console's run record.")
    progress.run, RUN_ACTION = begin_run(progress.path, RUN_KEY, progress.run, utc_now())
    if RUN_ACTION == "exit":
        print_wrapped(RETRY_SKIPPED_ERROR)
        dbutils.notebook.exit(json.dumps(progress.run))
    if RUN_ACTION == "fail":
        raise RuntimeError(RETRY_SKIPPED_ERROR)

def _job_run():
    # Called only from later cells: WORKSPACE_HOST and _notebook_auth_headers are set by the sign-in cell.
    resp = requests.get(WORKSPACE_HOST + "/api/2.2/jobs/runs/get", params={"run_id": JOB_RUN_ID},
                        headers=_notebook_auth_headers(), timeout=(2, 3))
    resp.raise_for_status()
    return resp.json()

# Polled while the load runs, since the platform's cancel can arrive late or not at all.
cancel_watch = CancelWatch(_job_run if JOB_RUN_ID else None, printer=print_wrapped, debug=DEBUG)

# COMMAND ----------

# MAGIC %md ## Connect to the index

# COMMAND ----------

# DBTITLE 1,Sign in and read the index
try:
    progress.at("sign_in")
    # Confirmed before any request, so an automatic retry can never find this run at "starting" once load may have been sent.
    progress.confirm(status="signing_in")
    WORKSPACE_HOSTNAME = spark.conf.get("spark.databricks.workspaceUrl")
    WORKSPACE_HOST = "https://" + WORKSPACE_HOSTNAME
    INDEX_PATH = urllib.parse.quote(INDEX_NAME, safe="")
    ENDPOINT_PATH = urllib.parse.quote(ENDPOINT_NAME, safe="")
    # Lets Databricks tell this notebook's requests apart in its request logs; it carries no names. Kept to 20
    # characters so logs that truncate user agents keep all of it.
    USER_AGENT_HEADERS = {"User-Agent": "ai-search-loadtest/1"}
    # Marks the notebook's AI Search calls as load-test traffic; not sent on token, jobs or serving-endpoint calls.
    SEARCH_SOURCE_HEADERS = {"X-Databricks-Vector-Search-Source": "AI_SEARCH_LOAD_TEST"}
    # Names used here and not defined below come from the cells above.
    INDEX_AUTHORIZATION = {"type": "unity_catalog_permission", "securable_type": "table",
                           "securable_object_name": INDEX_NAME, "operation": "ReadVectorIndex"}
    # Route-optimized embedding endpoints the token must also allow querying: their IDs (a name is rejected at sign-in),
    # keyed by name. Empty until a route-optimized endpoint is found after the index is read.
    token_serving_endpoints = {}
    # Get a new token well inside its one-hour lifetime so a rerun of the ramp doesn't fail with 401s.
    TOKEN_MAX_AGE_SECONDS = 45 * 60
    NOTEBOOK_USER_IDENTITY = "the notebook user"
    SP_IDENTITY = "service principal " + SP_CLIENT_ID

    def _authorization_details():
        return json.dumps([INDEX_AUTHORIZATION] + [
            {"type": "workspace_permission", "object_type": "serving-endpoints",
             "object_path": "/serving-endpoints/" + endpoint_id, "actions": ["query_inference_endpoint"]}
            for endpoint_id in token_serving_endpoints.values()], separators=(",", ":"))

    def _notebook_api_token():
        # Serverless and shared clusters may not offer the dbutils context used below, so this is tried first.
        try:
            from dbruntime.databricks_repl_context import get_context
            token = get_context().apiToken
            if token:
                return token
        except Exception:
            pass
        return dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

    def _notebook_auth_headers():
        """Headers for a workspace API call as the notebook itself; the load-test token is scoped to the index."""
        return {"Authorization": "Bearer " + _notebook_api_token(), **USER_AGENT_HEADERS}

    def _notebook_search_headers():
        """_notebook_auth_headers for an AI Search call, which also says it is load-test traffic."""
        return {**_notebook_auth_headers(), **SEARCH_SOURCE_HEADERS}

    def _token_from_response(resp):
        if resp.status_code != 200:
            return None
        try:
            return resp.json().get("access_token")
        except ValueError:
            return None

    class AuthorizationRefused(LoadTestError):
        """The sign-in turned down a permission the token asked for; unlike other failures, a retry won't help."""

        def __init__(self, message, cause="auth", next_step=False):
            super().__init__(message, cause, next_step)

    def _refused_authorization(resp):
        try:
            body = resp.json()
        except ValueError:
            return False
        # Covers invalid_authorization_details and invalid_authorization_details_permission_denied.
        return isinstance(body, dict) and str(body.get("error", "")).startswith("invalid_authorization_details")

    def _retriable(resp):
        return resp.status_code == 429 or resp.status_code >= 500

    class SignInUnavailable(LoadTestError):
        """The sign-in kept failing with 429s or 5xx: the workspace, not the identity, so signing in as another identity
        would only change whose access the run measures."""

    # Phases before the first query; the workspace_rate_limited next step says no load was sent yet.
    BEFORE_LOAD_PHASES = ("sign_in", "lookup")

    def _in_parentheses(message):
        """ (message), without the full stop a message (such as a reply from _reply_text) may end in, which would only
        clutter the punctuation after it; an ellipsis that says it was cut is kept."""
        return " (%s)" % (message[:-1] if message.endswith(".") and not message.endswith("...") else message)

    def _sign_in_unavailable(message, status):
        """The SignInUnavailable for a sign-in whose last reply had this 429 or 5xx status."""
        if progress.phase in BEFORE_LOAD_PHASES:
            return SignInUnavailable(message, "workspace_rate_limited" if status == 429 else "unknown")
        return SignInUnavailable("Could not renew the sign-in during the run%s. Rerun." % _in_parentheses(message),
                                 "unknown")

    last_exchange_failure = None

    def _exchange_notebook_token():
        """Returns (token, None) on success, else (None, whether a requested permission was refused)."""
        global last_exchange_failure
        # Only exception type names are printed, because messages can carry URLs or credentials.
        try:
            notebook_token = _notebook_api_token()
        except Exception as exc:
            print("Could not read the notebook identity token (%s)." % type(exc).__name__)
            return None, False
        failures = []
        # The status of each reply that failed after its retries, or None for a connection failure.
        statuses = []
        refused = retriable = False
        for attempt, subject_type in enumerate(("urn:databricks:params:oauth:token-type:personal-access-token",
                                                 "urn:ietf:params:oauth:token-type:access_token"), 1):
            # Token type names mean nothing to the reader, so failures are numbered instead.
            label = "try %d" % attempt
            try:
                resp = _with_workspace_retries(lambda: requests.post(
                    WORKSPACE_HOST + "/oidc/v1/token", headers=USER_AGENT_HEADERS, timeout=WORKSPACE_API_TIMEOUT_SECONDS,
                    data={
                        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
                        "subject_token": notebook_token,
                        "subject_token_type": subject_type,
                        "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
                        "scope": "all-apis",
                        "authorization_details": _authorization_details(),
                    }), SIGN_IN_ENDPOINT)
            except requests.RequestException as exc:
                retriable = True
                failures.append("%s: %s" % (label, type(exc).__name__))
                statuses.append(None)
                continue
            token = _token_from_response(resp)
            if token:
                last_exchange_failure = None
                return token, None
            refused = refused or _refused_authorization(resp)
            retriable = retriable or _retriable(resp)
            failures.append("%s: status %s" % (label, resp.status_code))
            statuses.append(resp.status_code)
            if _retriable(resp):
                # Its retries are spent; the other subject type would only wait them out again on the same limit.
                break
        message = "Could not sign in as the notebook user (%s)." % "; ".join(failures)
        # A retry that fails the same way adds nothing.
        if message != last_exchange_failure:
            print(message)
        last_exchange_failure = message
        turned_away = [status is None or status == 429 or status >= 500 for status in statuses]
        denied = refused or any(status in (401, 403) for status in statuses)
        # The workspace kept turning the sign-in away, which says nothing about permissions: the last try ended in a 429
        # or 5xx with no refusal, or every try ended in one of those or a connection failure.
        if not denied and turned_away[-1] and (statuses[-1] is not None or all(turned_away)):
            raise _sign_in_unavailable(message[:-1], statuses[-1])
        # A refusal next to a failure that may be transient could still go away on a retry.
        return None, refused and not retriable

    # Widest first, so the common all-apis secret signs in with one request; a secret refused a scope tries the next.
    SP_SCOPE_CHOICES = ("all-apis", "vector-search model-serving", "vector-search")
    # Found at the first sign-in and reused, so a refresh mid-run asks for exactly what worked.
    sp_scope = None

    def _scope_not_assigned(resp):
        """The refusal of a scope the secret lacks (OAuthScopeValidator.scala notAssignedScopeError), the only one worth
        trying a narrower scope for."""
        if resp.status_code not in (400, 401, 403):
            return False
        try:
            body = resp.json()
        except ValueError:
            return False
        return (isinstance(body, dict) and body.get("error") == "access_denied"
                and "are not assigned to the client" in str(body.get("error_description", "")))

    def _invalid_client(resp):
        """The sign-in's 401 for a client ID and secret that don't match a service principal in this workspace."""
        if resp.status_code != 401:
            return False
        try:
            body = resp.json()
        except ValueError:
            return False
        return isinstance(body, dict) and body.get("error") == "invalid_client"

    def _client_credentials_token():
        global sp_scope
        try:
            # The sign-in matches the secret's exact hash, and generated secrets have no whitespace, so a trailing newline
            # pasted into the secret would make a correct secret look wrong.
            secret = dbutils.secrets.get(scope=SP_SECRET_SCOPE, key=SP_SECRET_KEY).strip()
        except Exception as exc:
            # Only the exception type is shown, because messages can carry credentials.
            raise LoadTestError("Could not read the service principal's client secret from scope %s, key %s (%s)" % (
                SP_SECRET_SCOPE, SP_SECRET_KEY, type(exc).__name__), "service_principal_sign_in", next_step=True) from None
        scopes = (sp_scope,) if sp_scope else SP_SCOPE_CHOICES
        for index, scope in enumerate(scopes):
            try:
                resp = _with_workspace_retries(lambda: requests.post(
                    WORKSPACE_HOST + "/oidc/v1/token", auth=(SP_CLIENT_ID, secret), headers=USER_AGENT_HEADERS,
                    timeout=WORKSPACE_API_TIMEOUT_SECONDS, data={
                        "grant_type": "client_credentials",
                        "scope": scope,
                        "authorization_details": _authorization_details(),
                    }), SIGN_IN_ENDPOINT)
            except requests.RequestException as exc:
                raise LoadTestError("Service principal sign-in failed (%s)" % type(exc).__name__, "unreachable") from None
            token = _token_from_response(resp)
            if token:
                sp_scope = scope
                return token
            reply = _reply_text(resp.content, ERROR_TEXT_CHARS) if resp.status_code != 200 else ""
            detail = " " + reply if reply else ""
            message = "Service principal sign-in failed: status %s%s" % (resp.status_code, detail)
            if _retriable(resp):
                raise _sign_in_unavailable(message, resp.status_code)
            if sp_scope is None and _scope_not_assigned(resp):
                # Names the scopes this try asked for, which are the ones the reply refuses, so the reply isn't printed.
                names = scope.split()
                print_wrapped("The service principal's OAuth secret isn't allowed the %s scope%s%s." % (
                    " and ".join(names), "s" if len(names) > 1 else "",
                    "; trying " + " and ".join(scopes[index + 1].split()) if index + 1 < len(scopes) else ""))
                continue
            if _refused_authorization(resp):
                raise AuthorizationRefused(message)
            raise LoadTestError(message, "service_principal_sign_in", next_step=True,
                                detail="invalid_client" if _invalid_client(resp) else None)
        raise LoadTestError("Service principal sign-in failed: its OAuth secret has none of the needed scopes.",
                            "service_principal_scope", next_step=True)

    def _index_missing():
        """Whether the index doesn't exist, asked with the notebook's own token; False when that can't be told."""
        try:
            resp = requests.get(WORKSPACE_HOST + "/api/2.0/vector-search/indexes/" + INDEX_PATH,
                                headers=_notebook_search_headers(),
                                timeout=30)
            return (resp.json() or {}).get("error_code") == "RESOURCE_DOES_NOT_EXIST"
        except Exception:
            return False

    def _index_not_found(sign_in_error=None):
        return LoadTestError("Index %s was not found in this workspace. Check the index name.%s" % (
            INDEX_NAME, " Sign-in error: %s" % _redacted(sign_in_error) if sign_in_error else ""), "not_found")

    class RenewalFailed(LoadTestError):
        """A sign-in as an identity that already signed in failed; original is the failure, for callers that say on
        their own what the sign-in was for."""

        def __init__(self, message, original):
            super().__init__(message, "unknown")
            self.original = original

    def _renewed_service_principal_token():
        """A service principal that already signed in and then fails is unlikely to have a wrong client ID or secret
        scope, so the failure is "unknown", as SignInUnavailable is after load starts."""
        try:
            return _client_credentials_token()
        except (SignInUnavailable, AuthorizationRefused):
            raise
        except LoadTestError as exc:
            # Before load, the only renewal is the route-optimized re-sign-in.
            raise RenewalFailed(("Could not get a new service principal sign-in for the route-optimized embedding "
                                 "endpoint%s." if progress.phase in BEFORE_LOAD_PHASES else
                                 "Could not renew the service principal's sign-in during the run%s.")
                                % _in_parentheses(_redacted(exc)), exc) from None

    def _mint_token(identity=None):
        """Returns a token and who it signs in as. Given the identity of an earlier sign-in, it signs in only as that,
        so refreshing a token mid-run can't switch whose limits and permissions the run measures."""
        if identity == SP_IDENTITY:
            return _renewed_service_principal_token(), SP_IDENTITY
        # Whoever filled in the widgets chose the service principal; signing in as someone else would measure their access.
        if identity is None and SP_CONFIGURED:
            try:
                return _client_credentials_token(), SP_IDENTITY
            except SignInUnavailable:
                raise
            except RuntimeError as exc:
                sp_failure = exc
            # No sign-in can be scoped to an index that doesn't exist, so that is checked before blaming sign-in.
            if _index_missing():
                raise _index_not_found(sp_failure)
            if isinstance(sp_failure, AuthorizationRefused):
                raise AuthorizationRefused("The service principal can't get a token that queries this index. Sign-in "
                                           "error: %s" % _redacted(sp_failure), next_step=True)
            raise sp_failure
        token, refused = _exchange_notebook_token()
        if token:
            return token, NOTEBOOK_USER_IDENTITY
        if _index_missing():
            raise _index_not_found()
        # As for a service principal, a refresh that fails once load has started isn't advice about how to sign in.
        if identity == NOTEBOOK_USER_IDENTITY and not refused and progress.phase not in BEFORE_LOAD_PHASES:
            raise LoadTestError("Could not renew the sign-in during the run (%s). Rerun." % (
                last_exchange_failure or "Could not sign in as the notebook user.")[:-1], "unknown")
        # A console run takes its service principal from the console's form, so only a notebook run has these widgets.
        raise (AuthorizationRefused if refused else LoadTestError)(
            "Could not sign in with permission to query this index.%s" % (
                "" if CONSOLE_RUN else " To use a service principal, fill in the three service principal widgets."),
            "auth", next_step=True)

    def check_query_text_sign_in(embedding_endpoint, via_workspace, route_optimized):
        """Prints a warning before any load when the embedding endpoint may refuse this test's token, and returns it (None
        when there is none); embedding_endpoint is None for an index queried with vectors, and route_optimized None for an
        endpoint that couldn't be read. Only a warning: the gateway can embed without passing the token on (Model Serving's
        proxy, or the Serverless Platform query path), so the test query decides. Through the workspace the token isn't
        passed on at all, and route-optimized endpoints are assumed to accept the notebook user's (untested)."""
        if not embedding_endpoint or via_workspace:
            return None
        notes = []
        if route_optimized is None:
            notes.append("Could not read %s with the notebook user's token, so this test can't tell whether it is "
                         "route-optimized. If it is, queries may fail with 401; let the notebook user view it, then rerun."
                         % embedding_endpoint)
        if auth.identity == NOTEBOOK_USER_IDENTITY and not route_optimized:
            notes.append("%s may fail with 403 because %s; a service principal is recommended." % (
                "If it isn't, queries" if route_optimized is None else "Queries", NOTEBOOK_USER_CANNOT_EMBED % "usually "))
        if auth.identity == SP_IDENTITY and not _allows_model_serving(sp_scope):
            notes.append("The service principal's OAuth secret has the %s scope but not model-serving, so queries may fail "
                         "if the embedding endpoint needs the model-serving scope." % sp_scope)
        if not notes:
            return None
        warning = "WARNING: " + " ".join(notes)
        print_wrapped(warning)
        return warning

    def _failure_detail(exc, identity):
        # A notebook-user failure was printed by the token exchange, and its error only gives service principal advice.
        # A RenewalFailed's own message would repeat what the caller says the sign-in was for.
        return _in_parentheses(str(getattr(exc, "original", exc))[:ERROR_TEXT_CHARS]) if identity == SP_IDENTITY else ""

    class Auth:
        def __init__(self):
            self.headers = None
            self.minted_at = 0.0
            self.identity = None

        def fresh_headers(self):
            if self.headers is None or time.monotonic() - self.minted_at > TOKEN_MAX_AGE_SECONDS:
                try:
                    self.renew()
                except AuthorizationRefused as exc:
                    if not token_serving_endpoints:
                        raise
                    raise LoadTestError("Could not get a new token that can query the route-optimized embedding endpoint "
                                        "%s%s: %s may have lost Can Query on it or permission to read the index." % (
                                            ", ".join(token_serving_endpoints), _failure_detail(exc, self.identity),
                                            self.identity), "auth") from None
            return self.headers

        def renew(self):
            # Nothing changes until the sign-in succeeds, so a failed one keeps the current token.
            token, self.identity = _mint_token(self.identity)
            self.headers = {"Authorization": "Bearer " + token, "Content-Type": "application/json", **USER_AGENT_HEADERS,
                            **SEARCH_SOURCE_HEADERS}
            self.minted_at = time.monotonic()

    auth = Auth()

    # Reads of the index and endpoint, and sign-ins, retry 429s, 5xx and connection failures: shared workspaces hit their
    # organization's API request limit in bursts that usually pass within a minute.
    WORKSPACE_API_ATTEMPTS = 4
    WORKSPACE_API_BACKOFF_SECONDS = (5, 10, 20)
    # No retry starts once a read has spent this on its tries and waits.
    WORKSPACE_API_RETRY_BUDGET_SECONDS = 60
    # (connect, read) seconds per try. A try started just inside the budget can still end about 35 s past it, plus a
    # forced cancel check and any token refresh.
    WORKSPACE_API_TIMEOUT_SECONDS = (5, 30)
    # Each forced check is a Jobs API call, so a wait checks less often than the load stages do.
    WORKSPACE_API_CANCEL_CHECK_SECONDS = 15

    def workspace_api_retry_wait(failures, retry_after=None, elapsed=0.0):
        """Seconds to wait before trying again after `failures` failed tries, `elapsed` seconds into the read, or None to
        give up. A Retry-After in seconds replaces the backoff and gives up when longer than the budget left; its
        HTTP-date form is ignored."""
        if failures >= WORKSPACE_API_ATTEMPTS:
            return None
        try:
            wait = float(retry_after)
        except (TypeError, ValueError):
            wait = None
        remaining = WORKSPACE_API_RETRY_BUDGET_SECONDS - elapsed
        # Negative and NaN values fall back to the backoff; infinity gives up like any wait past the budget.
        if wait is not None and wait >= 0:
            # Waited out in full or not at all: a try before the reply's own wait is over would hit the same limit.
            wait = max(wait, 1.0)
            return wait if wait <= remaining else None
        wait = min(WORKSPACE_API_BACKOFF_SECONDS[min(failures, len(WORKSPACE_API_BACKOFF_SECONDS)) - 1], remaining)
        return wait if wait > 0 else None

    def _wait_checking_cancel(seconds):
        # A cancel from the console lands within one check interval instead of after the whole wait.
        end = time.monotonic() + seconds
        while True:
            cancel_watch.check(force=True)
            remaining = end - time.monotonic()
            if remaining <= 0:
                return
            time.sleep(min(remaining, WORKSPACE_API_CANCEL_CHECK_SECONDS))

    SIGN_IN_ENDPOINT = "The workspace's sign-in endpoint"

    def _with_workspace_retries(send, what):
        """Calls send() for a reply, trying again as workspace_api_retry_wait allows; what names the API in the retry
        line. Returns the last reply, or raises the last connection failure."""
        failures, started = 0, time.monotonic()
        while True:
            resp = failure = None
            try:
                resp = send()
            except requests.RequestException as exc:
                failure = exc
            if resp is not None and not _retriable(resp):
                return resp
            failures += 1
            wait = workspace_api_retry_wait(failures, None if resp is None else resp.headers.get("Retry-After"),
                                            time.monotonic() - started)
            if wait is None:
                if resp is None:
                    raise failure
                return resp
            print_wrapped("%s %s; trying again in %.0f s." % (
                what, "could not be reached (%s)" % type(failure).__name__ if resp is None
                else "returned status %s" % resp.status_code, wait))
            _wait_checking_cancel(wait)

    def _workspace_get(path, headers):
        """GETs a workspace API path with _with_workspace_retries; headers is called for each try."""
        return _with_workspace_retries(
            lambda: requests.get(WORKSPACE_HOST + path, headers=headers(), timeout=WORKSPACE_API_TIMEOUT_SECONDS),
            "The workspace API")

    def _reply_error_code(resp):
        """The reply's error_code (e.g. REQUEST_LIMIT_EXCEEDED), or None when it has none that is safe to show."""
        try:
            code = resp.json().get("error_code")
        except Exception:
            return None
        if isinstance(code, str) and 0 < len(code) <= 64 and all(char.isupper() or char == "_" for char in code):
            return code
        return None

    def _get_json(path):
        """The index's details, read as the identity the test signed in as, before any load is sent."""
        try:
            resp = _workspace_get(path, auth.fresh_headers)
        except requests.RequestException as exc:
            raise LoadTestError("Could not reach the workspace API to look up the index (%s), so no load was sent." % (
                type(exc).__name__), "unknown") from None
        if resp.status_code == 429:
            # The workspace's own limit, not the endpoint's, so the endpoint's capacity advice wouldn't apply.
            code = _reply_error_code(resp)
            # The code is already shown with the status, so the reply is shown without it.
            reply = _reply_text(resp.content, ERROR_TEXT_CHARS, with_error_code=False)
            raise LoadTestError("The workspace API rate limit was hit while looking up the index (status 429%s), so no load "
                                "was sent.%s" % (" " + code if code else "", " Reply: " + reply if reply else ""),
                                "workspace_rate_limited")
        if resp.status_code != 200:
            cause = reply_cause(resp.status_code, resp.content[:REPLY_PREFIX_BYTES])
            hint = ""
            if resp.status_code == 403 and auth.identity == SP_IDENTITY:
                hint = " Give the service principal permission to read it.%s" % (
                    "" if CONSOLE_RUN else " (Or clear its widgets to run as the notebook user, though usually not for an "
                    "index that computes its embeddings.)")
            reply = _reply_text(resp.content, ERROR_TEXT_CHARS)
            # An empty reply leaves no gap between the status and the hint.
            raise LoadTestError("Could not read the index as %s: status %s%s%s" % (
                auth.identity, resp.status_code, " " + reply if reply else "", hint),
                # Only fatal causes keep their class; a 5xx here is the workspace API, not the endpoint, so it becomes "unknown".
                cause if cause in FATAL_CAUSES else "unknown")
        return resp.json()

    def _endpoint_details(path):
        """The endpoint's details read as the notebook user (with the run's own token on a notebook-user run), or None after
        a WARNING. Querying an index needs no permission on its endpoint, so a failed read never stops the run: these
        details feed the scaling, endpoint-type and High QPS large-index warnings and the Target QPS the verdict compares
        against, which falls back to CONFIGURED_TARGET_QPS. A 429 that outlasts the retries is the workspace's limit and is
        no different."""
        try:
            # Read once, outside the retries, so this broad except can't swallow a cancel raised while they wait.
            headers = auth.fresh_headers() if auth.identity == NOTEBOOK_USER_IDENTITY else _notebook_search_headers()
        except Exception as exc:
            failure = "its token could not be read: %s" % type(exc).__name__
        else:
            try:
                resp = _workspace_get(path, lambda: headers)
            except requests.RequestException as exc:
                failure = "the workspace API could not be reached: %s" % type(exc).__name__
            else:
                details = None
                if resp.status_code == 200:
                    try:
                        details = resp.json()
                    except ValueError:
                        # Reported below as an unreadable reply, like a body that isn't an object.
                        pass
                if isinstance(details, dict):
                    return details
                code = _reply_error_code(resp)
                if resp.status_code == 200:
                    failure = "status 200, unreadable reply"
                elif code:
                    failure = "status %s %s" % (resp.status_code, code)
                else:
                    failure = "status %s" % resp.status_code
        print_wrapped("WARNING: Could not read the endpoint as %s (%s), so this test skips its checks for scaling in "
                      "progress, the endpoint type and large indexes at high QPS, and takes its Target QPS to be the one "
                      "this notebook was created with. Queries need no permission on the endpoint, so the test goes on."
                      % (NOTEBOOK_USER_IDENTITY, failure))
        return None

    def _allow_route_optimized_endpoint(endpoint_name, endpoint_id):
        """Signs in again so the token can also query the endpoint. Returns the endpoint's name if the sign-in refused
        that, else None."""
        global last_exchange_failure
        token_serving_endpoints[endpoint_name] = endpoint_id
        for attempt in range(2):
            try:
                auth.renew()
                return None
            except (AuthorizationRefused, SignInUnavailable) as exc:
                # A refusal won't change, and a sign-in turned away by the workspace has already used up its retries.
                error = exc
                break
            except RuntimeError as exc:
                error = exc
                # Other failures can be transient, so they get one retry.
                if attempt == 0:
                    time.sleep(2)
        # Only the retry's repeat is worth hiding; a later refresh failure is printed even if it looks the same.
        last_exchange_failure = None
        # The index-only token from before still works, and later refreshes stay index-only.
        token_serving_endpoints.clear()
        detail = _failure_detail(error, auth.identity)
        if isinstance(error, AuthorizationRefused):
            print_wrapped("Could not sign in with permission to query the route-optimized embedding endpoint %s%s, so "
                          "test queries will likely fail with 401. Give %s Can Query on %s, then rerun."
                          % (endpoint_name, detail, auth.identity, endpoint_name))
            return endpoint_name
        print_wrapped("Could not sign in with permission to query the route-optimized embedding endpoint %s%s, so test "
                      "queries may fail with 401. Rerun." % (endpoint_name, detail))
        return None


    # Signed in before the index is read, so a failure reading the index or endpoint isn't reported as a sign-in failure.
    auth.fresh_headers()
    print_wrapped("Signed in as %s." % auth.identity)
    progress.note(signed_in_as="service_principal" if auth.identity == SP_IDENTITY else "notebook_user")
    if auth.identity == NOTEBOOK_USER_IDENTITY:
        print_wrapped("Databricks recommends a service principal's OAuth token for high-QPS production workloads, so fill "
                      "in the service principal widgets for results that represent a production application.")

    progress.at("lookup")
    index_meta = _get_json("/api/2.0/vector-search/indexes/" + INDEX_PATH)
    if index_meta.get("endpoint_name") and index_meta["endpoint_name"] != ENDPOINT_NAME:
        raise LoadTestError("Index %s is on endpoint %s, not %s. Pick an index on this endpoint in the Index widget."
                            % (INDEX_NAME, index_meta["endpoint_name"], ENDPOINT_NAME), "not_found")

    index_url = (index_meta.get("status") or {}).get("index_url")
    if not index_url:
        raise LoadTestError("The index has no query URL yet.", "index_not_ready", next_step=True)
    if "://" not in index_url:
        index_url = "https://" + index_url
    if not index_url.startswith("https://"):
        raise RuntimeError("Refusing to send the OAuth token to a non-HTTPS index URL.")
    QUERY_URL = urllib.parse.urlsplit(index_url.rstrip("/") + "/query")
    QUERIES_VIA_WORKSPACE = QUERY_URL.hostname == WORKSPACE_HOSTNAME
    # True when the run fell back to the workspace route after the endpoint URL refused the notebook user's token.
    QUERY_ROUTE_FALLBACK = False

    INDEX_SPEC = index_meta.get("delta_sync_index_spec") or index_meta.get("direct_access_index_spec") or {}
    # Indexes that compute their embeddings take query text instead of vectors.
    TEXT_COLUMNS = INDEX_SPEC.get("embedding_source_columns") or []
    EMBEDDING_ENDPOINT = EMBEDDING_ENDPOINT_NAME = QUERY_ENDPOINT_NAME = SOURCE_ENDPOINT_NAME = None
    if TEXT_COLUMNS:
        # An index can embed queries with a different endpoint than the one it used to embed its source table. Queries
        # sent to the endpoint's own URL use the source table's; queries through the workspace use the one for queries.
        QUERY_ENDPOINT_NAME = TEXT_COLUMNS[0].get("model_endpoint_name_for_query")
        SOURCE_ENDPOINT_NAME = TEXT_COLUMNS[0].get("embedding_model_endpoint_name")
        EMBEDDING_ENDPOINT_NAME = ((QUERY_ENDPOINT_NAME or SOURCE_ENDPOINT_NAME) if QUERIES_VIA_WORKSPACE
                                   else (SOURCE_ENDPOINT_NAME or QUERY_ENDPOINT_NAME))
        EMBEDDING_ENDPOINT = ("the embedding endpoint " + EMBEDDING_ENDPOINT_NAME if EMBEDDING_ENDPOINT_NAME
                              else "its embedding endpoint")
        print_wrapped("This index computes its embeddings, so this test sends query text, and %s embeds every query."
                      % EMBEDDING_ENDPOINT)

    def _serving_endpoint(name):
        """The serving endpoint's details, or None if they can't be read."""
        # The notebook's own token: the load-test token is scoped to the index.
        try:
            resp = _workspace_get("/api/2.0/serving-endpoints/" + urllib.parse.quote(name, safe=""),
                                  _notebook_auth_headers)
            details = resp.json() if resp.status_code == 200 else None
        except (RunCancelled, RunTimedOut):
            raise
        except Exception:
            return None
        return details if isinstance(details, dict) else None

    EMBEDDING_SERVING_ENDPOINT = _serving_endpoint(EMBEDDING_ENDPOINT_NAME) if EMBEDDING_ENDPOINT_NAME else None
    # None when the endpoint couldn't be read: whether it is route-optimized is then unknown.
    EMBEDDING_ENDPOINT_ROUTE_OPTIMIZED = (None if EMBEDDING_SERVING_ENDPOINT is None
                                          else bool(EMBEDDING_SERVING_ENDPOINT.get("route_optimized")))
    # Noted before the test query so a run it stops still shows the console what kind of index it was.
    progress.note(query_route="workspace" if QUERIES_VIA_WORKSPACE else "endpoint",
                  embedding_endpoint=EMBEDDING_ENDPOINT_NAME, computes_embeddings=bool(TEXT_COLUMNS))
    check_query_text_sign_in(EMBEDDING_ENDPOINT, QUERIES_VIA_WORKSPACE, EMBEDDING_ENDPOINT_ROUTE_OPTIMIZED)

    # On the endpoint's own URL each query is embedded with the token the test sends, and a route-optimized endpoint
    # accepts that token only if it allows querying that endpoint. Through the workspace the token isn't passed on.
    ROUTE_OPTIMIZED_EMBEDDING_ENDPOINT = UNAUTHORIZED_ROUTE_OPTIMIZED_ENDPOINT = None
    if not QUERIES_VIA_WORKSPACE and EMBEDDING_ENDPOINT_ROUTE_OPTIMIZED and EMBEDDING_SERVING_ENDPOINT.get("id"):
        ROUTE_OPTIMIZED_EMBEDDING_ENDPOINT = EMBEDDING_ENDPOINT_NAME
        UNAUTHORIZED_ROUTE_OPTIMIZED_ENDPOINT = _allow_route_optimized_endpoint(
            EMBEDDING_ENDPOINT_NAME, EMBEDDING_SERVING_ENDPOINT["id"])

    def _external_embedding_model():
        """'<provider> model <name>' when the embedding endpoint serves an external model, else None (also if unreadable)."""
        entities = ((EMBEDDING_SERVING_ENDPOINT or {}).get("config") or {}).get("served_entities") or []
        for entity in entities:
            external = entity.get("external_model") or {}
            # This provider forwards to another Databricks model serving endpoint, so the Databricks advice applies.
            if external and external.get("provider") != "databricks-model-serving":
                return " ".join(part for part in (external.get("provider"), "model", external.get("name")) if part)
        return None

    EXTERNAL_EMBEDDING_MODEL = _external_embedding_model()

    def _embedding_endpoint_kind():
        """The kind of embedding endpoint: "pay_per_token" for a preconfigured pay-per-token Foundation Model APIs endpoint,
        "dedicated" for one that serves its own models, None when it can't be read or forwards elsewhere (external)."""
        config = (EMBEDDING_SERVING_ENDPOINT or {}).get("config")
        entities = (config.get("served_entities") if isinstance(config, dict) else None) or []
        # A reply of an unexpected shape leaves the kind unknown rather than stopping the notebook.
        if not isinstance(entities, list) or not all(isinstance(entity, dict) for entity in entities):
            return None
        foundation_models = [entity["foundation_model"] for entity in entities if entity.get("foundation_model")]
        # Tile endpoints also name a foundation_model, but without the endpoint type or the per-token price.
        if foundation_models and (EMBEDDING_SERVING_ENDPOINT.get("endpoint_type") == "FOUNDATION_MODEL_API" or any(
                isinstance(model, dict) and model.get("price_unit") for model in foundation_models)):
            return "pay_per_token"
        if not entities or foundation_models or any(entity.get("external_model") for entity in entities):
            return None
        return "dedicated"

    EMBEDDING_ENDPOINT_KIND = _embedding_endpoint_kind()
    if EXTERNAL_EMBEDDING_MODEL:
        print_wrapped("%s serves the external %s, so your rate limits and quota with that provider also limit this test."
                      % (_capitalized(EMBEDDING_ENDPOINT), EXTERNAL_EMBEDDING_MODEL))

    if QUERIES_VIA_WORKSPACE:
        print_wrapped("WARNING: Queries to this index don't use the network-optimized route, so each query takes longer and "
                      "they may be rate limited (HTTP 429) below your Target QPS. That route isn't available for this index "
                      "in this workspace right now; contact Databricks support for more details.")

    def _warn_if_proxied():
        """The load engine connects straight to QUERY_URL's host, not through a proxy (the test query may still use one)."""
        proxies = urllib.request.getproxies()
        if (proxies.get("https") or proxies.get("all")) and not urllib.request.proxy_bypass(QUERY_URL.hostname):
            print_wrapped("WARNING: This cluster has an HTTPS proxy set, but load requests connect directly to the "
                          "endpoint. If this cluster can only reach the endpoint through the proxy, load requests will "
                          "fail.")

    _warn_if_proxied()

    endpoint_meta = _endpoint_details("/api/2.0/vector-search/endpoints/" + ENDPOINT_PATH)
    ENDPOINT_UNREADABLE = endpoint_meta is None
    endpoint_meta = endpoint_meta or {}
    scaling_info = endpoint_meta.get("scaling_info") or {}
    throughput_info = endpoint_meta.get("throughput_info") or {}
    if (scaling_info.get("state") == "SCALING_CHANGE_IN_PROGRESS"
            or throughput_info.get("change_request_state") == "CHANGE_IN_PROGRESS"):
        print_wrapped("WARNING: The endpoint is still scaling, so results will understate its capacity. Rerun when it's "
                      "done.")
    # Read now, not when this notebook was created: the endpoint's Target QPS may have changed since.
    ENDPOINT_TARGET_QPS = int(scaling_info.get("requested_target_qps") or 0) or CONFIGURED_TARGET_QPS
    # On the workspace route the route itself may be the limit, where a higher Target QPS can't help (route_bottleneck).
    TARGET_QPS_ADVICE_NOTE = "If the endpoint falls behind, the verdict suggests a Target QPS.{}".format(
        " It doesn't if the route these queries take is the limit, since a higher Target QPS can't help then."
        if QUERIES_VIA_WORKSPACE else "")
    if TARGET_QPS > ENDPOINT_TARGET_QPS:
        print_wrapped("WARNING: Target QPS {:,} is above the endpoint's Target QPS {:,}. The endpoint provisions capacity "
                      "for its Target QPS and doesn't autoscale, so expect 429s or slowdowns above about {:,} QPS. "
                      "{}".format(TARGET_QPS, ENDPOINT_TARGET_QPS, ENDPOINT_TARGET_QPS, TARGET_QPS_ADVICE_NOTE))
    RUN_WARNING_CODES = [code for code, applies in (("queries_via_workspace", QUERIES_VIA_WORKSPACE),
                                                     ("above_endpoint_target_qps", TARGET_QPS > ENDPOINT_TARGET_QPS),
                                                     ("endpoint_unreadable", ENDPOINT_UNREADABLE)) if applies]
    # The console shows these beside the results, since it doesn't show the WARNING lines printed here.
    progress.note(warnings=RUN_WARNING_CODES, endpoint_target_qps=ENDPOINT_TARGET_QPS)
    ENDPOINT_TYPE = endpoint_meta.get("endpoint_type")
    INDEXED_ROW_COUNT = int((index_meta.get("status") or {}).get("indexed_row_count") or 0)
    # The top of the 10-100 range the best-practices docs recommend.
    MANY_RESULTS = 100
    # The endpoint sizes the capacity behind its Target QPS from a benchmark of queries returning this many results.
    SIZED_RESULTS = 10
    # The best-practices docs say 10 times more results can cut the QPS an endpoint serves by about 3 times. Other counts
    # use the same curve, which is an estimate.
    RESULTS_COST_EXPONENT = math.log10(3)
    # The endpoint page's High QPS warning fires above this many indexed rows times Target QPS.
    HIGH_QPS_LARGE_INDEX_ROW_QPS = 10000000000

    def _results_factor():
        """How many times a query with this many results costs over one with the count Target QPS is sized for."""
        # The docs' curve is for ANN; hybrid results are capped low and have no such figure, so none is estimated.
        if HYBRID_QUERIES:
            return 1.0
        return (NUM_RESULTS / float(SIZED_RESULTS)) ** RESULTS_COST_EXPONENT if NUM_RESULTS > SIZED_RESULTS else 1.0

    RESULTS_FACTOR = _results_factor()
    if RESULTS_FACTOR > 1.0:
        print_wrapped("WARNING: The endpoint's Target QPS provisions capacity for queries that return {} results, and this "
                      "test asks for {:,} per query, which costs more. The best-practices docs "
                      "(https://docs.databricks.com/aws/en/ai-search/best-practices) say 10 times more results can cut "
                      "the QPS an endpoint serves by about 3 times; for other counts this notebook estimates it from that "
                      "figure, here about {:.1f} times. {}".format(SIZED_RESULTS, NUM_RESULTS, RESULTS_FACTOR,
                                                                   TARGET_QPS_ADVICE_NOTE))
    if NUM_RESULTS > MANY_RESULTS:
        print_wrapped("WARNING: Results per query is {:,}. Use your production value, ideally 10-100.".format(NUM_RESULTS))
    if ENDPOINT_TYPE == "STANDARD" and INDEXED_ROW_COUNT * TARGET_QPS > HIGH_QPS_LARGE_INDEX_ROW_QPS:
        print_wrapped("WARNING: With this many indexed rows at this Target QPS, the endpoint may need more capacity than its "
                      "Target QPS provisions. Contact Databricks support to add it.")
    if ENDPOINT_TYPE and ENDPOINT_TYPE != "STANDARD":
        print_wrapped("WARNING: High QPS (the endpoint's Target QPS) is available for Standard endpoints only, so this %s "
                      "endpoint serves far less QPS. Expect high targets to fall short." % (
                          ENDPOINT_TYPE.lower().replace("_", "-")))
    if DEBUG:
        _not_reported = "not read" if ENDPOINT_UNREADABLE else "none reported"
        print("Endpoint scaling: %s" % (scaling_info or _not_reported))
        print("Endpoint throughput: %s" % (throughput_info or _not_reported))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# MAGIC %md ## Prepare queries

# COMMAND ----------

# DBTITLE 1,Query settings
try:
    progress.at("setup")
    progress.update_confirmed(status="preparing")
    # Every query is different, so repeated queries can't make latency look better than it is.
    MAX_SAMPLED_QUERIES = 20000
    PRIMARY_KEY = index_meta.get("primary_key")
    TOTAL_QUERIES = sum(int(math.ceil(SEND_RATE_FACTOR * qps * STAGE_SECONDS)) for qps in STAGE_QPS)
    # debug_level 1 makes the endpoint return its own handling time, shown as server latency.
    _query_fields = {"num_results": NUM_RESULTS, "debug_level": 1}
    if PRIMARY_KEY:
        _query_fields["columns"] = [PRIMARY_KEY]
    # Left out for ANN, the default, so those requests stay as they were.
    if HYBRID_QUERIES:
        _query_fields["query_type"] = "HYBRID"

    def _quoted_col(name):
        # Backticks make Spark read any name, even one with a dot, as a single column.
        return F.col("`" + name.replace("`", "``") + "`")

    def _source_rows(source_table, column, limit, value=None, keep=None, sample=True):
        """Up to limit values of one source table column, sampled with the run's seed when the table has more rows than
        limit (unless sample is False). value replaces the column as what is read, and keep replaces the non-null
        filter."""
        # The index's row count sizes the sample without scanning the table to count it.
        rows = spark.table(source_table).select(_quoted_col(column) if value is None else value).where(
            _quoted_col(column).isNotNull() if keep is None else keep)
        if sample and INDEXED_ROW_COUNT > limit:
            rows = rows.sample(fraction=min(1.0, 1.5 * limit / INDEXED_ROW_COUNT), seed=SEED)
        return rows.limit(limit).toLocalIterator()
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Query text (indexes that compute their embeddings)
try:
    cancel_watch.check(force=True)
    # Cut longer text to this many characters so request bodies stay small.
    MAX_QUERY_TEXT_CHARS = 2000
    # Every word is also a keyword to match, so hybrid sends a short, query-length prefix.
    HYBRID_QUERY_WORDS = 16
    # Later passes over the texts add " <pass number>", so no two queries send the same text.
    PASS_SUFFIX = re.compile(r" [0-9]+$")
    # Spark's error classes for a table this identity can't read or see; other read failures keep no cause.
    SOURCE_TABLE_ACCESS_ERRORS = ("INSUFFICIENT_PERMISSIONS", "PERMISSION_DENIED", "TABLE_OR_VIEW_NOT_FOUND")
    # Spark messages start with their error class, e.g. "[UNRESOLVED_COLUMN.WITH_SUGGESTION] ...". Only the class is kept
    # in the error, since the rest can quote table contents or paths.
    SPARK_ERROR_CLASS = re.compile(r"^\[([A-Z][A-Z0-9_]*(?:\.[A-Z0-9_]+)*)\]")

    def _hybrid_text(text):
        return " ".join(text.split()[:HYBRID_QUERY_WORDS])

    def _pass_tail(pass_number):
        # Closes the query text, after a later pass's number, so no two queries send the same text.
        return b' %d"}' % pass_number if pass_number else b'"}'

    def _source_read_error(exc, columns, table):
        """A LoadTestError when the table can't be read or seen, else a RuntimeError; columns is e.g. "column x"."""
        error_class = SPARK_ERROR_CLASS.match(str(exc).lstrip())
        read_failure = "Could not read %s of %s (%s%s)." % (
            columns, table, type(exc).__name__, ": " + error_class.group(1) if error_class else "")
        if not any(error in str(exc) for error in SOURCE_TABLE_ACCESS_ERRORS):
            return RuntimeError(read_failure)
        return LoadTestError(read_failure, "source_table_access", next_step=True)

    def _sample_texts(source_table, column, limit, sample=True):
        """Distinct non-blank texts from the column, cut for the query type, in a seeded random order."""
        distinct = {}
        # Cutting and dropping blank text in Spark keeps long text from being copied to the driver; the checks below
        # are a backstop.
        rows = _source_rows(source_table, column, limit,
                            value=F.substring(_quoted_col(column), 1, MAX_QUERY_TEXT_CHARS).alias(column),
                            keep=F.length(F.trim(_quoted_col(column))) > 0, sample=sample)
        for row in rows:
            if isinstance(row[0], str) and row[0].strip():
                text = row[0][:MAX_QUERY_TEXT_CHARS]
                # Cut before the checks below, so two texts the cut makes the same are sent once.
                distinct[_hybrid_text(text) if HYBRID_QUERIES else text] = None
        # A text that is another one plus a pass suffix would repeat one of that text's later queries.
        texts = [text for text in distinct if not (PASS_SUFFIX.search(text) and PASS_SUFFIX.sub("", text) in distinct)]
        return [texts[i] for i in np.random.default_rng(SEED).permutation(len(texts))]

    if TEXT_COLUMNS:
        _text_table, _text_column = INDEX_SPEC.get("source_table"), TEXT_COLUMNS[0].get("name")
        if index_meta.get("direct_access_index_spec") and not _text_table:
            raise LoadTestError("This Direct Access index has no source table to sample query text from.",
                                "unsupported_index", next_step=True)
        if not _text_table or not _text_column:
            raise LoadTestError("This index has no source table or text column to sample query text from.",
                                "unsupported_index", next_step=True)
        try:
            QUERY_TEXTS = _sample_texts(_text_table, _text_column, min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES))
            # The sample's share is sized from all indexed rows, so when few have text it can hold none: read the first
            # rows that do instead.
            if not QUERY_TEXTS and INDEXED_ROW_COUNT > min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES):
                QUERY_TEXTS = _sample_texts(_text_table, _text_column, min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES),
                                            sample=False)
        except Exception as exc:
            raise _source_read_error(exc, "column " + _text_column, _text_table) from exc
        if not QUERY_TEXTS:
            raise LoadTestError("Column %s of %s has no text to send as queries." % (_text_column, _text_table),
                                "no_query_text", next_step=True)
        print("Sampled %d texts from %s." % (len(QUERY_TEXTS), _text_table))
        # Heads omit the closing quote so a pass suffix can go inside the string; json.dumps escapes non-ASCII, so ASCII
        # encoding is safe.
        _text_body_heads = [(json.dumps(_query_fields)[:-1] + ', "query_text": ' + json.dumps(text)[:-1]).encode("ascii")
                            for text in QUERY_TEXTS]

        def query_body(query_index):
            pass_number, text_index = divmod(query_index, len(_text_body_heads))
            return _text_body_heads[text_index] + _pass_tail(pass_number)
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Query vectors (indexes with self-managed embeddings)
try:
    cancel_watch.check(force=True)
    PERTURBATION = 0.015
    EMBEDDING_DIMENSION = next((int(column["embedding_dimension"])
                                for column in INDEX_SPEC.get("embedding_vector_columns") or []
                                if column.get("embedding_dimension")), None)
    VECTOR_COLUMN = next((column.get("name") for column in INDEX_SPEC.get("embedding_vector_columns") or []
                          if column.get("name")), None)

    def _index_string_columns():
        """The index's string columns, or None if unreadable. Unity Catalog lists an index's columns as exactly those it
        syncs; in a live test the index GET didn't return columns_to_sync, even for an index created with it."""
        # The notebook's own token: the load-test token is scoped to querying the index.
        try:
            resp = _workspace_get("/api/2.1/unity-catalog/tables/" + INDEX_PATH, _notebook_auth_headers)
            columns = resp.json().get("columns") if resp.status_code == 200 else None
        except (RunCancelled, RunTimedOut):
            raise
        except Exception:
            return None
        # An index always has its key and vector columns, so an empty list says nothing about what it syncs.
        if not isinstance(columns, list) or not columns:
            return None
        return {column.get("name") for column in columns
                if isinstance(column, dict) and column.get("type_name") == "STRING"}

    def _hybrid_text_column():
        """The Text column widget's column as a Spark column, once checked to be a string column the index syncs."""
        source_table = INDEX_SPEC.get("source_table")
        if not source_table:
            raise LoadTestError("Hybrid queries take their text from the index's source table, and this index has none.",
                                "unsupported_index", next_step=True)
        if not TEXT_COLUMN:
            raise LoadTestError("Hybrid queries to an index with self-managed embeddings need a text column, and none was "
                                "given.", "hybrid_text_column", next_step=True)
        synced = _index_string_columns()
        if synced is None:
            print_wrapped("Could not read the index's columns from Unity Catalog, so whether it syncs %s wasn't checked."
                          % TEXT_COLUMN)
        elif TEXT_COLUMN not in synced:
            raise LoadTestError("The text column %s isn't a string column the index syncs." % TEXT_COLUMN,
                                "hybrid_text_column", next_step=True)
        try:
            field = next((field for field in spark.table(source_table).schema.fields if field.name == TEXT_COLUMN), None)
        except Exception as exc:
            raise _source_read_error(exc, "column " + TEXT_COLUMN, source_table) from exc
        # isinstance, so collated strings count too.
        if field is None or not isinstance(field.dataType, StringType):
            raise LoadTestError("%s has no string column %s." % (source_table, TEXT_COLUMN), "hybrid_text_column",
                                next_step=True)
        return _quoted_col(TEXT_COLUMN)

    HYBRID_TEXT = _hybrid_text_column() if HYBRID_QUERIES and not TEXT_COLUMNS else None
    # Recorded only once used: an index that embeds query text ignores the widget.
    if HYBRID_TEXT is not None:
        progress.note(text_column=TEXT_COLUMN)

    def _sample_vectors(limit, text=None, sample=True):
        """Returns (vectors, texts, None), or (None, None, why random vectors are used instead). texts is None, or with
        text (a Spark column) the hybrid query text of each vector's row. sample is passed to _source_rows."""
        source_table = INDEX_SPEC.get("source_table")
        column = VECTOR_COLUMN
        if not source_table or not column:
            return None, None, "this index has no source table to sample"
        vectors, texts, other_sizes = [], [], set()
        rows = (_source_rows(source_table, column, limit) if text is None else _source_rows(
            source_table, column, limit, value=[_quoted_col(column), F.substring(text, 1, MAX_QUERY_TEXT_CHARS)],
            keep=_quoted_col(column).isNotNull() & (F.length(F.trim(text)) > 0), sample=sample))
        for row in rows:
            vector = np.asarray(row[0], dtype=np.float32)
            if not np.isfinite(vector).all():
                continue
            if text is not None:
                if not (isinstance(row[1], str) and row[1].strip()):
                    continue
                row_text = _hybrid_text(row[1])
            # The index's dimension is what queries must match; without it, the first sampled vector sets the size.
            expected = EMBEDDING_DIMENSION or (len(vectors[0]) if vectors else len(vector))
            if len(vector) == expected:
                vectors.append(vector)
                if text is not None:
                    texts.append(row_text)
            else:
                other_sizes.add(len(vector))
        if not vectors:
            if other_sizes:
                return None, None, "vectors in %s.%s have %s values but the index expects %d" % (
                    source_table, column, "/".join(str(size) for size in sorted(other_sizes)), EMBEDDING_DIMENSION)
            return None, None, "%s has no usable vectors in %s" % (source_table, column)
        if text is None:
            unique, texts = np.unique(np.stack(vectors), axis=0), None
        else:
            # A vector seen twice keeps the text of its first row.
            unique, first_rows = np.unique(np.stack(vectors), axis=0, return_index=True)
            texts = [texts[row] for row in first_rows]
        print("Sampled %d vectors from %s." % (len(unique), source_table))
        order = np.random.default_rng(SEED).permutation(len(unique))
        return unique[order], None if texts is None else [texts[i] for i in order], None

    # Each query is a base vector (sampled, or random) plus noise from a fixed pool scaled to about PERTURBATION of that
    # vector's length; the first pass over the base vectors sends them unchanged. Values are sent as whole numbers with
    # one shared exponent (123456e-7, 6 significant digits) because formatting integers is about three times faster than
    # formatting floats, which matters at thousands of queries per second.
    NOISE_ROWS = 4096
    RANDOM_BASE_VECTORS = 4096

    def query_mantissas(query_index):
        pass_number, base_index = divmod(query_index, len(BASE_MANTISSAS))
        if not pass_number:
            return BASE_MANTISSAS[base_index]
        # Later passes add one noise row, then pairs of rows, which keeps queries unique even with few source vectors.
        gap, row = divmod(pass_number - 1, NOISE_ROWS)
        noise = NOISE[row]
        if gap:
            noise = noise + NOISE[(row + (gap - 1) % (NOISE_ROWS // 2 - 1) + 1) % NOISE_ROWS]
        return BASE_MANTISSAS[base_index] + np.rint(noise * NOISE_SCALES[base_index]).astype(np.int32)

    if not TEXT_COLUMNS:
        try:
            BASE_VECTORS, BASE_TEXTS, _random_reason = _sample_vectors(min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES), HYBRID_TEXT)
            # The sample's share is sized from all indexed rows, so when few have text it can hold no usable pair: read
            # the first rows that pass the filter instead. Retrying any empty result costs one bounded read on a run
            # that would stop anyway.
            if (BASE_VECTORS is None and HYBRID_TEXT is not None
                    and INDEXED_ROW_COUNT > min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES)):
                BASE_VECTORS, BASE_TEXTS, _random_reason = _sample_vectors(
                    min(TOTAL_QUERIES, MAX_SAMPLED_QUERIES), HYBRID_TEXT, sample=False)
        except Exception as exc:
            # Random vectors have no text to go with them, so a hybrid test needs the table.
            if HYBRID_TEXT is not None:
                raise _source_read_error(exc, "columns %s and %s" % (VECTOR_COLUMN, TEXT_COLUMN),
                                         INDEX_SPEC["source_table"]) from exc
            BASE_VECTORS, BASE_TEXTS, _random_reason = None, None, "the source table couldn't be sampled (%s)" % (
                type(exc).__name__)
        if BASE_VECTORS is None:
            if HYBRID_TEXT is not None:
                raise LoadTestError("No row of %s has both text in %s and a usable vector: %s." % (
                    INDEX_SPEC["source_table"], TEXT_COLUMN, _random_reason), "no_query_text", next_step=True)
            if not EMBEDDING_DIMENSION:
                raise RuntimeError("The embedding dimension is unknown and %s." % _random_reason)
            print_wrapped("WARNING: %s; using random vectors instead. Latency may differ from real queries."
                          % _capitalized(_random_reason))
        DIMENSION = EMBEDDING_DIMENSION or BASE_VECTORS.shape[1]
        _rng = np.random.default_rng(SEED)
        if BASE_VECTORS is None:
            BASE_VECTORS = _rng.standard_normal((RANDOM_BASE_VECTORS, DIMENSION)).astype(np.float32)
            BASE_VECTORS /= np.linalg.norm(BASE_VECTORS, axis=1, keepdims=True)
        _largest = float(np.abs(BASE_VECTORS).max()) * 1.1 or 1.0
        VECTOR_EXPONENT = 6 - int(math.ceil(math.log10(_largest)))
        _scale = 10.0 ** VECTOR_EXPONENT
        BASE_MANTISSAS = np.rint(BASE_VECTORS * _scale).astype(np.int32)
        # At least one unit of the last digit, so even a near-zero vector gets noise that changes it.
        NOISE_SCALES = np.maximum(
            PERTURBATION * np.linalg.norm(BASE_VECTORS, axis=1) / math.sqrt(DIMENSION) * _scale, 1.0).astype(np.float32)
        NOISE = _rng.standard_normal((NOISE_ROWS, DIMENSION)).astype(np.float32)
        VECTOR_FORMAT = "[" + ",".join(["%d" + ("e%d" % -VECTOR_EXPONENT if VECTOR_EXPONENT else "")] * DIMENSION) + "]"
        BODY_PREFIX = json.dumps(_query_fields)[:-1] + ', "query_vector": '

        if BASE_TEXTS is None:
            def query_body(query_index):
                return (BODY_PREFIX + VECTOR_FORMAT % tuple(query_mantissas(query_index).tolist()) + "}").encode("ascii")
        else:
            # Each vector's text, without its closing quote, so a later pass's number goes inside the string as for
            # text-only queries; for hybrid search that number is one more keyword, a rare one.
            _text_tails = [(', "query_text": ' + json.dumps(text)[:-1]).encode("ascii") for text in BASE_TEXTS]

            def query_body(query_index):
                pass_number, base_index = divmod(query_index, len(BASE_MANTISSAS))
                return ((BODY_PREFIX + VECTOR_FORMAT % tuple(query_mantissas(query_index).tolist())).encode("ascii")
                        + _text_tails[base_index] + _pass_tail(pass_number))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# MAGIC %md ## Check one query

# COMMAND ----------

# DBTITLE 1,Test query
try:
    cancel_watch.check(force=True)
    progress.at("preflight")
    # Two test queries on one connection: the second leaves out connection setup, so its time sizes the first stage
    # close to what it needs. If only the second fails, the first one's time is used rather than stopping.
    COLD_EMBEDDING_HINT = (" The embedding endpoint may be cold (scaled to zero); retry once it's warm."
                           if EMBEDDING_ENDPOINT else "")
    # Waits before each retry of the first test query. A busy endpoint can turn one query away (429, timeout, 5xx) and
    # still serve most of the ramp, so a query that keeps failing that way is recorded and the ramp runs anyway.
    PREFLIGHT_RETRY_WAITS_SECONDS = (3, 9)
    # Without a successful test query, the first stage is sized as for a query this fast; more connections open if
    # replies are slower.
    UNMEASURED_PROBE_SECONDS = 0.1

    # Woken every second; the Jobs API is asked at most every CANCEL_CHECK_SECONDS.
    TEST_QUERY_CANCEL_POLL_SECONDS = 1

    def _post_checking_cancel(session, **kwargs):
        """session.post to the index, run in a thread so a cancel from the console stops the wait; the abandoned request
        ends with its own timeout."""
        outcome = {}

        def post():
            try:
                outcome["response"] = session.post(QUERY_URL.geturl(), **kwargs)
            except BaseException as exc:
                outcome["error"] = exc

        sender = threading.Thread(target=post, daemon=True)
        sender.start()
        while True:
            sender.join(TEST_QUERY_CANCEL_POLL_SECONDS)
            if not sender.is_alive():
                break
            cancel_watch.check()
        if "error" in outcome:
            raise outcome["error"]
        return outcome["response"]

    def _query_headers():
        """The signed-in headers for a query. After a fallback its User-Agent says so, keeping the ai-search-loadtest
        prefix that request-log filters match."""
        headers = auth.fresh_headers()
        if not QUERY_ROUTE_FALLBACK:
            return headers
        return dict(headers, **{"User-Agent": USER_AGENT_HEADERS["User-Agent"] + " (workspace-fallback)"})

    # Shorter texts are left in: removing them would garble the reply, and they say little.
    MIN_SCRUBBED_QUERY_TEXT_CHARS = 4

    def _without_query_text(reply, body):
        """The reply (bytes) with the text of the query body it answers removed in any form a JSON reply could echo it,
        so sampled text never reaches an error message, an error sample or the run record."""
        # The ramp's bodies are memoryview slices, and a failed scrub must never end a stage, so any body is accepted.
        try:
            text = json.loads(bytes(body)).get("query_text")
        except Exception:
            return reply
        if not isinstance(text, str) or len(text) < MIN_SCRUBBED_QUERY_TEXT_CHARS:
            return reply
        reply = bytes(reply or b"")
        for form in {text, json.dumps(text)[1:-1], json.dumps(text, ensure_ascii=False)[1:-1]}:
            reply = reply.replace(form.encode("utf-8"), b"[query text]")
        return reply

    def _test_query(session, query_index):
        """(seconds, response, None) for a successful query, else (None, None, (HTTP status or None, cause, message))."""
        started = time.monotonic()
        try:
            response = _post_checking_cancel(session, headers=_query_headers(), data=query_body(query_index),
                                             timeout=60)
        except requests.RequestException as exc:
            # ConnectTimeout is also a Timeout, but a connection that never opens (e.g. a firewall) is unreachable.
            cause = ("timeout" if isinstance(exc, requests.Timeout) and not isinstance(exc, requests.ConnectionError)
                     else "unreachable")
            return None, None, (None, cause, "Could not reach the index (%s).%s" % (
                type(exc).__name__, COLD_EMBEDDING_HINT if cause == "timeout" else ""))
        if response.status_code == 200:
            return time.monotonic() - started, response, None
        cause = reply_cause(response.status_code, response.content[:REPLY_PREFIX_BYTES], bool(EMBEDDING_ENDPOINT))
        hint = ""
        refusal = EMBEDDING_ENDPOINT and embedding_token_refusal(
            response.status_code, response.content[:REPLY_PREFIX_BYTES], auth.identity,
            auth.identity == NOTEBOOK_USER_IDENTITY, QUERIES_VIA_WORKSPACE, EMBEDDING_ENDPOINT_ROUTE_OPTIMIZED,
            EMBEDDING_ENDPOINT_NAME, sp_scope)
        # A route-optimized endpoint the sign-in couldn't get Can Query for already has the more specific hint below.
        if refusal and not (refusal[0] == "auth" and UNAUTHORIZED_ROUTE_OPTIMIZED_ENDPOINT):
            cause, hint = refusal[0], " " + refusal[1]
        # The auth next step already covers the embedding endpoint; only a route-optimized one it couldn't sign in for is
        # worth naming.
        elif UNAUTHORIZED_ROUTE_OPTIMIZED_ENDPOINT and response.status_code in (401, 403):
            hint = (" %s also needs Can Query on the route-optimized endpoint %s."
                    % (_capitalized(auth.identity), UNAUTHORIZED_ROUTE_OPTIMIZED_ENDPOINT))
        elif cause == "timeout":
            hint = COLD_EMBEDDING_HINT
        reply = _reply_text(_without_query_text(response.content, query_body(query_index)), ERROR_TEXT_CHARS)
        # An empty reply leaves no gap between the status and the hint.
        return None, None, (response.status_code, cause, "The test query failed: status %s%s%s" % (
            response.status_code, " " + reply if reply else "", hint))

    def _test_queries(retry_waits):
        """(attempts, failure or None, [(seconds, response)]) for the test queries to QUERY_URL, retrying a failure that
        isn't fatal after each of retry_waits."""
        results = []
        with requests.Session() as session:
            for attempt, retry_wait in enumerate((0,) + tuple(retry_waits), 1):
                if retry_wait:
                    print_wrapped("The test query failed (%s); trying again in %d s." % (failure[1], retry_wait))
                    time.sleep(retry_wait)
                    cancel_watch.check(force=True)
                seconds, response, failure = _test_query(session, 0)
                if failure is None or failure[1] in FATAL_CAUSES:
                    break
            if failure is None:
                results.append((seconds, response))
                seconds, response, _ = _test_query(session, 1)
                if response is not None:
                    results.append((seconds, response))
        return attempt, failure, results

    probe_attempt, probe_failure, probe_results = _test_queries(PREFLIGHT_RETRY_WAITS_SECONDS)
    if probe_failure is not None and falls_back_to_workspace_route(
            probe_failure[1], auth.identity == NOTEBOOK_USER_IDENTITY, QUERIES_VIA_WORKSPACE,
            EMBEDDING_ENDPOINT_ROUTE_OPTIMIZED, (QUERY_ENDPOINT_NAME or SOURCE_ENDPOINT_NAME) == EMBEDDING_ENDPOINT_NAME):
        _endpoint_query_url = QUERY_URL
        # Built the way the index reports its URL when the endpoint's own route isn't available.
        QUERY_URL = urllib.parse.urlsplit(WORKSPACE_HOST + "/api/2.0/vector-search/indexes/" + INDEX_PATH + "/query")
        QUERIES_VIA_WORKSPACE = QUERY_ROUTE_FALLBACK = True
        print_wrapped("WARNING: %s Trying the workspace URL instead, which embeds query text without using your token."
                      % probe_failure[2])
        # Retried like a run that starts on the workspace URL: only a failure every query would hit puts it back.
        _fallback_attempts, _fallback_failure, _fallback_results = _test_queries(PREFLIGHT_RETRY_WAITS_SECONDS)
        if _fallback_failure is None or not (_fallback_failure[1] in FATAL_CAUSES or _fallback_failure[1] == "unreachable"):
            # Counted on the route the run goes on with, so a degraded run's "failed N times" is about that route.
            probe_attempt, probe_failure, probe_results = _fallback_attempts, _fallback_failure, _fallback_results
            # queries_via_workspace too, so a console that predates workspace_route_fallback still shows the route, though
            # with that code's contact-support advice; transitional, until such consoles are gone.
            RUN_WARNING_CODES = RUN_WARNING_CODES + ["queries_via_workspace", "workspace_route_fallback"]
            progress.note(query_route="workspace", warnings=RUN_WARNING_CODES)
            print_wrapped("WARNING: This test runs through the workspace URL because your sign-in can't be used to embed "
                          "query text on the endpoint's own URL. The workspace URL has a lower per-workspace rate limit, "
                          "so the results don't show the endpoint's full throughput. To test full throughput, run as a "
                          "service principal.")
            # The proxy check was for the endpoint host; repeat it for the workspace host.
            _warn_if_proxied()
            if TARGET_QPS > ENDPOINT_TARGET_QPS or RESULTS_FACTOR > 1.0:
                print_wrapped("The Target QPS warning above still applies. If the workspace URL turns out to be the "
                              "limit, the verdict won't suggest a higher Target QPS, since that can't help.")
        else:
            QUERY_URL, QUERIES_VIA_WORKSPACE, QUERY_ROUTE_FALLBACK = _endpoint_query_url, False, False
            if _fallback_failure[1] == "auth":
                # Refused on both routes: a missing permission (e.g. Can Query) is what to fix first.
                probe_attempt, probe_failure = _fallback_attempts, (
                    _fallback_failure[0], "auth", "Through the workspace URL: %s" % _fallback_failure[2])
            else:
                # The endpoint URL's refusal says what to fix; the workspace URL's failure is only printed.
                print_wrapped("The test query through the workspace URL also failed: %s" % _fallback_failure[2])
    if probe_failure is None:
        PROBE_SECONDS, probe = probe_results[-1]
        progress.note(preflight=preflight_record("ok", probe_attempt))
        print("Test query OK.")
    else:
        probe_status, probe_cause, probe_message = probe_failure
        # Every query would fail the same way, or none can reach the endpoint, so no stage is run.
        if probe_cause in FATAL_CAUSES or probe_cause == "unreachable":
            progress.note(preflight=preflight_record("failed", probe_attempt, probe_status, probe_cause,
                                                     _redacted(probe_message)))
            raise LoadTestError(probe_message, probe_cause, next_step=True)
        progress.note(preflight=preflight_record("degraded", probe_attempt, probe_status, probe_cause,
                                                 _redacted(probe_message)))
        # One query turned away says nothing about load, so no next step: the verdict says what to change.
        print_wrapped("WARNING: %s It failed %d times, so the ramp runs anyway and shows how many queries fail at each "
                      "stage." % (probe_message, probe_attempt))
        PROBE_SECONDS, probe = UNMEASURED_PROBE_SECONDS, None
    # The cells up to the ramp only define code, so a failure there is setup, not the test query.
    progress.at("setup")
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC ## Load engine
# MAGIC
# MAGIC No changes needed. How it works:
# MAGIC - Each query has a fixed send time. A query that can't be sent on time is skipped rather than sent late, so
# MAGIC   the rate never bursts above the stage's.
# MAGIC - Queries go over kept-alive HTTPS connections. More connections open as responses slow down; if new ones keep
# MAGIC   failing to open, it stops opening more.
# MAGIC - Up to one worker process per CPU core (at most 16) sends the queries, so one core doesn't cap the rate.
# MAGIC - It needs only libraries already on Databricks Runtime, so nothing needs installing. Queries go through a
# MAGIC   small built-in asyncio HTTP client rather than requests, so one process can keep hundreds of queries in
# MAGIC   flight.

# COMMAND ----------

# DBTITLE 1,HTTP connection
try:
    # Longer than the 30 seconds a Standard endpoint allows a query, so its own timeout reply can arrive first. A request
    # still waiting after this is counted as a timeout here.
    REQUEST_TIMEOUT_SECONDS = 40
    SSL_CONTEXT = ssl.create_default_context(cafile=requests_ca_bundle())

    class Connection:
        """One kept-alive HTTPS connection to the index, reopened when the server closes it."""

        def __init__(self):
            self.reader = self.writer = None
            # For the diagnostics: successful opens, and times the server ended this connection between replies.
            self.opened = 0
            self.server_closes = 0
            # A connection being set up counts toward a connection limit too.
            self.opening = False
            self.connect_error = None
            # When the request in flight times out (WorkerEngine.abort_late_requests enforces it), and whether it was
            # written yet: until then a slow open is a failed open, not a timed-out query.
            self.deadline = None
            self.timed_out = False
            self.request_sent = False
            # The last error reply's head and body, for the diagnostics' sample replies.
            self.error_reply = None
            # The last reply's body, for the diagnostics' sampled split of server time.
            self.reply = None

        def close(self):
            if self.writer is not None:
                self.writer.close()
            self.reader = self.writer = None

        async def _open(self):
            self.close()
            self.opening = True
            try:
                # Bounded by the request's deadline, so reopening can't stretch a request past its timeout.
                self.reader, self.writer = await asyncio.wait_for(asyncio.open_connection(
                    QUERY_URL.hostname, QUERY_URL.port or 443, ssl=SSL_CONTEXT), max(0.0, self.deadline - time.monotonic()))
            except ssl.SSLError:
                # Certificate and protocol faults keep their own kind, "TLS error".
                raise
            except (OSError, asyncio.TimeoutError, asyncio.IncompleteReadError) as exc:
                # For diagnostics: a refused connection and one reset during setup point at different limits.
                self.connect_error = type(exc).__name__
                if self.request_sent and time.monotonic() >= self.deadline:
                    self.timed_out = True
                raise ConnectFailed() from exc
            finally:
                self.opening = False
            self.opened += 1

        async def _read_head(self, request):
            """Sends the request and returns the response's status line and headers, or b"" if no reply arrived."""
            self.writer.write(request)
            self.request_sent = True
            try:
                return await self.reader.readuntil(b"\r\n\r\n")
            except asyncio.IncompleteReadError as exc:
                if exc.partial:
                    raise ResponseCutOff() from exc
                return b""
            except asyncio.LimitOverrunError:
                raise ValueError("response head too large") from None
            except ConnectionError:
                return b""

        async def _read_chunked_body(self):
            chunks = []
            while True:
                line = await self.reader.readline()
                if not line:
                    raise asyncio.IncompleteReadError(b"", None)
                size = int(line.split(b";")[0], 16)
                if size == 0:
                    # Optional trailer lines, then the blank line that ends the response.
                    while line != b"\r\n":
                        line = await self.reader.readline()
                        if not line:
                            raise asyncio.IncompleteReadError(b"", None)
                    return b"".join(chunks)
                chunks.append((await self.reader.readexactly(size + 2))[:-2])

        async def _exchange(self, request):
            reused = self.writer is not None and not self.reader.at_eof()
            if self.writer is not None and not reused:
                self.server_closes += 1
            if not reused:
                await self._open()
            head = await self._read_head(request)
            if not head and reused and not self.timed_out:
                # The server dropped this kept-alive connection before answering; queries are read-only, so the
                # request is sent once more on a new connection.
                await self._open()
                head = await self._read_head(request)
            if not head:
                raise ConnectionError("no reply")
            try:
                status = int(head.split(b"\r\n", 1)[0].split(b" ", 2)[1])
            except (IndexError, ValueError):
                raise ValueError("malformed status line") from None
            head = head.lower()
            length = _header(head, b"content-length")
            close = _header(head, b"connection") == b"close"
            try:
                if b"chunked" in (_header(head, b"transfer-encoding") or b""):
                    data = await self._read_chunked_body()
                elif length is not None:
                    data = await self.reader.readexactly(int(length))
                else:
                    # Without a length the body runs until the server closes the connection.
                    data = await self.reader.read()
                    close = True
            except (ConnectionError, asyncio.IncompleteReadError, ssl.SSLError) as exc:
                raise ResponseCutOff() from exc
            # 429 replies are kept too, since they show which rate limit refused the request.
            if status != 200:
                self.error_reply = (head, data)
            if close:
                # A body cut short by the timeout watchdog wasn't ended by the server.
                if not self.timed_out:
                    self.server_closes += 1
                self.close()
            return status, data

        async def send(self, request):
            """Returns (HTTP status, server time in ms); if the request failed, the status is the kind of failure."""
            self.deadline, self.timed_out, self.request_sent = time.monotonic() + REQUEST_TIMEOUT_SECONDS, False, False
            try:
                status, data = await self._exchange(request)
            except Exception as exc:
                status, data = _failure_kind(exc), None
            self.deadline = None
            if self.timed_out:
                # Aborted at its deadline, so any reply read up to then is incomplete.
                status = "timeout"
            if isinstance(status, str):
                self.close()
                return status, None
            if SERVER_TIME_SPLIT_EVERY:
                self.reply = data
            return status, _server_time(data) if status == 200 else None

        def abort_if_late(self, now):
            if self.deadline is not None and self.deadline < now and self.writer is not None:
                self.timed_out = True
                self.writer.transport.abort()

    def _header(head, name):
        """A header's value in a lowercased response head, or None."""
        start = head.find(b"\r\n" + name + b":")
        if start < 0:
            return None
        start += len(name) + 3
        return head[start:head.index(b"\r\n", start)].strip()

    # Response headers that show which layer answered an error; none of them carry credentials.
    SAMPLE_HEADERS = ("server", "via", "x-request-id")

    def _error_sample(status, head, data):
        """One error reply, shortened for the diagnostics: which layer answered and the start of its body."""
        headers = []
        for line in head.split(b"\r\n")[1:]:
            name, _, value = line.partition(b":")
            name = name.strip().decode("ascii", "replace")
            if name in SAMPLE_HEADERS or name.startswith("x-envoy-"):
                headers.append("%s=%s" % (name, _printable(value, 80)))
        return "HTTP %d: %s; body: %s" % (status, ", ".join(headers) or "no identifying headers",
                                          _printable(data, 200) or "empty")

    RESPONSE_TIME_PATTERN = re.compile(rb'"response_time"\s*:\s*(-?[0-9][0-9.eE+-]*)')

    def _server_time(data):
        """The endpoint's own handling time in ms from debug_info, or None."""
        start = data.rfind(b'"debug_info"')
        if start < 0:
            return None
        # A targeted search inside debug_info costs far less than parsing the whole response; json is the fallback.
        end = data.find(b"}", start)
        match = RESPONSE_TIME_PATTERN.search(data, start, end if end >= 0 else len(data))
        try:
            return float(match.group(1)) if match else float(json.loads(data)["debug_info"]["response_time"])
        except (ValueError, KeyError, TypeError):
            return None

    SERVER_TIME_SPLIT_PATTERNS = [re.compile(rb'"' + name + rb'"\s*:\s*(-?[0-9][0-9.eE+-]*)')
                                  for name in (b"embedding_gen_time", b"ann_time")]

    def _server_time_split(data):
        """The endpoint's embedding and search times in ms from debug_info; each is None if missing."""
        start = data.rfind(b'"debug_info"')
        if start < 0:
            return None, None
        end = data.find(b"}", start)
        times = []
        for pattern in SERVER_TIME_SPLIT_PATTERNS:
            match = pattern.search(data, start, end if end >= 0 else len(data))
            try:
                times.append(float(match.group(1)) if match else None)
            except ValueError:
                times.append(None)
        return times

    class ResponseCutOff(Exception):
        """The connection broke after the reply started, so this isn't a connection the server dropped for sitting idle."""

    class ConnectFailed(ConnectionError):
        """The connection failed or closed while being opened, before the request was sent."""

    def _failure_kind(exc):
        if isinstance(exc, ResponseCutOff):
            return "cut off"
        if isinstance(exc, ConnectFailed):
            return "connect failed"
        if isinstance(exc, ssl.SSLError):
            return "TLS error"
        # The connection closed or reset after the request was sent, before any reply arrived.
        if isinstance(exc, (ConnectionError, asyncio.IncompleteReadError)):
            return "connection closed"
        if isinstance(exc, ValueError):
            return "bad response"
        return type(exc).__name__

    def _request_head(headers):
        lines = ["POST %s HTTP/1.1" % QUERY_URL.path, "Host: %s" % QUERY_URL.netloc]
        lines += ["%s: %s" % item for item in headers.items()]
        return ("\r\n".join(lines) + "\r\n").encode("ascii")

    def _request(request_head, body):
        return b"".join((request_head, b"Content-Length: %d\r\n\r\n" % len(body), body))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Per-worker engine
try:
    # Measured from this many seconds into each stage, once responses flow steadily.
    MEASURE_AFTER_SECONDS = 3
    # Sending a query later than this after its time would burst above the stage rate, so it is skipped instead.
    SKIP_AFTER_SECONDS = 0.01
    MAX_CONNECTIONS_PER_WORKER = 256
    # Room above what the endpoint keeps in flight, so a slowdown doesn't use up connections before more open.
    CONNECTION_HEADROOM = 3
    # One or two failed opens can be chance; this many in a row suggests a connection limit.
    FAILED_OPENS_FOR_LIMIT = 3
    # The lowest limit one run of failed opens can set, so a short network blip can't shrink the pool to nothing.
    MIN_CONNECTION_LIMIT = 8
    # Weight of each request in the recent endpoint time, so a slowdown shows within about 100 requests.
    RECENT_WEIGHT = 0.01
    # Connections are added at most this often, so a burst of slow responses doesn't open dozens at once.
    GROW_SECONDS = 0.005
    # Each worker times a short sleep: waking late means ready work waited for CPU.
    LOOP_TICK_SECONDS = 0.005
    # An idle worker wakes well within a millisecond, so a recent delay above this means it is short of CPU.
    LOOP_BUSY_DELAY_SECONDS = SKIP_AFTER_SECONDS / 2
    # Weight of each tick in the recent delay, so it follows the last few dozen milliseconds.
    LOOP_DELAY_WEIGHT = 0.1
    # Error replies each worker keeps per stage for the diagnostics.
    ERROR_SAMPLES = 3
    # Reply classification lives in count_reply (Console run record cell) so the test query, ramp and verdict agree.
    EMBEDDING_INDEX = bool(EMBEDDING_ENDPOINT)
    # Failures where the connection closed before any reply, while opening or after sending: the "no reply" column.
    NO_REPLY_KINDS = ("connection closed", "connect failed")
    CONNECTION_COUNTERS = ("opened", "server_closes")
    # For an index that computes its embeddings, one reply in this many is read for how server time splits between
    # embedding and search (more often with diagnostics on); reading every reply would take CPU from sending.
    SERVER_TIME_SPLIT_EVERY = (10 if DEBUG else 50) if EMBEDDING_ENDPOINT else 0

    class StageStats:
        def __init__(self, measure_from):
            self.measure_from = measure_from
            # sent, rate_limited (429s), no_reply (NO_REPLY_KINDS) and errors (every other failure).
            self.counts = Counter()
            self.errors_by_kind = Counter()
            self.ok_in_window = 0
            self.skipped_in_window = 0
            # Total time requests held a connection, failures included; divided by requests it sizes the next stage. The
            # share sent in the measured window, divided by its length, is how many connections were in use on average.
            self.request_seconds = 0.0
            self.window_request_seconds = 0.0
            self.client_ms = array("d")
            self.server_ms = array("d")
            self.server_seconds = 0.0
            self.server_count = 0
            # Client minus server latency per request: network round trip plus any wait in this notebook.
            self.gap_ms = array("d")
            self.embedding_ms = array("d")
            self.search_ms = array("d")

        def split(self, scheduled_at, data):
            if scheduled_at >= self.measure_from:
                embedding_ms, search_ms = _server_time_split(data)
                if embedding_ms is not None:
                    self.embedding_ms.append(embedding_ms)
                if search_ms is not None:
                    self.search_ms.append(search_ms)

        def skip(self, scheduled_at):
            if scheduled_at >= self.measure_from:
                self.skipped_in_window += 1

        def record(self, status, scheduled_at, client_ms, server_ms):
            self.counts["sent"] += 1
            if status == 429:
                self.counts["rate_limited"] += 1
            elif status != 200:
                kind = status if isinstance(status, str) else "HTTP %d" % status
                self.counts["no_reply" if kind in NO_REPLY_KINDS else "errors"] += 1
                self.errors_by_kind[kind] += 1
            else:
                if server_ms is not None:
                    self.server_seconds += server_ms / 1000.0
                    self.server_count += 1
                # Counted by scheduled send time, so slow responses to the last sends of a stage still count.
                if scheduled_at >= self.measure_from:
                    self.ok_in_window += 1
                    self.client_ms.append(client_ms)
                    if server_ms is not None:
                        self.server_ms.append(server_ms)
                        self.gap_ms.append(client_ms - server_ms)

    class ConnectionLimit:
        """The most connections a worker holds (open or being opened) once new ones keep failing to open.

        FAILED_OPENS_FOR_LIMIT failures in a row set it to the connections held then, at least MIN_CONNECTION_LIMIT unless
        a limit is already set. It is lifted by an open that succeeds with more connections held than when it was set, so
        a network blip can't pin it low, or after a run of successful opens as long as the limit. That run grows fourfold
        each time the limit comes back, so reopening connections the endpoint closes only rarely tries to go past a real
        limit.
        """

        def __init__(self):
            self.value = None
            # What the worker actually held when the limit was set, which is what the report quotes.
            self.held_when_set = 0
            self.failures_in_row = 0
            self.opens_in_row = 0
            self.opens_to_lift = 0

        def allows(self, held):
            return self.value is None or held < self.value

        def opened(self, held):
            self.failures_in_row = 0
            self.opens_in_row += 1
            if self.value is not None and (held > self.held_when_set or self.opens_in_row >= self.opens_to_lift):
                self.value = None

        def failed(self, held):
            self.failures_in_row += 1
            self.opens_in_row = 0
            if self.failures_in_row >= FAILED_OPENS_FOR_LIMIT:
                lowest = MIN_CONNECTION_LIMIT if self.value is None else 1
                self.value, self.held_when_set, self.failures_in_row = max(lowest, held), held, 0
                self.opens_to_lift = max(4 * self.opens_to_lift, self.value)

    def _connection_counters(connections):
        return dict((key, sum(getattr(connection, key) for connection in connections)) for key in CONNECTION_COUNTERS)

    class WorkerEngine:
        """Runs inside one worker process; query n of every stage belongs to worker n % worker_count."""

        def __init__(self, worker_index, worker_count):
            self.worker_index = worker_index
            self.worker_count = worker_count
            self.connections = []
            self.bodies = None
            self.limit = ConnectionLimit()
            # The last limit found: warming and the next stage start from at most this many connections.
            self.start_limit = None
            # Failed opens and connection counters since the last stage ended, so each stage reports its warm-up too.
            self.connect_errors = Counter()
            self.counters_reported = _connection_counters([])
            # The last stage's mean loop delay, so the next stage's growth starts from it.
            self.loop_delay_seconds = 0.0
            # Until the endpoint reports its own time (debug_info), whole request times stand in for it.
            self.server_time_seen = False

        def open_connections(self):
            return sum(1 for connection in self.connections if connection.writer is not None or connection.opening)

        def note_open(self, connection, opened_before, status):
            if status == "connect failed":
                self.connect_errors[connection.connect_error] += 1
                self.limit.failed(self.open_connections())
            elif connection.opened > opened_before:
                self.limit.opened(self.open_connections())

        async def abort_late_requests(self):
            """Ends requests that run past REQUEST_TIMEOUT_SECONDS; one check a second costs far less than a timer per
            request."""
            while True:
                await asyncio.sleep(1)
                now = time.monotonic()
                for connection in self.connections:
                    connection.abort_if_late(now)

        async def prepare(self, send_rate, first_query, limit):
            """Builds up to limit of this worker's request bodies for the next stage before it starts, so sending them
            is cheap. They are packed into one buffer with offsets, which avoids per-body object overhead."""
            numbers = range(self.worker_index, int(math.ceil(send_rate * STAGE_SECONDS)), self.worker_count)[:limit]
            buffer, offsets = bytearray(), array("Q", [0])
            for number in numbers:
                buffer += query_body(first_query + number)
                offsets.append(len(buffer))
            self.bodies = (memoryview(buffer), offsets)

        async def warm(self, connection_count, rate, start, first_query, headers):
            """Opens the extra connections a stage needs before it starts, so connection setup isn't in its
            measured latency; each sends one query, paced at the stage rate to avoid a burst."""
            request_head = _request_head(headers)
            self.limit = ConnectionLimit()
            new_connections = [Connection() for _ in range(min(connection_count, self.start_limit or connection_count)
                                                             - len(self.connections))]
            # Listed before they open, so the limit counts each one as it opens.
            self.connections += new_connections

            async def warm_one(j, connection):
                number = self.worker_index + self.worker_count * j
                await asyncio.sleep(max(0.0, start + number / rate - time.monotonic()))
                if self.limit.allows(self.open_connections()):
                    status, _ = await connection.send(_request(request_head, query_body(first_query + number)))
                    self.note_open(connection, 0, status)

            # Connections that couldn't open stay listed; the stage starts on open ones.
            await asyncio.gather(*(warm_one(j, connection) for j, connection in enumerate(new_connections)))
            self.start_limit = self.limit.value or self.start_limit

        async def stage(self, qps, start, connection_count, first_query, headers, endpoint_seconds_hint,
                        network_seconds):
            request_head = _request_head(headers)
            end = start + STAGE_SECONDS
            stats = StageStats(start + MEASURE_AFTER_SECONDS)
            # This worker's queries in the stage, in send order.
            numbers = itertools.count(self.worker_index, self.worker_count)
            buffer, offsets = self.bodies if self.bodies else (None, ())
            prepared = len(offsets) - 1
            senders = set()
            # Connections with a running sender, and those dropped for the limit; live and peak count the senders.
            in_use, dropped = set(), []
            live = peak = busy = 0
            last_added = 0.0
            # The endpoint's recent time per request, so a mid-stage slowdown shows up within a fraction of a second.
            recent_endpoint_seconds = endpoint_seconds_hint
            error_samples = []
            # How late the event loop wakes from a sleep: recently, for growth, and summed over the measured window.
            loop_delay = {"recent": self.loop_delay_seconds, "total": 0.0, "ticks": 0}

            def worker_cpu_saturated():
                return loop_delay["recent"] > LOOP_BUSY_DELAY_SECONDS

            async def measure_loop_delay():
                while True:
                    before = time.monotonic()
                    await asyncio.sleep(LOOP_TICK_SECONDS)
                    delay = max(0.0, time.monotonic() - before - LOOP_TICK_SECONDS)
                    # Idle ticks before the stage starts would wash out the carried-over delay.
                    if before >= start:
                        loop_delay["recent"] += LOOP_DELAY_WEIGHT * (delay - loop_delay["recent"])
                    if stats.measure_from <= before < end:
                        loop_delay["total"] += delay
                        loop_delay["ticks"] += 1

            async def sample_window():
                """CPU time at the start and end of the measured window, which leaves out the idle wait before the stage
                and the wait for late replies after it."""
                await asyncio.sleep(max(0.0, stats.measure_from - time.monotonic()))
                first = (time.process_time(), time.monotonic())
                await asyncio.sleep(max(0.0, end - time.monotonic()))
                return first, (time.process_time(), time.monotonic())

            def add_sender(connection):
                nonlocal live, peak
                in_use.add(connection)
                live += 1
                peak = max(peak, live)
                senders.add(asyncio.ensure_future(sender(connection)))

            def drop(connection):
                nonlocal live
                in_use.discard(connection)
                self.connections.remove(connection)
                dropped.append(connection)
                live -= 1

            def over_limit(connection):
                """Whether this connection would have to open past the limit; the last sender always stays."""
                return connection.writer is None and live > 1 and not self.limit.allows(self.open_connections())

            def grow_if_needed():
                """Adds a sender before idle connections run out, so slower responses don't cost a skipped query."""
                nonlocal last_added
                idle = live - busy
                # A CPU-saturated worker is short of CPU, not connections. Delay from CPU used outside this notebook stops
                # growth too, which can lower its throughput.
                if (idle >= max(2, 0.1 * live) or live >= MAX_CONNECTIONS_PER_WORKER
                        or time.monotonic() - last_added < GROW_SECONDS or worker_cpu_saturated()):
                    return
                # Sized by the endpoint's time, not the request time, which also counts time responses wait for this
                # notebook's CPU and would grow the pool far past what the endpoint needs.
                in_flight_seconds = recent_endpoint_seconds + network_seconds
                if live >= CONNECTION_HEADROOM * qps / self.worker_count * in_flight_seconds:
                    return
                last_added = time.monotonic()
                spares = [connection for connection in self.connections if connection not in in_use]
                # An open spare first: it adds no connection, so it is fine even at the limit.
                spare = next((connection for connection in spares if connection.writer is not None), None)
                if spare is None:
                    if not self.limit.allows(self.open_connections()):
                        return
                    spare = spares[0] if spares else Connection()
                    if not spares:
                        self.connections.append(spare)
                add_sender(spare)

            async def sender(connection):
                nonlocal busy, recent_endpoint_seconds
                if over_limit(connection):
                    drop(connection)
                    return
                for number in numbers:
                    scheduled_at = start + number / qps
                    if scheduled_at >= end:
                        return
                    delay = scheduled_at - time.monotonic()
                    if delay > 0:
                        await asyncio.sleep(delay)
                    elif -delay > SKIP_AFTER_SECONDS:
                        stats.skip(scheduled_at)
                        # When the CPU keeps up, a skip means every connection was busy, so another may be needed.
                        grow_if_needed()
                        continue
                    index = number // self.worker_count
                    if index < prepared:
                        body = buffer[offsets[index]:offsets[index + 1]]
                    else:
                        body = query_body(first_query + number)
                    busy += 1
                    grow_if_needed()
                    sent_at = time.monotonic()
                    opened_before = connection.opened
                    status, server_ms = await connection.send(_request(request_head, body))
                    busy -= 1
                    finished_at = time.monotonic()
                    self.note_open(connection, opened_before, status)
                    stats.request_seconds += finished_at - sent_at
                    if stats.measure_from <= sent_at < end:
                        stats.window_request_seconds += finished_at - sent_at
                    if server_ms is not None:
                        self.server_time_seen = True
                    if server_ms is not None or not self.server_time_seen:
                        sample = finished_at - sent_at if server_ms is None else server_ms / 1000.0
                        recent_endpoint_seconds += RECENT_WEIGHT * (sample - recent_endpoint_seconds)
                    # Latency runs from the scheduled send time, so any delay in this notebook is counted too.
                    stats.record(status, scheduled_at, (finished_at - scheduled_at) * 1000.0, server_ms)
                    if SERVER_TIME_SPLIT_EVERY:
                        if status == 200 and index % SERVER_TIME_SPLIT_EVERY == 0:
                            stats.split(scheduled_at, connection.reply)
                        # Dropped once read, so idle connections don't each hold a reply body.
                        connection.reply = None
                    if isinstance(status, int) and status != 200:
                        count_reply(stats.counts, status, connection.error_reply[1][:REPLY_PREFIX_BYTES], EMBEDDING_INDEX)
                        if len(error_samples) < ERROR_SAMPLES:
                            error_samples.append(_error_sample(status, connection.error_reply[0],
                                                               _without_query_text(connection.error_reply[1], body)))
                    # Checked only after a send, so no scheduled query is lost with the connection.
                    if over_limit(connection):
                        drop(connection)
                        return

            ticker = asyncio.ensure_future(measure_loop_delay())
            window_sampler = asyncio.ensure_future(sample_window())
            try:
                # Open connections first, so the stage starts on them rather than on ones it would have to reopen.
                self.connections.sort(key=lambda connection: connection.writer is None)
                self.limit = ConnectionLimit()
                for connection in self.connections[:min(connection_count, self.start_limit or connection_count)]:
                    add_sender(connection)
                # Requests time out after REQUEST_TIMEOUT_SECONDS, so senders finish by then after the stage ends.
                while not all(task.done() for task in senders):
                    await asyncio.wait([task for task in senders if not task.done()])
            finally:
                # Otherwise the ticker would keep running into the next stage on this long-lived loop.
                ticker.cancel()
                await asyncio.gather(ticker, return_exceptions=True)
            (cpu_first, wall_first), (cpu_last, wall_last) = await window_sampler
            for task in senders:
                if task.exception() is not None:
                    raise task.exception()
            self.bodies = None
            self.loop_delay_seconds = loop_delay["total"] / max(1, loop_delay["ticks"])
            if self.limit.value is not None:
                self.start_limit = self.limit.value
            elif self.start_limit and peak > self.start_limit:
                # This stage held more connections without hitting a limit.
                self.start_limit = None
            counts = stats.counts
            for key, value in _connection_counters(self.connections + dropped).items():
                counts[key] = value - self.counters_reported[key]
            self.counters_reported = _connection_counters(self.connections)
            connect_errors, self.connect_errors = self.connect_errors, Counter()
            return {"counts": counts, "errors_by_kind": stats.errors_by_kind, "connect_errors": connect_errors,
                    "ok_in_window": stats.ok_in_window, "skipped_in_window": stats.skipped_in_window,
                    "server_seconds": stats.server_seconds, "server_count": stats.server_count,
                    "gap_ms": stats.gap_ms.tolist(), "prepared": prepared, "request_seconds": stats.request_seconds,
                    "loop_delay_ms": 1000.0 * self.loop_delay_seconds, "connections": peak,
                    "connection_limit": self.limit.held_when_set if self.limit.value is not None else None,
                    "cpu_seconds": cpu_last - cpu_first, "wall_seconds": wall_last - wall_first,
                    "window_request_seconds": stats.window_request_seconds,
                    "error_samples": error_samples,
                    "client_ms": stats.client_ms.tolist(), "server_ms": stats.server_ms.tolist(),
                    "embedding_ms": stats.embedding_ms.tolist(), "search_ms": stats.search_ms.tolist()}
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Worker processes
try:
    def _worker_main(worker_index, worker_count, pipe):
        # One event loop for the worker's whole life, so its connections stay open between stages.
        event_loop = asyncio.new_event_loop()
        engine = WorkerEngine(worker_index, worker_count)
        # Runs alongside every command on this loop.
        timeouts = event_loop.create_task(engine.abort_late_requests())
        while True:
            command = pipe.recv()
            if command is None:
                timeouts.cancel()
                return
            name, kwargs = command
            try:
                pipe.send(("ok", event_loop.run_until_complete(getattr(engine, name)(**kwargs))))
            except Exception as exc:
                pipe.send(("error", type(exc).__name__))

    # The driver's side of the load workers: starting them, waiting for their replies, and stopping them. Kept apart from
    # the engine so the stop-on-cancel path can be tested with real forked processes.
    import os
    import signal
    import time

    # The longest single wait for a worker's reply: a job cancel only lands between Python steps, not inside one wait.
    WAIT_SLICE_SECONDS = 1.0
    # How long close() waits for a worker to exit after each of: the stop command, SIGTERM, SIGKILL.
    CLOSE_WAIT_SECONDS = (2.0, 2.0, 1.0)
    # After a cancel or failure, stopping all workers takes at most about this long, the last KILL_WAIT_SECONDS of it for
    # killing any that ignored SIGTERM.
    STOP_DEADLINE_SECONDS = 3.0
    KILL_WAIT_SECONDS = 1.0


    class LoadWorker:
        """A forked process that sends its share of every stage, so the load isn't limited to one core. check is called
        between waits and may raise to abandon the wait (the notebook passes its cancel check)."""

        def __init__(self, context, target, worker_index, worker_count, check=lambda: None,
                     close_wait_seconds=CLOSE_WAIT_SECONDS):
            self.check = check
            self.close_wait_seconds = close_wait_seconds
            self.pipe, child_pipe = context.Pipe()
            self.process = context.Process(target=target, args=(worker_index, worker_count, child_pipe), daemon=True)
            self.process.start()
            child_pipe.close()

        def run(self, name, **kwargs):
            self.pipe.send((name, kwargs))

        def result(self, timeout):
            deadline = time.monotonic() + timeout
            try:
                ready = False
                while not ready and time.monotonic() < deadline:
                    self.check()
                    ready = self.pipe.poll(min(WAIT_SLICE_SECONDS, max(0.0, deadline - time.monotonic())))
                # A reply already queued counts even when no wait time was left.
                ready = ready or self.pipe.poll(0)
                status, payload = self.pipe.recv() if ready else (None, None)
            except (EOFError, OSError):
                status, payload = None, None
            if status == "error":
                raise RuntimeError("A load worker failed (%s)." % payload)
            if status != "ok":
                raise RuntimeError("A load worker stopped responding.")
            return payload

        def stop(self):
            """Ends the worker at once, so its load stops with it."""
            if self.process.is_alive():
                self.process.terminate()

        def close(self, deadline=None):
            """Asks the worker to exit, then terminates and finally kills it, so no worker outlives the run. Given a
            deadline (time.monotonic), every wait ends by it, and the waits before the kill end KILL_WAIT_SECONDS sooner."""
            stop_wait, terminate_wait, kill_wait = self.close_wait_seconds

            def wait(seconds, reserve=KILL_WAIT_SECONDS):
                if deadline is not None:
                    seconds = min(seconds, max(0.0, deadline - reserve - time.monotonic()))
                self.process.join(seconds)

            try:
                self.pipe.send(None)
            except (OSError, ValueError):
                pass
            wait(stop_wait)
            if self.process.is_alive():
                self.process.terminate()
                wait(terminate_wait)
            if self.process.is_alive():
                try:
                    os.kill(self.process.pid, signal.SIGKILL)
                except OSError:
                    pass
                wait(kill_wait, reserve=0.0)


    def _run_on_all(workers, name, timeout, **kwargs):
        for worker in workers:
            worker.run(name, **kwargs)
        return [worker.result(timeout) for worker in workers]


    class StoppingWorkers:
        """Closes every worker when the block ends. On any exception, a cancel included, it first stops all of them, since
        closing them one by one could leave the others sending load for seconds."""

        def __init__(self, workers):
            self.workers = workers

        def __enter__(self):
            return self.workers

        def __exit__(self, exc_type, exc, traceback):
            deadline = None
            try:
                if exc_type is not None:
                    for worker in self.workers:
                        worker.stop()
                    # One deadline for all of them, so stopping takes at most about that long however many there are.
                    deadline = time.monotonic() + STOP_DEADLINE_SECONDS
            finally:
                for worker in self.workers:
                    worker.close(deadline)
            return False
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Driver sizing
try:
    # Each worker process uses one core and sends up to about QPS_PER_WORKER, so the top stage gets enough of them with
    # CAPACITY_BUFFER to spare, but never more than the usable cores (or MAX_WORKERS).
    MAX_WORKERS = 16
    QPS_PER_WORKER = 200
    CAPACITY_BUFFER = 1.1
    # Connections a stage starts with: twice what its rate keeps in flight at the last stage's request time (at least
    # this), so it rarely has to open more while measuring.
    MIN_LATENCY_SECONDS = 0.05
    # Request bodies for a stage are built before it starts, in up to this share of free memory minus a margin.
    PREPARED_MEMORY_SHARE = 0.5
    PREPARED_MEMORY_MARGIN_BYTES = 512 * 1024 ** 2

    def _cpu_quota_cores():
        # Containers are often limited by a CPU quota rather than by the cores they can see.
        try:
            with open("/sys/fs/cgroup/cpu.max") as quota_file:
                quota, period = quota_file.read().split()[:2]
        except (OSError, ValueError):
            try:
                with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as quota_file:
                    quota = quota_file.read().strip()
                with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as period_file:
                    period = period_file.read().strip()
            except (OSError, ValueError):
                return None
        if quota in ("max", "-1") or int(period) <= 0:
            return None
        return max(1, int(int(quota) // int(period)))

    VISIBLE_CORES = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
    USABLE_CORES = min(VISIBLE_CORES, _cpu_quota_cores() or VISIBLE_CORES)
    WORKER_COUNT = max(1, min(int(math.ceil(CAPACITY_BUFFER * SEND_RATE_FACTOR * max(STAGE_QPS) / QPS_PER_WORKER)),
                              USABLE_CORES, MAX_WORKERS))

    def _connections_for(rate, latency_seconds, endpoint_seconds=None):
        """Connections per worker for a stage; once server timings exist, never more than the growth limit."""
        share = rate / float(WORKER_COUNT)
        planned = max(8, int(math.ceil(2 * share * max(latency_seconds, MIN_LATENCY_SECONDS))))
        if endpoint_seconds:
            planned = min(planned, max(8, int(math.ceil(CONNECTION_HEADROOM * share * endpoint_seconds))))
        return min(MAX_CONNECTIONS_PER_WORKER, planned)

    def _available_memory_bytes():
        """Free memory for this process: the smaller of the machine's available memory and any container limit."""
        limits = []
        try:
            with open("/proc/meminfo") as meminfo:
                limits += [int(line.split()[1]) * 1024 for line in meminfo if line.startswith("MemAvailable:")]
        except OSError:
            pass
        for limit_path, usage_path in (("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory.current"),
                                       ("/sys/fs/cgroup/memory/memory.limit_in_bytes",
                                        "/sys/fs/cgroup/memory/memory.usage_in_bytes")):
            try:
                with open(limit_path) as limit_file, open(usage_path) as usage_file:
                    limits.append(int(limit_file.read()) - int(usage_file.read()))
            except (OSError, ValueError):
                pass
        return min(limits) if limits else 0
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Diagnostics
try:
    # Printed with Show diagnostics on; the verdict also uses the CPU figures on every run. Every figure here is optional:
    # anything that can't be read is left out.
    # Container CPU throttling counters: cgroup v2, then v1, with the unit of each one's throttled time in seconds.
    CGROUP_CPU_STATS = (("/sys/fs/cgroup/cpu.stat", "throttled_usec", 1e-6),
                        ("/sys/fs/cgroup/cpu/cpu.stat", "throttled_time", 1e-9))

    def _driver_snapshot():
        """Machine-wide CPU ticks (total, idle, steal) and how many CPUs they cover, container throttling (periods,
        throttled, seconds) and network bytes (received, sent) on the driver; each is None if unreadable."""
        snapshot = {"time": time.monotonic(), "ticks": None, "cpus": None, "throttling": None, "network": None}
        try:
            with open("/proc/stat") as stat_file:
                lines = stat_file.readlines()
            # user nice system idle iowait irq softirq steal
            ticks = [int(value) for value in lines[0].split()[1:9]]
            snapshot["ticks"] = (sum(ticks), ticks[3] + ticks[4], ticks[7])
            snapshot["cpus"] = sum(1 for line in lines if line.startswith("cpu") and line[3:4].isdigit())
        except (OSError, ValueError, IndexError):
            pass
        for path, time_key, seconds_per_unit in CGROUP_CPU_STATS:
            try:
                with open(path) as stat_file:
                    fields = dict(line.split() for line in stat_file if line.strip())
                snapshot["throttling"] = (int(fields["nr_periods"]), int(fields["nr_throttled"]),
                                          int(fields[time_key]) * seconds_per_unit)
                break
            except (OSError, ValueError, KeyError):
                pass
        try:
            received = sent = 0
            with open("/proc/net/dev") as network_file:
                for line in network_file.readlines()[2:]:
                    name, _, fields = line.partition(":")
                    if name.strip() != "lo":
                        received += int(fields.split()[0])
                        sent += int(fields.split()[8])
            snapshot["network"] = (received, sent)
        except (OSError, ValueError, IndexError):
            pass
        return snapshot

    def _other_cpu_share(before, after, parts):
        """The share of the driver's CPU used by anything other than this notebook's workers, or None."""
        driver_busy = _driver_cpu_busy(before, after)
        if driver_busy is None or not after["cpus"]:
            return None
        return driver_busy - sum(_workers_busy(parts)) / float(after["cpus"])

    def _driver_cpu_busy(before, after):
        """The share of the driver's CPU time that was busy between two snapshots, or None."""
        if not (before["ticks"] and after["ticks"]):
            return None
        total, idle, _ = (late - early for late, early in zip(after["ticks"], before["ticks"]))
        return 1 - idle / float(max(1, total))

    def _workers_busy(parts):
        """Each worker's CPU time as a share of the stage's wall time."""
        return [part["cpu_seconds"] / max(1e-9, part["wall_seconds"]) for part in parts]

    def _cpu_summary(before, after, parts):
        """Whether this notebook's workers or the machine ran short of CPU in a stage."""
        sent = sum(part["counts"]["sent"] for part in parts)
        items = ["%.0f CPU us per request" % (1e6 * sum(part["cpu_seconds"] for part in parts) / max(1, sent)),
                 "workers busy " + ", ".join("%.0f%%" % (100.0 * busy) for busy in _workers_busy(parts))]
        driver_busy = _driver_cpu_busy(before, after)
        if driver_busy is not None:
            total, _, steal = (late - early for late, early in zip(after["ticks"], before["ticks"]))
            items.append("driver busy %.0f%%, taken by other VMs (steal) %.0f%%" % (
                100.0 * driver_busy, 100.0 * steal / float(max(1, total))))
        if before["throttling"] and after["throttling"]:
            periods, throttled, seconds = (late - early for late, early in zip(after["throttling"], before["throttling"]))
            items.append("held back by the container's CPU limit %.0f%% of the time (%.1f s)" % (
                100.0 * throttled / float(max(1, periods)), seconds))
        try:
            items.append("load average " + ", ".join("%.1f" % load for load in os.getloadavg()))
        except OSError:
            pass
        return "; ".join(items)

    def _ms_text(value):
        return "-" if value is None else "%.0f" % value

    def _by_kind(counts):
        return ", ".join("%s x%d" % item for item in sorted(counts.items())) or "none"

    def _print_diagnostics(fraction, stage, parts, connections_per_worker, prepare_seconds, cpu_before, cpu_after):
        print_wrapped("{}% stage (target {:,} QPS): {:,.0f} QPS completed, {:,.0f} QPS skipped".format(
            round(fraction * 100), stage["target_qps"], stage["achieved_qps"], stage["skipped_qps"]))
        client_ms, server_ms = ("/".join(_ms_text(value) for value in stage[key]) for key in ("client_ms", "server_ms"))
        print_wrapped("Latency p50/p90/p95/p99: client {} ms, server {} ms".format(client_ms, server_ms), "  ")
        if SERVER_TIME_SPLIT_EVERY:
            embedding_ms, search_ms = _ms_text(stage["embedding_ms"][0]), _ms_text(stage["search_ms"])
            print_wrapped("Server time p50 (from 1 in %d replies): embedding %s ms, search %s ms" % (
                SERVER_TIME_SPLIT_EVERY, embedding_ms, search_ms), "  ")
        # 429s are counted apart from errors_by_kind (the verdict treats them separately), but listing
        # "none" next to 429 samples reads as if nothing failed.
        failures_by_kind = dict(stage["errors_by_kind"])
        if stage["rate_limited"]:
            failures_by_kind["HTTP 429"] = stage["rate_limited"]
        print_wrapped("Failures by kind: %s" % _by_kind(failures_by_kind), "  ")
        limit = ("; stopped opening more at about {:,} open".format(stage["connection_limit"])
                 if stage["connection_limit"] is not None else "")
        failed_opens = sum(stage["connect_errors"].values())
        failed = " ({})".format(_by_kind(stage["connect_errors"])) if failed_opens else ""
        print_wrapped("Connections: {:,} at peak ({:,} planned per worker); {:,} opened; {:,} failed to open{}; {:,} "
                      "closed by the endpoint{}".format(stage["connections"], connections_per_worker,
                                                        stage["connections_opened"], failed_opens, failed,
                                                        stage["server_closed_connections"], limit), "  ")
        print_wrapped("CPU: %s" % _cpu_summary(cpu_before, cpu_after, parts), "  ")
        if cpu_before["network"] and cpu_after["network"]:
            seconds = max(1e-9, cpu_after["time"] - cpu_before["time"])
            received, sent = ((late - early) * 8e-6 / seconds for late, early in zip(cpu_after["network"],
                                                                                      cpu_before["network"]))
            print_wrapped("Driver network: %.0f Mbit/s in, %.0f Mbit/s out" % (received, sent), "  ")
        for sample in stage["error_samples"]:
            print_wrapped("Sample " + sample, "  ")
        print_wrapped("Loop delay by worker (how late each worker woke; above %.0f ms means it was short of CPU): %s ms" % (
            1000.0 * LOOP_BUSY_DELAY_SECONDS, ", ".join("%.1f" % delay for delay in stage["loop_delay_ms_by_worker"])),
                      "  ")
        # The fastest tenth of client minus server time; it grows when requests wait outside the endpoint's own time.
        print_wrapped("Client minus server time (fastest 10%%) about %.1f ms; %s request bodies prepared in %.1f s" % (
            1000.0 * stage["network_seconds"], "{:,}".format(stage["prepared_bodies"]), prepare_seconds), "  ")
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Ramp
try:
    # Above this share of a stage's requests failing, the ramp stops early: the stages after it would only add load to an
    # endpoint failing almost everything. Below it the ramp runs every stage, so each one shows how many queries fail.
    ALMOST_ALL_FAILED_SHARE = 90 / 100.0
    # Above this share of a stage's requests timing out or failing with 5xx errors, the ramp stops too: an endpoint failing
    # that way is struggling rather than shedding load with 429s, and more load can make it worse for its other traffic.
    DISTRESS_FAILED_SHARE = 25 / 100.0
    WORKER_REPLY_MARGIN_SECONDS = 30
    # Warm-up queries are numbered from here, far past any stage's, so they never repeat a measured query.
    WARM_QUERY_START = 10 ** 12
    def _percentiles(values, points=(0.5, 0.9, 0.95, 0.99)):
        values = sorted(values)
        return [values[min(len(values) - 1, int(len(values) * p))] if values else None for p in points]

    def _total(parts, key):
        total = Counter()
        for part in parts:
            total.update(part[key])
        return total

    def _stage_result(qps, parts):
        counts = _total(parts, "counts")
        window = float(STAGE_SECONDS - MEASURE_AFTER_SECONDS)
        loop_delay_ms = [part["loop_delay_ms"] for part in parts]
        return {"target_qps": qps,
                "achieved_qps": sum(part["ok_in_window"] for part in parts) / window,
                "skipped_qps": sum(part["skipped_in_window"] for part in parts) / window,
                "client_ms": _percentiles([value for part in parts for value in part["client_ms"]]),
                "server_ms": _percentiles([value for part in parts for value in part["server_ms"]]),
                "embedding_ms": _percentiles([value for part in parts for value in part["embedding_ms"]]),
                "search_ms": _percentiles([value for part in parts for value in part["search_ms"]], (0.5,))[0],
                "sent": counts["sent"], "rate_limited": counts["rate_limited"], "errors": counts["errors"],
                # Replies that named the embedding endpoint (indexes that compute their embeddings only).
                "embedding_rate_limited": counts["embedding_rate_limited"], "embedding_errors": counts["embedding_errors"],
                # 429s from Databricks' request rate limit rather than the endpoint.
                "request_limited": counts["request_limited"],
                # Replies saying the query ran past the endpoint's time limit, and 429s saying it turned load away.
                "timed_out": counts["timed_out"], "timed_out_429": counts["timed_out_429"],
                "overloaded": counts["overloaded"],
                "capacity_limited": counts["capacity_limited"],
                "no_reply": counts["no_reply"], "errors_by_kind": dict(_total(parts, "errors_by_kind")),
                # Warm-up before the stage included.
                "connections_opened": counts["opened"], "connect_errors": dict(_total(parts, "connect_errors")),
                "server_closed_connections": counts["server_closes"],
                "connections": sum(part["connections"] for part in parts),
                # Workers without a limit add their peak, so the total is what the driver held.
                "connection_limit": (sum(part["connections"] if part["connection_limit"] is None
                                         else part["connection_limit"] for part in parts)
                                     if any(part["connection_limit"] is not None for part in parts) else None),
                "loop_delay_ms_by_worker": loop_delay_ms,
                "loop_delay_high": max(loop_delay_ms) > 1000.0 * LOOP_BUSY_DELAY_SECONDS,
                "workers_busy": sum(_workers_busy(parts)) / len(parts),
                "error_samples": [sample for part in parts for sample in part["error_samples"]][:ERROR_SAMPLES],
                "prepared_bodies": sum(part["prepared"] for part in parts),
                "server_mean_seconds": (sum(part["server_seconds"] for part in parts)
                                        / max(1, sum(part["server_count"] for part in parts))
                                        if any(part["server_count"] for part in parts) else None),
                # A low percentile of client minus server time leaves out waits in this notebook: the network.
                "network_seconds": (_percentiles([value for part in parts for value in part["gap_ms"]], (0.1,))[0] or 0.0)
                                   / 1000.0,
                "mean_request_seconds": sum(part["request_seconds"] for part in parts) / max(1, counts["sent"]),
                "connections_in_use": sum(part["window_request_seconds"] for part in parts) / window}

    # Counts shared by the ramp and the verdict.
    def _failures(stage):
        return stage["rate_limited"] + stage["errors"] + stage["no_reply"]

    def _server_errors(stage):
        return sum(count for error, count in stage["errors_by_kind"].items() if error.startswith("HTTP 5"))

    def _share(stage, count):
        return count / float(max(1, stage["sent"]))

    def _limit(stage):
        # A limit set while nothing was open is no evidence of one.
        return stage["connection_limit"] or None

    def _stop_reason(stage):
        """(skip_reason, skip_cause) when the ramp stops after this stage, else (None, None). 429s alone keep the
        ALMOST_ALL_FAILED_SHARE bar, since they are the endpoint shedding load it can't take; timeouts and 5xx errors stop
        the ramp past DISTRESS_FAILED_SHARE."""
        causes = failures_by_cause(stage)
        if _share(stage, _failures(stage)) > ALMOST_ALL_FAILED_SHARE:
            return "almost_all_failed", main_cause(causes)
        if _share(stage, causes["timeout"] + causes["server_error"]) > DISTRESS_FAILED_SHARE:
            return "endpoint_distress", "timeout" if causes["timeout"] >= causes["server_error"] else "server_error"
        return None, None

    def _distress_reason(stage, lightest):
        """What an endpoint_distress stop says about its cause, as a clause to append."""
        causes = failures_by_cause(stage)
        if lightest and causes["timeout"] >= causes["server_error"]:
            # The lightest load already timed out, so the queries cost too much, not the load.
            return "; queries timed out even at the lightest load, so they are likely too expensive"
        named = stage.get("embedding_errors") or 0
        if EMBEDDING_ENDPOINT and causes["server_error"] > causes["timeout"] and 2 * named > causes["server_error"]:
            return "; most of them named {}, so it, not the AI Search endpoint, was struggling".format(EMBEDDING_ENDPOINT)
        return ", a sign the endpoint was struggling"

    # Stop causes that more load makes worse, so skipping the remaining stages protects the endpoint.
    ENDPOINT_LOAD_CAUSES = ("throttled", "embedding_rate_limited", "capacity_limited", "timeout", "server_error")

    def run_ramp():
        context = multiprocessing.get_context("fork")
        workers = []
        results = []
        with StoppingWorkers(workers):
            with warnings.catch_warnings():
                # Python warns about forking a process that has threads; the workers never use Spark or dbutils.
                warnings.simplefilter("ignore", DeprecationWarning)
                # Frozen objects are skipped by the workers' garbage collector, so a worker doesn't stall scanning the
                # notebook's objects.
                gc.collect()
                gc.freeze()
                try:
                    for index in range(WORKER_COUNT):
                        workers.append(LoadWorker(context, _worker_main, index, WORKER_COUNT, check=cancel_watch.check))
                finally:
                    gc.unfreeze()
            if DEBUG:
                print("%d worker processes; %d usable CPU cores of %d visible." % (
                    WORKER_COUNT, USABLE_CORES, VISIBLE_CORES))
            connections_per_worker = 0
            latency_seconds = PROBE_SECONDS
            network_seconds = None
            # Until a stage measures them, the test query stands in for the endpoint's time and the network.
            probe_server_ms = None if probe is None else _server_time(probe.content)
            endpoint_seconds_hint = PROBE_SECONDS if probe_server_ms is None else probe_server_ms / 1000.0
            probe_network_seconds = PROBE_SECONDS - endpoint_seconds_hint
            # bytearray grows in steps, so allow a little over the average body size.
            body_bytes = 1.15 * sum(len(query_body(index)) for index in range(1, 33)) / 32
            warm_query = WARM_QUERY_START
            first_query = 2  # queries 0 and 1 were the test queries
            for stage_index, (fraction, qps) in enumerate(zip(STAGE_FRACTIONS, STAGE_QPS)):
                progress.at("stage", stage_index)
                cancel_watch.check(force=True)
                send_rate = SEND_RATE_FACTOR * qps
                connection_count = _connections_for(
                    send_rate, latency_seconds, network_seconds is not None and endpoint_seconds_hint + network_seconds)
                memory_budget = PREPARED_MEMORY_SHARE * _available_memory_bytes() - PREPARED_MEMORY_MARGIN_BYTES
                prepare_limit = max(0, int(memory_budget / body_bytes / WORKER_COUNT))
                prepare_started = time.monotonic()
                _run_on_all(workers, "prepare", 600, send_rate=send_rate, first_query=first_query, limit=prepare_limit)
                prepare_seconds = time.monotonic() - prepare_started
                if connection_count > connections_per_worker:
                    # Warming is paced at the stage rate, and each connection's first query may take up to its timeout.
                    _run_on_all(workers, "warm", connection_count * WORKER_COUNT / send_rate + 3 * REQUEST_TIMEOUT_SECONDS,
                                connection_count=connection_count, rate=send_rate, start=time.monotonic() + 0.2,
                                first_query=warm_query, headers=_query_headers())
                    warm_query += connection_count * WORKER_COUNT
                    connections_per_worker = connection_count
                cpu_before = _driver_snapshot()
                parts = _run_on_all(workers, "stage", STAGE_SECONDS + REQUEST_TIMEOUT_SECONDS + WORKER_REPLY_MARGIN_SECONDS,
                                    qps=send_rate, start=time.monotonic() + 0.2, connection_count=connection_count,
                                    first_query=first_query, headers=_query_headers(),
                                    endpoint_seconds_hint=endpoint_seconds_hint,
                                    network_seconds=probe_network_seconds if network_seconds is None else network_seconds)
                cpu_after = _driver_snapshot()
                first_query += int(math.ceil(send_rate * STAGE_SECONDS))
                stage = _stage_result(qps, parts)
                stage["other_cpu_share"] = _other_cpu_share(cpu_before, cpu_after, parts)
                results.append(stage)
                skip_reason, stop_cause = _stop_reason(stage)
                skipped = [skipped_stage_record(index, STAGE_QPS[index], stop_cause, skip_reason)
                           for index in range(len(results), len(STAGE_QPS))] if skip_reason else []
                progress.update(completed_stages=len(results), stages=[stage_record(done) for done in results],
                                skipped_stages=skipped,
                                current_stage_target_qps=STAGE_QPS[len(results)] if len(results) < len(STAGE_QPS)
                                and not skipped else None)
                connections_per_worker = max([connections_per_worker] + [part["connections"] for part in parts])
                # The network round trip is taken once, from the first stage with server timings, before any load
                # could make this notebook queue requests.
                if network_seconds is None and stage["server_mean_seconds"] is not None:
                    network_seconds = max(0.0, stage["network_seconds"])
                if stage["server_mean_seconds"] is not None:
                    endpoint_seconds_hint = stage["server_mean_seconds"]
                    latency_seconds = stage["server_mean_seconds"] + (network_seconds or 0.0)
                elif stage["sent"]:
                    endpoint_seconds_hint = latency_seconds = stage["mean_request_seconds"]
                if DEBUG:
                    _print_diagnostics(fraction, stage, parts, connection_count, prepare_seconds, cpu_before, cpu_after)
                if skipped:
                    causes = failures_by_cause(stage)
                    why = ("{:.0f}% of its requests failed, the largest share {}".format(
                        100.0 * _share(stage, _failures(stage)), CAUSE_LABELS[stop_cause])
                           if skip_reason == "almost_all_failed" else
                           "{:.0f}% of its requests were timeouts or 5xx errors, more of them {}{}".format(
                               100.0 * _share(stage, causes["timeout"] + causes["server_error"]),
                               CAUSE_LABELS[stop_cause], _distress_reason(stage, len(results) == 1)))
                    # The verdict below says what to do next. Requests the request rate limit turned away or that got no
                    # reply never loaded the endpoint.
                    print_wrapped("Stopped after the {}% stage ({:,} QPS): {}, so the remaining {} skipped{}.".format(
                        round(fraction * 100), qps, why, "stage was" if len(skipped) == 1 else
                        "{} stages were".format(len(skipped)),
                        " to protect the endpoint" if stop_cause in ENDPOINT_LOAD_CAUSES else ""))
                    break
        progress.at("results")
        return results
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# MAGIC %md ## Run the ramp

# COMMAND ----------

# DBTITLE 1,Run
try:
    cancel_watch.check(force=True)
    progress.update(status="running", total_stages=len(STAGE_QPS), current_stage_target_qps=STAGE_QPS[0])
    # Starting the workers is setup; run_ramp marks each stage as it starts.
    progress.at("setup")
    results = run_ramp()
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC ## How to read the results
# MAGIC
# MAGIC **Columns**
# MAGIC - **achieved QPS**: successful responses per second, measured from 3 s into each stage. Each stage sends
# MAGIC   5% more than its target, so this can read slightly above it.
# MAGIC - **client** latency is measured in this notebook and includes the network; **server** latency is the
# MAGIC   endpoint's own handling time. Run this notebook in the endpoint's region for client latency close to what
# MAGIC   nearby applications see.
# MAGIC - **429**: requests that were rate limited: by Databricks' request rate limit in front of the endpoint, by the
# MAGIC   endpoint itself, by the endpoint's or index's capacity limit, or, for an index that computes its embeddings,
# MAGIC   by its embedding endpoint. The verdict counts each that the replies show, and says what to change. A query
# MAGIC   that runs past the endpoint's time limit may also get a 429; the verdict calls those timeouts.
# MAGIC - **errors**: error responses, timeouts, replies cut off partway or malformed, and TLS errors.
# MAGIC - **no reply**: requests whose connection closed before any reply, including new connections that failed to
# MAGIC   open.
# MAGIC - For an index that computes its embeddings, each query's text is first embedded by its model serving endpoint,
# MAGIC   and server latency includes that. Rate limits (429s), server errors (5xx) and slowdowns can then come from
# MAGIC   the embedding endpoint, not only the AI Search endpoint. The verdict names it when it may be the cause, and
# MAGIC   says how to get more QPS from it.
# MAGIC
# MAGIC **Verdict**
# MAGIC - The ramp runs every stage unless more than 90% of a stage's requests fail or are rate limited,
# MAGIC   or more than 25% time out or fail with 5xx errors (a sign the endpoint is struggling); then it skips
# MAGIC   the rest. Each stage's failures are listed by cause after the table, and the verdict says what to do next.
# MAGIC - **PASS**: the 100% stage completed at least 99% of your Target QPS, and at most
# MAGIC   1% of its requests failed (429s, errors and no reply all count).
# MAGIC - **REVIEW**: the 100% stage fell short because this notebook couldn't send the full rate, so this run can't
# MAGIC   judge the endpoint. The verdict says why and how to send more.
# MAGIC - **FAIL**: the ramp stopped before the 100% stage, or more than 1% of the 100% stage's
# MAGIC   requests failed, or the 100% stage fell short while this notebook sent the full rate.
# MAGIC - For each kind of failure, the verdict gives a count and what to do about it.
# MAGIC
# MAGIC **What this test shows, and what it doesn't**
# MAGIC - By default each stage is a short check of about 24 s. For a baseline, hold the 100% load for 5
# MAGIC   to 10 minutes: set the **Stages as shares of Target QPS** widget to 0.25,1.0 (keeping a light stage to compare
# MAGIC   with) and **Seconds per stage** to 300.
# MAGIC - Queries are approximate nearest neighbor (ANN) searches (text, for an index that computes its embeddings),
# MAGIC   or hybrid searches when **Query type** is HYBRID, with no filters or reranker, and they return only the
# MAGIC   primary key. Production queries with those, or with more columns, cost more. ANN tests don't show the cost
# MAGIC   of hybrid search, which uses about 4x the resources of ANN; the endpoint's Target QPS doesn't allow for
# MAGIC   that, so hybrid traffic may need a higher one.
# MAGIC - To get this QPS from an application, use a service principal's OAuth token and create the index object once
# MAGIC   and reuse it. Personal access tokens get a much lower per-workspace limit.
# MAGIC - Run it when no other index on this endpoint is syncing, since syncing competes with queries for the
# MAGIC   endpoint's resources.
# MAGIC - For an index with self-managed embeddings this test sends ready-made vectors, so it doesn't cover any
# MAGIC   embedding your application does before querying (for example, calling an external embedding service). Load
# MAGIC   test that service separately at the same QPS, and cache vectors for repeated query text where you can.
# MAGIC - More: [High QPS](https://docs.databricks.com/aws/en/ai-search/high-qps),
# MAGIC   [best practices](https://docs.databricks.com/aws/en/ai-search/best-practices) and
# MAGIC   [endpoint load testing](https://docs.databricks.com/aws/en/ai-search/endpoint-load-test).
# MAGIC
# MAGIC **Testing more than one notebook can send**
# MAGIC - How much one notebook can send depends on the driver's CPU cores, the query size and the endpoint's
# MAGIC   latency; the verdict says when this notebook was the limit.
# MAGIC - For more load, run several copies of this notebook at the same time on separate clusters (copies on one
# MAGIC   cluster share its CPU), each with Target QPS set to a share of the total.

# COMMAND ----------

# DBTITLE 1,Results table
try:
    def _ms(values):
        return "-" if values[0] is None else "/".join(_ms_text(value) for value in values)

    columns = ["stage", "stage target QPS", "achieved QPS", "client p50/p90/p95/p99 ms", "server p50/p90/p95/p99 ms", "429"]
    rows = [["%d%%" % round(fraction * 100), "{:,}".format(stage["target_qps"]), "{:,.0f}".format(stage["achieved_qps"]),
             _ms(stage["client_ms"]), _ms(stage["server_ms"]), str(stage["rate_limited"])]
            for fraction, stage in zip(STAGE_FRACTIONS, results)]
    for column, key in (("errors", "errors"), ("no reply", "no_reply")):
        if any(stage[key] for stage in results):
            columns.append(column)
            for row, stage in zip(rows, results):
                row.append(str(stage[key]))
    widths = [max(len(cell) for cell in column) for column in zip(columns, *rows)]
    for row in [columns] + rows:
        print("  ".join(cell.ljust(width) for cell, width in zip(row, widths)))

    # The latest connection limit, unless a later stage held more connections than it.
    connection_limit = None
    for stage in results:
        if _limit(stage) is not None:
            connection_limit = _limit(stage)
        elif connection_limit is not None and stage["connections"] > connection_limit:
            connection_limit = None
    if connection_limit is not None:
        print_wrapped("New connections to the endpoint kept failing to open once about {:,} were open, so this test kept to "
                      "that many.{}".format(connection_limit, "" if CONSOLE_RUN else
                                            " If the limit is per cluster, copies on separate clusters each get their own."))
    SKIP_REASON_TEXTS = {
        "almost_all_failed": "almost every request of the stage before it failed",
        "endpoint_distress": ("more than 25% of the requests of the stage before it timed out or failed "
                              "with 5xx errors")}
    for skipped_stage in progress.run["skipped_stages"]:
        print_wrapped("{}% stage (target {:,} QPS): skipped, since {}".format(
            round(STAGE_FRACTIONS[skipped_stage["stage_index"]] * 100), skipped_stage["target_qps"],
            SKIP_REASON_TEXTS[skipped_stage["skip_reason"]]))

    preflight = progress.run["preflight"]
    print_wrapped("Signed in as %s." % auth.identity)
    if preflight["status"] == "degraded":
        print_wrapped("The test query kept failing before the ramp (%s)." % CAUSE_LABELS[preflight["cause"]])
    # Failures by cause per stage, as the run record carries them; the verdict below says what to do about them.
    for fraction, stage in zip(STAGE_FRACTIONS, results):
        causes = {cause: count for cause, count in failures_by_cause(stage).items() if count}
        if causes:
            print_wrapped("{}% stage failures: {}".format(round(fraction * 100), ", ".join(
                "{} {:,}".format(CAUSE_LABELS[cause], count) for cause, count in causes.items())))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Verdict rules
try:
    # PASS needs the 100% stage to complete at least this share of Target QPS, with no more than MAX_FAILED_SHARE of
    # its requests failed; the extra 5% sent can't hide failures that way.
    PASS_FRACTION = 99 / 100.0
    MAX_FAILED_SHARE = 1 / 100.0
    # Names used here and not defined below come from the notebook's earlier cells (loadTestProgress.py and
    # loadTestStageCounts.py among them), which also set PASS_FRACTION and MAX_FAILED_SHARE just above.
    REQUIRED_QPS = PASS_FRACTION * TARGET_QPS
    # "Most" of a stage's failures, 429s or shortfall: more than this share of them, so a tie is not most.
    MOST_SHARE = 0.5
    # Skipped queries above this share of a failing stage's target rate are worth a sentence of their own.
    SKIPPED_NOTE_SHARE = 0.1
    # More copies than this would mean this notebook sent about a tenth of Target QPS or less, too little to size
    # copies from.
    MAX_SUGGESTED_COPIES = 10
    # Workers busier than this on average were short of CPU themselves; less busy, any delay came from other work.
    BUSY_WORKERS_SHARE = 0.8
    # Other processes using this share of the driver's CPU (beyond the workers) leave the workers waiting for it.
    OTHER_WORK_SHARE = 0.3
    # Error replies are only put down to load when they are at least this share of the 100% stage's requests, and at
    # least LOAD_ERROR_GROWTH times their share in the earlier stages; fewer say little about load.
    MIN_LOAD_ERROR_SHARE = 0.001
    LOAD_ERROR_GROWTH = 2
    # Embedding p50 growth below this share of server p50's growth rules the embedding endpoint out of a slowdown.
    EMBEDDING_GROWTH_NOTE_SHARE = 0.25
    # Embedding-named replies below this share of a stage's failures are too few to explain its timeouts.
    EMBEDDING_FAILURE_SHARE = 0.1
    # Connections count as all busy when the measured time per request kept at least this share of them in use.
    BUSY_CONNECTIONS_SHARE = 0.8
    # End-to-end time this many times the endpoint's time plus the network round trip means requests waited elsewhere.
    OUTSIDE_TIME_FACTOR = 2
    # A stage whose achieved QPS is at most this many times the one before's has stopped rising: the endpoint's limit,
    # rather than a test still ramping up its own sending.
    PLATEAU_GROWTH = 1.1
    # Server p50 at or above this at the lightest stage means queries are slow on their own, so more QPS can't help; well
    # below the endpoint's 30-second timeout, but far slower than a healthy index answers.
    SLOW_SERVER_P50_MS = 1000
    # Below this server p50 the verdict calls the endpoint fast, so a shortfall points away from it.
    FAST_SERVER_P50_MS = 100
    # Databricks' load-test targets (docs.databricks.com/aws/en/ai-search/endpoint-load-test), shown beside the verdict.
    TARGET_P95_MS = 500
    TARGET_P99_MS = 1000
    TARGET_SUCCESS_SHARE = 0.995
    OUTGROWN_CAPACITY = "if the index has grown a lot since it was created, contact Databricks support"
    # A console run has no notebook to scroll back through, and can't run copies of itself.
    CLIENT = "this test" if CONSOLE_RUN else "this notebook"
    DRIVER = "this test's compute" if CONSOLE_RUN else "the driver"
    # The gateway skips the embedding call for a query that carries its own vector (CommonQueryObject.needsEmbedding), so
    # the same index serves it.
    QUERY_VECTOR_ADVICE = ("skip embedding at query time: compute query embeddings in your application with the model "
                           "this index uses and send them as query_vector instead of query_text "
                           "(https://docs.databricks.com/api/ai-search/v1/query-index)")
    # Hybrid search still needs the text for its keyword search, so the vector goes beside it.
    if HYBRID_QUERIES:
        QUERY_VECTOR_ADVICE = ("skip embedding at query time: compute query embeddings in your application with the model "
                               "this index uses and send them as query_vector alongside query_text "
                               "(https://docs.databricks.com/api/ai-search/v1/query-index)")
    PROVISIONED_THROUGHPUT_DOCS = ("https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/"
                                   "deploy-prov-throughput-foundation-model-apis")
    HIGH_QPS_DOCS = "https://docs.databricks.com/aws/en/ai-search/high-qps"
    # The section the console's link for the hybrid_resources finding opens too.
    BEST_PRACTICES_ANN_HYBRID_DOCS = ("https://docs.databricks.com/aws/en/ai-search/best-practices"
                                      "#use-ann-for-efficiency-and-use-hybrid-when-necessary")
    # Codes the console shows its own text and docs link for (verdict finding_codes, one per finding, else None); keep in
    # step with the console's list.
    FINDING_CODES = ("hybrid_resources",)
    SERVING_OPTIMIZATION_DOCS = "https://docs.databricks.com/aws/en/machine-learning/model-serving/production-optimization"
    # A Target QPS scaled from a measured peak is rounded up to a multiple of this: the peak is approximate.
    TARGET_QPS_STEP = 50
    # A stage that completed less than this share of its own target fell short, which shows the endpoint hit a limit; when
    # none did, the stages only show what they were asked for.
    SHORT_STAGE_SHARE = 0.95
    # Names the generated notebook, since the console can't run copies itself.
    RUN_COPIES = "run {} copies of the generated load-test notebook at the same time on separate clusters, each with {}"
    # What a console run can change instead: more cores run more workers, each with its own connections.
    MORE_CORES = "use an existing cluster with more CPU cores (Advanced options > Compute)"
    ROUTE_UNAVAILABLE = ("that route isn't available for this index in this workspace right now; contact Databricks support "
                         "for more details")
    # Warning codes (RUN_WARNINGS) the verdict's own sentences already explain, so the run record leaves them out.
    covered_warnings = set()
    # Each recommendation code and the params its console message takes (LOAD_TEST_RECOMMENDATION_PARAMS in
    # loadTestRecommendations.ts, which a jest test keeps in step).
    RECOMMENDATION_PARAMS = {
        "set_target_qps": ("target_qps", "test_target_qps", "results_per_query", "results_cost_factor", "sized_results"),
        "raise_target_qps": ("endpoint_target_qps",),
        "return_fewer_results": (),
        "contact_support": (),
        "contact_support_if_unhelpful": (),
        "query_vector": (),
        "embedding_throughput": ("embedding_endpoint_name",),
        "embedding_provider_quota": ("embedding_endpoint_name",),
        "embedding_dedicated_capacity": ("embedding_endpoint_name",),
        "embedding_dedicated_endpoint": ("embedding_endpoint_name",),
        "embedding_scale_up": ("embedding_endpoint_name",),
        "route_contact_support": (),
        "use_service_principal": (),
        "request_limit_contact_support": (),
        "delete_data": (),
        "cheaper_queries": (),
        "outgrown_contact_support": (),
        "rerun_contact_support": (),
        "rerun_with_diagnostics": (),
        "check_permissions": ("computes_embeddings",),
        "check_network": (),
        "run_copies": ("copies", "copy_target_qps"),
        "more_cores": ("test_target_qps",),
        "rerun_after_fixing": (),
        "rerun_when_not_syncing": (),
    }
    # What each sentence that says what to do is made of: findings (text) and recommendations (_recommendation dicts), so
    # the console can show the recommendations apart. A sentence with no entry is one finding.
    sentence_parts = {}

    def _recommendation(code, said, text=None, **params):
        """said is the recommendation as its sentence words it; text, when given, says it on its own, for a list of next
        steps where the words around it (e.g. "it", "instead") would dangle."""
        return {"code": code, "params": params, "said": said, "text": text or said}

    def _said(sentence, *parts):
        """The sentence, noting the findings and recommendations it is made of."""
        sentence_parts[sentence] = list(parts)
        return sentence

    def _advice(code, said, text=None, **params):
        """A sentence that is one recommendation."""
        return _said(said, _recommendation(code, said, text, **params))

    def _or(recommendation, *alternatives):
        """The recommendation with what the reader can do instead of it, which is listed with it, not as another step."""
        for alternative in alternatives:
            alternative["alternative"] = True
        return dict(recommendation, alternatives=list(alternatives))

    ROUTE_UNAVAILABLE_FINDING, _route_step = ROUTE_UNAVAILABLE.split("; ")
    _said(ROUTE_UNAVAILABLE, ROUTE_UNAVAILABLE_FINDING, _recommendation(
        "route_contact_support", _route_step, "contact Databricks support about the network-optimized route, which isn't "
        "available for this index in this workspace right now"))
    # A run moved to the workspace route (QUERY_ROUTE_FALLBACK) took it only because of how it signed in.
    ROUTE_FALLBACK = ("{} took that route only because its sign-in can't be used to embed query text on the endpoint's "
                      "own URL; run it as a service principal to test the endpoint's full throughput").format(CLIENT)
    _route_fallback_finding, _route_fallback_step = ROUTE_FALLBACK.split("; ")
    _said(ROUTE_FALLBACK, _route_fallback_finding, _recommendation(
        "use_service_principal", _route_fallback_step, "run the test as a service principal to test the endpoint's full "
        "throughput on the endpoint's own URL"))
    ROUTE_ADVICE = ROUTE_FALLBACK if QUERY_ROUTE_FALLBACK else ROUTE_UNAVAILABLE
    # After a fallback the index can use the endpoint's own URL; only this test's queries left it.
    SLOWER_ROUTE = ("{}'s queries went through the workspace instead of the endpoint's own URL".format(CLIENT)
                    if QUERY_ROUTE_FALLBACK else "queries to this index don't use the network-optimized route")

    def _counted(count, singular, plural=None):
        """A count with its noun: 1 error, 2 errors."""
        return "{:,} {}".format(count, singular if count == 1 else (plural or singular + "s"))

    def _most(count, total):
        """Whether count is most of total: strictly more than MOST_SHARE of it, so a tie is not most."""
        return bool(total) and count > MOST_SHARE * total

    def _timeout_504s(stage):
        """Timeout replies counted among the 5xx errors rather than the 429s."""
        return stage["timed_out"] - stage["timed_out_429"]

    def _timeouts(stage):
        """Same count as failures_by_cause's timeout cause."""
        return stage_timeouts(stage)

    def _embedding_candidate_errors(stage):
        """5xx replies other than timeout replies, which the timeout sentences explain."""
        return max(0, _server_errors(stage) - _timeout_504s(stage))

    def _unnamed_server_errors(stage):
        """5xx replies that neither named the embedding endpoint nor said the query timed out."""
        return max(0, _embedding_candidate_errors(stage) - stage["embedding_errors"])

    def _denied(stage):
        return sum(stage["errors_by_kind"].get(code, 0) for code in ("HTTP 401", "HTTP 403"))

    def _endpoint_429s(stage):
        """429s from the AI Search endpoint itself: overload replies and any that no marker explains."""
        return max(0, stage["rate_limited"] - stage["request_limited"] - stage["embedding_rate_limited"]
                   - stage["capacity_limited"] - stage["timed_out_429"])

    def _from_embedding_endpoint(stage, key):
        """Whether most of a stage's 429s ("rate_limited") or 5xx ("errors") replies named the embedding endpoint."""
        count = stage["rate_limited"] if key == "rate_limited" else _embedding_candidate_errors(stage)
        named = stage["embedding_rate_limited"] if key == "rate_limited" else stage["embedding_errors"]
        return bool(EMBEDDING_ENDPOINT) and _most(named, count)

    def _embedding_named(stage):
        """Whether enough of a stage's failures named the embedding endpoint, and most of its 429s or 5xx did, that the
        embedding call may be what its timed-out queries waited for."""
        named = stage["embedding_rate_limited"] + stage["embedding_errors"]
        return bool(EMBEDDING_ENDPOINT) and named >= EMBEDDING_FAILURE_SHARE * _failures(stage) and (
            _from_embedding_endpoint(stage, "rate_limited") or _from_embedding_endpoint(stage, "errors"))

    def _timeouts_started_under_load(stage, results):
        """Whether a stage's timeouts grew well past the lightest stage's share, so they came with the load rather than
        with what each query costs."""
        return stage is not results[0] and (
            _share(stage, _timeouts(stage)) > LOAD_ERROR_GROWTH * _share(results[0], _timeouts(results[0])))

    def _failure_sources(stage, results):
        """Each kind of failure in the stage with its count, largest first; every failed request is in exactly one."""
        timeout_key = "timeout_under_load" if _timeouts_started_under_load(stage, results) else "timeout_query_cost"
        sources = [("request_limit", stage["request_limited"]),
                   ("embedding", stage["embedding_rate_limited"] + stage["embedding_errors"]),
                   ("capacity_limit", stage["capacity_limited"]),
                   ("endpoint_429", _endpoint_429s(stage)),
                   (timeout_key, _timeouts(stage)),
                   ("server_error", _unnamed_server_errors(stage)),
                   ("denied", _denied(stage)),
                   ("no_reply", stage["no_reply"])]
        sources.append(("other", max(0, _failures(stage) - sum(count for _, count in sources))))
        return sorted(((key, count) for key, count in sources if count), key=lambda source: -source[1])

    def _endpoint_owned(stage, results):
        """Failures the AI Search endpoint itself caused: its own 429s, and timeouts that grew with the load unless the
        embedding endpoint was failing too."""
        owned = _endpoint_429s(stage)
        if _timeouts_started_under_load(stage, results) and not _embedding_named(stage):
            owned += _timeouts(stage)
        return owned

    def _endpoint_turned_load_away(stage, results):
        """Whether the endpoint's own failures were most of the stage's, so its capacity is what fell short. Most is
        strictly more than half, so no other source can be as large, and a tie gets no Target QPS advice."""
        return _most(_endpoint_owned(stage, results), _failures(stage))

    SOURCE_SHORT_NAMES = {
        "request_limit": ("request turned away by the request rate limit", "requests turned away by the request rate limit"),
        "embedding": ("failure that named the embedding endpoint", "failures that named the embedding endpoint"),
        "capacity_limit": ("request refused at the storage capacity limit", "requests refused at the storage capacity limit"),
        "endpoint_429": ("429 from the endpoint", "429s from the endpoint"),
        "timeout_under_load": ("timeout", None), "timeout_query_cost": ("timeout", None),
        "server_error": ("5xx error", None), "denied": ("401 or 403 reply", "401 or 403 replies"),
        "no_reply": ("request with no reply", "requests with no reply"), "other": ("other failure", None)}

    def _embedding_step(finding, step, query_vector):
        """The finding, if any, then the step, with sending query vectors as the alternative to it when offered. Either
        one alone gets the embedding endpoint out of the way, so they are one step, not two; the alternative gets its own
        sentence, since the step can end in links."""
        findings = [finding] if finding else []
        said = "".join(text + ", so " for text in findings) + step["said"]
        if not query_vector:
            return [_said(said, *(findings + [step]))]
        vector = _recommendation("query_vector", "or, " + QUERY_VECTOR_ADVICE, QUERY_VECTOR_ADVICE)
        return [_said(said, *(findings + [_or(step, vector)])), _said(vector["said"], vector)]

    def _embedding_advice(full_advice, query_vector=True):
        """How to get more QPS from the embedding endpoint: in full when it is the largest source of trouble, else a
        short pointer. query_vector=False leaves out QUERY_VECTOR_ADVICE where capacity advice may follow."""
        if EXTERNAL_EMBEDDING_MODEL:
            # Provisioned throughput and scale to zero don't apply to an external provider.
            if not full_advice:
                return [_advice("embedding_provider_quota",
                                "to get more QPS from it, raise your rate limit or quota with that provider",
                                "to get more QPS from {}, raise your rate limit or quota with the provider of {}".format(
                                    EMBEDDING_ENDPOINT, EXTERNAL_EMBEDDING_MODEL),
                                embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME)]
            quota = ("to get more QPS from the embedding side, raise your rate limit or quota with {}, and the serving "
                     "endpoint's AI Gateway rate limits if any are set")
            return _embedding_step(
                "{} serves the external {}".format(EMBEDDING_ENDPOINT, EXTERNAL_EMBEDDING_MODEL),
                _recommendation("embedding_provider_quota", quota.format("that provider"),
                                quota.format("the provider of " + EXTERNAL_EMBEDDING_MODEL),
                                embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME), query_vector)
        if EMBEDDING_ENDPOINT_KIND == "pay_per_token":
            # Nothing on a shared endpoint can be raised, and updating an index changes only its payload logging, not its
            # embedding endpoint.
            if not full_advice:
                # Its own code, so the console doesn't show the full step, with its new index, for a possible cause.
                return [_advice("embedding_dedicated_endpoint",
                                "to get more QPS from it, serve its model on a dedicated endpoint, such as a provisioned "
                                "throughput endpoint",
                                "to get more QPS from {}, a pay-per-token endpoint, serve its model on a dedicated "
                                "endpoint, such as a provisioned throughput endpoint".format(EMBEDDING_ENDPOINT),
                                embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME)]
            def dedicated(model):
                return ("to get more QPS from the embedding side, serve {} on capacity of its own (a provisioned throughput "
                        "endpoint for the same model, or another dedicated Model Serving endpoint) and create a new index "
                        "that uses it, since an index's embedding endpoint is set when it is created; see {} and {}".format(
                            model, PROVISIONED_THROUGHPUT_DOCS, HIGH_QPS_DOCS))
            # The listed step names the endpoint itself, since the finding before it isn't there.
            named = ("the model of the pay-per-token endpoint " + EMBEDDING_ENDPOINT_NAME if EMBEDDING_ENDPOINT_NAME
                     else "the index's embedding model")
            return _embedding_step(
                "{} is a preconfigured pay-per-token Foundation Model APIs endpoint".format(EMBEDDING_ENDPOINT),
                _recommendation("embedding_dedicated_capacity", dedicated("the embedding model"), dedicated(named),
                                embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME), query_vector)
        if EMBEDDING_ENDPOINT_KIND == "dedicated":
            if not full_advice:
                return [_advice("embedding_scale_up",
                                "to get more QPS from it, raise its provisioned throughput or concurrency",
                                "to get more QPS from {}, raise its provisioned throughput or concurrency".format(
                                    EMBEDDING_ENDPOINT), embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME)]
            return _embedding_step(None, _recommendation(
                "embedding_scale_up",
                "to get more QPS from the embedding side, scale up {}: raise its provisioned throughput or provisioned "
                "concurrency, turn off scale to zero so it stays warm, and check its rate limits (see {})".format(
                    EMBEDDING_ENDPOINT, SERVING_OPTIMIZATION_DOCS),
                embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME), query_vector)
        if not full_advice:
            return [_advice("embedding_throughput", "to get more QPS from it, increase its throughput or check its rate limits",
                            "to get more QPS from {}, increase its throughput or check its rate limits".format(
                                EMBEDDING_ENDPOINT),
                            embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME)]
        return _embedding_step(None, _recommendation(
            "embedding_throughput",
            "to get more QPS from the embedding side, increase the throughput of {}, for example by serving the embedding "
            "model with provisioned throughput and turning off scale to zero so it stays warm, and check its rate "
            "limits".format(EMBEDDING_ENDPOINT),
            embedding_endpoint_name=EMBEDDING_ENDPOINT_NAME), query_vector)

    def _diagnostics_step():
        """Rerun advice for failures whose replies say more than the counts; a console run has no Show diagnostics."""
        if CONSOLE_RUN:
            return _advice("rerun_contact_support", "rerun, and if it repeats, contact Databricks support")
        return _advice("rerun_with_diagnostics", "rerun with Show diagnostics set to true to print sample replies, and if it "
                       "repeats, contact Databricks support")

    def _source_sentences(key, count, stage, results, most, full_embedding_advice, capacity_advised):
        """The sentences for one source: what its failures were, with the count, then this source's own next step. The
        endpoint's own 429s and timeouts get theirs once, after every source (the Target QPS advice or a rerun)."""
        if key == "request_limit" and QUERIES_VIA_WORKSPACE:
            covered_warnings.add("queries_via_workspace")
            label, detail = ("429s from the workspace's request limit for AI Search",
                             " (the replies say \"Too many requests\"); these requests never reached the endpoint")
            steps = [SLOWER_ROUTE + ", and the route they take allows far "
                     "fewer requests per second", ROUTE_ADVICE]
        elif key == "request_limit":
            label, detail = ("429s from Databricks' request rate limit in front of the AI Search endpoint",
                             " (the replies say \"Too many requests\"); these requests never reached the endpoint")
            steps = [_advice("request_limit_contact_support",
                             "contact Databricks support about this workspace's request limit for AI Search")]
        elif key == "embedding":
            parts = [_counted(number, noun) for number, noun in ((stage["embedding_rate_limited"], "429"),
                                                                 (stage["embedding_errors"], "5xx error")) if number]
            label = "replies that named {}".format(EMBEDDING_ENDPOINT)
            detail = (" ({})".format(_listed(parts)) if len(parts) > 1 else "") + ", not the AI Search endpoint"
            # query_vector advice says more endpoint capacity can't help, which would contradict capacity advice.
            steps = _embedding_advice(full_embedding_advice, query_vector=not capacity_advised)
        elif key == "capacity_limit":
            label, detail = ("429s saying the endpoint's or index's capacity limit was reached",
                             ", which is about the data stored, not the request rate")
            steps = [_advice("delete_data", "the replies say to delete some data, or contact Databricks support")]
        elif key == "endpoint_429":
            overloaded = min(stage["overloaded"], count)
            unclear = count - overloaded
            if not unclear:
                label, detail = ("429s in which the endpoint turned requests away under load",
                                 " (the replies say heavy load or too many requests)")
            elif not overloaded:
                label = "429s whose replies didn't say which limit refused them"
                if QUERIES_VIA_WORKSPACE:
                    covered_warnings.add("queries_via_workspace")
                    detail = ("; " + SLOWER_ROUTE + ", so the workspace's request "
                              "limit or the endpoint refused them, and if it was the workspace's limit, a higher Target QPS "
                              "won't help")
                elif EMBEDDING_ENDPOINT and not capacity_advised:
                    detail = "; they may have come from the AI Search endpoint or {}".format(EMBEDDING_ENDPOINT)
                else:
                    detail = "; they are counted as the endpoint's"
            else:
                label = "429s from the endpoint"
                detail = " ({} said heavy load or too many requests, and {:,} didn't say which limit refused them)".format(
                    _counted(overloaded, "reply", "replies"), unclear)
            steps = []
        elif key == "timeout_under_load":
            label = "timeouts, far fewer of which happened at the lightest stage"
            detail = ("; under this load {} couldn't keep up, so queries waited past the endpoint's time limit (30 seconds "
                      "on Standard endpoints)".format(
                          "the endpoint or its embedding endpoint"
                          if _embedding_named(stage) or _embedding_time_most(stage, results) else "the endpoint"))
            steps = []
        elif key == "timeout_query_cost":
            label = "timeouts, already common at the lightest stage"
            detail = ("; queries took longer than the endpoint allows (30 seconds on Standard endpoints), so each query "
                      "costs too much whatever the rate")
            steps = [_advice("cheaper_queries",
                             "to make queries cheaper, return fewer results (Results per query) and fewer columns"),
                     _slow_at_start_sentence(results) if _slow_at_start(results) else
                     _advice("contact_support_if_unhelpful", "if that doesn't help, contact Databricks support",
                             "if fewer results and columns don't help, contact Databricks support")]
        elif key == "server_error":
            # 504s that all said the query timed out are among the timeouts.
            codes = " or ".join(sorted(code for code, number in stage["errors_by_kind"].items()
                                       if code.startswith("HTTP 5")
                                       and not (code == "HTTP 504" and number <= _timeout_504s(stage)))) or "5xx"
            earlier = results[:_stage_index(stage, results)]
            earlier_share = (sum(_unnamed_server_errors(other) for other in earlier)
                             / float(max(1, sum(other["sent"] for other in earlier))))
            # For an index that computes its embeddings, a 5xx that names no endpoint may be the embedding call's.
            source = ("the endpoint, its embedding endpoint or something in front of them" if EMBEDDING_ENDPOINT else
                      "the endpoint or something in front of it")
            label = "{} replies".format(codes)
            # At the lightest stage there is no lighter load to compare with.
            detail = ("; they came from {}, which started turning requests away at this load".format(source)
                      if earlier and _share(stage, count) > LOAD_ERROR_GROWTH * earlier_share else
                      "; they came from {}".format(source))
            steps = [_diagnostics_step()]
        elif key == "denied":
            label, detail = "HTTP 401 and 403 replies", ", which mean requests weren't allowed"
            steps = [_advice("check_permissions",
                             "check that the identity this test signed in as can query this index{}, then rerun".format(
                                 " and its embedding endpoint" if EMBEDDING_ENDPOINT else ""),
                             computes_embeddings=bool(EMBEDDING_ENDPOINT))]
        elif key == "no_reply":
            connect_failed = stage["errors_by_kind"].get("connect failed", 0)
            label, detail = "requests with no reply", ""
            if connect_failed:
                detail = " ({} that failed to open{})".format(
                    _counted(connect_failed, "new connection"),
                    "" if _limit(stage) is None else " once about {:,} were open".format(_limit(stage)))
            connection_limit = "that can be a limit on connections from one cluster rather than on the endpoint"
            if _limit(stage) is not None and CONSOLE_RUN:
                repeats = "if it repeats, contact Databricks support"
                steps = [_said(connection_limit + "; " + repeats, connection_limit,
                               _recommendation("rerun_contact_support", repeats,
                                               "rerun, and if it repeats, contact Databricks support"))]
            elif _limit(stage) is not None:
                copies = RUN_COPIES.format("several", "a share of the Target QPS")
                steps = [_said(connection_limit + "; if it is, " + copies, connection_limit,
                               _recommendation("run_copies", copies, "if it is a connection limit, " + copies, copies=None,
                                               copy_target_qps=None))]
            else:
                network = ("connections that close or fail to open point at the network between this cluster and the "
                           "endpoint, such as a proxy or firewall")
                check = "check that this cluster can reach the endpoint directly, rerun, and if it repeats, contact Databricks support"
                steps = [_said(network + "; " + check, network, _recommendation("check_network", check))]
        else:
            known = ("timeout", "connect failed", "connection closed", "HTTP 401", "HTTP 403", "HTTP 408")
            kinds = sorted(code for code in stage["errors_by_kind"] if code not in known and not code.startswith("HTTP 5"))
            label, detail = "other failures", " ({})".format(", ".join(kinds)) if kinds else ""
            steps = [_diagnostics_step()]
        first = ("most failures were {}: {:,} of {:,}{}" if most else "{}: {:,}{}").format(
            *((label, count, _failures(stage), detail) if most else (label, count, detail)))
        return first, steps

    def _significant_sources(stage, results):
        """The sources with at least MIN_LOAD_ERROR_SHARE of the stage's requests, largest first."""
        return [(key, count) for key, count in _failure_sources(stage, results)
                if _share(stage, count) >= MIN_LOAD_ERROR_SHARE]

    ENDPOINT_SOURCES = ("endpoint_429", "timeout_under_load")

    def _failure_sentences(stage, results, capacity_advised, lead):
        """(breakdown, sentences): one sentence group per significant source. With lead, a source that is most of the
        failures says so; otherwise, when there are several, breakdown lists their counts for the summary sentence."""
        sources = _significant_sources(stage, results)
        if not sources:
            return None, []
        most_key = sources[0][0] if lead and _most(sources[0][1], _failures(stage)) else None
        breakdown = (_listed([_counted(count, *SOURCE_SHORT_NAMES[key]) for key, count in sources])
                     if lead and not most_key and len(sources) > 1 else None)
        # The embedding endpoint gets its full advice unless the AI Search endpoint's own failures were most of them.
        full_embedding_advice = sources[0][0] == "embedding" or not _endpoint_turned_load_away(stage, results)
        sentences, said = [], set()
        for key, count in sources:
            first, steps = _source_sentences(key, count, stage, results, key == most_key, full_embedding_advice,
                                             capacity_advised)
            if steps and all(step in said for step in steps):
                # Another source already gave this next step.
                sentences.append(_said(first + "; rerun as above", first))
            elif len(steps) == 1 and key != "embedding":
                # A one-sentence next step reads with its cause; longer advice keeps its own lines.
                sentences.append(_said(first + "; " + steps[0], first, *sentence_parts.get(steps[0], [steps[0]])))
            else:
                sentences += [first] + steps
            said.update(steps)
        others = [key for key, _ in sources if key not in ENDPOINT_SOURCES]
        if not capacity_advised and others and any(key in ENDPOINT_SOURCES for key, _ in sources):
            # One rerun note for the endpoint's own failures, naming the largest ones to fix first.
            names = [SOURCE_SHORT_NAMES[key][1] or SOURCE_SHORT_NAMES[key][0] + "s" for key in others]
            sentences.append(_advice("rerun_after_fixing", "once the {} are fixed, rerun to see whether the endpoint keeps up".format(
                _listed(names) if len(names) <= 2 else "other failures, starting with the {} and {},".format(*names[:2]))))
        return breakdown, sentences

    def _clean_qps(results):
        """The most QPS a stage completed with at most MAX_FAILED_SHARE of its requests failing."""
        clean = [stage["achieved_qps"] for stage in results if _share(stage, _failures(stage)) <= MAX_FAILED_SHARE]
        return max(clean) if clean else None

    def _stage_label(stage, results):
        return "%d%%" % round(STAGE_FRACTIONS[_stage_index(stage, results)] * 100)

    def _measured_peak_qps(results):
        """The endpoint's measured limit at its current Target QPS: the most QPS any stage completed, when some stage fell
        short of its own target (SHORT_STAGE_SHARE), some stage's requests were turned away by the endpoint itself (its
        429s, or timeouts that grew with the load), and that peak is below TARGET_QPS. None when there is no such limit."""
        if not any(_share(stage, _endpoint_owned(stage, results)) >= MIN_LOAD_ERROR_SHARE for stage in results):
            return None
        if not any(stage["achieved_qps"] < SHORT_STAGE_SHARE * stage["target_qps"] for stage in results):
            return None
        # Every stage counts, short or not: the endpoint served at least the best of them. achieved_qps counts successful
        # replies only, so failed requests don't inflate it.
        peak = max(stage["achieved_qps"] for stage in results)
        return peak if 0 < peak < TARGET_QPS else None

    def _recommended_target_qps(results):
        """(Target QPS for the endpoint to serve TARGET_QPS, measured peak or None). With a measured limit, the endpoint's
        Target QPS scaled by how far the peak fell short, rounded up to TARGET_QPS_STEP; otherwise an estimate from the
        results per query (RESULTS_FACTOR)."""
        peak = _measured_peak_qps(results)
        if peak is None:
            return int(math.ceil(TARGET_QPS * RESULTS_FACTOR)), None
        # Rounded first so float noise in an exact multiple doesn't add a step.
        steps = math.ceil(round(ENDPOINT_TARGET_QPS * TARGET_QPS / float(peak) / TARGET_QPS_STEP, 6))
        return int(steps) * TARGET_QPS_STEP, peak

    def _results_cost_text(peak):
        """What the result count adds to the advice: the estimated cost only when no measured peak replaces it."""
        # Without a measured peak the figure is just this test's QPS; no hybrid factor is known to scale it by.
        if HYBRID_QUERIES and peak is None:
            return (" of hybrid queries (this test's QPS as it is, since Target QPS sizing doesn't account for hybrid "
                    "search)")
        if RESULTS_FACTOR == 1.0:
            return ""
        if peak is not None:
            return " with {:,} results per query".format(NUM_RESULTS)
        return (" with {:,} results per query (an estimated {:.1f} times the cost of the {} results Target QPS is sized "
                "for)".format(NUM_RESULTS, RESULTS_FACTOR, SIZED_RESULTS))

    def _endpoint_target_qps_text():
        """The endpoint's Target QPS as the verdict quotes it; an unreadable endpoint leaves only the one at creation."""
        return "{:,}{}".format(ENDPOINT_TARGET_QPS, " (its value when this notebook was created)" if ENDPOINT_UNREADABLE
                               else "")

    def _more_capacity_sentences(stage, results):
        """What Target QPS the endpoint needs, when it fell behind at this stage (turned load away, slowed down or fell
        short) because it lacked capacity."""
        needed, peak = _recommended_target_qps(results)
        peak_found = ([] if peak is None else ["the endpoint completed at most about {:,.0f} QPS at its Target QPS of {}".format(
            peak, _endpoint_target_qps_text())])
        set_target_step = ("to serve {:,} QPS{}, set the endpoint's Target QPS to about {:,} and rerun once its scaling "
                           "change is applied".format(TARGET_QPS, _results_cost_text(peak), needed))
        estimated = RESULTS_FACTOR != 1.0 and peak is None
        set_target_recommendation = _recommendation(
            "set_target_qps", set_target_step, target_qps=needed, test_target_qps=TARGET_QPS,
            results_per_query=None if RESULTS_FACTOR == 1.0 else NUM_RESULTS,
            results_cost_factor=round(RESULTS_FACTOR, 1) if estimated else None,
            sized_results=SIZED_RESULTS if estimated else None)

        def set_target(*alternatives):
            """The Target QPS sentence, then its alternatives (recommendations) joined with "or"."""
            return _said("".join(text + ", so " for text in peak_found) + set_target_step
                         + "".join(", or " + alternative["said"] for alternative in alternatives),
                         *(peak_found + [_or(set_target_recommendation, *alternatives)] + list(alternatives)))
        if stage["target_qps"] > ENDPOINT_TARGET_QPS:
            served = _clean_qps(results)
            return ["the {} stage targeted {:,} QPS, above the endpoint's Target QPS of {}, which is the load its "
                    "capacity is provisioned for; the endpoint doesn't autoscale, so falling behind above it is "
                    "expected{}".format(
                        _stage_label(stage, results), stage["target_qps"], _endpoint_target_qps_text(),
                        "" if served is None else
                        ", and it completed about {:,.0f} QPS with almost no failures".format(served)),
                    set_target()]
        if RESULTS_FACTOR > 1.0 and stage["target_qps"] * RESULTS_FACTOR > ENDPOINT_TARGET_QPS:
            return ["the endpoint doesn't autoscale, and its Target QPS of {} provisions capacity for queries returning "
                    "{} results".format(_endpoint_target_qps_text(), SIZED_RESULTS),
                    set_target(_recommendation("return_fewer_results", "return fewer results",
                                               "return fewer results per query"))]
        first = ("the endpoint doesn't autoscale, and it fell behind at or below its own Target QPS of {}; the docs say "
                 "throughput scaling is best-effort and depends on index size, vector dimension and query "
                 "complexity".format(_endpoint_target_qps_text()))
        if peak is not None:
            return [first, set_target(_recommendation("contact_support", "contact Databricks support"))]
        raise_step = "raise the endpoint's Target QPS above {} and rerun once its scaling change is applied".format(
            _endpoint_target_qps_text())
        contact = _recommendation("contact_support", "contact Databricks support")
        return [first, _said(raise_step + ", or contact Databricks support",
                             _or(_recommendation("raise_target_qps", raise_step, endpoint_target_qps=ENDPOINT_TARGET_QPS),
                                 contact), contact)]

    def _slow_at_start(results):
        """Whether the lightest stage was already slow on the endpoint's own time."""
        base_server_p50 = results[0]["server_ms"][0] if results else None
        return base_server_p50 is not None and base_server_p50 >= SLOW_SERVER_P50_MS

    def _slow_at_start_sentence(results):
        slow = ("the endpoint was already slow at the lightest stage (server p50 {:,.0f} ms at {:,} QPS), so a higher "
                "Target QPS won't make each query faster".format(results[0]["server_ms"][0], results[0]["target_qps"]))
        return _said(slow + "; " + OUTGROWN_CAPACITY, slow, _recommendation("outgrown_contact_support", OUTGROWN_CAPACITY))

    def _fallback_next_step(stage):
        """What to do when no source of failures was large enough to say."""
        if not stage or _share(stage, _failures(stage)) <= MAX_FAILED_SHARE:
            # No failures to inspect: a slowdown or shortfall that may come from other work on the endpoint.
            return _advice("rerun_when_not_syncing", "rerun when no other index on this endpoint is syncing, since "
                           "syncing competes with queries; if it repeats, contact Databricks support")
        return _diagnostics_step()

    def _stage_index(stage, results):
        return next(i for i, other in enumerate(results) if other is stage)

    def _listed(items):
        return (", ".join(items[:-1]) + " and " if len(items) > 1 else "") + items[-1]

    def _request_limit_significant(stage, results):
        """Whether request-limit 429s were a significant source: a per-workspace limit, which copies of this notebook
        share."""
        return any(key == "request_limit" for key, _ in _significant_sources(stage, results))

    def _embedding_growth_share(stage, results):
        """Embedding p50's growth from the lightest stage as a share of server p50's growth, or None when either stage
        has no embedding sample or server p50 didn't grow."""
        embedding_p50, server_p50 = (stage.get("embedding_ms") or [None])[0], stage["server_ms"][0]
        base_embedding_p50, base_server_p50 = (results[0].get("embedding_ms") or [None])[0], results[0]["server_ms"][0]
        if not EMBEDDING_ENDPOINT or None in (embedding_p50, server_p50, base_embedding_p50, base_server_p50):
            return None
        server_growth = server_p50 - base_server_p50
        return (embedding_p50 - base_embedding_p50) / float(server_growth) if server_growth > 0 else None

    def _embedding_time_most(stage, results):
        """Whether embedding each query took most of the endpoint's time (server p50) at this stage, and its growth from
        the lightest stage is most of the server time's growth. A share that stays the same under load is what each query
        costs, not what held the endpoint back, so it doesn't count."""
        growth_share = _embedding_growth_share(stage, results)
        return growth_share is not None and growth_share > MOST_SHARE and _most(stage["embedding_ms"][0],
                                                                                stage["server_ms"][0])

    def _held_connections(stage):
        """About how many connections the stage kept open: all it opened, or as many as the connection limit let it."""
        return stage["connections"] if _limit(stage) is None else min(stage["connections"], _limit(stage))

    def _connections_busy(stage):
        """Whether the measured time per request kept most of the stage's connections in use."""
        return stage["connections_in_use"] >= BUSY_CONNECTIONS_SHARE * _held_connections(stage)

    def _waited_outside_endpoint(stage, results):
        """Whether the stage's mean time per request was more than OUTSIDE_TIME_FACTOR times the endpoint's mean time plus
        the lightest stage's network round trip, so requests waited somewhere else; None when the endpoint reported no
        time. On the workspace route that round trip already includes the route's own delay, so only the endpoint's time
        is compared there."""
        if stage["server_mean_seconds"] is None:
            return None
        network_seconds = 0 if QUERIES_VIA_WORKSPACE else results[0]["network_seconds"]
        return stage["mean_request_seconds"] > OUTSIDE_TIME_FACTOR * (stage["server_mean_seconds"] + network_seconds)

    def _endpoint_held_connections(stage, results):
        """Whether the stage's connections were busy because the endpoint got slower: most were in use, there were enough
        of them to send the required QPS at the lightest stage's time per request, server p50 more than doubled from the
        lightest stage to at least FAST_SERVER_P50_MS while achieved QPS stopped rising, the time per request was at most
        twice the endpoint's time plus the network round trip, and the driver's CPU didn't explain it. The queries skipped then follow from the slowdown.
        The workspace route keeps its own sentences."""
        index = _stage_index(stage, results)
        base, base_server_p50, server_p50 = results[0], results[0]["server_ms"][0], stage["server_ms"][0]
        if (QUERIES_VIA_WORKSPACE or index == 0 or stage["server_mean_seconds"] is None or not base_server_p50
                or not server_p50 or not base["mean_request_seconds"]):
            return False
        # Too few connections (a connection limit, say) would have fallen short at the lightest stage's speed too.
        enough_connections = _held_connections(stage) >= REQUIRED_QPS * base["mean_request_seconds"]
        # Still climbing, the shortfall may be this test's sending rather than the endpoint's limit.
        plateaued = stage["achieved_qps"] <= PLATEAU_GROWTH * results[index - 1]["achieved_qps"]
        # The same CPU test as the REVIEW's, so a run this lets through doesn't get the CPU advice there either.
        cpu_limited = stage["workers_busy"] >= BUSY_WORKERS_SHARE and WORKER_COUNT == USABLE_CORES
        # A fast endpoint can double its time and still be fast, so it isn't what held the connections.
        return (_connections_busy(stage) and enough_connections and plateaued
                and server_p50 > 2 * base_server_p50 and server_p50 >= FAST_SERVER_P50_MS
                and not _waited_outside_endpoint(stage, results)
                and not cpu_limited)

    def _workspace_route_sentences(stage, route_advised):
        """Why a run on the route that isn't network-optimized topped out with most connections busy: the route added time
        to each query, so the connections were all waiting on it. With route_advised, the request-limit sentences already
        say the route isn't available."""
        covered_warnings.add("queries_via_workspace")
        client_p50, server_p50 = stage["client_ms"][0], stage["server_ms"][0]
        embedding_p50 = (stage.get("embedding_ms") or [None])[0]
        if client_p50 is None or server_p50 is None or client_p50 <= server_p50:
            added = "which added time to each query"
        else:
            added = "which added about {:,.0f} ms per query on top of the endpoint's own {:,.0f} ms (server p50{})".format(
                client_p50 - server_p50, server_p50,
                "" if embedding_p50 is None else ", including {:,.0f} ms of embedding".format(embedding_p50))
        sentences = [SLOWER_ROUTE + ", {}, so the connections were all waiting "
                     "on it and {} couldn't send the full rate".format(added, CLIENT)]
        if not route_advised:
            sentences.append(ROUTE_ADVICE)
        return sentences

    def _target_check(name, value, target, met):
        return "%s %s, target %s: %s" % (name, value, target, "met" if met else "missed")
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Verdict
try:
    full = results[FULL_STAGE] if len(results) > FULL_STAGE else None
    # Judged on the 100% stage, or on the stage the ramp stopped at.
    judged = full or (results[-1] if results else None)
    sentences = []
    # Set when a REVIEW already blames CPU, so the CPU sentence below doesn't repeat it.
    cpu_explained = False
    slowed_down = False
    # Set when the endpoint's slowdown kept the connections busy, so the queries skipped didn't make it a REVIEW.
    endpoint_held = False
    # Set when the 100% stage fell short with few enough failures to pass, and not for lack of sending.
    fell_short = False
    # Set once a sentence says what to do next, so a FAIL never ends without one.
    next_step = False
    # Set when the route that isn't network-optimized held back what this notebook could send.
    route_limited = False
    if full:
        other_cpu_share = full.get("other_cpu_share")
        cpu_short = full["workers_busy"] >= BUSY_WORKERS_SHARE
        cpu_contended = (full["loop_delay_high"] and not cpu_short and other_cpu_share is not None
                         and other_cpu_share >= OTHER_WORK_SHARE)
    if full is None:
        kind = "FAIL"
        if results and _stop_reason(results[-1])[0] == "endpoint_distress":
            sentences.append("the ramp stopped before the 100% stage because more than {:.0f}% of a stage's requests timed "
                             "out or failed with 5xx errors{}".format(100 * DISTRESS_FAILED_SHARE,
                                                                     _distress_reason(results[-1], len(results) == 1)))
        else:
            sentences.append("the ramp stopped before the 100% stage because almost every request of a stage was failing")
    elif _share(full, _failures(full)) > MAX_FAILED_SHARE:
        kind = "FAIL"
    elif full["achieved_qps"] >= REQUIRED_QPS:
        kind = "PASS"
    else:
        base_server_p50, full_server_p50 = results[0]["server_ms"][0], full["server_ms"][0]
        endpoint_held = _endpoint_held_connections(full, results)
        # Skipped queries come first: a slower endpoint is no evidence when this notebook didn't send the full rate, unless
        # that slowdown is what kept its connections busy.
        if full["skipped_qps"] >= MOST_SHARE * (REQUIRED_QPS - full["achieved_qps"]) and not endpoint_held:
            # Most of the shortfall is queries this notebook skipped, so it can't be pinned on the endpoint.
            kind = "REVIEW"
            sustained = max([stage["achieved_qps"] for stage in results[:FULL_STAGE + 1] if stage["skipped_qps"]]
                            + [full["achieved_qps"]])
            # Sized by what this notebook sent without its extra 5%, since each copy adds its own.
            copies = max(2, int(math.ceil(TARGET_QPS / max(1.0, sustained / SEND_RATE_FACTOR))))
            held = _held_connections(full)
            request_ms = 1000.0 * full["mean_request_seconds"]
            server_ms = 1000.0 * full["server_mean_seconds"] if full["server_mean_seconds"] is not None else None
            # Copies share the workspace's request limit, so its 429s rule out copies; the request-limit sentences below
            # also say the route isn't available.
            request_limited = _request_limit_significant(full, results)
            detail = []
            if cpu_short and WORKER_COUNT == USABLE_CORES:
                cpu_explained = True
                reason = ", likely limited by the {} CPU {} of {}".format(
                    USABLE_CORES, "core" if USABLE_CORES == 1 else "cores", DRIVER)
            elif _connections_busy(full):
                reason = ", and most of its about {:,} connections were busy".format(held)
                waited_outside = bool(_waited_outside_endpoint(full, results))
                route_limited = QUERIES_VIA_WORKSPACE and waited_outside
                if route_limited:
                    detail = _workspace_route_sentences(full, request_limited)
                else:
                    detail.append("each query took about {:.0f} ms end to end{}".format(
                        request_ms, "" if server_ms is None else ", though the endpoint reported {:.0f} ms".format(
                            server_ms)))
                    if waited_outside:
                        detail.append("most of that time was outside the endpoint's reported handling time; this run "
                                      "can't tell where requests waited")
                    elif QUERIES_VIA_WORKSPACE and server_ms is None:
                        detail.append("the endpoint reported no handling time, so this run can't tell whether the route "
                                      "queries to this index take or the endpoint held it back")
                    elif QUERIES_VIA_WORKSPACE and server_ms >= MOST_SHARE * request_ms:
                        detail.append("most of it was the endpoint's own handling time{}, not time on the way to it".format(
                            ", including embedding each query" if EMBEDDING_ENDPOINT else ""))
            elif cpu_contended:
                cpu_explained = True
                reason = ", likely because other work on {} was using its CPU".format(DRIVER)
            else:
                reason = ""
            if not detail and full_server_p50 is not None:
                reason += ("; the endpoint stayed fast (server p50 %.0f ms)" if full_server_p50 < FAST_SERVER_P50_MS
                           else "; server p50 %.0f ms") % full_server_p50
            # No advice on sending more when the route or the request limit is the limit: more load hits it too. Copies
            # on separate clusters each bring their own connections and CPU; a console run can only use more cores.
            if request_limited or route_limited:
                sentences += ["{} completed about {:,} QPS{}".format(CLIENT, int(sustained), reason)] + detail
            elif copies > MAX_SUGGESTED_COPIES:
                sentences += ["{} completed only about {:,} QPS{}".format(CLIENT, int(sustained), reason)] + detail + [
                              "at about a tenth of the {:,} QPS target or less, this run can't judge the endpoint".format(
                                  TARGET_QPS),
                              _advice("more_cores", "instead, " + MORE_CORES, MORE_CORES + " instead", test_target_qps=None)
                              if CONSOLE_RUN else
                              _advice("run_copies", "instead, " + RUN_COPIES.format("several", "a share of the Target QPS"),
                                      RUN_COPIES.format("several", "a share of the Target QPS") + " instead", copies=None,
                                      copy_target_qps=None)]
            else:
                sentences += ["{} completed about {:,} QPS{}".format(CLIENT, int(sustained), reason)] + detail + [
                              _advice("more_cores", "to test {:,} QPS, ".format(TARGET_QPS) + MORE_CORES,
                                      test_target_qps=TARGET_QPS) if CONSOLE_RUN else
                              _advice("run_copies", "to test {:,} QPS, ".format(TARGET_QPS) + RUN_COPIES.format(
                                  copies, "Target QPS {:,}".format(int(math.ceil(TARGET_QPS / float(copies))))),
                                      copies=copies, copy_target_qps=int(math.ceil(TARGET_QPS / float(copies))))]
        elif base_server_p50 and full_server_p50 and full_server_p50 > 2 * base_server_p50:
            kind = "FAIL"
            slowed_down = True
            sentences.append("the endpoint slowed down: server p50 {:.0f} ms at {:,} QPS vs {:.0f} ms at {:,} QPS".format(
                full_server_p50, full["target_qps"], base_server_p50, results[0]["target_qps"]))
        else:
            kind = "FAIL"
            fell_short = True
            sentences.append("the 100% stage completed {:,.0f} QPS, below the {:,.0f} needed ({}% of {:,})".format(
                full["achieved_qps"], REQUIRED_QPS, round(PASS_FRACTION * 100), TARGET_QPS))
    # What failed, by source, each with its own next step. Capacity advice only when the endpoint itself fell behind: a
    # slowdown or shortfall without failures gets it only where the endpoint wasn't sized for the load, since otherwise
    # the run can't say capacity was the cause.
    endpoint_turned_load_away = bool(judged) and _endpoint_turned_load_away(judged, results)
    failing = kind == "FAIL" and bool(judged) and _share(judged, _failures(judged)) > MAX_FAILED_SHARE
    judged_sources = _significant_sources(judged, results) if judged else []
    top_source = judged_sources[0][0] if judged_sources else None
    # The embedding endpoint held the test back: its replies were the largest source of a failing stage's failures (unless
    # the endpoint's own were most), or it took most of the endpoint's time while nothing but timeouts failed much. More endpoint
    # capacity can't help then.
    embedding_bottleneck = kind == "FAIL" and bool(EMBEDDING_ENDPOINT) and bool(judged) and (
        (failing and top_source == "embedding" and not endpoint_turned_load_away)
        or (_embedding_time_most(judged, results)
            and (not failing or top_source in ("timeout_under_load", "timeout_query_cost"))))
    # The route that isn't network-optimized held the test back: its delay kept the connections waiting, or the workspace's
    # request limit on it was the largest source of a failing stage's failures (unless the endpoint's own were most). A
    # higher endpoint Target QPS can't help then. A few such 429s beside a slowdown or shortfall don't explain it.
    route_bottleneck = route_limited or (failing and QUERIES_VIA_WORKSPACE and top_source == "request_limit"
                                         and not endpoint_turned_load_away)
    if route_bottleneck:
        covered_warnings.add("above_endpoint_target_qps")
    capacity_advised = kind == "FAIL" and bool(judged) and not embedding_bottleneck and not route_bottleneck and (
        endpoint_turned_load_away
        or ((slowed_down or fell_short) and judged["target_qps"] * RESULTS_FACTOR > ENDPOINT_TARGET_QPS))
    # A few failures beside a REVIEW or a slowdown are listed without a lead, since they aren't why the stage fell short.
    breakdown, failure_sentences = (_failure_sentences(judged, results, capacity_advised, failing)
                                    if judged and kind != "PASS" else (None, []))
    if failing:
        summary = "{:.1f}% of requests at the {} stage failed".format(100.0 * _share(judged, _failures(judged)),
                                                                     _stage_label(judged, results))
        parts = [text for number, text in ((judged["rate_limited"], "{:,} with HTTP 429".format(judged["rate_limited"])),
                                           (judged["errors"], _counted(judged["errors"], "error")),
                                           (judged["no_reply"], "{:,} with no reply".format(judged["no_reply"])))
                 if number]
        sentences.append(summary + (": " + breakdown if breakdown else " ({})".format(", ".join(parts))))
    sentences += failure_sentences
    # Failures beside a slowdown or shortfall don't explain it, so that still needs its own next step.
    next_step = bool(failure_sentences) and failing
    if (failing or endpoint_held) and judged["skipped_qps"] >= SKIPPED_NOTE_SHARE * judged["target_qps"]:
        sentences.append("{} also skipped about {:,.0f} QPS it couldn't send on time, so the endpoint got less than this "
                         "stage's target load".format(CLIENT, judged["skipped_qps"]))
    if full:
        # Measured by how late the workers woke from short sleeps: ready work was waiting for CPU.
        if full["loop_delay_high"] and cpu_short and not cpu_explained:
            sentences.append("the CPU of {} was saturated during the 100% stage, so client latency includes time requests "
                             "waited for it".format(DRIVER))
        elif cpu_contended and not cpu_explained:
            sentences.append("other work on {} was using its CPU during the 100% stage, so client latency includes time "
                             "requests waited for it".format(DRIVER))
    # The endpoint embeds each query's text before searching, so a slowdown may be the embedding endpoint's, unless the
    # samples show embedding stayed about flat while server time grew.
    judged_embedding_growth = _embedding_growth_share(judged, results) if judged else None
    if embedding_bottleneck and _embedding_time_most(judged, results):
        next_step = True
        sentences.append("embedding each query with {} took most of the endpoint's time: {:,.0f} ms of its {:,.0f} ms "
                         "(server p50) at the {} stage".format(EMBEDDING_ENDPOINT, judged["embedding_ms"][0],
                                                               judged["server_ms"][0], _stage_label(judged, results)))
        sentences += _embedding_advice(True)
    elif embedding_bottleneck:
        next_step = True
    elif EMBEDDING_ENDPOINT and slowed_down and not (
            judged_embedding_growth is not None and judged_embedding_growth < EMBEDDING_GROWTH_NOTE_SHARE):
        next_step = True
        sentences.append("this index embeds every query with {}, so the slowdown may have come from that endpoint rather "
                         "than the AI Search endpoint".format(EMBEDDING_ENDPOINT))
        # Embedding is only a possible cause, so a short pointer, without the query_vector advice, which is for when
        # embedding held the run back.
        sentences += _embedding_advice(False)
    # The query-cost timeout sentences already say this when they appear.
    if (kind != "PASS" and _slow_at_start(results)
            and not any(key == "timeout_query_cost" for key, _ in (_significant_sources(judged, results) if judged else []))):
        sentences.append(_slow_at_start_sentence(results))
    # The code of each coded finding, by the finding as it is listed (capitalized, with its full stop).
    coded_findings = {}
    # Hybrid's extra cost (about 4x in Target QPS benchmarks) varies by query shape, so it's a finding, not part of a
    # suggested Target QPS.
    if (HYBRID_QUERIES and kind == "FAIL" and bool(judged) and not embedding_bottleneck and not route_bottleneck
            and (endpoint_turned_load_away or slowed_down or fell_short)):
        hybrid_resources = ("the endpoint's Target QPS is sized for ANN queries, and hybrid search uses about 4x the "
                            "resources of ANN (see the best-practices docs on ANN and hybrid: {})".format(
                                BEST_PRACTICES_ANN_HYBRID_DOCS))
        sentences.append(hybrid_resources)
        coded_findings[_capitalized(hybrid_resources) + "."] = "hybrid_resources"
    if capacity_advised:
        next_step = True
        sentences += _more_capacity_sentences(judged, results)
        # Only the above-target branch says the stage was above the endpoint's Target QPS; the others say it fell behind
        # at or below it, which doesn't explain the warning.
        if judged["target_qps"] > ENDPOINT_TARGET_QPS:
            covered_warnings.add("above_endpoint_target_qps")
    if kind == "FAIL" and not next_step:
        sentences.append(_fallback_next_step(judged))
    # Sources can share a next step, such as the rerun-and-contact-support one; it is said once.
    sentences = list(dict.fromkeys(sentences))

    def _listed_recommendation(part):
        return {"code": part["code"], "params": part["params"], "text": _capitalized(part["text"]) + ".",
                "alternatives": [_listed_recommendation(alternative) for alternative in part.get("alternatives", [])]}

    def _same_step(part, other):
        return (part["code"], part["params"]) == (other["code"], other["params"])

    def _listed_once(parts):
        """The recommendations among parts, each (code, params) once. Sources can share a next step worded apart, such as
        rerunning and then contacting support; the one that says the most is kept, where the first was said, with the
        alternatives of both. Alternatives go with the step they are an alternative to."""
        listed = []
        for part in parts:
            if not isinstance(part, dict) or part.get("alternative"):
                continue
            same = [index for index, other in enumerate(listed) if _same_step(part, other)]
            if not same:
                listed.append(part)
                continue
            kept = listed[same[0]]
            alternatives = list(kept.get("alternatives", []))
            alternatives += [alternative for alternative in part.get("alternatives", [])
                             if not any(_same_step(alternative, other) for other in alternatives)]
            listed[same[0]] = dict(part if len(part["text"]) > len(kept["text"]) else kept, alternatives=alternatives)
        return [_listed_recommendation(part) for part in listed]

    # The same sentences split into what happened and what to do, so the console can show the recommendations apart.
    verdict_parts = [part for sentence in sentences for part in sentence_parts.get(sentence, [sentence])]
    findings = list(dict.fromkeys(_capitalized(part) + "." for part in verdict_parts if not isinstance(part, dict)))
    recommendations = _listed_once(verdict_parts)

    # The verdict as printed below; the console's run record (loadTestProgress.py) keeps the same values. sentences is the
    # whole verdict, for consoles that read neither findings nor recommendations.
    verdict_result = {"kind": kind, "sentences": [_capitalized(sentence) + "." for sentence in sentences],
                      "targets": {"p95_ms": None, "p99_ms": None, "success_pct": None},
                      "findings": findings, "recommendations": recommendations,
                      "finding_codes": [coded_findings.get(finding) for finding in findings]}
    if judged:
        p95_ms, p99_ms = judged["client_ms"][2], judged["client_ms"][3]
        success_share = 1.0 - _share(judged, _failures(judged))
        verdict_result["targets"] = {"p95_ms": p95_ms, "p99_ms": p99_ms, "success_pct": 100.0 * success_share}

    # One sentence per line, so a long verdict stays readable; what to do comes last, numbered.
    print("VERDICT:", verdict_result["kind"])
    for sentence in verdict_result["findings"]:
        print_wrapped(sentence, "  ")
    if verdict_result["recommendations"]:
        print("Next steps:")
        for number, recommendation in enumerate(verdict_result["recommendations"], 1):
            print_wrapped("{}. {}".format(number, " Or: ".join(
                [recommendation["text"]] + [alternative["text"] for alternative in recommendation["alternatives"]])), "  ")

    if judged:
        # Skipped queries were never sent, so they are this notebook's limit, not failures.
        skipped_note = (" Queries this notebook skipped aren't sent, so the success rate doesn't count them."
                        if judged["skipped_qps"] else "")
        print_wrapped("Databricks' load-test targets at the %d%% stage: %s.%s" % (
            round(STAGE_FRACTIONS[FULL_STAGE if judged is full else len(results) - 1] * 100), "; ".join([
                _target_check("client p95 (successful requests)",
                              "-" if p95_ms is None else "{:,.0f} ms".format(p95_ms),
                              "under {:,} ms".format(TARGET_P95_MS), p95_ms is not None and p95_ms < TARGET_P95_MS),
                _target_check("client p99 (successful requests)",
                              "-" if p99_ms is None else "{:,.0f} ms".format(p99_ms),
                              "under {:,} ms".format(TARGET_P99_MS), p99_ms is not None and p99_ms < TARGET_P99_MS),
                _target_check("success", "%.2f%%" % (100.0 * success_share), "above %.1f%%" % (100.0 * TARGET_SUCCESS_SHARE),
                              success_share > TARGET_SUCCESS_SHARE)]), skipped_note))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise

# COMMAND ----------

# DBTITLE 1,Console run result
try:
    # The console shows these beside the verdict, so the ones its sentences already explain or rule out are left out.
    progress.update(warnings=[code for code in RUN_WARNING_CODES if code not in covered_warnings],
                    **succeeded_fields(verdict_result, len(results), len(STAGE_QPS)))
    if RUN_KEY:
        dbutils.notebook.exit(json.dumps(progress.run))
except BaseException as _cell_error:
    _record_cell_error(_cell_error)
    raise
