[ML-69867] Make the workspace walk pickers cancellable and stop blocking CLI exit - #800
Merged
xsh310 merged 3 commits intoSep 23, 2026
Merged
Conversation
…t (ML-69867) The `ug skills add`, `ug skills download`, and `ug mcp add services` pickers run a workspace wide catalog and schema walk on a background thread. That walk used a ThreadPoolExecutor whose worker threads are joined when the interpreter exits, so the CLI stalled for several seconds after its last output line while in flight HTTP probes finished. Closing the picker also did nothing to stop the walk. Replace the ThreadPoolExecutor in `walk_catalog_schemas` with daemon threads that are never joined, so process exit is never blocked even when a probe is still waiting on the network. Add a `threading.Event` that the picker sets the moment it closes, forwarded through the skills and mcp services loaders into the walk, so the walk stops as soon as the user finishes selecting or the 30 second deadline passes. All three walk backed picker loaders are updated; the new cancel_event parameters default to None, so non picker callers are unchanged. Shrinking the per probe HTTP timeout to the remaining deadline is left as a follow up. Co-authored-by: Isaac <no-reply@databricks.com>
…-69867) Address review feedback on the workspace walk change. * Replace the cancellation test so it trips the event from inside the probe, actually exercising probe phase cancellation. The old test set the event up front, so the walk returned during schema listing and never reached probing. * Add invariant tests for `_collect_concurrently`: a full success run collects every schema and returns None, and a run where some probes raise skips those results without stalling the drain loop. * Add generics, queue type parameters, and a short docstring to `_collect_concurrently`. Co-authored-by: Isaac <no-reply@databricks.com>
xsh310
marked this pull request as ready for review
September 23, 2026 03:20
xsh310
changed the title
sunishsheth2009
approved these changes
Sep 23, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What did you change, and why?
Change: The
ug skills add,ug skills download, andug mcp add servicespickers run a workspace wide catalog and schema walk on a background thread. That walk now runs on plain daemon threads that are never joined, and it stops early when the picker closes or the 30 second deadline passes.Why: The walk used a
ThreadPoolExecutorwhose worker threads are joined at interpreter exit, so the CLI stalled for several seconds after printing its last line while in flight HTTP probes drained. Closing the picker never stopped the walk. (ML-69867)Details:
walk_catalog_schemasno longer uses aThreadPoolExecutor. A small helper runs each probe ondaemon=Truethreads that are never joined, collecting results on the calling thread until a stop predicate trips, so process exit is never blocked even when a probe is still waiting on the network.threading.Eventowned by the picker is tripped when it closes (every exit path), forwarded through all three walk backed loaders (_skill_schema_background_loader,_mcp_services_background_loader,_skills_download_background_loader) into the walk. The walk stops on picker close or when the deadline passes.cancel_eventparameters default toNone, so non picker callers behave exactly as before.How do you know it works?
Testing: New unit tests in
tests/test_databricks.pycover that the walk returns promptly when the cancel event is set and when the deadline passes.uv run ty check src/passes, andruff checkplusruff format --checkpass on the changed files. Full suite: 2637 passed, 44 skipped; the only failures are two preexistingtest_e2e_user_agent.pycases that require capture server routing, unrelated to this change.Before:
Screen.Recording.2026-09-22.at.7.53.49.PM.mov
After:
Screen.Recording.2026-09-22.at.7.55.39.PM.mov
This pull request and its description were written by Isaac.