From 270fbae8ae0d032595d51564cfeb3da56074dd99 Mon Sep 17 00:00:00 2001
From: muhammed-shihebi <58932402+muhammed-shihebi@users.noreply.github.com>
Date: Wed, 1 May 2024 17:39:34 +0000
Subject: [PATCH] temp: update model remove worker endpoint

---
 .github/workflows/build-and-deploy.yml                |  4 ++--
 datastore-api/app/core/model_api.py                   | 10 ++++++++++
 model-manager/model_manager/app/routers/management.py |  8 +++++++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/build-and-deploy.yml b/.github/workflows/build-and-deploy.yml
index f63e98a0a..014de9944 100644
--- a/.github/workflows/build-and-deploy.yml
+++ b/.github/workflows/build-and-deploy.yml
@@ -31,7 +31,7 @@ jobs:
           "model-manager",
           "model-inference-transformer",
           "model-inference-onnx",
-          # "model-inference-sentence-transformer",
+          "model-inference-sentence-transformer",
           "model-inference-graph-transformer",
           "skill-manager",
           "skill-boolq",
@@ -48,7 +48,7 @@ jobs:
           "skill-multiple-choice-metaqa",
           "skill-kgqa-kqapro",
           "evaluator",
-          # "sensitivity",
+          "sensitivity",
         ]
         include:
           - build-args: ""
diff --git a/datastore-api/app/core/model_api.py b/datastore-api/app/core/model_api.py
index 64c0b545f..6ae9b99a7 100644
--- a/datastore-api/app/core/model_api.py
+++ b/datastore-api/app/core/model_api.py
@@ -178,6 +178,16 @@ async def predict(self, model_identifier, prediction_method, input_data):
         logger.debug(f"Request data: {input_data}")
         logger.debug(f"URL is: {self.square_api_url}/main/{model_identifier}/{prediction_method}")
 
+
+        # # check if the model is available
+        # req_result = requests.get(
+        #     f"{self.base_url}/{model_identifier}/health/heartbeat",
+        #     headers={"Authorization": f"Bearer {client_credentials()}"}
+        # )
+
+
+
+
         my_conn = aiohttp.TCPConnector()
         async with aiohttp.ClientSession(connector=my_conn) as session:
             async with session.post(
diff --git a/model-manager/model_manager/app/routers/management.py b/model-manager/model_manager/app/routers/management.py
index 4551954d9..42038db05 100644
--- a/model-manager/model_manager/app/routers/management.py
+++ b/model-manager/model_manager/app/routers/management.py
@@ -344,7 +344,13 @@ async def add_model_container(request: Request, identifier: str, num: int):
 
 
 @router.delete("/{identifier}/remove_worker/{num}")
-async def remove_model_container(request: Request, identifier: str, num: int):
+@router.delete("/{hf_username}/{identifier}/remove_worker/{num}")
+async def remove_model_container(
+    request: Request,
+    identifier: str,
+    num: int,
+    hf_username: str = None
+):
     """
     Remove a model worker from the platform
     """