#289 - Add register, update and delete model group functionality to support Model Access Control

opensearch_py_ml/ml_commons/model_access_control.py

@@ -14,6 +14,7 @@
     validate_create_model_group_parameters,
     validate_update_model_group_parameters,
     validate_delete_model_group_parameters,
+    validate_search_model_group_parameters
 )
 
 
@@ -49,31 +50,18 @@ def register_model_group(
 
     def update_model_group(
         self,
-        name: str,
-        description: str,
-        access_mode: str,
-        backend_roles: List[str],
-        add_all_backend_roles=False,
+        update_query: dict,
+        model_group_id: Optional[str]=None,
+        model_group_name: Optional[str]=None,
     ):
-        validate_update_model_group_parameters(
-            name, description, access_mode, backend_roles, add_all_backend_roles
-        )
-        body = {"name": name, "add_all_backend_roles": add_all_backend_roles}
-        if description:
-            body["description"] = description
-        if access_mode:
-            body["access_mode"] = access_mode
-        if backend_roles:
-            body["backend_roles"] = backend_roles
-
+
+        validate_update_model_group_parameters(update_query, model_group_id, model_group_name)
         return self.client.transport.perform_request(
-            method="PUT", url=f"{ML_BASE_URI}/{self.API_ENDPOINT}/_update", body=body
+            method="PUT", url=f"{ML_BASE_URI}/{self.API_ENDPOINT}/", body=update_query
         )
 
-    def search_model_group(self, query):
-        if isinstance(query, str):
-            query = json.loads(query)
-
+    def search_model_group(self, query: dict):
+        validate_search_model_group_parameters(query)
         return self.client.transport.perform_request(
             method="GET", url=f"{ML_BASE_URI}/{self.API_ENDPOINT}/_search", body=query
         )

opensearch_py_ml/ml_commons/validators/model_access_control.py

@@ -36,6 +36,14 @@ def _validate_model_group_add_all_backend_roles(add_all_backend_roles):
     if not isinstance(add_all_backend_roles, bool):
         raise ValueError("add_all_backend_roles should be a boolean")
 
+def _validate_model_group_query(query, operation=None):
+    if not isinstance(query, dict):
+        raise ValueError("query needs to be a dictionary")
+
+    if operation and not isinstance(operation, str):
+        raise ValueError("operation needs to be a string")
+
+
 
 def validate_create_model_group_parameters(
     name,
@@ -68,16 +76,20 @@ def validate_create_model_group_parameters(
             )
 
 
-def validate_update_model_group_parameters(
-    name,
-    description,
-    access_mode,
-    backend_roles,
-    add_all_backend_roles,
-):
-    validate_create_model_group_parameters(
-        name, description, access_mode, backend_roles, add_all_backend_roles
-    )
+def validate_update_model_group_parameters(update_query, model_group_id, model_group_name):
+    if model_group_id and model_group_name:
+        raise ValueError(
+            "You cannot specify both model_group_id and model_group_name at the same time"
+        )
+
+    if not isinstance(model_group_id, (NoneType, str)):
+        raise ValueError("Invalid model_group_id. model_group_id needs to be a string")
+
+    if not isinstance(model_group_name, (NoneType, str)):
+        raise ValueError("Invalid model_group_name. model_group_name needs to be a string")
+
+    if not isinstance(update_query, dict):
+        raise ValueError("Invalid update_query. update_query needs to be a dictionary")
 
 
 def validate_delete_model_group_parameters(model_group_id, model_group_name):
@@ -91,3 +103,7 @@ def validate_delete_model_group_parameters(model_group_id, model_group_name):
 
     if not isinstance(model_group_name, (NoneType, str)):
         raise ValueError("Invalid model_group_name. model_group_name needs to be a string")
+
+
+def validate_search_model_group_parameters(query):
+    _validate_model_group_query(query)

tests/ml_commons/test_model_access_control.py

@@ -11,8 +11,13 @@
 from opensearch_py_ml.ml_commons.model_access_control import ModelAccessControl
 from tests import OPENSEARCH_TEST_CLIENT
 from opensearchpy.exceptions import RequestError
+from packaging.version import parse as parse_version
+import os
+
+OPENSEARCH_VERSION = parse_version(os.environ['OPENSEARCH_VERSION'])
+MAC_MIN_VERSION = parse_version("2.8.0")
+MAC_UPDATE_MIN_VERSION = parse_version("2.11.0")
 
-OPENSEARCH_VERSION = OPENSEARCH_TEST_CLIENT.info()['version']['number']
 @pytest.fixture
 def client():
     return ModelAccessControl(OPENSEARCH_TEST_CLIENT)
@@ -21,7 +26,7 @@ def client():
 def test_model_group(client):
     model_group_name = "__test__model_group_1"
     client.delete_model_group(model_group_name=model_group_name)
-    # time.sleep(0.5)
+    time.sleep(0.5)
     client.register_model_group(
         name=model_group_name,
         description="test model group for opensearch-py-ml test cases",
@@ -31,6 +36,8 @@ def test_model_group(client):
     client.delete_model_group(model_group_name=model_group_name)
 
 
+
+@pytest.mark.skipif(OPENSEARCH_VERSION < MAC_MIN_VERSION, reason="Model groups are supported in OpenSearch 2.8.0 and above")
 def test_register_model_group(client):
 
     model_group_name1 = "__test__model_group_A"
@@ -80,24 +87,116 @@ def test_register_model_group(client):
         client.register_model_group(name=model_group_name2)
     assert exec_info.value.status_code == 400
     assert exec_info.match("The name you provided is already being used by a model group")
+
+
+@pytest.mark.skipif(OPENSEARCH_VERSION < MAC_UPDATE_MIN_VERSION, reason="Model groups updates are supported in OpenSearch 2.11.0 and above")
+def test_update_model_group(client, test_model_group):
+    # update model group name and description
+    update_query = {
+        "name": "__test__model_group_X",
+        "description": "updated description",
+    }
+    update_successful = False
+    try:
+        res = client.update_model_group(update_query, model_group_name=test_model_group)
+        assert isinstance(res, dict)
+        assert "status" in res
+        assert res['status'] == "Updated"
+        update_successful = True
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
 
-    for each in "ABC":
-        client.delete_model_group(model_group_name=f"__test__model_group_{each}")
+    if update_successful:
+        # revert model group name
+        update_query = {"name": test_model_group}
+        res = client.update_model_group(update_query)
+        assert isinstance(res, dict)
+        assert "status" in res
+        assert res['status'] == "Updated"
 
+
 
-def test_update_model_group():
-    import os
-    env_data = os.environ
-    print("Environ data = ", env_data)
-    assert False, "!@#"
 
-@pytest.mark.skipif(OPENSEARCH_VERSION >= "2.8.0")
-def test_delete_model_group(client):
-    assert False
 
-def test_search_model_group(client):
-    pass
+@pytest.mark.skipif(OPENSEARCH_VERSION < MAC_MIN_VERSION, reason="Model groups are supported in OpenSearch 2.8.0 and above")
+def test_search_model_group(client, test_model_group):
+    query1 = {"query": {"match": {"name": test_model_group}}, "size": 1}
+
+    try:
+        res = client.search_model_group(query1)
+        assert isinstance(res, dict)
+        assert "hits" in res and "hits" in res["hits"]
+        assert len(res['hits']['hits']) == 1
+        assert "_source" in res['hits']['hits']
+        assert "name" in res['hits']['hits']['_source']
+        assert test_model_group == res['hits']['hits']['_source']['name']
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
+
+    query2 = {"query": {"match": {"name": "test-unknown"}}, "size": 1}
+    try:
+        res = client.search_model_group(query2)
+        assert isinstance(res, dict)
+        assert "hits" in res and "hits" in res["hits"]
+        assert len(res['hits']['hits']) == 0
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
+
+@pytest.mark.skipif(OPENSEARCH_VERSION < MAC_MIN_VERSION, reason="Model groups are supported in OpenSearch 2.8.0 and above")
+def test_search_model_group_by_name(client, test_model_group):
+
+    try:
+        res = client.search_model_group_by_name(model_group_name=test_model_group)
+        assert isinstance(res, dict)
+        assert "hits" in res and "hits" in res["hits"]
+        assert len(res['hits']['hits']) == 1
+        assert "_source" in res['hits']['hits']
+        assert len(res['hits']['hits']['_source']) > 1
+        assert "name" in res['hits']['hits']['_source']
+        assert test_model_group == res['hits']['hits']['_source']['name']
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
+
+
+    try:
+        res = client.search_model_group_by_name(model_group_name=test_model_group, _source="name")
+        assert isinstance(res, dict)
+        assert "hits" in res and "hits" in res["hits"]
+        assert len(res['hits']['hits']) == 1
+        assert "_source" in res['hits']['hits']
+        assert len(res['hits']['hits']['_source']) == 1
+        assert "name" in res['hits']['hits']['_source']
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
+
+    try:
+        res = client.search_model_group_by_name(model_group_name="test-unknown")
+        assert isinstance(res, dict)
+        assert "hits" in res and "hits" in res["hits"]
+        assert len(res['hits']['hits']) == 0
+    except Exception as ex:
+        assert False,f"Failed to search model group due to unhandled error: {ex}"
 
 
-def test_search_model_group_by_name(client):
-    pass
+
+@pytest.mark.skipif(OPENSEARCH_VERSION < MAC_MIN_VERSION, reason="Model groups are supported in OpenSearch 2.8.0 and above")
+def test_delete_model_group(client, test_model_group):
+    # create a test model group
+
+    for each in "AB":
+        model_group_name = f"__test__model_group_{each}"
+        res = client.delete_model_group(model_group_name=model_group_name)
+        assert res is None or isinstance(res, dict)
+        if isinstance(res, dict):
+            assert 'result' in res
+            assert res['result'] in ['not_found', 'deleted']
+
+    res = client.delete_model_group(model_group_id="test-unknown")
+    assert isinstance(res, dict)
+    assert "result" in res
+    assert res['result'] == 'not_found'
+
+    res = client.delete_model_group(model_group_name=test_model_group)
+    assert isinstance(res, dict)
+    assert "result" in res
+    assert res['result'] == 'deleted'