setup:
  - skip:
      features: headers
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model:
        model_id: "test_model"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "pass_through": {
                "tokenization": {
                  "bert": {
                    "with_special_tokens": false
                  }
                }
              }
            }
          }
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model_vocabulary:
        model_id: "test_model"
        body: >
          { "vocabulary": ["[PAD]","[UNK]","these", "are", "my", "words"] }
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model_definition_part:
        model_id: "test_model"
        part: 0
        body: >
          {
            "total_definition_length":1630,
            "definition": "UEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAAUAA4Ac2ltcGxlbW9kZWwvZGF0YS5wa2xGQgoAWlpaWlpaWlpaWoACY19fdG9yY2hfXwpTdXBlclNpbXBsZQpxACmBfShYCAAAAHRyYWluaW5ncQGIdWJxAi5QSwcIXOpBBDQAAAA0AAAAUEsDBBQACAgIAAAAAAAAAAAAAAAAAAAAAAAdAEEAc2ltcGxlbW9kZWwvY29kZS9fX3RvcmNoX18ucHlGQj0AWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWnWOMWvDMBCF9/yKI5MMrnHTQsHgjt2aJdlCEIp9SgWSTpykFvfXV1htaYds0nfv473JqhjhkAPywbhgUbzSnC02wwZAyqBYOUzIUUoY4XRe6SVr/Q8lVsYbf4UBLkS2kBk1aOIPxbOIaPVQtEQ8vUnZ/WlrSxTA+JCTNHMc4Ig+Eles+Jod+iR3N/jDDf74wxu4e/5+DmtE9mUyhdgFNq7bZ3ekehbruC6aTxS/c1rom6Z698WrEfIYxcn4JGTftLA7tzCnJeD41IJVC+U07kumUHw3E47Vqh+xnULeFisYLx064mV8UTZibWFMmX0p23wBUEsHCE0EGH3yAAAAlwEAAFBLAwQUAAgICAAAAAAAAAAAAAAAAAAAAAAAJwA5AHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5LmRlYnVnX3BrbEZCNQBaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWrWST0+DMBiHW6bOod/BGS94kKpo2Mwyox5x3pbgiXSAFtdR/nQu3IwHiZ9oX88CaeGu9tL0efq+v8P7fmiGA1wgTgoIcECZQqe6vmYD6G4hAJOcB1E8NazTm+ELyzY4C3Q0z8MsRwF+j4JlQUPEEo5wjH0WB9hCNFqgpOCExZY5QnnEw7ME+0v8GuaIs8wnKI7RigVrKkBzm0lh2OdjkeHllG28f066vK6SfEypF60S+vuYt4gjj2fYr/uPrSvRv356TepfJ9iWJRN0OaELQSZN3FRPNbcP1PTSntMr0x0HzLZQjPYIEo3UaFeiISRKH0Mil+BE/dyT1m7tCBLwVO1MX4DK3bbuTlXuy8r71j5Aoho66udAoseOnrdVzx28UFW6ROuO/lT6QKKyo79VU54emj9QSwcInsUTEDMBAAAFAwAAUEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAAZAAYAc2ltcGxlbW9kZWwvY29uc3RhbnRzLnBrbEZCAgBaWoACKS5QSwcIbS8JVwQAAAAEAAAAUEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAATADsAc2ltcGxlbW9kZWwvdmVyc2lvbkZCNwBaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaMwpQSwcI0Z5nVQIAAAACAAAAUEsBAgAAAAAICAAAAAAAAFzqQQQ0AAAANAAAABQAAAAAAAAAAAAAAAAAAAAAAHNpbXBsZW1vZGVsL2RhdGEucGtsUEsBAgAAFAAICAgAAAAAAE0EGH3yAAAAlwEAAB0AAAAAAAAAAAAAAAAAhAAAAHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5UEsBAgAAFAAICAgAAAAAAJ7FExAzAQAABQMAACcAAAAAAAAAAAAAAAAAAgIAAHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5LmRlYnVnX3BrbFBLAQIAAAAACAgAAAAAAABtLwlXBAAAAAQAAAAZAAAAAAAAAAAAAAAAAMMDAABzaW1wbGVtb2RlbC9jb25zdGFudHMucGtsUEsBAgAAAAAICAAAAAAAANGeZ1UCAAAAAgAAABMAAAAAAAAAAAAAAAAAFAQAAHNpbXBsZW1vZGVsL3ZlcnNpb25QSwYGLAAAAAAAAAAeAy0AAAAAAAAAAAAFAAAAAAAAAAUAAAAAAAAAagEAAAAAAACSBAAAAAAAAFBLBgcAAAAA/AUAAAAAAAABAAAAUEsFBgAAAAAFAAUAagEAAJIEAAAAAA==",
            "total_parts": 1
          }
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model:
        model_id: "another_test_model"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "pass_through": {
                "tokenization": {
                  "bert": {
                    "with_special_tokens": false
                  }
                }
              }
            }
          }
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model_vocabulary:
        model_id: "another_test_model"
        body: >
          { "vocabulary": ["[PAD]","[UNK]","these", "are", "my", "words"] }
  - do:
      headers:
        Authorization: "Basic eF9wYWNrX3Jlc3RfdXNlcjp4LXBhY2stdGVzdC1wYXNzd29yZA==" # run as x_pack_rest_user, i.e. the test setup superuser
      ml.put_trained_model_definition_part:
        model_id: "another_test_model"
        part: 0
        body: >
          {
            "total_definition_length":1630,
            "definition": "UEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAAUAA4Ac2ltcGxlbW9kZWwvZGF0YS5wa2xGQgoAWlpaWlpaWlpaWoACY19fdG9yY2hfXwpTdXBlclNpbXBsZQpxACmBfShYCAAAAHRyYWluaW5ncQGIdWJxAi5QSwcIXOpBBDQAAAA0AAAAUEsDBBQACAgIAAAAAAAAAAAAAAAAAAAAAAAdAEEAc2ltcGxlbW9kZWwvY29kZS9fX3RvcmNoX18ucHlGQj0AWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWnWOMWvDMBCF9/yKI5MMrnHTQsHgjt2aJdlCEIp9SgWSTpykFvfXV1htaYds0nfv473JqhjhkAPywbhgUbzSnC02wwZAyqBYOUzIUUoY4XRe6SVr/Q8lVsYbf4UBLkS2kBk1aOIPxbOIaPVQtEQ8vUnZ/WlrSxTA+JCTNHMc4Ig+Eles+Jod+iR3N/jDDf74wxu4e/5+DmtE9mUyhdgFNq7bZ3ekehbruC6aTxS/c1rom6Z698WrEfIYxcn4JGTftLA7tzCnJeD41IJVC+U07kumUHw3E47Vqh+xnULeFisYLx064mV8UTZibWFMmX0p23wBUEsHCE0EGH3yAAAAlwEAAFBLAwQUAAgICAAAAAAAAAAAAAAAAAAAAAAAJwA5AHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5LmRlYnVnX3BrbEZCNQBaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWrWST0+DMBiHW6bOod/BGS94kKpo2Mwyox5x3pbgiXSAFtdR/nQu3IwHiZ9oX88CaeGu9tL0efq+v8P7fmiGA1wgTgoIcECZQqe6vmYD6G4hAJOcB1E8NazTm+ELyzY4C3Q0z8MsRwF+j4JlQUPEEo5wjH0WB9hCNFqgpOCExZY5QnnEw7ME+0v8GuaIs8wnKI7RigVrKkBzm0lh2OdjkeHllG28f066vK6SfEypF60S+vuYt4gjj2fYr/uPrSvRv356TepfJ9iWJRN0OaELQSZN3FRPNbcP1PTSntMr0x0HzLZQjPYIEo3UaFeiISRKH0Mil+BE/dyT1m7tCBLwVO1MX4DK3bbuTlXuy8r71j5Aoho66udAoseOnrdVzx28UFW6ROuO/lT6QKKyo79VU54emj9QSwcInsUTEDMBAAAFAwAAUEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAAZAAYAc2ltcGxlbW9kZWwvY29uc3RhbnRzLnBrbEZCAgBaWoACKS5QSwcIbS8JVwQAAAAEAAAAUEsDBAAACAgAAAAAAAAAAAAAAAAAAAAAAAATADsAc2ltcGxlbW9kZWwvdmVyc2lvbkZCNwBaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaWlpaMwpQSwcI0Z5nVQIAAAACAAAAUEsBAgAAAAAICAAAAAAAAFzqQQQ0AAAANAAAABQAAAAAAAAAAAAAAAAAAAAAAHNpbXBsZW1vZGVsL2RhdGEucGtsUEsBAgAAFAAICAgAAAAAAE0EGH3yAAAAlwEAAB0AAAAAAAAAAAAAAAAAhAAAAHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5UEsBAgAAFAAICAgAAAAAAJ7FExAzAQAABQMAACcAAAAAAAAAAAAAAAAAAgIAAHNpbXBsZW1vZGVsL2NvZGUvX190b3JjaF9fLnB5LmRlYnVnX3BrbFBLAQIAAAAACAgAAAAAAABtLwlXBAAAAAQAAAAZAAAAAAAAAAAAAAAAAMMDAABzaW1wbGVtb2RlbC9jb25zdGFudHMucGtsUEsBAgAAAAAICAAAAAAAANGeZ1UCAAAAAgAAABMAAAAAAAAAAAAAAAAAFAQAAHNpbXBsZW1vZGVsL3ZlcnNpb25QSwYGLAAAAAAAAAAeAy0AAAAAAAAAAAAFAAAAAAAAAAUAAAAAAAAAagEAAAAAAACSBAAAAAAAAFBLBgcAAAAA/AUAAAAAAAABAAAAUEsFBgAAAAAFAAUAagEAAJIEAAAAAA==",
            "total_parts": 1
          }
---
"Test getting and putting Fill Mask with two mask tokens, as well as exceptions caused by requests with the wrong token":
  - do:
      ml.put_trained_model:
        model_id: "bert_fill_mask_model"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "bert": {
                    "with_special_tokens": false
                  }
                }
              }
            }
          }

  - do:
      ml.put_trained_model:
        model_id: "roberta_fill_mask_model"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "roberta": {
                    "with_special_tokens": false
                  }
                }
              }
            }
          }
  - do:
      ml.put_trained_model:
        model_id: "with_correct_mask_token"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "bert": {
                    "with_special_tokens": false
                  }
                },
                "mask_token": "[MASK]"
              }
            }
          }
  - do:
      ml.put_trained_model:
        model_id: "with_other_correct_mask_token"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "roberta": {
                    "with_special_tokens": false
                  }
                },
                "mask_token": "<mask>"
              }
            }
          }
  - do:
      ml.get_trained_models:
        model_id: "bert_fill_mask_model"
  - match: { trained_model_configs.0.inference_config.fill_mask.mask_token: "[MASK]" }
  - do:
      ml.get_trained_models:
        model_id: "roberta_fill_mask_model"
  - match: { trained_model_configs.0.inference_config.fill_mask.mask_token: "<mask>" }
  - do:
      catch: /IllegalArgumentException. Mask token requested was \[<mask>\] but must be \[\[MASK\]\] for this model/
      ml.put_trained_model:
        model_id: "incorrect_mask_token"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "bert": {
                    "with_special_tokens": false
                  }
                },
                "mask_token": "<mask>"
              }
            }
          }
  - do:
      catch: /IllegalArgumentException. Mask token requested was \[\[MASK\]\] but must be \[<mask>\] for this model/
      ml.put_trained_model:
        model_id: "incorrect_mask_token"
        body: >
          {
            "description": "simple model for testing",
            "model_type": "pytorch",
            "inference_config": {
              "fill_mask": {
                "tokenization":{
                  "roberta": {
                    "with_special_tokens": false
                  }
                },
                "mask_token": "[MASK]"
              }
            }
          }
---
"Test start deployment fails with missing model definition":

  - do:
      ml.put_trained_model:
        model_id: distilbert-finetuned-sst
        body: >
          {
            "description": "distilbert-base-uncased-finetuned-sst-2-english.pt",
            "model_type": "pytorch",
            "inference_config": {
              "ner": { }
            }
          }

  - do:
      catch: /Could not find trained model definition \[distilbert-finetuned-sst\]/
      ml.start_trained_model_deployment:
        model_id: distilbert-finetuned-sst

---
"Test start deployment fails while model download in progress":
  - do:
      ml.put_trained_model:
        model_id: .elser_model_2
        body: >
          {
            "input": {
          	"field_names": ["text_field"]
            }
          }
  # Set a low timeout so the test doesn't actually wait
  # for the model download to complete
  - do:
      catch: /Model download task is currently running\. Wait for trained model \[.elser_model_2\] download task to complete then try again/
      ml.start_trained_model_deployment:
        model_id: .elser_model_2
        timeout: 1s
  - do:
      ml.delete_trained_model:
        model_id: .elser_model_2
  - do:
      catch: /No known trained model with model_id \[.elser_model_2\]/
      ml.start_trained_model_deployment:
        model_id: .elser_model_2
---
"Test start and stop deployment with no cache":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        cache_size: 0
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.cache_size: "0" }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }
---
"Test start and stop deployment with cache":
  - skip:
      features: allowed_warnings

  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        cache_size: 10kb
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.cache_size: 10kb }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "are" }
            ]
          }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.get_trained_models_stats:
        model_id: "test_model"
  - match: { count: 1 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_count: 3 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_cache_hit_count: 1 }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }
---
"Test start and stop deployment with low priority":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.priority: low }
  - match: { assignment.task_parameters.number_of_allocations: 1 }
  - match: { assignment.task_parameters.threads_per_allocation: 1 }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }
---
"Test start deployment with low priority and multiple allocations":
  - do:
      catch: /\[number_of_allocations\] must be 1 when \[priority\] is low/
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        number_of_allocations: 3
---
"Test start deployment with low priority and multiple threads per allocation":
  - do:
      catch: /\[threads_per_allocation\] must be 1 when \[priority\] is low/
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        threads_per_allocation: 4
---
"Test start deployment with low priority and max allocations more than one":
  - do:
      catch: /\[max_number_of_allocations\] must be 1 when \[priority\] is low/
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        body:
          adaptive_allocations:
            enabled: true
            min_number_of_allocations: 0
            max_number_of_allocations: 2
---
"Test update deployment":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      # We update to the same value of 1 as if the test runs on a node with just 1 processor it would fail otherwise
      ml.update_trained_model_deployment:
        model_id: test_model
        body: >
          {
            "number_of_allocations": 1
          }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }

  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_deployment
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      # We update to the same value of 1 as if the test runs on a node with just 1 processor it would fail otherwise
      ml.update_trained_model_deployment:
        model_id: test_model_deployment
        body: >
          {
            "number_of_allocations": 1
          }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  # Update with a query parameter
  - do:
      # We update to the same value of 1 as if the test runs on a node with just 1 processor it would fail otherwise
      ml.update_trained_model_deployment:
        model_id: test_model_deployment
        number_of_allocations: 1
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      catch: /\[number_of_allocations\] must be a positive integer/
      ml.update_trained_model_deployment:
        model_id: test_model_deployment
        number_of_allocations: -1
---
"Test clear deployment cache":
  - skip:
      features: allowed_warnings

  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_deployment_cache_test
        cache_size: 10kb
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.cache_size: 10kb }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "are" }
            ]
          }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.get_trained_models_stats:
        model_id: "test_model"
  - match: { count: 1 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_count: 3 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_cache_hit_count: 1 }


  - do:
      ml.clear_trained_model_deployment_cache:
        model_id: test_model_deployment_cache_test
  - match: { cleared: true }

  - do:
      ml.infer_trained_model:
        model_id: "test_model"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.get_trained_models_stats:
        model_id: "test_model"
  - match: { count: 1 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_count: 4 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_cache_hit_count: 1 }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }
---
"Test update deployment with low priority to multiple allocations":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        number_of_allocations: 1
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.priority: low }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      ml.update_trained_model_deployment:
        model_id: test_model
        body:
          number_of_allocations: 3
      catch: /\[number_of_allocations\] must be 1 when \[priority\] is low/

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model
  - match: { stopped: true }
---
"Test update deployment with low priority and max allocations more than one":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        priority: "low"
        body:
          adaptive_allocations:
            enabled: true
            min_number_of_allocations: 0
            max_number_of_allocations: 1
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.priority: low }
  - match: { assignment.task_parameters.number_of_allocations: 1 }

  - do:
      catch: /\[max_number_of_allocations\] must be 1 when \[priority\] is low/
      ml.update_trained_model_deployment:
        model_id: test_model
        body:
          adaptive_allocations:
            max_number_of_allocations: 2
---
"Test put model alias on pytorch model":
  - do:
      ml.put_trained_model_alias:
        model_alias: "pytorch"
        model_id: "test_model"
  - do:
      ml.get_trained_models:
        model_id: "pytorch"
  - match: { count: 1 }
  - length: { trained_model_configs: 1 }
  - match: { trained_model_configs.0.model_id: "test_model" }

  - do:
      ml.put_trained_model_alias:
        model_alias: "pytorch"
        model_id: "another_test_model"
        reassign: true
  - do:
      ml.get_trained_models:
        model_id: "pytorch"
  - match: { count: 1 }
  - length: { trained_model_configs: 1 }
  - match: { trained_model_configs.0.model_id: "another_test_model" }
---
"Test update model alias on pytorch model to undeployed model":
  - skip:
      features: allowed_warnings
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }

  - do:
      ml.put_trained_model_alias:
        model_alias: "pytorch"
        model_id: "test_model"

  - do:
      ml.infer_trained_model:
        model_id: "pytorch"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }
  - do:
      ml.get_trained_models_stats:
        model_id: "test_model"
  - match: { count: 1 }
  - match: { trained_model_stats.0.deployment_stats.nodes.0.inference_count: 1 }

  - do:
      catch: /cannot reassign model_alias \[pytorch\] to model \[another_test_model\] from model \[test_model\] as it is not yet deployed/
      ml.put_trained_model_alias:
        model_alias: "pytorch"
        model_id: "another_test_model"
        reassign: true

---
"Test include model definition status":
  - do:
      ml.get_trained_models:
        model_id: test_model
        include: definition_status
  - match: { trained_model_configs.0.fully_defined: true }

---
"Test stop deployments with allow_no_match":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_for_search
        wait_for: started
  - match: { assignment.assignment_state: started }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model_for_search
  - match: { stopped: true }

  - do:
      catch: missing
      ml.stop_trained_model_deployment:
        allow_no_match: false
        model_id: test_model_for_search
  - match: { error.reason: "No known model deployment with id [test_model_for_search]" }

  - do:
      ml.stop_trained_model_deployment:
        allow_no_match: true
        model_id: test_model_for_search
  - match: { stopped: true }

---
"Test start and stop multiple deployments":
  - skip:
      features: allowed_warnings
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_for_search
        priority: low
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.deployment_id: test_model_for_search }
  - do:
      ml.infer_trained_model:
        model_id: "test_model_for_search"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }

  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_for_ingest
        priority: low
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.deployment_id: test_model_for_ingest }
  - do:
      ml.infer_trained_model:
        model_id: "test_model_for_ingest"
        body: >
          {
            "docs": [
              { "input": "words" }
            ]
          }
  - do:
      ml.get_trained_models_stats:
        model_id: test_model
  - match: { count: 1 } # one model matched
  - match: { trained_model_stats.0.model_id: test_model }
  - match: { trained_model_stats.0.deployment_stats.deployment_id: test_model_for_ingest }
  - match: { trained_model_stats.1.model_id: test_model }
  - match: { trained_model_stats.1.deployment_stats.deployment_id: test_model_for_search }

  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model_for_search
  - match: { stopped: true }
  - do:
      ml.stop_trained_model_deployment:
        model_id: test_model_for_ingest
  - match: { stopped: true }

---
"Test cannot start 2 deployments with the same Id":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_deployment
        wait_for: started
  - match: { assignment.assignment_state: started }

  - do:
      catch: /Could not start model deployment because an existing deployment with the same id \[test_model_deployment\] exist/
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: test_model_deployment
        wait_for: started

---
"Test cannot start when deployment Id matches a different model":
  - do:
      catch: /Deployment id \[another_test_model\] is the same as an another model which is not the model being deployed. Deployment id can be the same as the model being deployed but cannot match a different model/
      ml.start_trained_model_deployment:
        model_id: test_model
        deployment_id: another_test_model

---
"Test start deployment with default Id":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        wait_for: started
  - match: { assignment.assignment_state: started }
  - match: { assignment.task_parameters.model_id: test_model }
  - match: { assignment.task_parameters.deployment_id: test_model }

---
"Test cannot create model with a deployment Id":
  - do:
      ml.start_trained_model_deployment:
        model_id: test_model
        wait_for: started
        deployment_id: test_model_deployment
  - match: { assignment.assignment_state: started }

  - do:
      catch: /Cannot create model \[test_model_deployment\] the model id is the same as the deployment id of a current model deployment/
      ml.put_trained_model:
        model_id: test_model_deployment
        body: >
          {
            "model_type": "pytorch",
            "inference_config": {
              "ner": { }
            }
          }

---
"Test put model config with Japanese tokenizer":
  - do:
      ml.put_trained_model:
        model_id: j_bert
        body: >
          {
            "description": "model config with Japanese tokenizer",
            "model_type": "pytorch",
            "inference_config": {
              "pass_through": {
                "tokenization": {
                  "bert_ja": {
                    "with_special_tokens": false,
                    "max_sequence_length": 512,
                    "truncate": "first",
                    "span": -1
                  }
                }
              }
            }
          }

---
"Test put model config with prefix strings":
  - do:
      ml.put_trained_model:
        model_id: model_with_prefixes
        body: >
          {
            "model_type": "pytorch",
            "inference_config": {
              "text_embedding": { }
            },
            "prefix_strings": {
              "search": "this is a query",
              "ingest": "this is a passage"
            }
          }
  - match: { prefix_strings.search: "this is a query" }
  - match: { prefix_strings.ingest: "this is a passage" }

  - do:
      ml.get_trained_models:
        model_id: model_with_prefixes
  - match: { trained_model_configs.0.prefix_strings.search: "this is a query" }
  - match: { trained_model_configs.0.prefix_strings.ingest: "this is a passage" }


  - do:
      ml.put_trained_model:
        model_id: model_with_search_prefix
        body: >
          {
            "model_type": "pytorch",
            "inference_config": {
              "text_embedding": { }
            },
            "prefix_strings": {
              "search": "this is a query"
            }
          }
  - match: { prefix_strings.search: "this is a query" }
  - is_false: prefix_strings.ingest

  - do:
      ml.put_trained_model:
        model_id: model_with_ingest_prefix
        body: >
          {
            "model_type": "pytorch",
            "inference_config": {
              "text_embedding": { }
            },
            "prefix_strings": {
              "ingest": "this is a passage"
            }
          }
  - is_false: prefix_strings.search
  - match: { prefix_strings.ingest: "this is a passage" }
