Replicas limit (#465)

* try to mock * add tests * use read_env_vars_and_defaults_ints * add new a test * add new a test
marqo-ai · May 8, 2023 · 15f3093 · 15f3093
1 parent 8dcfb3d
commit 15f3093
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 9 deletions.
diff --git a/src/marqo/tensor_search/configs.py b/src/marqo/tensor_search/configs.py
@@ -59,6 +59,7 @@ def default_env_vars() -> dict:
         EnvVars.MARQO_EF_CONSTRUCTION_MAX_VALUE: 4096,
         EnvVars.MARQO_MAX_VECTORISE_BATCH_SIZE: 16,
         EnvVars.MARQO_MAX_DELETE_DOCS_COUNT: 10000,
-        EnvVars.MARQO_MAX_SEARCHABLE_TENSOR_ATTRIBUTES: None
+        EnvVars.MARQO_MAX_SEARCHABLE_TENSOR_ATTRIBUTES: None,
+        EnvVars.MARQO_MAX_NUMBER_OF_REPLICAS: 1,
     }
 
diff --git a/src/marqo/tensor_search/enums.py b/src/marqo/tensor_search/enums.py
@@ -107,6 +107,7 @@ class EnvVars:
     MARQO_MAX_VECTORISE_BATCH_SIZE = "MARQO_MAX_VECTORISE_BATCH_SIZE"
     MARQO_MAX_SEARCHABLE_TENSOR_ATTRIBUTES = "MARQO_MAX_SEARCHABLE_TENSOR_ATTRIBUTES"
     MARQO_MAX_DELETE_DOCS_COUNT = "MARQO_MAX_DELETE_DOCS_COUNT"
+    MARQO_MAX_NUMBER_OF_REPLICAS = "MARQO_MAX_NUMBER_OF_REPLICAS"
 
 
 class RequestType:

diff --git a/src/marqo/tensor_search/models/settings_object.py b/src/marqo/tensor_search/models/settings_object.py
@@ -1,6 +1,6 @@
 from marqo.tensor_search import enums as ns_enums
 from marqo.tensor_search.enums import IndexSettingsField as NsFields, EnvVars
-from marqo.tensor_search.utils import read_env_vars_and_defaults
+from marqo.tensor_search.utils import read_env_vars_and_defaults, read_env_vars_and_defaults_ints
 
 settings_schema = {
     "$schema": "https://json-schema.org/draft/2019-09/schema",
@@ -128,7 +128,7 @@
                                 NsFields.hnsw_ef_construction: {
                                     "type": "integer",
                                     "minimum": 1,
-                                    "maximum": int(read_env_vars_and_defaults(EnvVars.MARQO_EF_CONSTRUCTION_MAX_VALUE)),
+                                    "maximum": read_env_vars_and_defaults_ints(EnvVars.MARQO_EF_CONSTRUCTION_MAX_VALUE),
                                     "examples": [
                                         128
                                     ]
@@ -192,6 +192,7 @@
         NsFields.number_of_replicas: {
             "type": "integer",
             "minimum": 0,
+            "maximum": read_env_vars_and_defaults_ints(EnvVars.MARQO_MAX_NUMBER_OF_REPLICAS),
             "examples": [
                 1
             ]

diff --git a/tests/tensor_search/test_create_index.py b/tests/tensor_search/test_create_index.py
@@ -1,7 +1,6 @@
 import pprint
 from typing import Any, Dict
-import pytest
-import os
+from unittest.mock import patch
 import requests
 from marqo.tensor_search.enums import IndexSettingsField, EnvVars
 from marqo.errors import MarqoApiError, MarqoError, IndexNotFoundError
@@ -11,6 +10,7 @@
 from marqo.tensor_search.enums import IndexSettingsField as NsField
 from unittest import mock
 from marqo import errors
+from marqo.errors import InvalidArgError
 
 class TestCreateIndex(MarqoTestCase):
 
@@ -324,24 +324,134 @@ def test_set_number_of_shards(self):
         assert intended_shard_count == int(resp.json()[self.index_name_1]['settings']['index']['number_of_shards'])
 
     def test_set_number_of_replicas(self):
-        """ does it work if other params are filled?"""
         intended_replicas_count = 4
-        res_0 = tensor_search.create_vector_index(
+        from marqo.tensor_search.models.settings_object import settings_schema
+        with patch.dict(settings_schema["properties"][NsField.number_of_replicas], maximum=10):
+            res_0 = tensor_search.create_vector_index(
+                index_name=self.index_name_1, config=self.config,
+                index_settings={
+                    "index_defaults": {
+                        "treat_urls_and_pointers_as_images": True,
+                        "model": "ViT-B/32",
+                    },
+                    NsField.number_of_replicas: intended_replicas_count
+                }
+            )
+        resp = requests.get(
+            url=self.authorized_url + f"/{self.index_name_1}",
+            headers=self.generic_header,
+            verify=False
+        )
+        assert intended_replicas_count == int(resp.json()[self.index_name_1]['settings']['index']['number_of_replicas'])
+
+    def test_configurable_max_number_of_replicas(self):
+        maximum_number_of_replicas = 5
+        large_intended_replicas_count = 10
+        small_intended_replicas_count = 3
+        from marqo.tensor_search.models.settings_object import settings_schema
+
+        with patch.dict(settings_schema["properties"][NsField.number_of_replicas], maximum=maximum_number_of_replicas):
+            # a large value exceeding limits should not work
+            try:
+                res_0 = tensor_search.create_vector_index(
+                    index_name=self.index_name_1, config=self.config,
+                    index_settings={
+                        "index_defaults": {
+                            "treat_urls_and_pointers_as_images": True,
+                            "model": "ViT-B/32",
+                        },
+                        NsField.number_of_replicas: large_intended_replicas_count
+                    }
+                )
+                raise AssertionError
+            except InvalidArgError as e:
+                pass
+
+            try:
+                tensor_search.delete_index(config=self.config, index_name=self.index_name_1)
+            except IndexNotFoundError:
+                pass
+
+            # a small value should work
+            res_1 = tensor_search.create_vector_index(
+                index_name=self.index_name_1, config=self.config,
+                index_settings={
+                    "index_defaults": {
+                        "treat_urls_and_pointers_as_images": True,
+                        "model": "ViT-B/32",
+                    },
+                    NsField.number_of_replicas: small_intended_replicas_count
+                }
+            )
+            resp = requests.get(
+                url=self.authorized_url + f"/{self.index_name_1}",
+                headers=self.generic_header,
+                verify=False
+            )
+            assert small_intended_replicas_count == int(
+                resp.json()[self.index_name_1]['settings']['index']['number_of_replicas'])
+
+            try:
+                tensor_search.delete_index(config=self.config, index_name=self.index_name_1)
+            except IndexNotFoundError:
+                pass
+
+            # the same number should also work
+            res_1 = tensor_search.create_vector_index(
+                index_name=self.index_name_1, config=self.config,
+                index_settings={
+                    "index_defaults": {
+                        "treat_urls_and_pointers_as_images": True,
+                        "model": "ViT-B/32",
+                    },
+                    NsField.number_of_replicas: maximum_number_of_replicas
+                }
+            )
+            resp = requests.get(
+                url=self.authorized_url + f"/{self.index_name_1}",
+                headers=self.generic_header,
+                verify=False
+            )
+            assert maximum_number_of_replicas == int(
+                resp.json()[self.index_name_1]['settings']['index']['number_of_replicas'])
+
+    def test_default_max_number_of_replicas(self):
+        large_intended_replicas_count = 2
+        small_intended_replicas_count = 0
+        # a large value exceeding limits should not work
+        try:
+            res_0 = tensor_search.create_vector_index(
+                index_name=self.index_name_1, config=self.config,
+                index_settings={
+                    "index_defaults": {
+                        "treat_urls_and_pointers_as_images": True,
+                        "model": "ViT-B/32",
+                    },
+                    NsField.number_of_replicas: large_intended_replicas_count
+                }
+            )
+            raise AssertionError
+        except InvalidArgError as e:
+            pass
+
+        # a small value should work
+        res_1 = tensor_search.create_vector_index(
             index_name=self.index_name_1, config=self.config,
             index_settings={
                 "index_defaults": {
                     "treat_urls_and_pointers_as_images": True,
                     "model": "ViT-B/32",
                 },
-                NsField.number_of_replicas: intended_replicas_count
+                NsField.number_of_replicas: small_intended_replicas_count
             }
         )
         resp = requests.get(
             url=self.authorized_url + f"/{self.index_name_1}",
             headers=self.generic_header,
             verify=False
         )
-        assert intended_replicas_count == int(resp.json()[self.index_name_1]['settings']['index']['number_of_replicas'])
+        assert small_intended_replicas_count == int(
+            resp.json()[self.index_name_1]['settings']['index']['number_of_replicas'])
 
     def test_field_limits(self):
         index_limits = [1, 5, 10, 100, 1000]