From f34ab6a9c499822090bbfcf4366acba5f23f2020 Mon Sep 17 00:00:00 2001
From: sunby <sunbingyi1992@gmail.com>
Date: Tue, 8 Sep 2026 11:15:00 +0800
Subject: [PATCH] feat: add load collection warmup option

Signed-off-by: sunby <sunbingyi1992@gmail.com>
---
 pymilvus/client/prepare.py     |  7 +++++++
 tests/unit/test_load_warmup.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 tests/unit/test_load_warmup.py

diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
index ace1f4411..0fb163da9 100644
--- a/pymilvus/client/prepare.py
+++ b/pymilvus/client/prepare.py
@@ -2228,6 +2228,13 @@ def load_collection(
             priority = kwargs.get("priority")
             req.load_params["load_priority"] = priority
 
+        if "warmup" in kwargs:
+            if kwargs["warmup"] != "sync":
+                raise ParamError(message="load parameter warmup only supports sync")
+            if req.refresh:
+                raise ParamError(message="load parameter warmup cannot be used with refresh")
+            req.load_params["warmup"] = "sync"
+
         return req
 
     @classmethod
diff --git a/tests/unit/test_load_warmup.py b/tests/unit/test_load_warmup.py
new file mode 100644
index 000000000..326cb28f6
--- /dev/null
+++ b/tests/unit/test_load_warmup.py
@@ -0,0 +1,30 @@
+import pytest
+from pymilvus.client.prepare import Prepare
+from pymilvus.exceptions import ParamError
+
+
+def test_load_sync_warmup_preserves_other_options():
+    request = Prepare.load_collection(
+        "warmup_test", replica_number=2, warmup="sync", load_fields=["id"], priority="HIGH"
+    )
+    assert request.load_params["warmup"] == "sync"
+    assert request.load_params["load_priority"] == "HIGH"
+    assert request.replica_number == 2
+    assert list(request.load_fields) == ["id"]
+
+
+def test_load_without_warmup_keeps_existing_semantics():
+    assert "warmup" not in Prepare.load_collection("warmup_test").load_params
+    assert Prepare.load_collection("warmup_test", refresh=True).refresh
+
+
+@pytest.mark.parametrize("warmup", ["disable", "async", "", None, True, 1])
+def test_load_rejects_unsupported_warmup(warmup):
+    with pytest.raises(ParamError, match="only supports sync"):
+        Prepare.load_collection("warmup_test", warmup=warmup)
+
+
+@pytest.mark.parametrize("refresh_key", ["refresh", "_refresh"])
+def test_load_sync_warmup_rejects_refresh(refresh_key):
+    with pytest.raises(ParamError, match="cannot be used with refresh"):
+        Prepare.load_collection("warmup_test", warmup="sync", **{refresh_key: True})