Skip to content

Commit ecabe2c

Browse files
committed
Fix CI test for test_worker.py
1 parent 7cebef6 commit ecabe2c

File tree

2 files changed

+13
-1
lines changed

2 files changed

+13
-1
lines changed

xinference/core/tests/test_worker.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,13 @@ def __init__(
2929
cuda_devices: List[int],
3030
):
3131
super().__init__(supervisor_address, main_pool, cuda_devices)
32+
self._gpu_memory_info = {}
33+
for gpu_idx in cuda_devices:
34+
self._gpu_memory_info[gpu_idx] = {
35+
"total": 24000,
36+
"used": 0,
37+
"available": 24000
38+
}
3239

3340
async def __post_create__(self):
3441
pass

xinference/core/worker.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -538,8 +538,13 @@ def allocate_devices(self, model_uid: str, n_gpu: int) -> List[int]:
538538
remaining_needed = n_gpu - len(selected_devices)
539539

540540
# Get GPUs sorted by available memory (most available first)
541+
# Exclude GPUs that are already allocated by user_specified models
541542
candidate_gpus = [
542-
dev for dev in self._total_gpu_devices if dev not in selected_devices
543+
dev
544+
for dev in self._total_gpu_devices
545+
if dev not in selected_devices
546+
and dev not in self._gpu_to_model_uid
547+
and dev not in user_specified_allocated_devices
543548
]
544549

545550
gpu_memory_list = []

0 commit comments

Comments
 (0)