Rewrite the GPUPool in Rust (#2243) by YaxinCheng · Pull Request #2277 · cocoindex-io/cocoindex

YaxinCheng · 2026-07-14T05:37:48Z

Summary

Rewrite GPUPool in Rust and expose it as PyObj
Support >1 integer values to acquire multiple fully available GPUs
Introduce temp-env dev dependency for

Motivation

Previous pr (#2224) introduced the GPUPool to the system, but it was written in Python. Rewriting it in Rust would enhance the reusability for future Rust SDK and would improve performance for this most intensive logic.

Changes

Refactored: `GPUPool` (`gpu_pool.rs`)

Completely rewrote GPUPool in Python (runner.py) to GPUPool (gpu_pool.rs) allows acquiring and releasing a fraction of GPU resources from and to a GPU.
- acquiring and releasing logics are updated now to be more optimized
Rewrote related helper functions such as detect_num_gpus in Rust, which reads different environment variables to determine the number of GPUs.
Add a new function that accepts an integer value n where n >= 1 and acquires n number of fully available / non-occupied / capacity = 1.0 GPUs. When there are not enough GPUs, no GPU would be acquired until all are available to be acquired

New `PyGPUPool`

To use the GPUPool in Python, a new bridge object PyGPUPool is created. It exposes all the public functions from GPUPool to the Python end.

New temp_env dev dependency for tests

Added temp_env dev dependency to properly test functions that depends on environment variables.

Calling std::env::set_var is unsafe and it messes up the actual environment. For tests, using temp_env is more manageable and safer.

Tests (`gpu_pool.rs`)

22 tests:

Rewrote all python tests for GPUPool and detect_num_gpus in Rust.
New tests testing acquire_full_gpus:
- Acquire when there are enough GPUs
- Acquire when GPUs are not enough, and can only succeed after others release
- No partial acquiring test
- Attempting to acquire more GPUs than available

Since the detect_num_gpus can be calling nvidia-smi external command to acquire number of GPUs, in the test this behaviour is stubbed / mocked using echo to directly return a given value passed in. This stub can be controlled through a few test-only environment variables using temp-env. For example, MOCK_NVIDIA_SMI_STDOUT controls what is returned from the echo command, and MOCK_NVIDIA_SMI_EXIT_CODE controls what exit code is returned from running the command.

Backward Compatibility

The GPUPool is bridged into Python to replace the original Python GPUPool. The callers have been updated to use the new function names if there are any.

badmonster0 · 2026-07-15T04:54:19Z

thanks a lot @YaxinCheng , @georgeh0 can help take a look!

…herd problem

…The current strategy is to find GPUs with the shortest waitlist

…ed in, and take in NonZeroUsize when needed. Affected: * `acquire` * `release` * `new`

georgeh0 · 2026-07-23T03:34:12Z

+        exclude_gpus: &HashSet<usize>,
+    ) -> (usize, oneshot::Receiver<()>) {
+        let (sender, recv) = oneshot::channel();
+        let reserved_gpu = Self::find_shortest_queue(&state.reserved, exclude_gpus);


I think using shortest queue may still end up with sub-optimal scheduling, e.g. if there's task running super long, tasks will be pending in the queue while others are idle (without either running or reserved tasks).

In my mind a better approach will be:

One shared queue for all waiting tasks.

Each GPU is reserved for at most 1 task at any time, as there's no point to reserve a CPU for a second task before the existing one is scheduled.

Fractional (<1 GPU): When capacity opens up on a reserved GPU, the reserved task either absorbs the freed capacity or acquires enough total capacity to be scheduled.

Full GPU(s): Claims intact GPUs. While waiting for remaining GPUs, its already-reserved GPUs are held and not given to others.

Queue processing: When a GPU without a reserved task has capacity, inspect the head of the shared queue. The head task reserves/allocates available GPUs; once all its required GPUs are reserved, it pops off the queue.

Let me know if it makes sense. Thanks!

So if I understand correctly, when a GPU frees up, it goes to the shared queue to get the next task. What will happen if the capacity of this GPU is not enough to handle the requirement of the next task? Say, the GPU has a capacity of 0.6 after releasing, but the top task in the queue required 0.7? Does it just wait until the top task is handled by another GPU, then check the top of the queue again?

IIUC, acquire would then have two behaviours:

if we can find a GPU with enough capacity, we assign to the GPU and deduct the capacity

if not, we enqueue the task to the shared queue and wait

In this case, for release, let's represent it in pseudocode:

fn release(capacity, gpu) { capacities[gpu] += capacity loop { if !has_reservation(gpu) && !task_queue.is_empty() { add_reservation(gpu, task_queue.pop_left()) } if let Some(reserved_task) = get_reservation(gpu) { if reversed_task.demand_capacity <= capacities[gpu] { capacities[gpu] -= reversed_task.demand_capacity remove_reservation(gpu) reserved_task.notify_processing() } } if has_reservation(gpu) || task_queue.is_empty() { break } } }

is this what you meant?

The pseudocode almost aligns with the direction in my mind (and it already answers your question in the previous comment), with one revision: if the task at the front of the queue requires multiple GPUs, it's only popped out of the queue once all required GPUs are reserved for it.

e.g. if there're 3 GPUs (G1, G2, G3) all without available capacity, 2 tasks in the queue, T1 requires 2 GPU and T2 requires 0.5 GPU

GPUs: G1 (capacity=0), G2(capacity=0), G3(capacity=0) Queue: T1(req=2) T2(req=0.5)

Now G1 is released with 0.5 capacity, we can reserve it for T1, but T1 remains at the front of the queue:

GPUs: G1 (capacity=0.5, reserved_for=T1), G2(capacity=0), G3(capacity=0) Queue: T1(req=2, reserved=[G1]) T2(req=0.5)

Later, G2 is released with 0.5 capacity, we can reserve it for T1, now T2 is reserved with 2 GPUs, so it can be popped from the queue front

GPUs: G1 (capacity=0.5, reserved_for=T1), G2(capacity=0.5, reserved_for=T1), G3(capacity=0) Queue: T2(req=0.5) Fully reserved tasks: T1(req=2, reserved=[G1, G2])

From now on, if any available capacity is released for G3, it can be reserved for T2.

WDYT?

Hi @georgeh0 , thanks for the explanation.

I have a draft implementation on this, but I feel like there are a few more scenarios that need clarifications.

Should we randomly reserve or keep the task in the queue?

GPUs: G1 (cap=0.5), G2 (cap=0.6) Queue: T1 (req=0.7)

My concern is, if we reserve, say G1, then maybe G2 will be freed up next and we cannot change it

Following case 1, if we do not reserve, and a new task comes in:

GPUs: G1 (cap=0.5), G2 (cap=0.6) Queue: T1 (req=0.7), T2 (req = 0.1)

this will get stuck

Which path should we pick?

Hi @YaxinCheng ,

Good question.

I think we may just reserve either G1 or G2 for T1. Probably G2 is better as the available cap is higher. Yes, if G1 is freed up sooner, G1 will be idle which is a waste. But since we only reserve a GPU for at most 1 task at a time, the waste is bounded. So I think it's acceptable. If we go with 2 (i.e. not reserve), there'll be risk of starvation, which is worse.

Besides, there's an alternative way in my mind now: instead of reserving a fixed GPU determined at reserving time, we may reserve N GPUs with most available capacity, dynamically.

In your example, when T1 comes, we reserve "the GPU with most available cap" for it, which may change over time. Later, if G1 is freed up, now G1 has most capacity, so it's reserved one, and T1 can be assigned to it.

If multiple tasks have reserved GPUs, e.g. T1 and T2 each has 1 reserved GPU, if T1 comes earlier, then the GPU with most cap is reserved for T1, the GPU with the 2nd most cap is reserved for T2, etc.

This will make sure each task always get more available GPUs than tasks coming later, but later tasks still get a chance to be scheduled earlier if there're extra GPUs that fits it.

This alternative way might be slightly better, but may be more complex to implement.

WDYT?

If multiple tasks have reserved GPUs, e.g. T1 and T2 each has 1 reserved GPU, if T1 comes earlier, then the GPU with most cap is reserved for T1, the GPU with the 2nd most cap is reserved for T2, etc.

This will have another issue when T1 requires more resources than T2, and the most available one has enough for T2 but not T1. Due to the ordering, T2 will have to wait.

Or in a different case,

GPUs: G1 (cap=0.4, reserved=T1), G2 (cap = 0.3, reserved=T2) Tasks: T1 (req=0.6), T2 (req=0.4)

If we release 0.2 to G2, it will become 0.5, which is enough for T2 but not T1. However, since T1 reserves the most available one, now it will have G2, and T2 will have G1

GPUs: G1 (cap=0.4, reserved=T2), G2 (cap = 0.5, reserved=T1) Tasks: T1 (req=0.6), T2 (req=0.4)

In this case, the two tasks will still be stuck. Is this actually okay?

Hi @georgeh0 , the current version does the basic acquiring and releasing:

when it acquires (including acquire_full), it acquires, or reserves, or get put into the queue

when it releases, it will add back the capacity, then keep popping the head of the queue if it can be fulfilled or reserved

I am thinking maybe we can have a separate PR for the ordering solution. But before that, please have a look at this version and let me know. Thank you

Hi @georgeh0 , I was able to update the reserving logic to the flex strategy we discussed. The first reserved task will use the GPU with the most capacity, even when the capacities change, the most available will be reserved for this task. Please have a look, I have reduced the logic, so it should be slightly easier to review. Thank you

I have taken a look at the 1200 lines of code again, most of them are test related code or comments. Sorry for bloating up the CL, but it may be the most concise it can be at this stage. Thanks for taking the time to review it

…just >= the required amount

…e boundary check

…just >= the required amount

georgeh0

Thanks! Haven't gone through all code yet. Sent out what I have now.

georgeh0 · 2026-07-31T06:35:41Z

+            return Err(anyhow::anyhow!(
+                "Acquired fraction must be between 0.0 and 1.0, got 0"
+            ));


Let's use error types and helper functions/macros defined in https://github.com/cocoindex-io/cocoindex/blob/main/rust/utils/src/error.rs

I did not realize the error exists. I have switched. Let me know if this is not the right way to use the macros

georgeh0 · 2026-07-31T06:36:53Z

+        }
+        let receiver = {
+            let mut task_queue = self.task_queue.lock().await;
+            if task_queue.len() < self.num_gpus {


This condition seems not taking tasks requiring multiple GPUs into consideration.

For tasks taking multiple GPUs, we will have multiple tasks consecutively. So this should be fine

* group pool and capacities into a single struct avoid lock acquiring order limitation * create a struct for PendingTask for better readabilities * switch to std::sync::Mutex * use error kind

georgeh0 · 2026-08-02T05:04:03Z

+    ///
+    fn process_task_queue(pool: &mut PoolState) {
+        while let Some(pending_task) = pool.task_queue.front()
+            && let Some(gpu_id) = pool.capacities.find(&pending_task.demand)


Here seems not strictly FIFO.

If the 1st task in the queue requires 1 but the GPU with maximum capacity only has 0.5, and the 2nd one only requires 0.5, this will assign the GPU to the second one. This may end up with starvation for the 1st task.

If the first task is in queue, say it requires 0.6, then it reserves the one with the max capacity, say 0.5. Every other GPU should have a capacity less than or equal to 0.5 in this case. The second task comes in requires 0.5, and it will only be assigned if there is another GPU with 0.5 capacity (the max one is reserved by the first task in queue), or it will be put into the queue and wait for the first task to be processed. This should be fine, right?

Sorry, looking at the code again, I misunderstood that. There won't be starvation.

But it has another problem, e.g.

Available GPUs: 0.5 0.4

Tasks demad: 0.6, 0.3

In this case, we reserve the 1st GPU for the 1st task, and the 2nd GPU can be assigned to the 2nd task.

Note that "blocking others until the first is scheduled" is a simple and feasible solution. And this is the original solution we discussed earlier this this thread. But its implementation can be much simpler: just a single queue without a reservation system. Since we're already using a reservation system, we should achieve its benefit: it reserves a dedicated GPU to the earlier task, at the same time, allocations for the remaining GPUs are not blocked by it.

I am sorry, I don't get this. In the example you provided, the second task (0.3) won't even be in the queue. Because when the task comes in, the situation is:

GPUs: G1 (0.5, reserved), G2 (0.4) Queue: T1 (0.6)

The task 2 shows up, T2(0.3), and it will directly be assigned, and won't even reach the queue. Everything in the queue is strictly Fifo. Even if we managed to work around and put the T2 into the queue, then we have

GPUs: G1 (0.5, reserved), G2 (0.4) Queue: T1 (0.6), T2 (0.3)

In this case, the queue is frozen. Even if T2 can technically be handled by G2, it won't. Because in the code, the process_acquisition_queue strictly follows FIFO, and it won't work on T2 until T1 is done.

As mentioned here:

instead of reserving a fixed GPU determined at reserving time, we may reserve N GPUs with most available capacity, dynamically
... when T1 comes, we reserve "the GPU with most available cap" for it, which may change over time. Later, if G1 is freed up, now G1 has most capacity, so it's reserved one, and T1 can be assigned to it...

What is the issue that we need to fix here? Are you saying, if T1 reserves G1, and we cannot assign any task until T1 is finished, even though there are enough GPUs available to process T2?

In the example you provided, the second task (0.3) won't even be in the queue. Because when the task comes in, the situation is:

GPUs: G1 (0.5, reserved), G2 (0.4) Queue: T1 (0.6)

The task 2 shows up, T2(0.3), and it will directly be assigned, and won't even reach the queue.

e.g. when T2 shows up, G2 may have less capacity (e.g. 0.2) than its demand, so T2 will be inserted into the queue:

GPUs: G1 (0.5, reserved), G2 (0.2, reserved) Queue: T1 (0.6), T2 (0.3)

Later, G2 has another 0.2 capacity becoming available, now we have,

GPUs: G1 (0.5, reserved), G2 (0.4, reserved) Queue: T1 (0.6), T2 (0.3)

Please let me know if it's not how the current code is working.

What is the issue that we need to fix here?

The problem we want to solve is: T1 already has its reserved GPU, so it shouldn't prevent other GPUs to be assigned to other tasks. This is holding more than the necessary GPU(s) for a task, and will cause inefficiency.

In this case, the queue is frozen. Even if T2 can technically be handled by G2, it won't. Because in the code, the process_acquisition_queue strictly follows FIFO, and it won't work on T2 until T1 is done.

As mentioned here:

instead of reserving a fixed GPU determined at reserving time, we may reserve N GPUs with most available capacity, dynamically
... when T1 comes, we reserve "the GPU with most available cap" for it, which may change over time. Later, if G1 is freed up, now G1 has most capacity, so it's reserved one, and T1 can be assigned to it...

To clarify, the problem I'm trying to point out in my latest comment is not about FIFO unfollowed (this was what my first comment is saying, which was misleading and based on my misunderstanding for potential starvation; I corrected in my next comment).

FIFO should be applied to tasks before they're either reserved or assigned. Once they already get their reserved GPU, later tasks get a chance to be reserved or assigned, no more FIFO here. If we follow "strict FIFO" even for reserved tasks, we shouldn't even directly assign non-reserved GPU to new task when it already has enough capacity, right? We don't have to be subject to "strict FIFO across all tasks including reserved ones".

Also the paragraph quoted above ("... when T1 comes ...") is just using a simple example to explain how dynamic adjusted reservation works for one task. It says nothing about freezing later GPUs/tasks in the queue.

e.g. when T2 shows up, G2 may have less capacity (e.g. 0.2) than its demand, so T2 will be inserted into the queue:

GPUs: G1 (0.5, reserved), G2 (0.2, reserved) Queue: T1 (0.6), T2 (0.3)

Later, G2 has another 0.2 capacity becoming available, now we have,

GPUs: G1 (0.5, reserved), G2 (0.4, reserved) Queue: T1 (0.6), T2 (0.5)

Yes, this is correct. It will be like this. But I think we previously agreed that the queue should be strict FIFO. In this case, G2 should not be handled, because it's blocked by G1. Do we want the G2 to be handled in this case instead?

Strict FIFO is one acceptable option I discussed in this early thread, but the implementation should be much simpler and don't need the "reserve" mechanism.

Since we already introduce the "reserve" mechanism, it's no longer strict FIFO. As discussed in that thread, the main motivation of introducing the "reserve" mechanism is to avoid the head task block others when there're more available GPUs. If we want to go with strict FIFO, "reserve" is not needed at all.

I see. I think I have misunderstood it. I have just pushed a commit, which should allow T2 in this case gets assigned to G2. The queue is no longer strictly FIFO. PHAL

@georgeh0 can you please have a look and let me know if there are any other issues? Thank you

* Optimize acquire_full * Change term `task` to `acquisition`

georgeh0 · 2026-08-03T06:45:07Z

+    ///
+    fn process_task_queue(pool: &mut PoolState) {
+        while let Some(pending_task) = pool.task_queue.front()
+            && let Some(gpu_id) = pool.capacities.find(&pending_task.demand)


Sorry, looking at the code again, I misunderstood that. There won't be starvation.

But it has another problem, e.g.

Available GPUs: 0.5 0.4

Tasks demad: 0.6, 0.3

In this case, we reserve the 1st GPU for the 1st task, and the 2nd GPU can be assigned to the 2nd task.

Note that "blocking others until the first is scheduled" is a simple and feasible solution. And this is the original solution we discussed earlier this this thread. But its implementation can be much simpler: just a single queue without a reservation system. Since we're already using a reservation system, we should achieve its benefit: it reserves a dedicated GPU to the earlier task, at the same time, allocations for the remaining GPUs are not blocked by it.

georgeh0 · 2026-08-21T06:49:38Z

+            .collect::<Vec<_>>();
+            (acquired_gpus, receivers)
+        };
+        let gpu_ids = futures::future::try_join_all(receivers).await?;


We need to consider cancellation safety. e.g. if the coroutine is cancelled during this await in the middle, we may leak some GPU capacities: they may never be refilled. We may consider add a drop guard.

Updated the code to add an error handling branch. Now it should release the already acquired GPUs when the waiting is cancelled. PHAL

YaxinCheng added 6 commits July 12, 2026 00:41

Add GPUPool in the core rust code

829279b

Switch python to use rust GPUPool

5365787

Updated unit tests for GPUPool

a05e77e

Add function to acquire integer number of gpus

bba983b

Add back comments

953cf0d

Add clarification on the acquire full gpus

89c4646

badmonster0 requested a review from georgeh0 July 15, 2026 04:54

georgeh0 reviewed Jul 15, 2026

View reviewed changes

Comment thread rust/utils/src/gpu_pool.rs Outdated

Comment thread rust/utils/src/gpu_pool.rs Outdated

YaxinCheng added 2 commits July 15, 2026 17:30

release.notify_one() instead of notify_all() to avoid the thundering …

527bc4f

…herd problem

Expose acquire_full to python

4d0dbab

YaxinCheng force-pushed the gpu_pool branch from 48d2834 to 4d0dbab Compare July 16, 2026 04:02

georgeh0 reviewed Jul 16, 2026

View reviewed changes

Comment thread rust/utils/src/gpu_pool.rs Outdated

Comment thread rust/utils/src/gpu_pool.rs Outdated

YaxinCheng added 7 commits July 16, 2026 19:21

acquire_full returns error instead of crash

2f85c53

Updated the comment

b231cf7

Explicit waiting queue for GPUPool

95422da

Add a Waiter struct for readability

73e4e51

Introduce GPU reservation for GPUs that are not currently available. …

81af87f

…The current strategy is to find GPUs with the shortest waitlist

Use try_join_all to better handle error

ed40a2d

Update the functions to return errors when unexpected values are pass…

d8a0c04

…ed in, and take in NonZeroUsize when needed. Affected: * `acquire` * `release` * `new`

georgeh0 reviewed Jul 23, 2026

View reviewed changes

YaxinCheng added 8 commits July 23, 2026 17:09

find_available no longer finds the max, but finds the amount that is …

e6eb8d4

…just >= the required amount

Introduce GPUFraction type which handles f32 to usize, and handles th…

567c2ad

…e boundary check

find_available no longer finds the max, but finds the amount that is …

70cd5d2

…just >= the required amount

Add a few traits to GPUFraction, and also make SCALE private

9043f87

Revamped queuing and handling strategy

4379ee0

Code cleanup to improve readability

1c2c325

Updated scheduling and reserving strategy

808f869

Add comments and improve code readability

dd3ba9f

georgeh0 reviewed Jul 31, 2026

View reviewed changes

YaxinCheng added 2 commits July 31, 2026 19:08

Fix code based on comments

6e7f512

* group pool and capacities into a single struct avoid lock acquiring order limitation * create a struct for PendingTask for better readabilities * switch to std::sync::Mutex * use error kind

releas is now sync

a7acfd0

georgeh0 reviewed Aug 2, 2026

View reviewed changes

Comment thread rust/utils/src/gpu_pool.rs Outdated

YaxinCheng added 2 commits August 2, 2026 19:27

Fix code based on comments

f7b088d

* Optimize acquire_full * Change term `task` to `acquisition`

Clean up

95c229e

georgeh0 reviewed Aug 3, 2026

View reviewed changes

YaxinCheng added 5 commits August 3, 2026 19:32

Merge find excluding

eeba17d

Updated tests

b574039

Update function names

2251496

The queue is no longer strict FIFO

53365a1

Use the right length

54031e6

georgeh0 reviewed Aug 21, 2026

View reviewed changes

Release acquired GPUs when reservation is cancelled

ee61291

Uh oh!

Conversation

YaxinCheng commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Changes

Refactored: GPUPool (gpu_pool.rs)

New PyGPUPool

New temp_env dev dependency for tests

Tests (gpu_pool.rs)

Backward Compatibility

Uh oh!

badmonster0 commented Jul 15, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

georgeh0 Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

georgeh0 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

georgeh0 Aug 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YaxinCheng Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

YaxinCheng commented Jul 14, 2026 •

edited

Loading

Refactored: `GPUPool` (`gpu_pool.rs`)

New `PyGPUPool`

Tests (`gpu_pool.rs`)

georgeh0 Jul 23, 2026 •

edited

Loading

georgeh0 Aug 2, 2026 •

edited

Loading

YaxinCheng Aug 4, 2026 •

edited

Loading