Skip to content

fix(plugins): memoise BigQuery analytics readiness across re-initialisation - #7020

Open
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-7017-bqaa-readiness-reruns
Open

fix(plugins): memoise BigQuery analytics readiness across re-initialisation#7020
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-7017-bqaa-readiness-reruns

Conversation

@chelsealong

Copy link
Copy Markdown
Contributor

Summary

Closes #7017.

BigQueryAgentAnalyticsPlugin re-ran its full table/view readiness pass on every re-initialisation instead of once per process. Each pass issues one CREATE OR REPLACE VIEW statement per _EVENT_VIEW_DEFS entry (25 at HEAD), and it is awaited from before_run_callback, so the DDL sits on the agent request path.

Any deployment that shares one plugin instance across short-lived Runners (Runner.close()PluginManager.close()plugin.close() clears _started) re-runs all the view DDL per request. The reporter measured ~19k view-DDL jobs/hour, per-table BigQuery quota exhausted in ~2h, and ~25s added to median latency — a 2.3.0 → 2.8.0 regression.

In 2.3.0 the readiness pass was gated behind the schema cache, so it ran at most once per process. At HEAD that gate was deliberately removed so a failed first attempt would still retry — but the comment's load-bearing assumption ("once _started is True the steady state pays no extra RPC") only holds when the plugin is initialised once per process, which is false whenever it is closed and reused.

Fix

  • Memoise readiness success in a new _schema_ready flag that, like _schema, survives close()/shutdown(), and gate _ensure_schema_exists on it in _lazy_setup.
    • A failed attempt raises before the flag is set, so it is still retried on the next setup (preserves the 2.8.0 intent).
    • A successful attempt is memoised, so it is not repeated per re-initialisation (restores the 2.3.0 cost profile). No new steady-state RPC.
  • Secondary: emit exactly one WARNING when a generation mismatch aborts an otherwise-successful setup, so a plugin churning through full setups is no longer completely silent (the issue observed 12 log lines against ~2000 setup runs).

Scope is limited to the readiness memoisation and the one diagnostic log line. The view SQL, _EVENT_VIEW_DEFS, and the _ensure_started backoff are untouched (the issue's non-goals).

Testing Plan

Added TestReadinessMemoisedAcrossReinit with two tests:

  • test_readiness_pass_runs_once_across_reinit_cycles — 3 _ensure_started() / shutdown() cycles over one shared plugin issue the view DDL once total (client.query called len(_EVENT_VIEW_DEFS) times), not once per cycle.
  • test_failed_readiness_is_retried_and_not_memoised — a first readiness attempt that raises leaves _schema_ready False and is retried on the next setup, which then succeeds and memoises.

Evidence the tests prove the bug: with the source change stashed (test kept), both new tests fail (AttributeError: ... _schema_ready / would-be 75 vs 25 DDL statements); with the fix restored they pass.

# without the fix (source change stashed):
2 failed in 24.20s

# with the fix:
tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py  505 passed, 6 skipped in 482.90s

Lint/format on the changed files: pyink (2 files unchanged), isort (exit 0), ruff check (all checks passed).

AI assistance disclosure

This change was prepared with AI assistance (Claude). A human reviewed the diff, the tests, and the failing-without-fix evidence before submission.

…sation

BigQueryAgentAnalyticsPlugin re-ran its full table/view readiness pass on
every re-initialisation instead of once per process. Each pass issues one
CREATE OR REPLACE VIEW statement per _EVENT_VIEW_DEFS entry, awaited from
before_run_callback, so the DDL sat on the request path. A host that builds
a short-lived Runner per request over one shared plugin closes the plugin
after every request (Runner.close() -> PluginManager.close() ->
plugin.close()), which clears _started, so the next request re-ran all the
view DDL. In production this exhausted the per-table BigQuery quota and
added ~25s to median latency (2.3.0 -> 2.8.0 regression).

Remember that readiness succeeded in a flag (_schema_ready) that, like
_schema, survives close()/shutdown(), and gate the readiness pass on it. A
*failed* attempt raises before the flag is set and is still retried on the
next setup (the 2.8.0 intent); a *successful* one is not repeated
(restoring the 2.3.0 cost profile). Also log one WARNING when a generation
mismatch aborts an otherwise-successful setup, so a plugin churning through
full setups is no longer silent.

Closes google#7017
@lupuletic

Copy link
Copy Markdown
Contributor

Reporter of #7017 here. Thanks for picking this up so quickly — the fix matches what we landed as a local mitigation, and putting _schema_ready next to _schema reads better than where we had it.

One gap worth closing before merge: __getstate__ does not reset _schema_ready, so the memo crosses the pickle boundary.

__setstate__ has the setdefault backfill, so legacy pickles are safe. But __getstate__ resets every other piece of runtime truth — _started, _init_pid, client, _executor, parser, _loop_state_by_loop — and _schema_ready is the one assertion in that group that is about a remote resource existing rather than about local data. Verified against this branch (4196a06):

__getstate__ carries _schema_ready = True
after pickle round-trip, _schema_ready = True
_started reset to                     = False
_init_pid reset to                    = 0

The plugin has a custom __getstate__ precisely because it does get pickled and restored into another runtime. On a first deploy the destination dataset/table does not exist yet, so a restored plugin with _schema_ready = True skips _ensure_schema_exists entirely and starts appending rows to a missing table. That is the exact failure the comment being replaced was guarding against ("mark the plugin started against a missing/unready table"), reintroduced through the pickle path rather than the retry path.

_schema can safely survive pickling because it is just column definitions. Readiness cannot, because it is a claim about the destination.

Suggested one-liner, alongside the existing resets:

     state["_started"] = False
+    state["_schema_ready"] = False
     state["_startup_error"] = None

Cost is one readiness pass per process, which is the 2.3.0 profile this PR is restoring anyway — it does not weaken the fix, since the whole problem was re-verifying many times within one process.

A test, if useful — it slots onto the existing TestForkSafety.test_getstate_resets_pid, which already asserts the _started reset:

  def test_getstate_resets_pid(self):
    """Pickle state should have _init_pid = 0 to force re-init."""
    plugin = self._make_plugin()
    plugin._schema_ready = True
    state = plugin.__getstate__()
    assert state["_init_pid"] == 0
    assert state["_started"] is False
    # Readiness is a claim about the destination dataset, so a pickle
    # restored elsewhere must verify it again rather than assume it.
    assert state["_schema_ready"] is False

Two smaller notes, neither blocking:

  • The new WARNING is a real improvement — that silent path is why we saw 12 log lines against roughly 2,000 setup runs. Worth noting it fires on a benign race that is most likely under exactly the host shape this PR fixes, so it should get quieter after the memo lands rather than becoming noise.
  • Out of scope here, but the residual after this PR is that the first pass in each process still runs 25 client.query(sql).result() calls serially on the request path, via before_run_callback. Starting the jobs and then collecting results would make that one round of latency instead of 25. Happy to open that separately if you want it.

The middleware side of the trigger is filed at ag-ui-protocol/ag-ui#2642, so the per-request teardown that drives the re-init gets fixed there too.

@lupuletic lupuletic left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified this branch (4196a06) against the reported production failure. The fix works and the approach is right.

I reproduced the host shape from #7017 — one shared plugin instance, _ensure_started() then close() per request, which is what before_run_callback and Runner.close() do — and counted the CREATE OR REPLACE VIEW statements over 5 cycles:

view DDL statements
origin/main 125 (25 per cycle)
this branch 25 (once, total)

That is exactly the intended shape: linear in re-initialisations before, constant after. The existing suite is also green on the branch — 511 passed in tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py.

I also like two choices here more than what we had in our own local mitigation:

  • Putting _schema_ready immediately after self._schema in __init__ makes the "pure data, survives close" grouping self-evident.
  • Driving the tests through real shutdown() cycles rather than poking _started directly. A test that sets the private flag would keep passing if a future close() started clearing the memo, so it would not actually guard the regression. Yours does.

The one thing I would still fix before merge is the __getstate__ reset I described in my earlier comment — the memo currently crosses the pickle boundary while every other runtime field is reset, so a plugin restored into an environment whose dataset does not exist yet would skip table creation. One line, and it does not weaken the fix, since the problem being solved is re-verification within one process.

Everything else reads good to me. Thanks for turning this around so fast.

For anyone landing here from the same symptom: the middleware-side trigger — a per-request Runner closing plugins owned by a long-lived App — is filed at ag-ui-protocol/ag-ui#2642. This PR makes the plugin resilient to that; that one stops the churn at source.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] BigQueryAgentAnalyticsPlugin re-runs all view DDL on every re-initialisation, on the request path (2.3.0 -> 2.8.0 regression)

3 participants