ML-Commons train api functionality

conftest.py

@@ -0,0 +1,66 @@
+import pytest
+import time
+import os
+from opensearchpy import OpenSearch, helpers
+from sklearn.datasets import load_iris
+
+@pytest.fixture
+def opensearch_client():
+    opensearch_host = os.environ.get("OPENSEARCH_HOST", "https://localhost:9200")
+    opensearch_admin_user = os.environ.get("OPENSEARCH_ADMIN_USER", "admin")
+    opensearch_admin_password = os.environ.get("OPENSEARCH_ADMIN_PASSWORD", "admin")
+    client = OpenSearch(
+        hosts=[opensearch_host],
+        http_auth=(opensearch_admin_user, opensearch_admin_password),
+        verify_certs=False,
+    )
+    yield client
+
+    # tear down
+    client.transport.close()
+
+
+@pytest.fixture
+def iris_index_client(opensearch_client: OpenSearch):
+    index_name = "test__index__iris_data"
+    index_mapping = {
+        "mappings": {
+            "properties": {
+                "sepal_length": {"type": "float"},
+                "sepal_width": {"type": "float"},
+                "petal_length": {"type": "float"},
+                "petal_width": {"type": "float"},
+                "species": {"type": "keyword"}
+            }
+        }
+    }
+
+    if opensearch_client.indices.exists(index=index_name):
+        opensearch_client.indices.delete(index=index_name)
+    opensearch_client.indices.create(index=index_name, body=index_mapping)
+
+    iris = load_iris()
+    iris_data = iris.data
+    iris_target = iris.target
+    iris_species = [iris.target_names[i] for i in iris_target]
+
+    actions = [
+        {   '_index': index_name,
+            "_source":{
+                "sepal_length": sepal_length,
+                "sepal_width": sepal_width,
+                "petal_length": petal_length,
+                "petal_width": petal_width,
+                "species": species
+            }
+        }
+        for (sepal_length, sepal_width, petal_length, petal_width), species in zip(iris_data, iris_species)
+    ]
+
+    helpers.bulk(opensearch_client, actions)
+    # without the sleep, test is failing.
+    time.sleep(2)
+
+    yield opensearch_client, index_name
+
+    opensearch_client.indices.delete(index=index_name)

opensearch_py_ml/ml_commons/__init__.py

@@ -9,6 +9,7 @@
 
 from opensearch_py_ml.ml_commons.ml_commons_client import MLCommonClient
 from opensearch_py_ml.ml_commons.model_execute import ModelExecute
+from opensearch_py_ml.ml_commons.model_train import ModelTrain
 from opensearch_py_ml.ml_commons.model_uploader import ModelUploader
 
-__all__ = ["MLCommonClient", "ModelExecute", "ModelUploader"]
+__all__ = ["MLCommonClient", "ModelExecute", "ModelUploader", "ModelTrain"]

opensearch_py_ml/ml_commons/ml_commons_client.py

@@ -8,7 +8,7 @@
 
 import json
 import time
-from typing import Any, List, Union
+from typing import Any, List, Optional, Union
 
 from deprecated.sphinx import deprecated
 from opensearchpy import OpenSearch
@@ -22,6 +22,7 @@
     TIMEOUT,
 )
 from opensearch_py_ml.ml_commons.model_execute import ModelExecute
+from opensearch_py_ml.ml_commons.model_train import ModelTrain
 from opensearch_py_ml.ml_commons.model_uploader import ModelUploader
 
 
@@ -35,6 +36,7 @@ def __init__(self, os_client: OpenSearch):
         self._client = os_client
         self._model_uploader = ModelUploader(os_client)
         self._model_execute = ModelExecute(os_client)
+        self._model_train = ModelTrain(os_client)
 
     def execute(self, algorithm_name: str, input_json: dict) -> dict:
         """
@@ -580,3 +582,12 @@ def delete_task(self, task_id: str) -> object:
             method="DELETE",
             url=API_URL,
         )
+
+    def train_model(
+        self, algorithm_name: str, input_json: dict, is_async: Optional[bool] = False
+    ) -> dict:
+        """
+        This method trains an ML model
+        """
+
+        return self._model_train._train(algorithm_name, input_json, is_async)

opensearch_py_ml/ml_commons/model_train.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Any modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+import json
+from typing import Optional
+
+from opensearchpy import OpenSearch
+
+from opensearch_py_ml.ml_commons.ml_common_utils import ML_BASE_URI
+
+
+class ModelTrain:
+    """
+    Class for training models using ML Commons train API.
+    """
+
+    API_ENDPOINT = "models/_train"
+
+    def __init__(self, os_client: OpenSearch):
+        self._client = os_client
+
+    def _train(
+        self, algorithm_name: str, input_json: dict, is_async: Optional[bool] = True
+    ) -> dict:
+        """
+        This method trains an ML model
+        """
+
+        params = {}
+        if not isinstance(input_json, dict):
+            input_json = json.loads(input_json)
+        if is_async:
+            params["async"] = "true"
+
+        return self._client.transport.perform_request(
+            method="POST",
+            url=f"{ML_BASE_URI}/_train/{algorithm_name}",
+            body=input_json,
+            params=params,
+        )

tests/ml_commons/test_model_train.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Any modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+
+from opensearchpy import OpenSearch
+
+from opensearch_py_ml.ml_commons import MLCommonClient, ModelTrain
+
+
+def test_init(opensearch_client):
+    ml_client = MLCommonClient(opensearch_client)
+    assert isinstance(ml_client._client, OpenSearch)
+    assert isinstance(ml_client._model_train, ModelTrain)
+
+
+def test_train(iris_index_client):
+    client, test_index_name = iris_index_client
+    ml_client = MLCommonClient(client)
+    algorithm_name = "kmeans"
+    input_json_sync = {
+        "parameters": {"centroids": 3, "iterations": 10, "distance_type": "COSINE"},
+        "input_query": {
+            "_source": ["petal_length", "petal_width"],
+            "size": 10000,
+        },
+        "input_index": [test_index_name],
+    }
+    response = ml_client.train_model(algorithm_name, input_json_sync)
+    assert isinstance(response, dict)
+    assert "model_id" in response
+    assert "status" in response
+    assert response["status"] == "COMPLETED"
+
+    input_json_async = {
+        "parameters": {"centroids": 3, "iterations": 10, "distance_type": "COSINE"},
+        "input_query": {
+            "_source": ["petal_length", "petal_width"],
+            "size": 10000,
+        },
+        "input_index": [test_index_name],
+    }
+    response = ml_client.train_model(algorithm_name, input_json_async, is_async=True)
+
+    assert isinstance(response, dict)
+    assert "task_id" in response
+    assert "status" in response
+    assert response["status"] == "CREATED"