ML-Commons train api functionality

CHANGELOG.md

@@ -8,6 +8,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add workflow and scripts for automating model listing updating process by @thanawan-atc in ([#210](https://github.com/opensearch-project/opensearch-py-ml/pull/210))
 - Add script to trigger ml-models-release jenkins workflow with generic webhook by @thanawan-atc in ([#211](https://github.com/opensearch-project/opensearch-py-ml/pull/211))
 - Add example notebook for tracing and registering a CLIPTextModel to OpenSearch with the Neural Search plugin by @patrickbarnhart in ([#283](https://github.com/opensearch-project/opensearch-py-ml/pull/283))
+- Add support for train api functionality by @rawwar in ([#310](https://github.com/opensearch-project/opensearch-py-ml/pull/310))
 
 ### Changed
 - Modify ml-models.JenkinsFile so that it takes model format into account and can be triggered with generic webhook by @thanawan-atc in ([#211](https://github.com/opensearch-project/opensearch-py-ml/pull/211))

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 = "_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}/{ModelTrain.API_ENDPOINT}/{algorithm_name}",
+            body=input_json,
+            params=params,
+        )

tests/ml_commons/test_ml_commons_client.py

@@ -513,6 +513,3 @@ def test_integration_model_train_register_full_cycle():
                 except:  # noqa: E722
                     raised = True
                 assert raised == False, "Raised Exception in deleting model"
-
-
-test_integration_model_train_register_full_cycle()

tests/ml_commons/test_model_train.py

@@ -0,0 +1,103 @@
+# 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 time
+
+import pytest
+from opensearchpy import OpenSearch, helpers
+from sklearn.datasets import load_iris
+
+from opensearch_py_ml.ml_commons import MLCommonClient, ModelTrain
+from tests import OPENSEARCH_TEST_CLIENT
+
+ml_client = MLCommonClient(OPENSEARCH_TEST_CLIENT)
+
+
+@pytest.fixture
+def iris_index():
+    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 ml_client._client.indices.exists(index=index_name):
+        ml_client._client.indices.delete(index=index_name)
+    ml_client._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(ml_client._client, actions)
+    # without the sleep, test is failing.
+    time.sleep(2)
+
+    yield index_name
+
+    ml_client._client.indices.delete(index=index_name)
+
+
+def test_init():
+    assert isinstance(ml_client._client, OpenSearch)
+    assert isinstance(ml_client._model_train, ModelTrain)
+
+
+def test_train(iris_index):
+    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": [iris_index],
+    }
+    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": [iris_index],
+    }
+    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"