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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user