mirror of
https://github.com/bigscience-workshop/petals
synced 2024-10-31 09:20:41 +00:00
523a7cad33
1. Added `from petals.client import *` to `petals/__init__.py`, so you can write just that:
```python
from petals import DistributedBloomForCausalLM
```
I didn't do the same with server, since its classes are supposed to by used by `petals.cli.run_server`, not end-users. Though it's still possible to do `from petals.server.smth import smth` if necessary.
2. Fixed one more logging issue: log lines from hivemind were shown twice due to a bug in #156.
3. Removed unused `runtime.py`, since the server actually uses `hivemind.moe.Runtime`, and `runtime.py` has no significant changes comparing to it.
72 lines
2.8 KiB
Python
72 lines
2.8 KiB
Python
import multiprocessing as mp
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
from hivemind.moe.server.runtime import Runtime
|
|
|
|
from petals.server.task_pool import PrioritizedTaskPool
|
|
|
|
|
|
@pytest.mark.forked
|
|
def test_priority_pools():
|
|
outputs_queue = mp.SimpleQueue()
|
|
results_valid = mp.Event()
|
|
|
|
def dummy_pool_func(x):
|
|
time.sleep(0.1)
|
|
y = x**2
|
|
outputs_queue.put((x, y))
|
|
return (y,)
|
|
|
|
class DummyBackend:
|
|
def __init__(self, pools):
|
|
self.pools = pools
|
|
|
|
def get_pools(self):
|
|
return self.pools
|
|
|
|
pools = (
|
|
PrioritizedTaskPool(dummy_pool_func, name="A", max_batch_size=1),
|
|
PrioritizedTaskPool(dummy_pool_func, name="B", max_batch_size=1),
|
|
)
|
|
|
|
runtime = Runtime({str(i): DummyBackend([pool]) for i, pool in enumerate(pools)}, prefetch_batches=0)
|
|
runtime.start()
|
|
|
|
def process_tasks():
|
|
futures = []
|
|
futures.append(pools[0].submit_task(torch.tensor([0]), priority=1))
|
|
futures.append(pools[0].submit_task(torch.tensor([1]), priority=1))
|
|
time.sleep(0.01)
|
|
futures.append(pools[1].submit_task(torch.tensor([2]), priority=1))
|
|
futures.append(pools[0].submit_task(torch.tensor([3]), priority=2))
|
|
futures.append(pools[0].submit_task(torch.tensor([4]), priority=10))
|
|
futures.append(pools[0].submit_task(torch.tensor([5]), priority=0))
|
|
futures.append(pools[0].submit_task(torch.tensor([6]), priority=1))
|
|
futures.append(pools[1].submit_task(torch.tensor([7]), priority=11))
|
|
futures.append(pools[1].submit_task(torch.tensor([8]), priority=1))
|
|
for i, f in enumerate(futures):
|
|
assert f.result()[0].item() == i**2
|
|
results_valid.set()
|
|
|
|
proc = mp.Process(target=process_tasks)
|
|
proc.start()
|
|
proc.join()
|
|
assert results_valid.is_set()
|
|
|
|
ordered_outputs = []
|
|
while not outputs_queue.empty():
|
|
ordered_outputs.append(outputs_queue.get()[0].item())
|
|
|
|
assert ordered_outputs == [0, 5, 1, 2, 6, 8, 3, 4, 7]
|
|
# 0 - first batch is loaded immediately, before everything else
|
|
# 5 - highest priority task overall
|
|
# 1 - first of several tasks with equal lowest priority (1)
|
|
# 2 - second earliest task with priority 1, fetched from pool B
|
|
# 6 - third earliest task with priority 1, fetched from pool A again
|
|
# 8 - last priority-1 task, pool B
|
|
# 3 - task with priority 2 from pool A
|
|
# 4 - task with priority 10 from pool A
|
|
# 7 - task with priority 11 from pool B
|