forked from extern/easydiffusion
493 lines
19 KiB
Python
493 lines
19 KiB
Python
"""task_manager.py: manage tasks dispatching and render threads.
|
|
Notes:
|
|
render_threads should be the only hard reference held by the manager to the threads.
|
|
Use weak_thread_data to store all other data using weak keys.
|
|
This will allow for garbage collection after the thread dies.
|
|
"""
|
|
import json
|
|
import traceback
|
|
import logging
|
|
|
|
TASK_TTL = 15 * 60 # seconds, Discard last session's task timeout
|
|
|
|
import torch
|
|
import queue, threading, time, weakref
|
|
from typing import Any, Hashable
|
|
|
|
from sd_internal import TaskData, device_manager
|
|
from diffusionkit.types import GenerateImageRequest
|
|
|
|
log = logging.getLogger()
|
|
|
|
THREAD_NAME_PREFIX = ''
|
|
ERR_LOCK_FAILED = ' failed to acquire lock within timeout.'
|
|
LOCK_TIMEOUT = 15 # Maximum locking time in seconds before failing a task.
|
|
# It's better to get an exception than a deadlock... ALWAYS use timeout in critical paths.
|
|
|
|
DEVICE_START_TIMEOUT = 60 # seconds - Maximum time to wait for a render device to init.
|
|
|
|
class SymbolClass(type): # Print nicely formatted Symbol names.
|
|
def __repr__(self): return self.__qualname__
|
|
def __str__(self): return self.__name__
|
|
class Symbol(metaclass=SymbolClass): pass
|
|
|
|
class ServerStates:
|
|
class Init(Symbol): pass
|
|
class LoadingModel(Symbol): pass
|
|
class Online(Symbol): pass
|
|
class Rendering(Symbol): pass
|
|
class Unavailable(Symbol): pass
|
|
|
|
class RenderTask(): # Task with output queue and completion lock.
|
|
def __init__(self, req: GenerateImageRequest, task_data: TaskData):
|
|
task_data.request_id = id(self)
|
|
self.render_request: GenerateImageRequest = req # Initial Request
|
|
self.task_data: TaskData = task_data
|
|
self.response: Any = None # Copy of the last reponse
|
|
self.render_device = None # Select the task affinity. (Not used to change active devices).
|
|
self.temp_images:list = [None] * req.num_outputs * (1 if task_data.show_only_filtered_image else 2)
|
|
self.error: Exception = None
|
|
self.lock: threading.Lock = threading.Lock() # Locks at task start and unlocks when task is completed
|
|
self.buffer_queue: queue.Queue = queue.Queue() # Queue of JSON string segments
|
|
async def read_buffer_generator(self):
|
|
try:
|
|
while not self.buffer_queue.empty():
|
|
res = self.buffer_queue.get(block=False)
|
|
self.buffer_queue.task_done()
|
|
yield res
|
|
except queue.Empty as e: yield
|
|
@property
|
|
def status(self):
|
|
if self.lock.locked():
|
|
return 'running'
|
|
if isinstance(self.error, StopAsyncIteration):
|
|
return 'stopped'
|
|
if self.error:
|
|
return 'error'
|
|
if not self.buffer_queue.empty():
|
|
return 'buffer'
|
|
if self.response:
|
|
return 'completed'
|
|
return 'pending'
|
|
@property
|
|
def is_pending(self):
|
|
return bool(not self.response and not self.error)
|
|
|
|
# Temporary cache to allow to query tasks results for a short time after they are completed.
|
|
class DataCache():
|
|
def __init__(self):
|
|
self._base = dict()
|
|
self._lock: threading.Lock = threading.Lock()
|
|
def _get_ttl_time(self, ttl: int) -> int:
|
|
return int(time.time()) + ttl
|
|
def _is_expired(self, timestamp: int) -> bool:
|
|
return int(time.time()) >= timestamp
|
|
def clean(self) -> None:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.clean' + ERR_LOCK_FAILED)
|
|
try:
|
|
# Create a list of expired keys to delete
|
|
to_delete = []
|
|
for key in self._base:
|
|
ttl, _ = self._base[key]
|
|
if self._is_expired(ttl):
|
|
to_delete.append(key)
|
|
# Remove Items
|
|
for key in to_delete:
|
|
(_, val) = self._base[key]
|
|
if isinstance(val, RenderTask):
|
|
log.debug(f'RenderTask {key} expired. Data removed.')
|
|
elif isinstance(val, SessionState):
|
|
log.debug(f'Session {key} expired. Data removed.')
|
|
else:
|
|
log.debug(f'Key {key} expired. Data removed.')
|
|
del self._base[key]
|
|
finally:
|
|
self._lock.release()
|
|
def clear(self) -> None:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.clear' + ERR_LOCK_FAILED)
|
|
try: self._base.clear()
|
|
finally: self._lock.release()
|
|
def delete(self, key: Hashable) -> bool:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.delete' + ERR_LOCK_FAILED)
|
|
try:
|
|
if key not in self._base:
|
|
return False
|
|
del self._base[key]
|
|
return True
|
|
finally:
|
|
self._lock.release()
|
|
def keep(self, key: Hashable, ttl: int) -> bool:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.keep' + ERR_LOCK_FAILED)
|
|
try:
|
|
if key in self._base:
|
|
_, value = self._base.get(key)
|
|
self._base[key] = (self._get_ttl_time(ttl), value)
|
|
return True
|
|
return False
|
|
finally:
|
|
self._lock.release()
|
|
def put(self, key: Hashable, value: Any, ttl: int) -> bool:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.put' + ERR_LOCK_FAILED)
|
|
try:
|
|
self._base[key] = (
|
|
self._get_ttl_time(ttl), value
|
|
)
|
|
except Exception as e:
|
|
log.error(traceback.format_exc())
|
|
return False
|
|
else:
|
|
return True
|
|
finally:
|
|
self._lock.release()
|
|
def tryGet(self, key: Hashable) -> Any:
|
|
if not self._lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('DataCache.tryGet' + ERR_LOCK_FAILED)
|
|
try:
|
|
ttl, value = self._base.get(key, (None, None))
|
|
if ttl is not None and self._is_expired(ttl):
|
|
log.debug(f'Session {key} expired. Discarding data.')
|
|
del self._base[key]
|
|
return None
|
|
return value
|
|
finally:
|
|
self._lock.release()
|
|
|
|
manager_lock = threading.RLock()
|
|
render_threads = []
|
|
current_state = ServerStates.Init
|
|
current_state_error:Exception = None
|
|
tasks_queue = []
|
|
session_cache = DataCache()
|
|
task_cache = DataCache()
|
|
weak_thread_data = weakref.WeakKeyDictionary()
|
|
idle_event: threading.Event = threading.Event()
|
|
|
|
class SessionState():
|
|
def __init__(self, id: str):
|
|
self._id = id
|
|
self._tasks_ids = []
|
|
@property
|
|
def id(self):
|
|
return self._id
|
|
@property
|
|
def tasks(self):
|
|
tasks = []
|
|
for task_id in self._tasks_ids:
|
|
task = task_cache.tryGet(task_id)
|
|
if task:
|
|
tasks.append(task)
|
|
return tasks
|
|
def put(self, task, ttl=TASK_TTL):
|
|
task_id = id(task)
|
|
self._tasks_ids.append(task_id)
|
|
if not task_cache.put(task_id, task, ttl):
|
|
return False
|
|
while len(self._tasks_ids) > len(render_threads) * 2:
|
|
self._tasks_ids.pop(0)
|
|
return True
|
|
|
|
def thread_get_next_task():
|
|
from sd_internal import renderer
|
|
if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT):
|
|
log.warn(f'Render thread on device: {renderer.context.device} failed to acquire manager lock.')
|
|
return None
|
|
if len(tasks_queue) <= 0:
|
|
manager_lock.release()
|
|
return None
|
|
task = None
|
|
try: # Select a render task.
|
|
for queued_task in tasks_queue:
|
|
if queued_task.render_device and renderer.context.device != queued_task.render_device:
|
|
# Is asking for a specific render device.
|
|
if is_alive(queued_task.render_device) > 0:
|
|
continue # requested device alive, skip current one.
|
|
else:
|
|
# Requested device is not active, return error to UI.
|
|
queued_task.error = Exception(queued_task.render_device + ' is not currently active.')
|
|
task = queued_task
|
|
break
|
|
if not queued_task.render_device and renderer.context.device == 'cpu' and is_alive() > 1:
|
|
# not asking for any specific devices, cpu want to grab task but other render devices are alive.
|
|
continue # Skip Tasks, don't run on CPU unless there is nothing else or user asked for it.
|
|
task = queued_task
|
|
break
|
|
if task is not None:
|
|
del tasks_queue[tasks_queue.index(task)]
|
|
return task
|
|
finally:
|
|
manager_lock.release()
|
|
|
|
def thread_render(device):
|
|
global current_state, current_state_error
|
|
|
|
from sd_internal import renderer, model_manager
|
|
try:
|
|
renderer.init(device)
|
|
except Exception as e:
|
|
log.error(traceback.format_exc())
|
|
weak_thread_data[threading.current_thread()] = {
|
|
'error': e
|
|
}
|
|
return
|
|
|
|
weak_thread_data[threading.current_thread()] = {
|
|
'device': renderer.context.device,
|
|
'device_name': renderer.context.device_name,
|
|
'alive': True
|
|
}
|
|
|
|
current_state = ServerStates.LoadingModel
|
|
model_manager.load_default_models(renderer.context)
|
|
|
|
current_state = ServerStates.Online
|
|
|
|
while True:
|
|
session_cache.clean()
|
|
task_cache.clean()
|
|
if not weak_thread_data[threading.current_thread()]['alive']:
|
|
log.info(f'Shutting down thread for device {renderer.context.device}')
|
|
model_manager.unload_all(renderer.context)
|
|
return
|
|
if isinstance(current_state_error, SystemExit):
|
|
current_state = ServerStates.Unavailable
|
|
return
|
|
task = thread_get_next_task()
|
|
if task is None:
|
|
idle_event.clear()
|
|
idle_event.wait(timeout=1)
|
|
continue
|
|
if task.error is not None:
|
|
log.error(task.error)
|
|
task.response = {"status": 'failed', "detail": str(task.error)}
|
|
task.buffer_queue.put(json.dumps(task.response))
|
|
continue
|
|
if current_state_error:
|
|
task.error = current_state_error
|
|
task.response = {"status": 'failed', "detail": str(task.error)}
|
|
task.buffer_queue.put(json.dumps(task.response))
|
|
continue
|
|
log.info(f'Session {task.task_data.session_id} starting task {id(task)} on {renderer.context.device_name}')
|
|
if not task.lock.acquire(blocking=False): raise Exception('Got locked task from queue.')
|
|
try:
|
|
def step_callback():
|
|
global current_state_error
|
|
|
|
if isinstance(current_state_error, SystemExit) or isinstance(current_state_error, StopAsyncIteration) or isinstance(task.error, StopAsyncIteration):
|
|
renderer.context.stop_processing = True
|
|
if isinstance(current_state_error, StopAsyncIteration):
|
|
task.error = current_state_error
|
|
current_state_error = None
|
|
log.info(f'Session {task.task_data.session_id} sent cancel signal for task {id(task)}')
|
|
|
|
current_state = ServerStates.LoadingModel
|
|
model_manager.resolve_model_paths(task.task_data)
|
|
model_manager.set_vram_optimizations(renderer.context, task.task_data)
|
|
model_manager.reload_models_if_necessary(renderer.context, task.task_data)
|
|
|
|
current_state = ServerStates.Rendering
|
|
task.response = renderer.make_images(task.render_request, task.task_data, task.buffer_queue, task.temp_images, step_callback)
|
|
# Before looping back to the generator, mark cache as still alive.
|
|
task_cache.keep(id(task), TASK_TTL)
|
|
session_cache.keep(task.task_data.session_id, TASK_TTL)
|
|
except Exception as e:
|
|
task.error = e
|
|
task.response = {"status": 'failed', "detail": str(task.error)}
|
|
task.buffer_queue.put(json.dumps(task.response))
|
|
log.error(traceback.format_exc())
|
|
continue
|
|
finally:
|
|
# Task completed
|
|
task.lock.release()
|
|
task_cache.keep(id(task), TASK_TTL)
|
|
session_cache.keep(task.task_data.session_id, TASK_TTL)
|
|
if isinstance(task.error, StopAsyncIteration):
|
|
log.info(f'Session {task.task_data.session_id} task {id(task)} cancelled!')
|
|
elif task.error is not None:
|
|
log.info(f'Session {task.task_data.session_id} task {id(task)} failed!')
|
|
else:
|
|
log.info(f'Session {task.task_data.session_id} task {id(task)} completed by {renderer.context.device_name}.')
|
|
current_state = ServerStates.Online
|
|
|
|
def get_cached_task(task_id:str, update_ttl:bool=False):
|
|
# By calling keep before tryGet, wont discard if was expired.
|
|
if update_ttl and not task_cache.keep(task_id, TASK_TTL):
|
|
# Failed to keep task, already gone.
|
|
return None
|
|
return task_cache.tryGet(task_id)
|
|
|
|
def get_cached_session(session_id:str, update_ttl:bool=False):
|
|
if update_ttl:
|
|
session_cache.keep(session_id, TASK_TTL)
|
|
session = session_cache.tryGet(session_id)
|
|
if not session:
|
|
session = SessionState(session_id)
|
|
session_cache.put(session_id, session, TASK_TTL)
|
|
return session
|
|
|
|
def get_devices():
|
|
devices = {
|
|
'all': {},
|
|
'active': {},
|
|
}
|
|
|
|
def get_device_info(device):
|
|
if device == 'cpu':
|
|
return {'name': device_manager.get_processor_name()}
|
|
|
|
mem_free, mem_total = torch.cuda.mem_get_info(device)
|
|
mem_free /= float(10**9)
|
|
mem_total /= float(10**9)
|
|
|
|
return {
|
|
'name': torch.cuda.get_device_name(device),
|
|
'mem_free': mem_free,
|
|
'mem_total': mem_total,
|
|
}
|
|
|
|
# list the compatible devices
|
|
gpu_count = torch.cuda.device_count()
|
|
for device in range(gpu_count):
|
|
device = f'cuda:{device}'
|
|
if not device_manager.is_device_compatible(device):
|
|
continue
|
|
|
|
devices['all'].update({device: get_device_info(device)})
|
|
|
|
devices['all'].update({'cpu': get_device_info('cpu')})
|
|
|
|
# list the activated devices
|
|
if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('get_devices' + ERR_LOCK_FAILED)
|
|
try:
|
|
for rthread in render_threads:
|
|
if not rthread.is_alive():
|
|
continue
|
|
weak_data = weak_thread_data.get(rthread)
|
|
if not weak_data or not 'device' in weak_data or not 'device_name' in weak_data:
|
|
continue
|
|
device = weak_data['device']
|
|
devices['active'].update({device: get_device_info(device)})
|
|
finally:
|
|
manager_lock.release()
|
|
|
|
return devices
|
|
|
|
def is_alive(device=None):
|
|
if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('is_alive' + ERR_LOCK_FAILED)
|
|
nbr_alive = 0
|
|
try:
|
|
for rthread in render_threads:
|
|
if device is not None:
|
|
weak_data = weak_thread_data.get(rthread)
|
|
if weak_data is None or not 'device' in weak_data or weak_data['device'] is None:
|
|
continue
|
|
thread_device = weak_data['device']
|
|
if thread_device != device:
|
|
continue
|
|
if rthread.is_alive():
|
|
nbr_alive += 1
|
|
return nbr_alive
|
|
finally:
|
|
manager_lock.release()
|
|
|
|
def start_render_thread(device):
|
|
if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('start_render_thread' + ERR_LOCK_FAILED)
|
|
log.info(f'Start new Rendering Thread on device: {device}')
|
|
try:
|
|
rthread = threading.Thread(target=thread_render, kwargs={'device': device})
|
|
rthread.daemon = True
|
|
rthread.name = THREAD_NAME_PREFIX + device
|
|
rthread.start()
|
|
render_threads.append(rthread)
|
|
finally:
|
|
manager_lock.release()
|
|
timeout = DEVICE_START_TIMEOUT
|
|
while not rthread.is_alive() or not rthread in weak_thread_data or not 'device' in weak_thread_data[rthread]:
|
|
if rthread in weak_thread_data and 'error' in weak_thread_data[rthread]:
|
|
log.error(f"{rthread}, {device}, error: {weak_thread_data[rthread]['error']}")
|
|
return False
|
|
if timeout <= 0:
|
|
return False
|
|
timeout -= 1
|
|
time.sleep(1)
|
|
return True
|
|
|
|
def stop_render_thread(device):
|
|
try:
|
|
device_manager.validate_device_id(device, log_prefix='stop_render_thread')
|
|
except:
|
|
log.error(traceback.format_exc())
|
|
return False
|
|
|
|
if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('stop_render_thread' + ERR_LOCK_FAILED)
|
|
log.info(f'Stopping Rendering Thread on device: {device}')
|
|
|
|
try:
|
|
thread_to_remove = None
|
|
for rthread in render_threads:
|
|
weak_data = weak_thread_data.get(rthread)
|
|
if weak_data is None or not 'device' in weak_data or weak_data['device'] is None:
|
|
continue
|
|
thread_device = weak_data['device']
|
|
if thread_device == device:
|
|
weak_data['alive'] = False
|
|
thread_to_remove = rthread
|
|
break
|
|
if thread_to_remove is not None:
|
|
render_threads.remove(rthread)
|
|
return True
|
|
finally:
|
|
manager_lock.release()
|
|
|
|
return False
|
|
|
|
def update_render_threads(render_devices, active_devices):
|
|
devices_to_start, devices_to_stop = device_manager.get_device_delta(render_devices, active_devices)
|
|
log.debug(f'devices_to_start: {devices_to_start}')
|
|
log.debug(f'devices_to_stop: {devices_to_stop}')
|
|
|
|
for device in devices_to_stop:
|
|
if is_alive(device) <= 0:
|
|
log.debug(f'{device} is not alive')
|
|
continue
|
|
if not stop_render_thread(device):
|
|
log.warn(f'{device} could not stop render thread')
|
|
|
|
for device in devices_to_start:
|
|
if is_alive(device) >= 1:
|
|
log.debug(f'{device} already registered.')
|
|
continue
|
|
if not start_render_thread(device):
|
|
log.warn(f'{device} failed to start.')
|
|
|
|
if is_alive() <= 0: # No running devices, probably invalid user config.
|
|
raise EnvironmentError('ERROR: No active render devices! Please verify the "render_devices" value in config.json')
|
|
|
|
log.debug(f"active devices: {get_devices()['active']}")
|
|
|
|
def shutdown_event(): # Signal render thread to close on shutdown
|
|
global current_state_error
|
|
current_state_error = SystemExit('Application shutting down.')
|
|
|
|
def render(render_req: GenerateImageRequest, task_data: TaskData):
|
|
current_thread_count = is_alive()
|
|
if current_thread_count <= 0: # Render thread is dead
|
|
raise ChildProcessError('Rendering thread has died.')
|
|
|
|
# Alive, check if task in cache
|
|
session = get_cached_session(task_data.session_id, update_ttl=True)
|
|
pending_tasks = list(filter(lambda t: t.is_pending, session.tasks))
|
|
if current_thread_count < len(pending_tasks):
|
|
raise ConnectionRefusedError(f'Session {task_data.session_id} already has {len(pending_tasks)} pending tasks out of {current_thread_count}.')
|
|
|
|
new_task = RenderTask(render_req, task_data)
|
|
if session.put(new_task, TASK_TTL):
|
|
# Use twice the normal timeout for adding user requests.
|
|
# Tries to force session.put to fail before tasks_queue.put would.
|
|
if manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT * 2):
|
|
try:
|
|
tasks_queue.append(new_task)
|
|
idle_event.set()
|
|
return new_task
|
|
finally:
|
|
manager_lock.release()
|
|
raise RuntimeError('Failed to add task to cache.')
|