Skip to content

Cut ug startup cost: fast --version entry point and lazy heavy imports - #1018

Merged
david-siqi-liu merged 7 commits into
databricks:mainfrom
yaoharry:harry-yao/port-startup-imports
Oct 13, 2026
Merged

Cut ug startup cost: fast --version entry point and lazy heavy imports#1018
david-siqi-liu merged 7 commits into
databricks:mainfrom
yaoharry:harry-yao/port-startup-imports

Conversation

@yaoharry

@yaoharry yaoharry commented Oct 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • Point both console scripts at ucode.entrypoint:main. A bare --version or
    -V reads installed distribution metadata without importing the full CLI.
    Other arguments still delegate to ucode.cli.main, including mixed version
    arguments. Existing ucode.cli:main and python -m ucode.cli remain supported.
  • Defer the Databricks SDK OAuth import until the SDK backend is used,
    gateway_proxy until a relayed Claude or traced Codex launch, and
    codex_interposer until Codex smart routing, before its app-server is spawned.
  • Verify console scripts with --help on the legacy distribution-migration
    path, since the fast version path no longer proves the CLI can import.
  • Add fresh-interpreter import guards and component tests for version suffixes,
    distribution-name normalization, search-path precedence, zip/egg fallback,
    incomplete metadata, and unchanged argument dispatch.

Command registration, agent SPEC registration, questionary UI imports,
telemetry imports on normal launches, and the skills executor remain eager.
The port keeps the original production patch's intent and adds regression
coverage and coverage documentation.

Startup Timings

Median wall time of 15 interleaved before/after runs after one warm-up, using
CPython 3.12.3 on this Linux VM with frozen dependencies and cached bytecode.
Baseline is main at 4d8ff8d; after is this port.

Measurement Before After
ug --version 1,053 ms 13 ms
ug --help 1,066 ms 255 ms
ug claude --help 1,067 ms 261 ms
python -c 'import ucode.cli' 1,146 ms 216 ms

Stdout and stderr were byte-identical for every compared invocation. Comparisons
were taken before the port commit, when both installations reported the same
distribution version; post-commit console checks verify the new installed version.
These are local startup measurements, not live agent launch latency.

Verification

  • just lint: Ruff check, Ruff format check, and ty check src/ passed.
  • Startup and installed-console-script checks: 33 passed.
  • Full ordinary suite at the current 800b75e base, excluding
    tests/test_e2e.py: 3,402 passed, six failed, 13 skipped.
    Unchanged main at that base in the same environment:
    3,373 passed, the same six failures, 13 skipped. JUnit failure sets match
    exactly; there are no new failures.
  • After rebasing onto 800b75e, lint and all 33 startup/console-script checks
    passed again.
  • Full-suite runs removed ambient bearer/token overrides and limited PATH
    to the checkout virtualenv plus /usr/bin:/bin, avoiding real agent
    execution. The six baseline failures depend on Databricks CLI discovery or
    authentication that is not available in this isolated environment.
  • Frozen dependency installation used the internal Python package proxy;
    uv.lock is unchanged.
  • Independent diff review found no correctness or regression issues.

Installed-Package Verification

  • Built one immutable PR wheel and ran a four-cell fresh-consumer matrix:
    primary entry point ug/ucode, with MCP 2.2.0/default resolution or
    MCP 1.28.0 pinned. All 20 selected package/protocol tests passed. Each run
    used separate runtime/test virtualenvs, installed real Claude 2.1.280 and
    Codex 0.154.0, and verified site-packages provenance.
  • The five selected tests cover public help/version, fresh-home status,
    actionable unconfigured auth, token-only helper stdout, and real local
    MCP initialization/tool-list exchanges through both command names.
    Below-floor upgrade cases were explicitly excluded; these are not
    live inference or backend proxy tests.
  • A real uv tool install --reinstall transition from the 800b75e
    baseline wheel to the PR wheel regenerated both scripts with the new
    entry point. All 28 before/after help, mixed-version-argument, and invalid
    argument comparisons matched, accounting only for the installed version.
  • Actual installed-script import profiles showed no CLI/Typer/metadata or
    heavy-module imports for bare version requests, and no deferred heavy
    modules on the help path. Median version wall times in the reinstall
    probe were 1,258 -> 17 ms (ug) and 1,414 -> 13 ms (ucode).

Live E2E Limits

No authenticated model-inference E2E or live relay/tracing/interposer test has
run. This host has existing OS-managed Claude/Codex settings, so the live runner
refuses it. API-only workspace/config/catalog reads succeeded using the explicitly
selected Isaac profile, but the clean-Arca attempt was blocked by SSH/dbcert
authentication. No agent inference ran. The ambient
Databricks token was not used and developer/OS-managed settings were not changed.

The SDK-backed custom OAuth path has no existing live integration coverage.
It additionally needs a registered client and legitimate browser consent or
cache. The public ug upgrade command currently hardcodes upstream, so the
pre-merge candidate transition was tested through the documented uv reinstall
path, not a redirected or mocked upgrade source.

Review Follow-Up

faec297 documents --help as the startup-import check and adds four component
cases showing that intact version metadata can still print successfully while
help fails on a missing startup dependency. The focused startup, console-script,
upgrade, and integration-contract selection passed all 92 tests; just lint
passed again. No dependency or lock-file changes were included.

The companion Isaac fix is
https://github.com/databricks-eng/universe/pull/2760310. It requires a bounded
ucode --help check before accepting a matching installed pin, including Isaac's
metadata shortcut, and after installation. Its unit and installer suites passed,
including real offline uv repair of a deleted dependency at the same pin.
The Isaac wheel pin is unchanged; land the companion before the pin bump that
includes this Gateway change.

Harry Yao @timothytamm-db
Cut ug startup cost with a fast version entry point and lazy imports
a68bc97 
Port the startup-performance patch from universe PR #2741961 onto current main. Keep ordinary CLI dispatch unchanged, defer command-specific dependencies, and verify upgraded commands with --help.

Add regression coverage for import boundaries, version lookup and fallback, search-path precedence, and argument compatibility.

Co-authored-by: Timothy Tamm <timothy.tamm@databricks.com>
Signed-off-by: Harry Yao <harry.yao@databricks.com>
@yaoharry
yaoharry marked this pull request as ready for review October 6, 2026 23:55
@lilly-luo lilly-luo added the ug-review Run the automated UG review label Oct 7, 2026
@github-actions

github-actions Bot commented Oct 7, 2026 •
edited
Loading

Copy link
Copy Markdown

UG review

The change introduces a metadata-only version path and defers heavy imports until relevant launches. The metadata scan misses an extracted-egg case that can report the wrong installed version.

Major

  • Preserve extracted egg precedence — src/ucode/entrypoint.py:35: can we fall back for extracted .egg directories on sys.path? their EGG-INFO is skipped, so --version can report a later dist-info version instead of the first installed copy.

Automated advisory review of 0f4019acd634 using Lilly's UG review rubric. It does not approve or block this PR.

david-siqi-liu
david-siqi-liu approved these changes Oct 7, 2026
View reviewed changes

@david-siqi-liu david-siqi-liu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM, nice startup win. One rollout note inline about Isaac's install health check; please land that alongside the Isaac pin bump.

Comment thread src/ucode/entrypoint.py
Harry Yao
Clarify CLI health probes after startup-import review
faec297 
Document --help as the startup-import smoke check and cover intact metadata with missing startup dependencies for both command names and version flags.

Signed-off-by: Harry Yao <harry.yao@databricks.com>
This was referenced Oct 8, 2026
Closed
Merged
Harry Yao and others added 3 commits October 8, 2026 23:04
Harry Yao
Merge origin/main into harry-yao/port-startup-imports
d2fbf05 
Signed-off-by: Harry Yao <harry.yao@databricks.com>
Harry Yao
Merge remote-tracking branch 'origin/main' into codex/pr1018
0f4019a 
Signed-off-by: Harry Yao <harry.yao@databricks.com>
@david-siqi-liu
Merge remote-tracking branch 'ugh/main' into pr-1018-update
268bc4f 
# Conflicts:
#	tests/integration/README.md
david-siqi-liu added a commit to yaoharry/unity-gateway that referenced this pull request Oct 12, 2026
@david-siqi-liu Isaac
Move fork-PR credentialed CI off pull_request_target (databricks#1044)
b349075 
After this PR, no workflow uses `pull_request_target`. Fork PRs get
credentialed CI when an org member (`MEMBER` or `OWNER`) comments a
command on the PR. Every comment-triggered run uses the default branch's
workflow definitions. Design doc (option B):
https://docs.google.com/document/d/1v-6znVBCb8YfGxr12SBd80w3i9n4U0fMEYohyz9Aoyw/edit

## Commands

Org members comment these on a PR:

- `/integration-test` runs the full Integration workflow with secrets on
the PR's current head commit.
- `/integration-test <sha>` does the same, but only if the head is still
`<sha>` (7 to 40 hex characters). If the fork pushed since, it refuses
and replies with the new head.
- `/ug-review` runs the UG rubric review.
- `/user-journey-check` runs the user-journey gate. For in-repo PRs it
also runs automatically, so the command is mostly for forks.

Each command posts a commit status and a comment with the run link on
the PR. The statuses are `Fork integration` and `User Journey Test
Required`. Each new push needs a new command. Anyone else who comments
`/integration-test` gets a reply saying only org members can run
integration tests.

## Before and after

Live integration on fork PRs:
- Before: every credentialed job skips, because fork PRs get no secrets
on `pull_request`. That covers the E2E CUJs, workspace validation, and
the smoke, full, and managed journeys.
- After: a member reviews, then comments `/integration-test` or
`/integration-test <sha>`. It runs `integration.yml` on the pinned
commit with secrets.

Live integration on in-repo PRs:
- Before: runs automatically.
- After: unchanged. `/integration-test` also works as an on-demand
rerun, and it never cancels the automatic run.

UG review:
- Before: adding the `ug-review` label, via `pull_request_target`.
- After: a member comments `/ug-review`. It still runs only
default-branch code and never executes the PR.

User Journey Test Required:
- Before: `pull_request_target` on every PR.
- After: `pull_request` for in-repo PRs, and `/user-journey-check` for
fork PRs. It still reads the PR only through the API. The gate checks
the head is unchanged across its diff reads and reports the fork-PR
status only on that evaluated commit, so a push during the check never
inherits the old verdict.

`pull_request_target`:
- Before: used by `ug-review.yml` and `user-journey-test-required.yml`.
- After: not used anywhere, and a contract test fails if it comes back.
No admin event-policy exception is needed before GitHub blocks it on
public repos on 2026-11-02.

## Fork safety in `integration.yml`

- It takes an optional `ref` input. When the input is empty, existing
callers behave exactly as before.
- Every checkout pins `ref` and sets `persist-credentials: false`.
- setup-uv skips its cache for ref-pinned runs, so fork builds cannot
poison it.
- Concurrency groups key on the PR number, so runs for different PRs
never cancel each other. Concurrency is set per job, so ordinary PR
comments never cancel a run.

`AGENTS.md` now says workflows must not use `pull_request_target`. The
new `fork-integration.yml` call to `integration.yml` carries the same
inline `zizmor: ignore[secrets-inherit]` as `ci.yml` (databricks#1089).

## Testing

- Contract tests cover the triggers, the member-only gate, and the
authorize script, run against a fake `gh`. They also check the
non-member reply, the pinned credential-free checkouts, the disabled
caches, and the status reports.
- `issue_comment` workflows only run from the default branch, so the
first live run happens after merge. The plan is to comment
`/integration-test` on a fork PR such as databricks#1017 or databricks#1018, and
`/ug-review` on any PR.

OIDC note: under `issue_comment`, the OIDC subject is the default
branch, not a PR ref. If the JFrog trust policy for the Windows
installation lane keys on PR refs, that one lane fails on fork runs. The
other lanes use secrets, not OIDC.

This pull request and its description were written by Isaac.

---------

Co-authored-by: Isaac <no-reply@databricks.com>
david-siqi-liu and others added 2 commits October 13, 2026 01:17
@david-siqi-liu
Merge remote-tracking branch 'ugh/main' into pr-1018-update
1b46a92
@david-siqi-liu Isaac
Skip the Windows console-script launcher in the fast version lookup
8739e7a 
On Windows, sys.path[0] is the launcher file itself (Scripts\ug.exe, Scripts\pytest.exe). os.listdir on it raises NotADirectoryError, which the reader treated as an unhandled entry and returned None, so the fast path never ran on Windows and test_installed_version_matches_importlib_metadata failed once Windows unit tests ran for fork PRs. Skip .exe entries; other files may be zip archives and still fall back.

Co-authored-by: Isaac <no-reply@databricks.com>
@david-siqi-liu

david-siqi-liu commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

/integration-test 8739e7a

@github-actions

github-actions Bot commented Oct 13, 2026

Copy link
Copy Markdown

Fork integration success on 8739e7a: https://github.com/databricks/unity-gateway/actions/runs/38333455051

@david-siqi-liu
david-siqi-liu added this pull request to the merge queue Oct 13, 2026
Merged via the queue into databricks:main with commit b5bb40e Oct 13, 2026
36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@david-siqi-liu david-siqi-liu david-siqi-liu approved these changes

Assignees

No one assigned

Labels

ug-review Run the automated UG review

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@yaoharry @david-siqi-liu @lilly-luo