Fixed an error in DeleleMlModel runner. (#537)

Signed-off-by: Vesa Pehkonen <[email protected]>

osbenchmark/worker_coordinator/runner.py

@@ -2400,6 +2400,11 @@ def __repr__(self, *args, **kwargs):
 class DeleteMlModel(Runner):
     @time_func
     async def __call__(self, opensearch, params):
+        async def _is_deployed(model_id):
+            resp = await opensearch.transport.perform_request('GET', '_plugins/_ml/models/' + model_id)
+            state = resp.get('model_state')
+            return state in ('PARTIALLY_DEPLOYED', 'DEPLOYED')
+
         body= {
             "query": {
                 "match_phrase": {
@@ -2422,8 +2427,16 @@ async def __call__(self, opensearch, params):
                     model_ids.add(id)
 
         for model_id in model_ids:
-            resp=await opensearch.transport.perform_request('POST', '/_plugins/_ml/models/' + model_id + '/_undeploy')
-            resp=await opensearch.transport.perform_request('DELETE', '/_plugins/_ml/models/' + model_id)
+            await opensearch.transport.perform_request('POST', '/_plugins/_ml/models/' + model_id + '/_undeploy')
+
+        for model_id in model_ids:
+            timeout = params.get('undeploy-timeout', 10)
+            end = time.time() + timeout
+            while await _is_deployed(model_id):
+                await asyncio.sleep(1)
+                if time.time() > end:
+                    raise TimeoutError("Timeout when undeploying ml-model.")
+            await opensearch.transport.perform_request('DELETE', '/_plugins/_ml/models/' + model_id)
 
     def __repr__(self, *args, **kwargs):
         return "delete-ml-model"