Added get_cached_task to replace task_cache.tryGet in server.py

Now updated cache TTL on /stream and temp images endpoints.
Keep images alive longer when browser keeps reading the endpoints.
This commit is contained in:
Marc-Andre Ferland 2022-10-22 13:52:13 -04:00
parent 46a46877ed
commit 364e364429
2 changed files with 11 additions and 4 deletions

View File

@ -326,6 +326,13 @@ def thread_render(device):
print(f'Session {task.request.session_id} task {id(task)} completed.') print(f'Session {task.request.session_id} task {id(task)} completed.')
current_state = ServerStates.Online current_state = ServerStates.Online
def get_cached_task(session_id:str, update_ttl:bool=False):
# By calling keep before tryGet, wont discard if was expired.
if update_ttl and not task_cache.keep(session_id, TASK_TTL):
# Failed to keep task, already gone.
return None
return task_cache.tryGet(session_id)
def is_first_cuda_device(device): def is_first_cuda_device(device):
from . import runtime # When calling runtime from outside thread_render DO NOT USE thread specific attributes or functions. from . import runtime # When calling runtime from outside thread_render DO NOT USE thread specific attributes or functions.
return runtime.is_first_cuda_device(device) return runtime.is_first_cuda_device(device)

View File

@ -251,7 +251,7 @@ def ping(session_id:str=None):
# Alive # Alive
response = {'status': str(task_manager.current_state)} response = {'status': str(task_manager.current_state)}
if session_id: if session_id:
task = task_manager.task_cache.tryGet(session_id) task = task_manager.get_cached_task(session_id)
if task: if task:
response['task'] = id(task) response['task'] = id(task)
if task.lock.locked(): if task.lock.locked():
@ -302,7 +302,7 @@ def render(req : task_manager.ImageRequest):
@app.get('/image/stream/{session_id:str}/{task_id:int}') @app.get('/image/stream/{session_id:str}/{task_id:int}')
def stream(session_id:str, task_id:int): def stream(session_id:str, task_id:int):
#TODO Move to WebSockets ?? #TODO Move to WebSockets ??
task = task_manager.task_cache.tryGet(session_id) task = task_manager.get_cached_task(session_id, update_ttl=True)
if not task: raise HTTPException(status_code=410, detail='No request received.') # HTTP410 Gone if not task: raise HTTPException(status_code=410, detail='No request received.') # HTTP410 Gone
if (id(task) != task_id): raise HTTPException(status_code=409, detail=f'Wrong task id received. Expected:{id(task)}, Received:{task_id}') # HTTP409 Conflict if (id(task) != task_id): raise HTTPException(status_code=409, detail=f'Wrong task id received. Expected:{id(task)}, Received:{task_id}') # HTTP409 Conflict
if task.buffer_queue.empty() and not task.lock.locked(): if task.buffer_queue.empty() and not task.lock.locked():
@ -320,7 +320,7 @@ def stop(session_id:str=None):
raise HTTPException(status_code=409, detail='Not currently running any tasks.') # HTTP409 Conflict raise HTTPException(status_code=409, detail='Not currently running any tasks.') # HTTP409 Conflict
task_manager.current_state_error = StopAsyncIteration('') task_manager.current_state_error = StopAsyncIteration('')
return {'OK'} return {'OK'}
task = task_manager.task_cache.tryGet(session_id) task = task_manager.get_cached_task(session_id)
if not task: raise HTTPException(status_code=404, detail=f'Session {session_id} has no active task.') # HTTP404 Not Found if not task: raise HTTPException(status_code=404, detail=f'Session {session_id} has no active task.') # HTTP404 Not Found
if isinstance(task.error, StopAsyncIteration): raise HTTPException(status_code=409, detail=f'Session {session_id} task is already stopped.') # HTTP409 Conflict if isinstance(task.error, StopAsyncIteration): raise HTTPException(status_code=409, detail=f'Session {session_id} task is already stopped.') # HTTP409 Conflict
task.error = StopAsyncIteration('') task.error = StopAsyncIteration('')
@ -328,7 +328,7 @@ def stop(session_id:str=None):
@app.get('/image/tmp/{session_id}/{img_id:int}') @app.get('/image/tmp/{session_id}/{img_id:int}')
def get_image(session_id, img_id): def get_image(session_id, img_id):
task = task_manager.task_cache.tryGet(session_id) task = task_manager.get_cached_task(session_id, update_ttl=True)
if not task: raise HTTPException(status_code=410, detail=f'Session {session_id} has not submitted a task.') # HTTP410 Gone if not task: raise HTTPException(status_code=410, detail=f'Session {session_id} has not submitted a task.') # HTTP410 Gone
if not task.temp_images[img_id]: raise HTTPException(status_code=425, detail='Too Early, task data is not available yet.') # HTTP425 Too Early if not task.temp_images[img_id]: raise HTTPException(status_code=425, detail='Too Early, task data is not available yet.') # HTTP425 Too Early
try: try: