feat: add priority flag (#106)

* jobs submitted with priority high will be added to the start of the redis queue

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

run.sh

@@ -3,7 +3,9 @@
 if nvcc --version
 then
   export CUDA_VISIBLE_DEVICES=0
+  export LLAMA_N_CTX=44000
+else
+  export LLAMA_N_CTX=8182
 fi
 
-export LLAMA_N_CTX=44000
 poetry run python -m uvicorn skynet.main:app --reload

skynet/modules/ttt/summaries/jobs.py

@@ -19,7 +19,7 @@
 
 from .persistence import db
 from .processor import process, process_azure, process_open_ai
-from .v1.models import DocumentMetadata, DocumentPayload, Job, JobId, JobStatus, JobType, Processors
+from .v1.models import DocumentMetadata, DocumentPayload, Job, JobId, JobStatus, JobType, Priority, Processors
 
 log = get_logger(__name__)
 
@@ -101,7 +101,11 @@ async def create_job(job_type: JobType, payload: DocumentPayload, metadata: Docu
 
     log.info(f"Created job {job.id}.")
 
-    await db.rpush(PENDING_JOBS_KEY, job_id)
+    if payload.priority == Priority.HIGH:
+        await db.lpush(PENDING_JOBS_KEY, job_id)
+    else:
+        await db.rpush(PENDING_JOBS_KEY, job_id)
+
     await update_summary_queue_metric()
 
     return JobId(id=job_id)

skynet/modules/ttt/summaries/v1/models.py

@@ -10,14 +10,27 @@ class HintType(Enum):
     TEXT = 'text'
 
 
+class Priority(Enum):
+    NORMAL = 'normal'
+    HIGH = 'high'
+
+
 class DocumentPayload(BaseModel):
     text: str
     hint: HintType = HintType.TEXT
+    priority: Priority = Priority.NORMAL
     prompt: str | None = None
 
     model_config = {
         'json_schema_extra': {
-            'examples': [{'text': 'Your text here', 'hint': 'text', 'prompt': 'Summarize the following text {text}'}]
+            'examples': [
+                {
+                    'text': 'Your text here',
+                    'hint': 'text',
+                    'priority': 'normal',
+                    'prompt': 'Summarize the following text {text}',
+                }
+            ]
         }
     }