Optimize k8s launcher and modify k8s test cases

LazyAGI · Dec 25, 2024 · 06f0d87 · 06f0d87
1 parent 1f960cb
commit 06f0d87
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 169 deletions.
diff --git a/lazyllm/launcher.py b/lazyllm/launcher.py
@@ -563,17 +563,19 @@ def _wait_for_gateway_update(self, timeout, check_interval):
         def _create_httproute(self):
             custom_api = k8s.client.CustomObjectsApi()
 
+            httproute_name = f"httproute-{self.deployment_name}"
             httproute_spec = {
                 "apiVersion": "gateway.networking.k8s.io/v1beta1",
                 "kind": "HTTPRoute",
                 "metadata": {
-                    "name": f"httproute-{self.deployment_name}",
+                    "name": httproute_name,
                     "namespace": self.namespace
                 },
                 "spec": {
                     "parentRefs": [{
                         "name": self.gateway_name,
-                        "port": self.deployment_port
+                        "port": self.deployment_port,
+                        "sectionName": httproute_name
                     }],
                     "rules": [{
                         "matches": [{
@@ -601,7 +603,7 @@ def _create_httproute(self):
                     plural="httproutes",
                     body=httproute_spec
                 )
-                LOG.info(f"Kubernetes HTTPRoute 'httproute-{self.deployment_name}' created successfully.")
+                LOG.info(f"Kubernetes HTTPRoute '{httproute_name}' created successfully.")
             except k8s.client.rest.ApiException as e:
                 LOG.error(f"Exception when creating HTTPRoute: {e}")
                 raise

diff --git a/tests/k8s_tests/test_example.py b/tests/k8s_tests/test_example.py
diff --git a/tests/k8s_tests/test_pipeline_k8s.py b/tests/k8s_tests/test_pipeline_k8s.py
@@ -0,0 +1,54 @@
+import lazyllm
+from lazyllm import launchers, pipeline
+
+def demo1(input): return input * 2
+
+def demo2(input): return input * 3
+
+def demo3(input): return input * 4
+
+def demo4(input): return input * 5
+
+class TestPipelineK8s(object):
+    def test_single_pipeline(self):
+        with pipeline() as ppl:
+            ppl.m1 = lazyllm.ServerModule(demo1, launcher=launchers.k8s()).start()
+            ppl.m2 = lazyllm.ServerModule(demo2, launcher=launchers.k8s()).start()
+            ppl.m3 = lazyllm.ServerModule(demo3, launcher=launchers.k8s()).start()
+            ppl.m4 = lazyllm.ServerModule(demo4, launcher=launchers.k8s()).start()
+
+        assert ppl(2) == 240
+        lazyllm.launcher.cleanup()
+
+    def test_pipeline_server(self):
+        with pipeline() as p1:
+            p1.m1 = demo1
+            p1.m2 = demo2
+        module1 = lazyllm.ServerModule(p1, launcher=launchers.k8s())
+
+        with pipeline() as p2:
+            p2.m1 = module1
+            p2.m2 = demo3
+
+        module2 = lazyllm.ServerModule(p2, launcher=launchers.k8s())
+
+        with pipeline() as p3:
+            p3.m1 = module2
+            p3.m2 = demo4
+
+        module3 = lazyllm.ServerModule(p3, launcher=launchers.k8s())
+        module3.start()
+        assert module3(2) == 240
+        lazyllm.launcher.cleanup()
+
+    def test_nesting_pipeline(self):
+        with pipeline() as p:
+            with pipeline() as p.m1:
+                with pipeline() as p.m1.mm1:
+                    p.m1.mm1.m1 = lazyllm.ServerModule(demo1, launcher=launchers.k8s()).start()
+                    p.m1.mm1.m2 = lazyllm.ServerModule(demo2, launcher=launchers.k8s()).start()
+                p.m1.mm2 = lazyllm.ServerModule(demo3, launcher=launchers.k8s()).start()
+            p.m2 = lazyllm.ServerModule(demo4, launcher=launchers.k8s()).start()
+
+        assert p(2) == 240
+        lazyllm.launcher.cleanup()