fix: context-length fallback logging, batch trajectory durability, pool cleanup

Salvage of #6629 by aaronlab (kshitijk4poor reworked against current main).

Three concerns from the original PR, reworked to address review feedback:

1. Context-length fallback diagnostic (agent/model_metadata.py):
   get_model_context_length() silently returned 256K when all 9 detection
   methods failed. Users with small-context models (8K, 32K) would get 256K
   silently, causing hard-to-debug API context-length errors. Added a
   warning log at the step 9 fallback with model name, base_url, and the
   correct config override hint (model.context_length, not context_length).
   The token-estimation ceiling-division fix from the original PR already
   landed on main (5c2ecdec) with CJK handling — not duplicated here.

2. Fsync for batch trajectory writes (batch_runner.py):
   Trajectory entries were written without flush/fsync, but the checkpoint
   immediately marked them as completed. A crash between write and disk
   sync would leave the checkpoint claiming completion with no trajectory
   data on disk. Added flush() + os.fsync() before checkpoint update.

3. Pool cleanup on interruption (batch_runner.py):
   Ctrl+C during pool.imap_unordered() relied on context manager cleanup
   which can hang. Added explicit pool.terminate() + pool.join() for both
   KeyboardInterrupt and Exception paths. The original PR used
   pool.join(timeout=10) which is invalid — CPython's Pool.join() takes
   no timeout parameter. Fixed to use pool.join() without arguments.

Tests:
  - test_warning_emitted_on_fallback: verifies warning fires at step 9
  - test_no_warning_when_cached: verifies no false warning when cache hits
  - test_trajectory_entry_is_synced_to_disk: verifies os.fsync is called
  - test_pool_terminate_called_on_exception: verifies cleanup on RuntimeError
  - test_pool_terminate_called_on_keyboard_interrupt: verifies cleanup on Ctrl+C
  - test_pool_join_called_without_timeout: verifies no timeout arg to join()
  - test_real_pool_join_accepts_no_timeout: integration check on CPython API

Co-authored-by: Aaron Lab <aaronlab@users.noreply.github.com>
This commit is contained in:
kshitijk4poor
2026-08-01 14:10:29 +05:30
committed by kshitij
parent 34c11fa689
commit a1ff62a139
4 changed files with 258 additions and 1 deletions
+9
View File
@@ -485,6 +485,8 @@ def _process_batch_worker(args: Tuple) -> Dict[str, Any]:
# Append to batch output file
with open(batch_output_file, 'a', encoding='utf-8') as f:
f.write(json.dumps(trajectory_entry, ensure_ascii=False) + "\n")
f.flush()
os.fsync(f.fileno())
# Aggregate tool statistics
for tool_name, stats in result.get("tool_stats", {}).items():
@@ -978,8 +980,15 @@ class BatchRunner:
except Exception as ckpt_err:
# Don't fail the run if checkpoint write fails
print(f"⚠️ Warning: Failed to save incremental checkpoint: {ckpt_err}")
except KeyboardInterrupt:
print("\n⚠️ Interrupted — terminating batch workers...")
pool.terminate()
pool.join()
raise
except Exception as e:
logger.error("Batch worker failed: %s", e, exc_info=True)
pool.terminate()
pool.join()
raise
finally:
root_logger.setLevel(original_level)