[tinker] Retry external sample requests through the lazy adapter-load race

[john-hahn]

[tinker] Retry external sample requests through the lazy adapter-load race

Summary

ExternalInferenceClient._forward_to_engine extracts a LoRA checkpoint to
external_inference_lora_base, references it by name, and relies on the lora_filesystem_resolver
plugin to load it lazily on the engine's first reference to that name. The first /completions
request for a freshly extracted adapter can be issued before the engine has registered it and come
back as a 404 (or a 400 reporting the model does not exist), which currently fails the whole sample
request.

This wraps the completion call in a bounded retry. On the "adapter not loaded yet" signal — a 404, a
400 whose body reports a missing model, or a transport error — it proactively asks the engine to
load the adapter (POST /load_lora_adapter) and retries with exponential backoff. Other client
errors (e.g. an over-length request) and server errors are surfaced immediately.

Changes

• Add _is_retriable_sample_error and retry constants to external_inference.py.
• Add _request_adapter_load and _post_completion methods on ExternalInferenceClient; route the
  completion call in _forward_to_engine through _post_completion.
• Add tests/tinker/test_external_inference.py.

Design notes

• Retries are scoped to the adapter-load race: 404, 400-with-model-missing body, and transport
  errors. A 4xx that is a genuine client error and any 5xx are re-raised on the first occurrence, so
  real failures are not masked or delayed.
• Bounded: 6 attempts, exponential backoff capped at 2.0s, so a persistent failure adds ~4s before
  surfacing.
• The proactive load is only attempted for LoRA requests (base_model is None) and is best-effort —
  if the endpoint is disabled or runtime LoRA updating is off, the retry alone still gives the
  filesystem resolver time to register the adapter. POST /load_lora_adapter resolves under the
  client's /v1 base, matching vLLM's dynamic-LoRA endpoint.
• No config or signature change; the success path is unchanged (one request, no backoff).

Testing

• Unit (CPU): uv run --extra dev --extra jax pytest tests/tinker/test_external_inference.py —
  passing (10 cases: retry predicate for 404 / 400-missing / 400-other / 403 / 500 / transport;
  happy path, self-heal-then-succeed, retry-exhaustion, no-retry-on-5xx, base-model-never-loads).

Checklist

• Follows Google Python style; comments describe what the code does.
• CPU unit test added; the patched file is unchanged upstream since [tinker] Add deterministic routing support for sample requests #1720, so this rebases cleanly.
• No paths/naming changes affecting .claude/ docs; no example-script or doc changes needed.
• Added lines verified <= 120 cols and ast-parse clean; bash format.sh (ruff + black) to be run in-repo before opening.

[gemini-code-assist]

Code Review

This pull request introduces a self-healing retry mechanism for the external inference client to handle lazy-loading of LoRA adapters. When a sample request fails with a 404 or a specific 400 error, the client proactively requests the engine to load the adapter and retries the request. Unit tests are also added to verify this behavior. The review feedback suggests setting an explicit timeout on the proactive adapter-load request to prevent blocking the retry loop, and avoiding retries on 404/400 errors when sampling from a base model, as these represent permanent configuration errors rather than transient lazy-load races.

[gemini-code-assist]

The proactive load request uses the default timeout of the http_client, which is configured to be 300 seconds (5 minutes). If the /load_lora_adapter endpoint hangs or responds very slowly, it will block the entire retry loop for a long time, defeating the purpose of a fast retry/fail. Since this proactive load is a best-effort call, we should set a shorter, explicit timeout (e.g., 10 seconds).

Suggested change

	try:
	await http_client.post(
	"/load_lora_adapter",
	json={"lora_name": model_name, "lora_path": str(self.lora_base_dir / model_name)},
	)
	try:
	await http_client.post(
	"/load_lora_adapter",
	json={"lora_name": model_name, "lora_path": str(self.lora_base_dir / model_name)},
	timeout=10.0,
	)

[gemini-code-assist]

When sampling from a base model (base_model is not None), a 404 or 400 indicating that the model does not exist is a permanent configuration error (e.g., a misspelled model name), not a transient lazy-load race. Retrying these errors for base models wastes time and delays error reporting. We should only retry 404/400 errors when sampling a LoRA adapter (base_model is None), while still retrying transient transport errors for both.

Suggested change

	def _is_retriable_sample_error(exc: Exception) -> bool:
	"""Return whether a failed sample request should be retried after a proactive adapter load.
	The engine loads LoRA adapters lazily, so a freshly extracted adapter can be sampled before it is
	registered. Retry transient transport errors and the engine's "adapter not loaded yet" responses
	(a 404, or a 400 whose body reports the model does not exist). Malformed requests and server
	errors are not retried.
	"""
	if isinstance(exc, httpx.TransportError):
	return True
	if isinstance(exc, httpx.HTTPStatusError):
	status_code = exc.response.status_code
	if status_code == 404:
	return True
	if status_code == 400:
	body = exc.response.text.lower()
	return "does not exist" in body or "not found" in body
	return False
	def _is_retriable_sample_error(exc: Exception, *, is_lora: bool = True) -> bool:
	"""Return whether a failed sample request should be retried after a proactive adapter load.
	The engine loads LoRA adapters lazily, so a freshly extracted adapter can be sampled before it is
	registered. Retry transient transport errors and the engine's "adapter not loaded yet" responses
	(a 404, or a 400 whose body reports the model does not exist). Malformed requests and server
	errors are not retried.
	"""
	if isinstance(exc, httpx.TransportError):
	return True
	if is_lora and isinstance(exc, httpx.HTTPStatusError):
	status_code = exc.response.status_code
	if status_code == 404:
	return True
	if status_code == 400:
	body = exc.response.text.lower()
	return "does not exist" in body or "not found" in body
	return False

[gemini-code-assist]

Pass the is_lora flag to _is_retriable_sample_error to avoid retrying permanent 404/400 errors for base models.

Suggested change

	for attempt in range(_MAX_SAMPLE_RETRIES):
	try:
	response = await http_client.post("/completions", json=payload, headers=headers)
	response.raise_for_status()
	return response.json()
	except (httpx.HTTPStatusError, httpx.TransportError) as exc:
	if attempt == _MAX_SAMPLE_RETRIES - 1 or not _is_retriable_sample_error(exc):
	raise
	if base_model is None:
	await self._request_adapter_load(http_client, model_name)
	await asyncio.sleep(min(_RETRY_MAX_DELAY_SEC, _RETRY_BASE_DELAY_SEC * (2**attempt)))
	is_lora = base_model is None
	for attempt in range(_MAX_SAMPLE_RETRIES):
	try:
	response = await http_client.post("/completions", json=payload, headers=headers)
	response.raise_for_status()
	return response.json()
	except (httpx.HTTPStatusError, httpx.TransportError) as exc:
	if attempt == _MAX_SAMPLE_RETRIES - 1 or not _is_retriable_sample_error(exc, is_lora=is_lora):
	raise
	if is_lora:
	await self._request_adapter_load(http_client, model_name)
	await asyncio.sleep(min(_RETRY_MAX_DELAY_SEC, _RETRY_BASE_DELAY_SEC * (2**attempt)))

[gemini-code-assist]

Update the base model test to verify that transient transport errors are retried, but permanent 404 errors are not retried.

Suggested change

	def test_base_model_never_loads_adapter(monkeypatch):
	"""Base-model sampling retries transient errors but never calls the adapter-load endpoint."""
	scripted = _ScriptedClient([404, 200])
	result = _run(_make_client(), scripted, monkeypatch, base_model="base")
	assert result == {"choices": []}
	assert "/load_lora_adapter" not in scripted.calls
	def test_base_model_retries_transport_error(monkeypatch):
	"""Base-model sampling retries transient transport errors but never calls the adapter-load endpoint.
	"""
	scripted = _ScriptedClient([httpx.ConnectError("boom"), 200])
	result = _run(_make_client(), scripted, monkeypatch, base_model="base")
	assert result == {"choices": []}
	assert "/load_lora_adapter" not in scripted.calls
	def test_base_model_does_not_retry_404(monkeypatch):
	"""Base-model sampling does not retry on 404 (model not found)."""
	scripted = _ScriptedClient([404, 200])
	with pytest.raises(httpx.HTTPStatusError):
	_run(_make_client(), scripted, monkeypatch, base_model="base")