Skip to content

Commit c48b0fb

Browse files
authored
enhance: Enable resource group api in milvus client (#2513) (#2514)
issue: milvus-io/milvus#38739 pr: #2513 this PR enable those API in milvus client: - create_resource_group - update_resource_groups - drop_resource_group - describe_resource_group - list_resource_groups - transfer_replica cause the resource group prefer to use declarative api, so the transfer_node API is deprecated in milvus client, it's recommand to use update_resource_group to manage node num in resource group. Signed-off-by: Wei Liu <[email protected]>
1 parent 292390d commit c48b0fb

File tree

2 files changed

+205
-0
lines changed

2 files changed

+205
-0
lines changed

examples/resource_group.py

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
from pymilvus import (
2+
MilvusClient,
3+
DataType,
4+
)
5+
from pymilvus.client.constants import DEFAULT_RESOURCE_GROUP
6+
7+
from pymilvus.client.types import (
8+
ResourceGroupConfig,
9+
)
10+
11+
fmt = "\n=== {:30} ===\n"
12+
dim = 8
13+
collection_name = "hello_milvus"
14+
milvus_client = MilvusClient("http://localhost:19530")
15+
16+
17+
## create collection and load collection
18+
print("create collection and load collection")
19+
collection_name = "hello_milvus"
20+
has_collection = milvus_client.has_collection(collection_name, timeout=5)
21+
if has_collection:
22+
milvus_client.drop_collection(collection_name)
23+
24+
schema = milvus_client.create_schema(enable_dynamic_field=True)
25+
schema.add_field("id", DataType.INT64, is_primary=True)
26+
schema.add_field("embeddings", DataType.FLOAT_VECTOR, dim=dim)
27+
schema.add_field("title", DataType.VARCHAR, max_length=64)
28+
milvus_client.create_collection(collection_name, schema=schema, consistency_level="Strong")
29+
index_params = milvus_client.prepare_index_params()
30+
index_params.add_index(field_name = "embeddings", metric_type="L2")
31+
index_params.add_index(field_name = "title", index_type = "Trie", index_name="my_trie")
32+
milvus_client.create_index(collection_name, index_params)
33+
milvus_client.load_collection(collection_name)
34+
35+
36+
## create resource group
37+
print("create resource group")
38+
milvus_client.create_resource_group("rg1")
39+
milvus_client.create_resource_group("rg2")
40+
41+
## update resource group
42+
configs = {
43+
"rg1": ResourceGroupConfig(
44+
requests={"node_num": 1},
45+
limits={"node_num": 5},
46+
transfer_from=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
47+
transfer_to=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
48+
),
49+
"rg2": ResourceGroupConfig(
50+
requests={"node_num": 4},
51+
limits={"node_num": 4},
52+
transfer_from=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
53+
transfer_to=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
54+
),
55+
}
56+
milvus_client.update_resource_groups(configs)
57+
58+
## describe resource group
59+
print("describe rg1")
60+
result = milvus_client.describe_resource_group("rg1")
61+
print(result)
62+
63+
print("describe rg2")
64+
result = milvus_client.describe_resource_group("rg2")
65+
print(result)
66+
67+
## list resource group
68+
print("list resource group")
69+
result = milvus_client.list_resource_groups()
70+
print(result)
71+
72+
## transfer replica
73+
print("transfer replica to rg1")
74+
milvus_client.transfer_replica(DEFAULT_RESOURCE_GROUP, "rg1", collection_name, 1)
75+
print("describe rg1 after transfer replica in")
76+
result = milvus_client.describe_resource_group("rg1")
77+
print(result)
78+
79+
milvus_client.transfer_replica("rg1", DEFAULT_RESOURCE_GROUP, collection_name, 1)
80+
print("describe rg1 after transfer replica out")
81+
result = milvus_client.describe_resource_group("rg1")
82+
print(result)
83+
84+
## drop resource group
85+
print("drop resource group")
86+
# create resource group
87+
configs = {
88+
"rg1": ResourceGroupConfig(
89+
requests={"node_num": 0},
90+
limits={"node_num": 0},
91+
transfer_from=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
92+
transfer_to=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
93+
),
94+
"rg2": ResourceGroupConfig(
95+
requests={"node_num": 0},
96+
limits={"node_num": 0},
97+
transfer_from=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
98+
transfer_to=[{"resource_group": DEFAULT_RESOURCE_GROUP}],
99+
),
100+
}
101+
milvus_client.update_resource_groups(configs)
102+
milvus_client.drop_resource_group("rg1")
103+
milvus_client.drop_resource_group("rg2")
104+
105+

pymilvus/milvus_client/milvus_client.py

+100
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
ExtraList,
1212
LoadState,
1313
OmitZeroDict,
14+
ResourceGroupConfig,
1415
construct_cost_extra,
1516
)
1617
from pymilvus.client.utils import is_vector_type
@@ -1525,3 +1526,102 @@ def remove_privileges_from_group(
15251526
"""
15261527
conn = self._get_connection()
15271528
conn.remove_privileges_from_group(group_name, privileges, timeout=timeout, **kwargs)
1529+
1530+
def create_resource_group(self, name: str, timeout: Optional[float] = None, **kwargs):
1531+
"""Create a resource group
1532+
It will success whether or not the resource group exists.
1533+
1534+
Args:
1535+
name: The name of the resource group.
1536+
Raises:
1537+
MilvusException: If anything goes wrong.
1538+
"""
1539+
conn = self._get_connection()
1540+
return conn.create_resource_group(name, timeout, **kwargs)
1541+
1542+
def update_resource_groups(
1543+
self,
1544+
configs: Dict[str, ResourceGroupConfig],
1545+
timeout: Optional[float] = None,
1546+
):
1547+
"""Update resource groups.
1548+
This function updates the resource groups based on the provided configurations.
1549+
1550+
Args:
1551+
configs: A mapping of resource group names to their configurations.
1552+
timeout: The timeout value in seconds. Defaults to None.
1553+
Raises:
1554+
MilvusException: If anything goes wrong.
1555+
"""
1556+
conn = self._get_connection()
1557+
return conn.update_resource_groups(configs, timeout)
1558+
1559+
def drop_resource_group(
1560+
self,
1561+
name: str,
1562+
timeout: Optional[float] = None,
1563+
):
1564+
"""Drop a resource group
1565+
It will success if the resource group is existed and empty, otherwise fail.
1566+
1567+
Args:
1568+
name: The name of the resource group.
1569+
timeout: The timeout value in seconds. Defaults to None.
1570+
Raises:
1571+
MilvusException: If anything goes wrong.
1572+
"""
1573+
conn = self._get_connection()
1574+
return conn.drop_resource_group(name, timeout)
1575+
1576+
def describe_resource_group(self, name: str, timeout: Optional[float] = None):
1577+
"""Drop a resource group
1578+
It will success if the resource group is existed and empty, otherwise fail.
1579+
1580+
Args:
1581+
name: The name of the resource group.
1582+
timeout: The timeout value in seconds. Defaults to None.
1583+
Returns:
1584+
ResourceGroupInfo: The detail info of the resource group.
1585+
Raises:
1586+
MilvusException: If anything goes wrong.
1587+
"""
1588+
conn = self._get_connection()
1589+
return conn.describe_resource_group(name, timeout)
1590+
1591+
def list_resource_groups(self, timeout: Optional[float] = None):
1592+
"""list all resource group names
1593+
1594+
Args:
1595+
timeout: The timeout value in seconds. Defaults to None.
1596+
Returns:
1597+
list[str]: all resource group names
1598+
Raises:
1599+
MilvusException: If anything goes wrong.
1600+
"""
1601+
conn = self._get_connection()
1602+
return conn.list_resource_groups(timeout)
1603+
1604+
def transfer_replica(
1605+
self,
1606+
source_group: str,
1607+
target_group: str,
1608+
collection_name: str,
1609+
num_replicas: int,
1610+
timeout: Optional[float] = None,
1611+
):
1612+
"""transfer num_replica from source resource group to target resource group
1613+
1614+
Args:
1615+
source_group: source resource group name
1616+
target_group: target resource group name
1617+
collection_name: collection name which replica belong to
1618+
num_replicas: transfer replica num
1619+
timeout: The timeout value in seconds. Defaults to None.
1620+
1621+
Raises:
1622+
MilvusException: If anything goes wrong.
1623+
"""
1624+
conn = self._get_connection()
1625+
return conn.transfer_replica(
1626+
source_group, target_group, collection_name, num_replicas, timeout
1627+
)

0 commit comments

Comments
 (0)