feat: restart server on job timeout (#104)

Co-authored-by: Avram Tudor <[email protected]>

Dockerfile

@@ -64,6 +64,9 @@ ENV \
  # https://docs.python.org/3/using/cmdline.html#envvar-PYTHONDONTWRITEBYTECODE
  PYTHONDONTWRITEBYTECODE=1 \
  PYTHONPATH=/app \
+ OUTLINES_CACHE_DIR=/app/vllm/outlines \
+ VLLM_CONFIG_ROOT=/app/vllm/config \
+ HF_HOME=/app/hf \
  LLAMA_PATH="/models/Llama-3.1-8B-Instruct-Q8_0.gguf"
 
 VOLUME [ "/models" ]

skynet/modules/ttt/openai_api/app.py

@@ -1,3 +1,4 @@
+import os
 import subprocess
 
 from skynet import http_client
@@ -71,12 +72,12 @@ def destroy():
 
 
 def restart():
- log.info('Restarting OpenAI API server...')
+ log.info('Restarting Skynet...')
 
  OPENAI_API_RESTART_COUNTER.inc()
 
- destroy()
- initialize()
+ # rely on the supervisor to restart the process
+ os._exit(1)
 
 
 __all__ = ['destroy', 'initialize', 'restart']

skynet/modules/ttt/summaries/jobs.py

@@ -36,7 +36,7 @@ def can_run_next_job() -> bool:
  return 'summaries:executor' in modules and (current_task is None or current_task.done())
 
 
-def get_job_processor(customer_id: str) -> str:
+def get_job_processor(customer_id: str) -> Processors:
  options = get_credentials(customer_id)
  secret = options.get('secret')
  api_type = options.get('type')
@@ -125,6 +125,29 @@ async def run_job(job: Job) -> None:
  exit_task.cancel()
 
 
+async def update_done_job(job: Job, result: str, processor: Processors, has_failed: bool = False) -> None:
+ should_expire = not has_failed or processor != Processors.LOCAL
+
+ updated_job = await update_job(
+ expires=redis_exp_seconds if should_expire else None,
+ job_id=job.id,
+ end=time.time(),
+ status=JobStatus.ERROR if has_failed else JobStatus.SUCCESS,
+ result=result,
+ )
+
+ if not should_expire:
+ await db.rpush(ERROR_JOBS_KEY, job.id)
+ SUMMARY_ERROR_COUNTER.inc()
+
+ await db.lrem(RUNNING_JOBS_KEY, 0, job.id)
+
+ SUMMARY_DURATION_METRIC.observe(updated_job.computed_duration)
+ SUMMARY_INPUT_LENGTH_METRIC.observe(len(updated_job.payload.text))
+
+ log.info(f"Job {updated_job.id} duration: {updated_job.computed_duration} seconds")
+
+
 async def _run_job(job: Job) -> None:
  has_failed = False
  result = None
@@ -177,26 +200,7 @@ async def _run_job(job: Job) -> None:
  has_failed = True
  result = str(e)
 
- should_expire = not has_failed or processor != Processors.LOCAL
-
- updated_job = await update_job(
- expires=redis_exp_seconds if should_expire else None,
- job_id=job.id,
- end=time.time(),
- status=JobStatus.ERROR if has_failed else JobStatus.SUCCESS,
- result=result,
- )
-
- if not should_expire:
- await db.rpush(ERROR_JOBS_KEY, job.id)
- SUMMARY_ERROR_COUNTER.inc()
-
- await db.lrem(RUNNING_JOBS_KEY, 0, job.id)
-
- SUMMARY_DURATION_METRIC.observe(updated_job.computed_duration)
- SUMMARY_INPUT_LENGTH_METRIC.observe(len(updated_job.payload.text))
-
- log.info(f"Job {updated_job.id} duration: {updated_job.computed_duration} seconds")
+ await update_done_job(job, result, processor, has_failed)
 
 
 def create_run_job_task(job: Job) -> asyncio.Task:
@@ -239,6 +243,8 @@ async def restart_on_timeout(job: Job) -> None:
 
  log.warning(f"Job {job.id} timed out after {job_timeout} seconds, restarting...")
 
+ await update_done_job(job, "Job timed out", Processors.LOCAL, has_failed=True)
+
  restart_openai_api()