Add priority to the API

requirements.txt

@@ -104,4 +104,4 @@ uvicorn==0.12.2
 Werkzeug==1.0.1
 zipp==3.4.0
 starlette_exporter==0.7.0
-git+git://github.com/airenas/smart_load_balancer.git@v0.1.26#egg=smart_load_balancer
+git+git://github.com/airenas/smart_load_balancer.git@v0.2.32#egg=smart_load_balancer

service/api/api.py

@@ -18,6 +18,7 @@ class Input(BaseModel):
     text: str
     voice: Optional[str] = None
     speed_control_alpha: Optional[float] = Field(alias='speedAlpha')
+    priority: Optional[int] = 0
 
 
 class Output(BaseModel):

service/routers/model.py

@@ -18,7 +18,7 @@ def calculate(inp: api.Input, request: Request):
     if not inp.voice:
         raise HTTPException(status_code=400, detail="No voice")
 
-    res_data = request.app.calculate(inp.text, inp.voice, inp.speed_control_alpha)
+    res_data = request.app.calculate(inp.text, inp.voice, inp.speed_control_alpha, inp.priority)
     res = api.Output(data=res_data)
     logger.info("Response ready")
     return res

service/service.py

@@ -100,10 +100,10 @@ def calc_model(voice: str, in_data: ModelData, workers_data):
         with app.metrics.calc_metric.labels(voice).time():
             return model.calculate(in_data.text, in_data.speed_control_alpha)
 
-    def calc(text, voice, speed_control_alpha):
+    def calc(text, voice, speed_control_alpha, priority: int = 0):
         if not voice:
             raise Exception("no voice")
-        work = Work(name=voice, data=ModelData(text, speed_control_alpha), work_func=calc_model)
+        work = Work(name=voice, data=ModelData(text, speed_control_alpha), work_func=calc_model, priority=priority)
         app.balancer.add_wrk(work)
         res = work.wait()
         if res.err is not None:

service/test_service.py

@@ -88,10 +88,11 @@ def test_calculate_fail_empty():
 def test_calculate():
     client, app = init_test_app()
 
-    def test_calc(text, model, speedAlpha):
+    def test_calc(text, model, speedAlpha, priority):
         assert text == "in text"
         assert model == "m"
         assert speedAlpha is None
+        assert priority == 0
         return "olia"
 
     app.calculate = test_calc
@@ -103,7 +104,7 @@ def test_calc(text, model, speedAlpha):
 def test_calculate_pass_speed():
     client, app = init_test_app()
 
-    def test_calc(text, model, speedAlpha):
+    def test_calc(text, model, speedAlpha, priority):
         assert text == "in text"
         assert model == "m"
         assert speedAlpha == 1.2
@@ -115,6 +116,21 @@ def test_calc(text, model, speedAlpha):
     assert response.json() == {"data": "olia", "error": None}
 
 
+def test_calculate_pass_priority():
+    client, app = init_test_app()
+
+    def test_calc(text, model, speedAlpha, priority):
+        assert text == "in text"
+        assert model == "m"
+        assert priority == 1000
+        return "olia"
+
+    app.calculate = test_calc
+    response = client.post("/model", json={"text": "in text", "voice": "m", "speedAlpha": 1.2, "priority": 1000})
+    assert response.status_code == 200
+    assert response.json() == {"data": "olia", "error": None}
+
+
 def test_environment():
     os.environ["CONFIG_FILE"] = "/m1/c.yaml"
     os.environ["DEVICE"] = "cuda"