fix: improve Dial reliability for slow exporters#873
Open
evakhoni wants to merge 5 commits into
Open
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds configurable exporter dial timeouts across configuration and ChangesDial timeout configuration and CLI wiring
Lease dialing and tunnel setup
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/packages/jumpstarter/jumpstarter/client/lease.py (1)
95-95: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winLease dataclass default still 30s; config default is now 60s.
Lease.dial_timeoutfield default remains30.0whileClientConfigV1Alpha1Lease.dial_timeoutwas bumped to60.0. Any directLease(...)construction that omitsdial_timeoutwill silently use the stale 30s value.🩹 Proposed fix
- dial_timeout: float = field(default=30.0) # Timeout in seconds for Dial retry loop when exporter not ready + dial_timeout: float = field(default=60.0) # Timeout in seconds for Dial retry loop when exporter not ready🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease.py` at line 95, The Lease dataclass still hardcodes an outdated 30.0 default for dial_timeout while the config default has moved to 60.0. Update the default in Lease so direct Lease(...) construction matches ClientConfigV1Alpha1Lease and does not silently use the stale timeout; keep the change localized to the Lease definition and ensure any related default wiring stays consistent.
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/client/lease_test.py (1)
578-630: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the remaining
_dial_with_retrybranches.Only the UNAVAILABLE retry/timeout paths are tested. The FAILED_PRECONDITION "not ready" retry path, the "permission denied" →
lease_transferredpath, and the generic "connection lost" immediate-raise path are untested.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease_test.py` around lines 578 - 630, The TestDialWithRetry coverage only exercises the UNAVAILABLE retry behavior in Lease._dial_with_retry, so add tests for the remaining branches: FAILED_PRECONDITION should retry until the lease becomes ready, PERMISSION_DENIED should set lease_transferred and stop retrying, and non-retryable connection-lost errors should raise immediately. Extend TestDialWithRetry using _make_lease_for_dial and mock Lease.controller.Dial to simulate each grpc.StatusCode path and assert the expected retry or state-change behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Around line 399-422: The per-connection dial in serve_unix_async is not being
retried or translated, so a transient controller.Dial failure can escape as a
raw AioRpcError and break TemporaryUnixListener handling. Update the inner
_tunnel_handler in lease.py to catch the Dial call failure and map exporter
UNAVAILABLE cases to ExporterUnreachableError, or route the connection through
the existing retry helper used by _dial_with_retry before calling
connect_router_stream. Keep the change localized to serve_unix_async and
_tunnel_handler so each accepted connection gets the same retry/translation
behavior as the shell path.
---
Outside diff comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease.py`:
- Line 95: The Lease dataclass still hardcodes an outdated 30.0 default for
dial_timeout while the config default has moved to 60.0. Update the default in
Lease so direct Lease(...) construction matches ClientConfigV1Alpha1Lease and
does not silently use the stale timeout; keep the change localized to the Lease
definition and ensure any related default wiring stays consistent.
---
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease_test.py`:
- Around line 578-630: The TestDialWithRetry coverage only exercises the
UNAVAILABLE retry behavior in Lease._dial_with_retry, so add tests for the
remaining branches: FAILED_PRECONDITION should retry until the lease becomes
ready, PERMISSION_DENIED should set lease_transferred and stop retrying, and
non-retryable connection-lost errors should raise immediately. Extend
TestDialWithRetry using _make_lease_for_dial and mock Lease.controller.Dial to
simulate each grpc.StatusCode path and assert the expected retry or state-change
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: de0b5284-fefe-47a2-8e87-253407fb69dd
📒 Files selected for processing (7)
python/packages/jumpstarter-cli/jumpstarter_cli/common.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/config/client.pypython/packages/jumpstarter/jumpstarter/config/client_config_test.pypython/packages/jumpstarter/jumpstarter/config/env.py
evakhoni
force-pushed
the
fix/remove-spurious-dial-probe
branch
from
July 9, 2026 21:44
5e4ced6 to
a9bae27
Compare
mangelajo
reviewed
Jul 10, 2026
_wait_for_ready_connection triggered a spurious handle_async instance via TemporaryUnixListener. When the probe's Dial retry timed out, structured concurrency cancelled the real handler mid-flight. The probe is unnecessary: socket.listen() happens synchronously, gRPC has lazy connect, and DirectLease works without it. Also fixes duplicate Dial messages on slow exporters. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Adds a readiness-check Dial with retry in `serve_unix_async` before accepting connections. Previously, the per-connection Dial and the HTTP/2 handshake raced — if Dial took >18s retrying on a slow exporter, the handshake timed out and killed the connection. Now the exporter is confirmed ready before the first gRPC client connects. Also extracts `_dial_with_retry` from `handle_async` and raises `ExporterUnreachableError` on timeout instead of raw `AioRpcError`, so the shell's retry loop can distinguish retryable failures. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Raises the default dial_timeout from 30s to 60s to handle slow exporters that need 30-90s to finish cleanup and start their Listen stream. Also exposes --dial-timeout as a CLI flag (matching --retry-timeout and --acquisition-timeout) so users with even slower hardware can override it. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Adds test coverage for error paths in _dial_with_retry and the serve_unix_async connection flow to protect against regressions. New tests in TestDialWithRetry: - FAILED_PRECONDITION timeout path (exporter not ready) - Permission denied path (lease transferred to another client) - Generic/unknown error path (catch-all handler) New TestServeUnixAsync class: - Verifies readiness check happens once before accepting connections - Verifies each connection gets its own Dial for unique router tunnels - Verifies connect_router_stream called with correct args Test cleanup: - Strengthened assertion in unknown error test (check for "lost" not just "exporter") - Moved ExporterUnreachableError and anyio imports to file top - Fixed indentation of lease_async mock functions in shell_test.py - Added pragma: no cover to client.py dial_timeout resolution (matches retry_timeout pattern) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
evakhoni
force-pushed
the
fix/remove-spurious-dial-probe
branch
from
July 12, 2026 14:51
c993232 to
5356e9b
Compare
Member
Author
|
had a merge conflict, just solved |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/source/getting-started/configuration/files.md`:
- Around line 49-52: Update the `leases.dial_timeout` comment to describe it as
the deadline for the dial retry loop while waiting for an exporter that is not
yet ready, rather than only the connection timeout. Preserve the existing
default value and units.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b1755feb-e3cf-4625-8436-fefa4fb63ed7
📒 Files selected for processing (10)
docs/source/getting-started/configuration/files.mddocs/source/getting-started/configuration/loading-order.mdpython/packages/jumpstarter-cli/jumpstarter_cli/common.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/config/client.pypython/packages/jumpstarter/jumpstarter/config/client_config_test.pypython/packages/jumpstarter/jumpstarter/config/env.py
🚧 Files skipped from review as they are similar to previous changes (7)
- python/packages/jumpstarter/jumpstarter/config/env.py
- python/packages/jumpstarter-cli/jumpstarter_cli/common.py
- python/packages/jumpstarter/jumpstarter/config/client_config_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter/jumpstarter/config/client.py
- python/packages/jumpstarter/jumpstarter/client/lease_test.py
- python/packages/jumpstarter/jumpstarter/client/lease.py
Adds documentation for lease timeout configuration and other previously undocumented environment variables. Client configuration updates: - leases YAML config section (acquisition_timeout, dial_timeout, retry_timeout) - JMP_RETRY_TIMEOUT, JMP_DIAL_TIMEOUT timeout env vars - JMP_OIDC_CALLBACK_PORT (useful for SSH tunneling) - JMP_GRPC_PASSPHRASE (passphrase-protected exporters) - Shell session variables (JMP_LEASE, JMP_EXPORTER, JMP_EXPORTER_LABELS) automatically set by jmp shell Exporter configuration updates: - JMP_DISABLE_COMPRESSION (disable stream compression) Updated loading-order.md to include timeout flags (--retry-timeout, --dial-timeout) in the configuration hierarchy examples. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
evakhoni
force-pushed
the
fix/remove-spurious-dial-probe
branch
from
July 12, 2026 16:54
5356e9b to
d3368b7
Compare
mangelajo
approved these changes
Jul 13, 2026
mangelajo
requested changes
Jul 13, 2026
mangelajo
left a comment
mangelajo
left a comment
Member
There was a problem hiding this comment.
hold on, I need to verify something, I remember we were blocking the unix socket until beforeLease hooks end for different reasons.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/packages/jumpstarter/jumpstarter/client/lease_test.py (1)
730-740: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winFixed
anyio.sleep(0.1)makes handler completion racy.The assertions on
dial_call_count/router_stream_callsrely on the background_tunnel_handlerfinishing within a hard-coded 0.1s window. Under CI load this can flake. Prefer signaling completion from the mock and awaiting it before asserting.♻️ Event-based synchronization
router_stream_calls = [] + handler_done = anyio.Event() `@asynccontextmanager` async def mock_connect_router_stream(endpoint, token, stream, tls_config, grpc_options): router_stream_calls.append((endpoint, token, tls_config, grpc_options)) yield + handler_done.set() with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): with patch("jumpstarter.client.lease.connect_router_stream", side_effect=mock_connect_router_stream): async with lease.serve_unix_async() as socket_path: assert readiness_check_called async with await anyio.connect_unix(socket_path): - # Give the handler time to process - await anyio.sleep(0.1) + with anyio.fail_after(5): + await handler_done.wait()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/packages/jumpstarter/jumpstarter/client/lease_test.py` around lines 730 - 740, Replace the fixed anyio.sleep delay in the serve_unix_async test with event-based synchronization: have the mock handler signal an AnyIO event after completing its Dial/router_stream work, then await that event before asserting dial_call_count and router_stream_calls. Keep the existing socket connection and readiness checks unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python/packages/jumpstarter/jumpstarter/client/lease_test.py`:
- Around line 730-740: Replace the fixed anyio.sleep delay in the
serve_unix_async test with event-based synchronization: have the mock handler
signal an AnyIO event after completing its Dial/router_stream work, then await
that event before asserting dial_call_count and router_stream_calls. Keep the
existing socket connection and readiness checks unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 76713a13-c1b6-4738-8a64-f62fa1d8ad39
📒 Files selected for processing (10)
docs/source/getting-started/configuration/files.mddocs/source/getting-started/configuration/loading-order.mdpython/packages/jumpstarter-cli/jumpstarter_cli/common.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell.pypython/packages/jumpstarter-cli/jumpstarter_cli/shell_test.pypython/packages/jumpstarter/jumpstarter/client/lease.pypython/packages/jumpstarter/jumpstarter/client/lease_test.pypython/packages/jumpstarter/jumpstarter/config/client.pypython/packages/jumpstarter/jumpstarter/config/client_config_test.pypython/packages/jumpstarter/jumpstarter/config/env.py
🚧 Files skipped from review as they are similar to previous changes (9)
- docs/source/getting-started/configuration/loading-order.md
- docs/source/getting-started/configuration/files.md
- python/packages/jumpstarter/jumpstarter/config/client_config_test.py
- python/packages/jumpstarter-cli/jumpstarter_cli/common.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell.py
- python/packages/jumpstarter/jumpstarter/config/client.py
- python/packages/jumpstarter/jumpstarter/config/env.py
- python/packages/jumpstarter/jumpstarter/client/lease.py
- python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py
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.
Problem
When connecting to slow hardware exporters (e.g., NVIDIA Jetson AGX Orin),
jmp shellfrequently failed or required multiple lease attempts before establishing a connection. Three compounding issues were identified:Spurious connection probe: Before opening the shell, the client sent a test connection to the Unix socket to verify it was ready. This probe triggered the same connection handler as the real shell connection, spawning a second Dial attempt in the same task group. When the probe's connection was closed (having served its purpose), its handler failed — and because both handlers shared a task group, the failure cancelled the entire group, killing the real connection along with it.
gRPC handshake timeout: Each connection to the exporter requires a Dial call to the controller, which retries with exponential backoff while the exporter starts its Listen stream. The Dial retries and the gRPC HTTP/2 handshake share the same connection — while Dial is retrying, the gRPC client on the other end is waiting for an HTTP/2 SETTINGS frame. If Dial retries take longer than gRPC's internal handshake timeout (~18-20 seconds), the handshake times out and the connection is dropped. This affects any exporter where the Dial retry phase exceeds the handshake window, though slow exporters make it far more likely.
Insufficient dial timeout: The default Dial timeout of 30 seconds was too short for exporters that need more than 30 seconds to finish cleanup from a previous session and restart their Listen stream. When the timeout was exceeded, the client released the lease and acquired a new one, which forced the exporter to restart its setup from scratch — creating a churn cycle that made the problem worse with each iteration.
Fix
Probe removal
The
_wait_for_ready_connectionmethod and its invocation inserve_unix_asyncare removed entirely. The probe was originally added to block until the Unix socket was accepting connections, but this is unnecessary — the socket is created byTemporaryUnixListenerbeforeserve_unix_asyncyields, so it is always ready by the time the shell connects. Theconnect_uniximport that was only used by the probe is also removed.Pre-Dial readiness check
The connection handler (
handle_async) previously did both Dial-with-retry and router stream setup in a single method. This meant the Dial retries and the gRPC handshake were coupled on the same connection, creating the timeout race.The fix splits these concerns:
serve_unix_asyncnow calls a new_dial_with_retrymethod before accepting any connections. This readiness check ensures the exporter has its Listen stream active. The Dial response from this check is intentionally discarded — each Unix socket connection still gets its own per-connection Dial (since each connection needs a unique router tunnel), but since the exporter is already confirmed ready, these per-connection Dials succeed immediately (typically under 1 second), well within the gRPC handshake timeout._dial_with_retryis extracted from the oldhandle_asyncmethod. It uses exponential backoff (0.3s base, 2s max) bounded bydial_timeout, and raisesExporterUnreachableErroron timeout instead of the rawAioRpcError. This allows the shell's outer retry loop to properly distinguish retryable exporter failures from other gRPC errors.Dial timeout increase and CLI flag
The default
dial_timeoutis raised from 30s to 60s to accommodate exporters that need more time to become ready. The field is also exposed for user configuration:--dial-timeoutonjmp shell(e.g.,--dial-timeout 90s,--dial-timeout 2m)JMP_DIAL_TIMEOUTleases.dial_timeout(previously excluded from serialization, now persisted)This matches the existing pattern for
--retry-timeout/JMP_RETRY_TIMEOUTand--acquisition-timeout.Test plan
make pkg-test-jumpstarter— 560 passedmake lint-fix— all checks passedjmp shell -n <exporter>connects reliably with default 60s timeoutjmp shell --dial-timeout 10triggers retry/release as expected