NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-19 12:57:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix infinite loop in wait_for_sem by updating semaphore values inside loop

Browse source
This commit is contained in:
Alexey Gorbatovski 2025-07-06 00:27:45 +03:00
parent 2c8340bece
commit 35c542328a
1 changed files with 24 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

43
atroposlib/envs/server_handling/server_manager.py
View file

@ -141,32 +141,37 @@ class ServerManager:
for server in self.servers:
await server.update_weight(weight)
async def wait_for_sem(self, is_training):
async def wait_for_sem(self, is_training: bool):
"""
Wait for a server to be available. This is used to prevent the client from
overwhelming the server with requests.
"""
if is_training:
eval_vals = [
(
def get_available_slots():
if is_training:
eval_vals = [
(
max(0, server.eval_sem._value - server.eval_sem.min_val())
if server.eval_sem._value != server.eval_sem.max_val
else 0
)
for server in self.servers
]
return [
max(0, (server.sem._value - server.sem.min_val()) - eval_val)
for server, eval_val in zip(self.servers, eval_vals)
]
else:
return [
max(0, server.eval_sem._value - server.eval_sem.min_val())
if server.eval_sem._value != server.eval_sem.max_val
else 0
)
for server in self.servers
]
sem_vals = [
max(0, (server.sem._value - server.sem.min_val()) - eval_val)
for server, eval_val in zip(self.servers, eval_vals)
]
else:
sem_vals = [
max(0, server.eval_sem._value - server.eval_sem.min_val())
for server in self.servers
]
while all([sem_val <= 0 for sem_val in sem_vals]):
for server in self.servers
]
sem_vals = get_available_slots()
while all(sem_val <= 0 for sem_val in sem_vals):
# None available... wait
await asyncio.sleep(1)
sem_vals = get_available_slots()
async def chat_completion(self, **kwargs) -> ChatCompletion:
n = kwargs.get("n", 1)