瀏覽代碼

fix renaming missprint

Pavel Samygin 3 年之前
父節點
當前提交
45caeefb0d
共有 1 個文件被更改,包括 3 次插入3 次删除
  1. 3 3
      src/server/backend.py

+ 3 - 3
src/server/backend.py

@@ -11,11 +11,11 @@ import torch
 from hivemind import use_hivemind_log_handler
 from hivemind.moe.server.module_backend import ModuleBackend
 from hivemind.moe.server.task_pool import Task, TaskPool
-from hivemind.utils import InvalidStateError, MPFuture, get_logger
+from hivemind.utils import InvalidStateError, get_logger
 
 from src.bloom.from_pretrained import BloomBlock
 from src.server.cache import MemoryCache
-from src.server.task_broker import SimpleBroker, TaskBrokerBase
+from src.server.task_broker import DustBrokerBase, SimpleBroker
 
 use_hivemind_log_handler("in_root_logger")
 logger = get_logger(__file__)
@@ -30,7 +30,7 @@ class PrioritizedTask:
 
 
 class PrioritizedTaskPool(TaskPool):
-    def __init__(self, *args, broker: TaskBrokerBase = SimpleBroker(), **kwargs):
+    def __init__(self, *args, broker: DustBrokerBase = SimpleBroker(), **kwargs):
         super().__init__(*args, **kwargs)
         self.broker = broker
         self.dust_queue = mp.Queue(maxsize=self.tasks.maxsize)