Applied suggestions from code review.

Signed-off-by: Nathalie Jonathan <[email protected]>

tests/plugins/ml/models/create_metadata.yaml → tests/plugins/ml/models/chunk.yaml

@@ -1,29 +1,19 @@
 $schema: ../../../../json_schemas/test_story.schema.yaml
 
-description: Test deprecated meta and chunk APIs for model creation and chunk upload.
+description: Test deprecated chunk API for chunk upload.
 distributions:
   excluded:
     - amazon-managed
     - amazon-serverless
-warnings:
-  multiple-paths-detected: false
 prologues:
   - path: /_cluster/settings
     method: PUT
     request:
       payload:
         persistent:
           plugins.ml_commons.allow_registering_model_via_local_file: true
-epilogues:
-  - path: /_plugins/_ml/models/{model_id}
-    parameters:
-      model_id: ${model_metadata.model_id}
-    method: DELETE
-    status: [200, 404]
-chapters:
-  - synopsis: Create model metadata.
+  - path: /_plugins/_ml/models/meta
     id: model_metadata
-    path: /_plugins/_ml/models/meta
     method: POST
     request:
       payload:
@@ -37,10 +27,15 @@ chapters:
           model_type: bert
           embedding_dimension: 384
           framework_type: sentence_transformers
-    response:
-      status: 200
     output:
       model_id: payload.model_id
+epilogues:
+  - path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${model_metadata.model_id}
+    method: DELETE
+    status: [200, 404]
+chapters:
   - synopsis: Upload chunk 0.
     path: /_plugins/_ml/models/{model_id}/chunk/{chunk_number}
     method: POST

tests/plugins/ml/models/meta.yaml

@@ -0,0 +1,41 @@
+$schema: ../../../../json_schemas/test_story.schema.yaml
+
+description: Test deprecated meta API for model creation.
+distributions:
+  excluded:
+    - amazon-managed
+    - amazon-serverless
+prologues:
+  - path: /_cluster/settings
+    method: PUT
+    request:
+      payload:
+        persistent:
+          plugins.ml_commons.allow_registering_model_via_local_file: true
+epilogues:
+  - path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${model_metadata.model_id}
+    method: DELETE
+    status: [200, 404]
+chapters:
+  - synopsis: Create model metadata.
+    id: model_metadata
+    path: /_plugins/_ml/models/meta
+    method: POST
+    request:
+      payload:
+        name: huggingface/sentence-transformers/msmarco-distilbert-base-tas-b
+        version: 1.0.1
+        model_format: TORCH_SCRIPT
+        model_content_hash_value: 123abc
+        total_chunks: 3
+        model_config:
+          all_config: path
+          model_type: bert
+          embedding_dimension: 384
+          framework_type: sentence_transformers
+    response:
+      status: 200
+    output:
+      model_id: payload.model_id

tests/plugins/ml/models/predict.yaml

@@ -5,8 +5,6 @@ distributions:
   excluded:
     - amazon-managed
     - amazon-serverless
-warnings:
-  multiple-paths-detected: false
 prologues:
   - path: /_cluster/settings
     method: PUT
@@ -55,6 +53,19 @@ prologues:
       model_id: ${get_completed_register_model_task.model_id}
     output:
       task_id: payload.task_id
+  - path: /_plugins/_ml/tasks/{task_id}
+    id: get_completed_deploy_model_task
+    method: GET
+    parameters:
+      task_id: ${deploy_model.task_id}
+    retry:
+      count: 3
+      wait: 10000
+      until:
+        - path: payload.state
+          equal: COMPLETED
+    output:
+      model_id: payload.model_id
 epilogues:
   - path: /_plugins/_ml/models/{model_id}/_undeploy
     method: POST
@@ -72,21 +83,6 @@ epilogues:
     parameters:
       model_group_id: ${create_model_group.test_model_group_id}
 chapters:
-  - synopsis: Wait to get completed task.
-    id: get_completed_deploy_model_task
-    path: /_plugins/_ml/tasks/{task_id}
-    method: GET
-    parameters:
-      task_id: ${deploy_model.task_id}
-    response:
-      status: 200
-      payload:
-        state: COMPLETED
-    output:
-      model_id: payload.model_id
-    retry:
-      count: 3
-      wait: 10000
   - synopsis: Predict new data.
     id: predict_data
     path: /_plugins/_ml/models/{model_id}/_predict
@@ -95,18 +91,16 @@ chapters:
       model_id: ${get_completed_register_model_task.model_id}
     request:
       payload:
-        query_text: today is sunny
+        query_text: The best selling book series in history is Harry Potter
         text_docs:
-          - how are you
-          - it is winter
-          - today is july fifth
-          - today is sunny
+          - Harry Potter is written by J.K. Rowling
+          - The Great Gatsby is a story of wealth and tragedy
+          - The Lord of the Rings is an epic high fantasy novel
+          - The best selling book series in history is Harry Potter
     response:
       status: 200
       payload:
         inference_results:
           - output:
               - data:
                   - 101
-    output:
-      prediction_results: payload.inference_results[0].output[0].data

tests/plugins/ml/models/register_meta.yaml

@@ -0,0 +1,42 @@
+$schema: ../../../../json_schemas/test_story.schema.yaml
+
+description: Test the creation of model metadata.
+distributions:
+  excluded:
+    - amazon-managed
+    - amazon-serverless
+version: '>= 2.11'
+prologues:
+  - path: /_cluster/settings
+    method: PUT
+    request:
+      payload:
+        persistent:
+          plugins.ml_commons.allow_registering_model_via_local_file: true
+epilogues:
+  - path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${model_metadata.model_id}
+    method: DELETE
+    status: [200, 404]
+chapters:
+  - synopsis: Register model metadata.
+    id: model_metadata
+    path: /_plugins/_ml/models/_register_meta
+    method: POST
+    request:
+      payload:
+        name: huggingface/sentence-transformers/msmarco-distilbert-base-tas-b
+        version: 1.0.1
+        model_format: TORCH_SCRIPT
+        model_content_hash_value: 123abc
+        total_chunks: 3
+        model_config:
+          all_config: path
+          model_type: bert
+          embedding_dimension: 384
+          framework_type: sentence_transformers
+    response:
+      status: 200
+    output:
+      model_id: payload.model_id

tests/plugins/ml/models/undeploy_specific_models.yaml → tests/plugins/ml/models/undeploy.yaml

@@ -1,14 +1,18 @@
 $schema: ../../../../json_schemas/test_story.schema.yaml
 
-description: Test the undeployment of specific models from specific nodes.
+description: Test the undeployment of specific models from specific nodes and a model from all nodes.
 distributions:
   excluded:
     - amazon-managed
     - amazon-serverless
 version: '>= 2.11'
-warnings:
-  multiple-paths-detected: false
 prologues:
+  - path: /_cluster/settings
+    method: PUT
+    request:
+      payload:
+        persistent:
+          plugins.ml_commons.allow_custom_deployment_plan: true
   - path: /_plugins/_ml/models/_register
     id: register_model
     method: POST
@@ -19,59 +23,67 @@ prologues:
         model_format: TORCH_SCRIPT
     output:
       task_id: payload.task_id
-epilogues:
-  - path: /_plugins/_ml/models/{model_id}
-    parameters:
-      model_id: ${get_completed_register_model_task.model_id}
-    method: DELETE
-    status: [200, 404]
-chapters:
-  - synopsis: Wait to get completed task.
+  - path: /_plugins/_ml/tasks/{task_id}
     id: get_completed_register_model_task
-    path: /_plugins/_ml/tasks/{task_id}
     method: GET
     parameters:
       task_id: ${register_model.task_id}
-    response:
-      status: 200
-      payload:
-        state: COMPLETED
-    output:
-      model_id: payload.model_id
     retry:
       count: 3
       wait: 10000
-  - synopsis: Deploy a model.
-    id: deploy_model
-    path: /_plugins/_ml/models/{model_id}/_deploy
+      until:
+        - path: payload.state
+          equal: COMPLETED
+    output:
+      model_id: payload.model_id
+  - path: /_plugins/_ml/models/{model_id}/_deploy
+    id: deploy_model   
     method: POST
     parameters:
       model_id: ${get_completed_register_model_task.model_id}
     output:
       task_id: payload.task_id
-    response:
-      status: 200
-  - synopsis: Wait to get completed task.
+  - path: /_plugins/_ml/tasks/{task_id}
     id: get_completed_deploy_model_task
-    path: /_plugins/_ml/tasks/{task_id}
     method: GET
     parameters:
       task_id: ${deploy_model.task_id}
-    response:
-      status: 200
-      payload:
-        state: COMPLETED
-    output:
-      model_id: payload.model_id
     retry:
       count: 3
-      wait: 10000      
-  - synopsis: Undeploy a model.
+      wait: 10000
+      until:
+        - path: payload.state
+          equal: COMPLETED
+    output:
+      node_id: payload.worker_node[0]
+      model_id: payload.model_id
+epilogues:
+  - path: /_plugins/_ml/models/{model_id}/_undeploy
+    method: POST
+    parameters:
+      model_id: ${get_completed_register_model_task.model_id}
+    status: [200, 404]
+  - path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${get_completed_register_model_task.model_id}
+    method: DELETE
+    status: [200, 404]
+chapters:  
+  - synopsis: Undeploy specific models from specific nodes.
     path: /_plugins/_ml/models/_undeploy
     method: POST
     request:
       payload:
+        node_ids:
+          - ${get_completed_deploy_model_task.node_id}
         model_ids: 
           - ${get_completed_register_model_task.model_id}
     response:
       status: 200
+  - synopsis: Undeploy a model from all nodes.
+    path: /_plugins/_ml/models/{model_id}/_undeploy
+    method: POST
+    parameters:
+      model_id: ${get_completed_register_model_task.model_id}
+    response:
+      status: 200