mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-12 07:10:52 +08:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import concurrent.futures
|
|
from typing import Callable
|
|
|
|
from pebble import ProcessPool, ProcessFuture
|
|
|
|
from ..component_model.executor_types import Executor, ExecutePromptArgs
|
|
|
|
|
|
class ProcessPoolExecutor(ProcessPool, Executor):
|
|
def __init__(self, max_workers: int = 1):
|
|
super().__init__(max_workers=1)
|
|
|
|
|
|
def shutdown(self, wait=True, *, cancel_futures=False):
|
|
if cancel_futures:
|
|
raise NotImplementedError("cannot cancel futures in this implementation")
|
|
if wait:
|
|
self.close()
|
|
else:
|
|
self.stop()
|
|
return
|
|
|
|
def schedule(self, function: Callable,
|
|
args: list = (),
|
|
kwargs=None,
|
|
timeout: float = None) -> ProcessFuture:
|
|
# todo: restart worker when there is insufficient VRAM or the workflows are sufficiently different
|
|
# try:
|
|
# args: ExecutePromptArgs
|
|
# prompt, prompt_id, client_id, span_context, progress_handler, configuration = args
|
|
#
|
|
# except ValueError:
|
|
# pass
|
|
if kwargs is None:
|
|
kwargs = {}
|
|
return super().schedule(function, args, kwargs, timeout)
|
|
|
|
def submit(self, fn, /, *args, **kwargs) -> concurrent.futures.Future:
|
|
return self.schedule(fn, args=list(args), kwargs=kwargs, timeout=None)
|