#

Signed-off-by: Theo Truong <[email protected]>
nhtruong · Aug 12, 2024 · 49eee04 · 49eee04
1 parent 788d0f4
commit 49eee04
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 7 deletions.
diff --git a/spec/namespaces/ml.yaml b/spec/namespaces/ml.yaml
@@ -147,7 +147,7 @@ components:
       content:
         application/json:
           schema:
-            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.register_model@200:
       content:
         application/json:
@@ -157,7 +157,7 @@ components:
       content:
         application/json:
           schema:
-            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.get_task@200:
       content:
         application/json:

diff --git a/spec/schemas/ml._common.yaml b/spec/schemas/ml._common.yaml
@@ -108,10 +108,10 @@ components:
       properties:
         model_id:
           type: string
-          description: The model ID.
-        state:
+        task_id:
+          type: string
+        status:
           type: string
-          description: The state.
           enum:
             - CANCELLED
             - COMPLETED
@@ -124,6 +124,7 @@ components:
           description: Task type.
           enum:
             - DEPLOY_MODEL
+            - REGISTER_MODEL
         function_name:
           type: string
         worker_node:
@@ -139,4 +140,4 @@ components:
         is_async:
           type: boolean
       required:
-        - state
+        - status
diff --git a/tests/default/ml/models.yaml b/tests/default/ml/models.yaml
@@ -11,6 +11,8 @@ prologues:
           plugins:
             ml_commons:
               only_run_on_ml_node: false
+  - path: /_plugins/_ml/tasks/*
+    method: DELETE
 chapters:
   - synopsis: Create model.
     id: create_model
@@ -26,6 +28,7 @@ chapters:
     output:
       task_id: payload.task_id
   - synopsis: Wait for task.
+    id: register_model
     path: /_plugins/_ml/tasks/{task_id}
     method: GET
     warnings:
@@ -36,6 +39,13 @@ chapters:
       status: 200
       payload:
         state: COMPLETED
+    output:
+      model_id: payload.model_id
     retry:
-      count: 3
+      count: 5
       wait: 30000
+  - synopsis: Delete model.
+    path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${register_model.model_id}
+    method: DELETE