Repository: allenai/olmocr
File: olmocr/pipeline.py (Python)

Code before:
        logger.info(f"Calculated items_per_group: {items_per_group} based on average pages per PDF: {avg_pages_per_pdf:.2f}")

        # Now call populate_queue
        await work_queue.populate_queue(pdf_work_paths, items_per_group)

    if args.stats:
        print_stats(args)
        return

    if args.beaker:
        submit_beaker_job(args)
        return

    # If you get this far, then you are doing inference and need a GPU
    check_sglang_version()
    check_torch_gpu_available()

    logger.info(f"Starting pipeline with PID {os.getpid()}")

    # Download the model before you do anything else
    await download_model(args.model)

    # Initialize the work queue
    qsize = await work_queue.initialize_queue()

    if qsize > 0:
        # Create a semaphore to control worker access
        # We only allow one worker to move forward with requests, until the server has no more requests in its queue
        # This lets us get full utilization by having many workers, but also to be outputting dolma docs as soon as possible
        # As soon as one worker is no longer saturating the gpu, the next one can start sending requests
        semaphore = asyncio.Semaphore(1)

        sglang_server = asyncio.create_task(sglang_server_host(args, semaphore))

        await sglang_server_ready()

        metrics_task = asyncio.create_task(metrics_reporter(work_queue))

        # Create worker tasks to process the queue concurrently.
        worker_tasks = []
        for i in range(args.workers):
            task = asyncio.create_task(worker(args, work_queue, semaphore, worker_id=i))
            worker_tasks.append(task)

        # Wait for all worker tasks to finish
        await asyncio.gather(*worker_tasks)

        # Wait for server to stop
        process_pool.shutdown(wait=False)

        sglang_server.cancel()

Review comment (suggestion):
Looks good, instead of indenting all this code, can you please instead just check:

```
if qsize == 0:
 logger.info("No work to do, exiting")
 return
```

Code after:
        logger.info(f"Calculated items_per_group: {items_per_group} based on average pages per PDF: {avg_pages_per_pdf:.2f}")

        # Now call populate_queue
        await work_queue.populate_queue(pdf_work_paths, items_per_group)

    if args.stats:
        print_stats(args)
        return

    if args.beaker:
        submit_beaker_job(args)
        return

    # If you get this far, then you are doing inference and need a GPU
    check_sglang_version()
    check_torch_gpu_available()

    logger.info(f"Starting pipeline with PID {os.getpid()}")

    # Download the model before you do anything else
    await download_model(args.model)

    # Initialize the work queue
    qsize = await work_queue.initialize_queue()

    if qsize == 0:
        logger.info("No work to do, exiting")
        return
    # Create a semaphore to control worker access
    # We only allow one worker to move forward with requests, until the server has no more requests in its queue
    # This lets us get full utilization by having many workers, but also to be outputting dolma docs as soon as possible
    # As soon as one worker is no longer saturating the gpu, the next one can start sending requests
    semaphore = asyncio.Semaphore(1)

    sglang_server = asyncio.create_task(sglang_server_host(args, semaphore))

    await sglang_server_ready()

    metrics_task = asyncio.create_task(metrics_reporter(work_queue))

    # Create worker tasks to process the queue concurrently.
    worker_tasks = []
    for i in range(args.workers):
        task = asyncio.create_task(worker(args, work_queue, semaphore, worker_id=i))
        worker_tasks.append(task)

    # Wait for all worker tasks to finish
    await asyncio.gather(*worker_tasks)

    # Wait for server to stop
    process_pool.shutdown(wait=False)