Bug: assign ReadTimeout can leave a live orphan assignment; reconcile before retrying

[hcsolakoglu]

I found an ambiguous-outcome failure in assignment creation on current main (e25c2655).

Client.assign() first checks the assignment for a notebook hash and then calls _post_assignment(). The current code handles ColabRequestError, but it does not reconcile a requests.exceptions.ReadTimeout from the assignment POST.

The problematic sequence is:

GET /tun/m/assign?...nbh=<hash>
POST /tun/m/assign?...nbh=<same hash>
→ backend may commit the assignment
→ client times out waiting for the POST response
→ ReadTimeout propagates
→ local SessionState is never created
→ the server-side assignment can remain live

During an earlier local audit, this was recorded as an observed provisioning failure: the assignment POST timed out and a later assignment query showed a live server-side assignment. I no longer have an independent raw transaction log for that incident, so I am treating that observation as historical evidence; the ambiguous-outcome behavior itself is independently confirmed by the current source and focused regression tests.

This makes an immediate retry unsafe. A new colab new call generates a new notebook hash, so blindly retrying allocation after an ambiguous POST outcome can request another runtime while the first assignment is still alive. I am not claiming that every retry necessarily creates a duplicate, only that the first POST outcome is unknown and should be reconciled before issuing another allocation request.

Suggested recovery invariant

After a read timeout on the POST, do not re-POST. Re-query using the same notebook hash from the original request.

A single immediate GET can race backend materialization, so a short bounded reconciliation window is safer. My reference implementation uses:

0s → 1s → 2s → 4s

If an Assignment appears, return/adopt it normally. If nothing materializes within the bounded window, re-raise the timeout and show an actionable message telling the user to run colab sessions before retrying.

I also validated the load-bearing assumption independently: querying the same notebook hash after a successful assignment returns the same server-side endpoint.

Reference implementation:

• branch: fix/assign-timeout-reconcile
• commit: b44c7c947eaa9c17c104ba135c89953f7d1739cf

Focused regression checks rerun:

test_cli_new_timeout_no_orphan_friendly_message
test_client_assign_post_timeout_adopts_orphan
test_client_assign_post_timeout_adopts_late_orphan
test_client_assign_post_timeout_no_orphan_reraises

All 4 passed.

The practical risk is lifecycle/resource leakage: a server-side assignment that the CLI failed to bind locally may continue to exist and consume resources until the user discovers and cleans it up.

I am using a Discussion rather than opening an external PR because the repository's current CONTRIBUTING.md says external PRs are not being accepted and directs issues/pain points to Discussions.

[lw-liuwang]

I encountered this same allocation-timeout problem while provisioning a TPU v6e1 with google-colab-cli==0.7.4 on Linux. I have a small implementation covering configurable HTTP timeouts and bounded allocation recovery, shared here in accordance with the repository's current contribution policy.

Personal fork and patch:

• Fork: https://github.com/lw-liuwang/google-colab-cli
• Branch: https://github.com/lw-liuwang/google-colab-cli/tree/fix-assignment-timeout-recovery
• Commit: lw-liuwang@a41650a
• Diff: lw-liuwang/google-colab-cli@main...fix-assignment-timeout-recovery

Follow-up tracking: issue #171
and PR #172, submitted from
my personal fork. I understand the current external-contribution policy; these
provide a reviewable implementation alongside this Discussion.

Observed behavior

On the installed CLI, allocation POSTs hit the approximately 120-second read timeout. With the recovery patch applied locally, a subsequent GET using the original notebook hash returned the allocated TPU runtime after three timeouts. Allocation took approximately 7 minutes 18 seconds overall, and a subsequent remote JAX device query confirmed a TPU device.

This was a real run of the locally patched installed package. It is historical evidence for the recovery approach, not a claim that the complete allocation flow has been independently rerun after publishing this fork commit.

Patch behavior

• COLAB_CLI_REQUEST_TIMEOUT overrides the HTTP timeout in seconds; unset preserves the underlying session default. An explicit caller-supplied timeout= takes precedence.
• COLAB_CLI_ASSIGN_RETRIES bounds additional allocation POST attempts (default 3; 0 disables retries).
• COLAB_CLI_RETRY_BACKOFF controls exponential backoff (default 10 seconds, giving 10/20/40-second waits).
• Only POST ReadTimeout and HTTP 502/503/504 are retried. HTTP 412 retains the existing TooManyAssignmentsError behavior; other HTTP errors propagate.
• Before another POST, GET queries the original notebook hash. If an Assignment is returned, it is adopted directly. Otherwise the next POST reuses that same notebook hash and the refreshed XSRF token.
• An initial or recovery GET failure propagates. The implementation does not provide unlimited waiting or alter kernel execution timeouts, runtime lifetime, or quotas.

Difference from the reconciliation-only proposal

Your proposed bounded GET-only reconciliation followed by re-raising the timeout is more conservative than this patch. My implementation can re-POST after a GET still reports no assignment; reusing the notebook hash avoids introducing a new request identity, but I have not established a general backend idempotency guarantee or ruled out delayed materialization races.

I would be interested in maintainer guidance on keeping recovery GET-only after ambiguous timeouts, while handling explicit transient HTTP errors with a separate bounded retry policy. The personal branch is an implementation to discuss, not evidence that these lifecycle edge cases are fully resolved.

Validation

27 focused client/consumption tests passed, covering:

• HTTP timeout override, default preservation, and explicit-argument precedence.
• Recovery using the same notebook hash after read timeouts and 502/503/504 responses.
• Refreshed tokens, eventual POST success, retry exhaustion, and 10/20/40-second backoff.
• Nonretryable errors, disabling retries, and existing allocation behavior.

The modified client also completed a real read-only assignment-list request. Mocked timeout tests do not substitute for a fresh end-to-end allocation/recovery test; that remains outstanding for the published branch.

Trying the branch

git clone --branch fix-assignment-timeout-recovery --single-branch \
  https://github.com/lw-liuwang/google-colab-cli.git
cd google-colab-cli
uv sync --locked --python 3.13

COLAB_CLI_REQUEST_TIMEOUT=300 \
COLAB_CLI_ASSIGN_RETRIES=3 \
COLAB_CLI_RETRY_BACKOFF=10 \
  uv run colab --auth oauth2 new -s timeout-recovery-test --tpu v6e1

Creating a runtime allocates real resources. After checking the session and retrieving any results, release it with:

uv run colab --auth oauth2 stop -s timeout-recovery-test
uv run colab --auth oauth2 sessions

If allocation fails ambiguously, inspect sessions before requesting another runtime; this patch does not guarantee that every failure leaves no server-side assignment.