yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-11-04 12:07:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Hopefully a cleaner pipeline

Browse Source
This commit is contained in:
Jake Poznanski 2025-07-23 16:48:56 +00:00
parent b4c5913772
commit ce32ceb83f
1 changed files with 9 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
olmocr/pipeline.py
View File

@ -669,6 +669,10 @@ async def vllm_server_task(model_name_or_path, args, semaphore):
except asyncio.CancelledError: except asyncio.CancelledError:
logger.info("Got cancellation request for VLLM server") logger.info("Got cancellation request for VLLM server")
proc.terminate() proc.terminate()
try:
await asyncio.wait_for(proc.wait(), timeout=10.0)
except asyncio.TimeoutError:
logger.warning("VLLM server did not terminate within 10 seconds")
raise raise
timeout_task.cancel() timeout_task.cancel()
@ -1209,6 +1213,9 @@ async def main():
vllm_server.cancel() vllm_server.cancel()
metrics_task.cancel() metrics_task.cancel()
# Wait for cancelled tasks to complete
await asyncio.gather(vllm_server, metrics_task, return_exceptions=True)
# Output final metrics summary # Output final metrics summary
metrics_summary = metrics.get_metrics_summary() metrics_summary = metrics.get_metrics_summary()
logger.info("=" * 80) logger.info("=" * 80)
@ -1233,7 +1240,8 @@ async def main():
) )
# Output finished_on_attempt statistics # Output finished_on_attempt statistics
logger.info("\nPages finished by attempt number:") logger.info("")
logger.info("Pages finished by attempt number:")
total_finished = sum(total_metrics.get(f"finished_on_attempt_{i}", 0) for i in range(args.max_page_retries)) total_finished = sum(total_metrics.get(f"finished_on_attempt_{i}", 0) for i in range(args.max_page_retries))
cumulative = 0 cumulative = 0