Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixed an error in DeleleMlModel runner. #537

Merged
merged 14 commits into from
Jul 1, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions osbenchmark/worker_coordinator/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand All @@ -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)
vpehkone marked this conversation as resolved.
Show resolved Hide resolved
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"
Expand Down
Loading