fix(ai): update ai-video selection suspension

[ad-astra-video]

What does this pull request do? Explain your changes. (required)

Suspension was not working because the penalty was always 3. This logic was a carryover from transcoding where the suspender always started at a refresh count of 0 because a new session manager was created with each stream. For AI, we are reusing the session manager and the suspender so the refresh count does not reset between requests. The fix to suspension is to consider the current refresh count when calculating the penalty so it is 3 more than the current refresh count in the suspender.

There was also an issue where the discoveryPoolSize was always 100 and with limited orchestrators providing models a refresh of sessions was being done with every request. I added an initialPoolSize field to track the last refresh pool size to use with the shouldRefreshSessions logic rather than 100. This stabilizes the suspender to allow more orchestrators to be tried with each Select call.

Last update was moving the signalRefresh() for the suspender that increments the refresh counter in the suspender to the Refresh function makes it more stable that every time we refresh sessions we add to the suspender refresh count

Happy to segregate some of these changes to separate PRs. The suspension fixes can be added separately without dependency on ai-worker PR.

Specific updates (required)

• Updates suspender to use the current refresh count of the suspender in the selector.
• Moves penalty to the AISessionSelector to make it easier to update and available for calculations on the suspension needed
• releases all Os when there are none in the warm and cold pool
• Adds option to not use managed containers.

How did you test each of these updates (required)

I have been running these updates on my gateway. Tested 1-200 requests with 5-10 workers sending to gateway. All completed with 1-2 orchestrators providing Bytedance model.

Does this pull request close any open issues?

Checklist:

• Read the contribution guide
• make runs successfully
• All tests in ./test.sh pass
• README and other documentation updated
• Pending changelog updated

[victorges]

Feel like I don't have context to officially approve this, but left some comments. Only nits tho, the implementation makes sense for the PR description.

[leszko]

I'm a little lost with this suspension logic. So, I see that:

• the pool.suspender.count is increased every time signalRefresh() is called
• pool.penalty is always set to 3
• last_count is always set to suspender.count + 3

So, that logic would mean that we're not taking the suspended orchestrator until 3 times the signalRefresh() is called. Is this the idea of this suspension mechanism? That we don't allow the given O to get selected in the 3 refresh sessions?

[ad-astra-video]

Yes thats right. A small penalty to take the orchestrator out of the working set for short period. With transcoding the refresh count starts at 0 for each new session. For AI the session pools are used across requests and are not reset. If we don't add the current suspender refresh count it wont suspend any orchestrators.

[leszko]

release all orchestrators

shouldn't we then just remove them from the suspender.list() rather than calling signalRefresh()? My understanding is that if penalty = 3, then we would need to call signalRefresh() 3 times in order to "release all orchestrators from suspension".

[ad-astra-video]

I did it this way thinking that there could be more than 3 orchestrators suspended so it would be less loops to just signalRefresh() 3 times. An alternative would be to just create a new suspender for the selector to clear it or kick all the orchs out of the suspended list (will require new function to do second option).

[rickstaa]

@leszko Are we still planning to merge these fixes? They help alleviate some of the orchestrator stickiness issues we've observed over the past few months.

[leszko]

@leszko Are we still planning to merge these fixes? They help alleviate some of the orchestrator stickiness issues we've observed over the past few months.

Well...if it helps Orchestrators, then yes. Why not? 🙃 Please address the PR review comments and re-request review :)

[ad-astra-video]

@rickstaa i have updated this PR.

I also think adding something like this would be beneficial for Gateway operators to see the current AI session pool state. WYDT? ad-astra-video@9d42251

EDIT: we included the /getAISessionsPoolInfo endpoint on the gateway cli webserver to provide data about AI pools.

[codecov]

Codecov Report

Attention: Patch coverage is 15.62500% with 81 lines in your changes missing coverage. Please review.

Project coverage is 32.15960%. Comparing base (390af43) to head (af7434f).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
server/handlers.go	2.27273%	42 Missing and 1 partial ⚠️
server/ai_session.go	0.00000%	35 Missing ⚠️
server/ai_process.go	81.25000%	3 Missing ⚠️

Additional details and impacted files

@@                 Coverage Diff                 @@
##              master       #3033         +/-   ##
===================================================
- Coverage   32.18835%   32.15960%   -0.02875%     
===================================================
  Files            147         147                 
  Lines          40670       40753         +83     
===================================================
+ Hits           13091       13106         +15     
- Misses         26807       26874         +67     
- Partials         772         773          +1     

Files with missing lines	Coverage Δ
server/webserver.go	95.87629% <100.00000%> (+0.04296%)	⬆️
server/ai_process.go	1.66945% <81.25000%> (+1.07723%)	⬆️
server/ai_session.go	2.33333% <0.00000%> (-0.18466%)	⬇️
server/handlers.go	52.19092% <2.27273%> (-1.77991%)	⬇️

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 390af43...af7434f. Read the comment docs.

Files with missing lines	Coverage Δ
server/webserver.go	95.87629% <100.00000%> (+0.04296%)	⬆️
server/ai_process.go	1.66945% <81.25000%> (+1.07723%)	⬆️
server/ai_session.go	2.33333% <0.00000%> (-0.18466%)	⬇️
server/handlers.go	52.19092% <2.27273%> (-1.77991%)	⬇️

[ad-astra-video]

@rickstaa this is ready for review. Titan did some testing and the suspension was working as expected.

pool_info_1000.json

pool_info_1_orch.json

[leszko]

Added two comments. My suggestion is also to maybe split this PR to make it smaller, because I think it includes some unrelated changes.

1. Wrt reserving capacity inside CheckAICapacity(), it looks counter-intuitive, but I may not have the full context. It raises a lot of questions in my head, like what if we reserved the capacity, but the request failed at some point. Won't O stop working? Maybe we need to have some timeout to release it?
2. Wrt: penalty mechanism, any chance it's possible to send this change in its own separate PR (or is it related to reserve capacity change?).

[rickstaa]

Added two comments. My suggestion is also to maybe split this PR to make it smaller, because I think it includes some unrelated changes.

1. Wrt reserving capacity inside CheckAICapacity(), it looks counter-intuitive, but I may not have the full context. It raises a lot of questions in my head, like what if we reserved the capacity, but the request failed at some point. Won't O stop working? Maybe we need to have some timeout to release it?
2. Wrt: penalty mechanism, any chance it's possible to send this change in its own separate PR (or is it related to reserve capacity change?).

@ad-astra-video I agree with @leszko if we can have the suspense fix in a seperate pull request we can move quicker since that part is more battle tested and reviewd 🙏🏻.

[ad-astra-video]

I have split the PR. If want split another way please tell me specifically which parts want in separate PRs. To me, all the remaining changes are needed to stablize orchestrator pools for batch jobs so I have left in this PR.

I retested all batch pipelines after splitting PR.

[leszko]

I have split the PR. If want split another way please tell me specifically which parts want in separate PRs. To me, all the remaining changes are needed to stablize orchestrator pools for batch jobs so I have left in this PR.

I retested all batch pipelines after splitting PR.

Where is the other PR? Could you share the link?

[leszko]

LGTM

[leszko]

@rickstaa do you want to review or should I merge?

[rickstaa]

@ad-astra-video, I made some small code improvements in ad-astra-video#32. Once those are merged, we're good to go. I thoroughly battle-tested it, trying to break it, but couldn't. 🚀