Optmize scale helm env (#7263)

Signed-off-by: NThreadDao <yufen.zong@zilliz.com>

tests/python_client/scale/helm_env.py

@@ -6,6 +6,8 @@ from utils.util_log import test_log as log
 from common import common_func as cf
 from scale import scale_common as sc
 
+milvus_chart_path = sc.get_milvus_chart_env_var()
+
 
 class HelmEnv:
     def __init__(self, release_name=None, **kwargs):
@@ -37,14 +39,11 @@ class HelmEnv:
                       f'--set queryNode.replicas={self.query_node} ' \
                       f'{self.release_name} . '
         log.debug(f'install_cmd: {install_cmd}')
-        log.debug(f'MILVUS CHART: {sc.get_milvus_chart_env_var()}')
         try:
-            os.system(f'cd {sc.get_milvus_chart_env_var()} && {install_cmd}')
-        except Exception as e:
-            raise
-        # raise Exception("Failed to deploy cluster milvus")
-        #     todo
-        #   return svc ip
+            os.system(f'cd {milvus_chart_path} && {install_cmd}')
+        except Exception:
+            raise Exception("Failed to deploy cluster milvus")
+        return self.get_service_ip()
 
     def helm_upgrade_cluster_milvus(self, **kwargs):
         """
@@ -57,19 +56,14 @@ class HelmEnv:
         data_node = kwargs.get(constants.DATA_NODE, self.data_node)
         index_node = kwargs.get(constants.INDEX_NODE, self.index_node)
         query_node = kwargs.get(constants.QUERY_NODE, self.query_node)
-        upgrade_cmd = f'helm upgrade --install ' \
-                      f'--set image.all.repository={constants.IMAGE_REPOSITORY} ' \
-                      f'--set image.all.tag={constants.IMAGE_TAG} ' \
-                      f'--set cluster.enabled=true ' \
-                      f'--set service.type=LoadBalancer ' \
+        upgrade_cmd = f'helm upgrade {self.release_name} . ' \
                       f'--set proxy.replicas={proxy} ' \
                       f'--set dataNode.replicas={data_node} ' \
                       f'--set indexNode.replicas={index_node} ' \
                       f'--set queryNode.replicas={query_node} ' \
-                      f'{self.release_name} . '
+                      f'--reuse-values'
         log.debug(f'upgrade_cmd: {upgrade_cmd}')
-        log.debug(f'MILVUS CHART: {sc.get_milvus_chart_env_var()}')
-        if os.system(f'cd {sc.get_milvus_chart_env_var()} && {upgrade_cmd}'):
+        if os.system(f'cd {milvus_chart_path} && {upgrade_cmd}'):
             raise Exception(f'Failed to upgrade cluster milvus with {kwargs}')
 
     def helm_uninstall_cluster_milvus(self):
@@ -87,7 +81,7 @@ class HelmEnv:
         # delete plusar
         # delete_pvc_plusar_cmd = "kubectl delete pvc scale-test-milvus-pulsar"
 
-    def get_svc_external_ip(self):
+    def get_service_ip(self):
         from kubernetes import client, config
         # from kubernetes.client.rest import ApiException
         config.load_kube_config()
@@ -120,12 +114,11 @@ class HelmEnv:
 
 if __name__ == '__main__':
     # default deploy q replicas
-    release_name = "milvus-chaos"
+    release_name = "scale-proxy"
     env = HelmEnv(release_name=release_name)
-    # host = env.get_svc_external_ip()
-    # log.debug(host)
     # env.helm_install_cluster_milvus()
+    # host = env.get_service_ip()
     # env.helm_upgrade_cluster_milvus(queryNode=2)
-    env.helm_uninstall_cluster_milvus()
+    # env.helm_uninstall_cluster_milvus()
+    env.export_all_logs()
     # sleep(5)
-    # env.export_all_logs()

tests/python_client/scale/scale_common.py

@@ -23,21 +23,23 @@ def get_milvus_chart_env_var(var=constants.MILVUS_CHART_ENV):
     return milvus_helm_chart
 
 
-def e2e_milvus(host, c_name):
+def e2e_milvus(host, c_name, collection_exist=False):
     # connect
     connections.add_connection(default={"host": host, "port": 19530})
     connections.connect(alias='default')
 
     # create
-    # c_name = cf.gen_unique_str(prefix)
     collection_w = ApiCollectionWrapper()
-    collection_w.init_collection(name=c_name, schema=cf.gen_default_collection_schema())
-    # collection_w.init_collection(name=c_name)
+    if collection_exist:
+        collection_w.init_collection(name=c_name)
+    else:
+        collection_w.init_collection(name=c_name, schema=cf.gen_default_collection_schema())
 
     # insert
     data = cf.gen_default_list_data(ct.default_nb)
     mutation_res, _ = collection_w.insert(data)
     assert mutation_res.insert_count == ct.default_nb
+    log.debug(collection_w.num_entities)
 
     # create index
     collection_w.create_index(ct.default_float_vec_field_name, ct.default_index)
@@ -50,6 +52,7 @@ def e2e_milvus(host, c_name):
     search_res, _ = collection_w.search(data[-1][:ct.default_nq], ct.default_float_vec_field_name,
                                         ct.default_search_params, ct.default_limit)
     assert len(search_res[0]) == ct.default_limit
+    log.debug(search_res[0][0].id)
 
     # query
     ids = search_res[0].ids[0]

tests/python_client/scale/test_data_node_scale.py

@@ -27,10 +27,9 @@ class TestDataNodeScale:
         expected: two collection create and insert op are both correctly
         """
         # deploy all nodes one pod cluster milvus with helm
-        release_name = "scale-test"
+        release_name = "scale-data"
         env = HelmEnv(release_name=release_name)
-        env.helm_install_cluster_milvus()
-        host = env.get_svc_external_ip()
+        host = env.helm_install_cluster_milvus()
 
         # connect
         connections.add_connection(default={"host": host, "port": 19530})
@@ -71,12 +70,12 @@ class TestDataNodeScale:
         method: 1.create collection and insert df 2. shrink dataNode 3.insert df
         expected: verify the property of collection which channel on shrink pod
         """
-        release_name = "scale-test"
+        release_name = "scale-data"
         env = HelmEnv(release_name=release_name, dataNode=2)
-        env.helm_install_cluster_milvus(image_pull_policy=constants.IF_NOT_PRESENT)
+        host = env.helm_install_cluster_milvus(image_pull_policy=constants.IF_NOT_PRESENT)
 
         # connect
-        connections.add_connection(default={"host": '10.98.0.8', "port": 19530})
+        connections.add_connection(default={"host": host, "port": 19530})
         connections.connect(alias='default')
 
         c_name = "data_scale_one"

tests/python_client/scale/test_index_node_scale.py

@@ -28,8 +28,7 @@ class TestIndexNodeScale:
         """
         release_name = "scale-index"
         env = HelmEnv(release_name=release_name)
-        env.helm_install_cluster_milvus()
-        host = env.get_svc_external_ip()
+        host = env.helm_install_cluster_milvus()
 
         # connect
         connections.add_connection(default={"host": host, "port": 19530})
@@ -82,10 +81,10 @@ class TestIndexNodeScale:
         """
         release_name = "scale-index"
         env = HelmEnv(release_name=release_name, indexNode=2)
-        env.helm_install_cluster_milvus()
+        host = env.helm_install_cluster_milvus()
 
         # connect
-        connections.add_connection(default={"host": '10.98.0.8', "port": 19530})
+        connections.add_connection(default={"host": host, "port": 19530})
         connections.connect(alias='default')
 
         data = cf.gen_default_dataframe_data(nb)
@@ -112,8 +111,7 @@ class TestIndexNodeScale:
         collection_w.drop_index()
         assert not collection_w.has_index()[0]
 
-        # expand indexNode from 1 to 2
-        # pdb.set_trace()
+        # expand indexNode from 2 to 1
         env.helm_upgrade_cluster_milvus(indexNode=1)
 
         start = datetime.datetime.now()

tests/python_client/scale/test_proxy_scale.py

 import pytest
+from pymilvus import connections
 
 from scale.helm_env import HelmEnv
 from common import common_func as cf
@@ -25,9 +26,7 @@ class TestProxyScale:
         # deploy all nodes one pod cluster milvus with helm
         release_name = "scale-proxy"
         env = HelmEnv(release_name=release_name)
-        env.helm_install_cluster_milvus()
-        host = env.get_svc_external_ip()
-        # host = "10.98.0.8"
+        host = env.helm_install_cluster_milvus()
 
         c_name = cf.gen_unique_str(prefix)
         sc.e2e_milvus(host, c_name)
@@ -35,8 +34,8 @@ class TestProxyScale:
         # scale proxy
         env.helm_upgrade_cluster_milvus(proxy=2)
 
-        c_name_2 = cf.gen_unique_str(prefix)
-        sc.e2e_milvus(host, c_name_2)
+        # c_name_2 = cf.gen_unique_str(prefix)
+        sc.e2e_milvus(host, c_name, collection_exist=True)
 
     def test_shrink_proxy(self):
         """
@@ -50,9 +49,7 @@ class TestProxyScale:
         # deploy all nodes one pod cluster milvus with helm
         release_name = "scale-proxy"
         env = HelmEnv(release_name=release_name, proxy=2)
-        env.helm_install_cluster_milvus()
-        host = env.get_svc_external_ip()
-        # host = "10.98.0.8"
+        host = env.helm_install_cluster_milvus()
 
         c_name = cf.gen_unique_str(prefix)
         sc.e2e_milvus(host, c_name)
@@ -60,5 +57,5 @@ class TestProxyScale:
         # scale proxy
         env.helm_upgrade_cluster_milvus(proxy=1)
 
-        c_name_2 = cf.gen_unique_str(prefix)
-        sc.e2e_milvus(host, c_name_2)
\ No newline at end of file
+        # c_name_2 = cf.gen_unique_str(prefix)
+        sc.e2e_milvus(host, c_name, collection_exist=True)

tests/python_client/scale/test_query_node_scale.py

@@ -26,8 +26,7 @@ class TestQueryNodeScale:
     def test_expand_query_node(self):
         release_name = "scale-query"
         env = HelmEnv(release_name=release_name)
-        env.helm_install_cluster_milvus()
-        host = env.get_svc_external_ip()
+        host = env.helm_install_cluster_milvus()
 
         # connect
         connections.add_connection(default={"host": host, "port": 19530})
@@ -78,10 +77,10 @@ class TestQueryNodeScale:
         # deploy
         release_name = "scale-query"
         env = HelmEnv(release_name=release_name, queryNode=2)
-        env.helm_install_cluster_milvus(image_pull_policy=constants.IF_NOT_PRESENT)
+        host = env.helm_install_cluster_milvus(image_pull_policy=constants.IF_NOT_PRESENT)
 
         # connect
-        connections.add_connection(default={"host": '10.98.0.8', "port": 19530})
+        connections.add_connection(default={"host": host, "port": 19530})
         connections.connect(alias='default')
 
         # collection one
@@ -108,7 +107,6 @@ class TestQueryNodeScale:
         assert res2[0].ids[0] == data[0][0]
 
         # scale queryNode pod
-        pdb.set_trace()
         env.helm_upgrade_cluster_milvus(queryNode=1)
 
         # search