Direct model loading from object storage with Runai Model Streamer (#17948)

Signed-off-by: Noa Neria <noa@run.ai>

test/registered/model_loading/test_runai_model_loader.py

@@ -0,0 +1,50 @@
+import unittest
+
+import sglang as sgl
+from sglang.srt.environ import temp_set_env
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, suite="stage-b-test-1-gpu-small")
+
+TEST_GCS_MODEL = "gs://vertex-model-garden-public-us/codegemma/codegemma-2b/"
+
+PROMPTS = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+
+class TestRunaiModelLoader(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        with temp_set_env(
+            GOOGLE_CLOUD_PROJECT="fake-project",
+            RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS="true",
+            CLOUD_STORAGE_EMULATOR_ENDPOINT="https://storage.googleapis.com",
+        ):
+            cls.engine = sgl.Engine(
+                model_path=TEST_GCS_MODEL,
+                load_format="runai_streamer",
+                cuda_graph_max_bs=1,
+                max_total_tokens=64,
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "engine") and cls.engine:
+            cls.engine.shutdown()
+
+    def test_generate_produces_output(self):
+        outputs = self.engine.generate(PROMPTS)
+        self.assertEqual(len(outputs), len(PROMPTS))
+        for i, output in enumerate(outputs):
+            text = output["text"]
+            self.assertIsInstance(text, str)
+            self.assertGreater(len(text), 0, f"Prompt {i} produced empty output")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_runai_utils.py

@@ -0,0 +1,57 @@
+import unittest
+from pathlib import Path
+
+from sglang.srt.configs.load_config import LoadFormat
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
+
+
+class TestRunaiUtils(CustomTestCase):
+    def test_is_runai_obj_uri_s3(self):
+        self.assertTrue(is_runai_obj_uri("s3://bucket/model/"))
+        self.assertTrue(is_runai_obj_uri("S3://Bucket/Model/"))
+
+    def test_is_runai_obj_uri_gs(self):
+        self.assertTrue(is_runai_obj_uri("gs://bucket/model/"))
+        self.assertTrue(is_runai_obj_uri("GS://Bucket/Model/"))
+
+    def test_is_runai_obj_uri_az(self):
+        self.assertTrue(is_runai_obj_uri("az://container/model/"))
+        self.assertTrue(is_runai_obj_uri("AZ://Container/Model/"))
+
+    def test_is_runai_obj_uri_local_paths(self):
+        self.assertFalse(is_runai_obj_uri("/path/to/model"))
+        self.assertFalse(is_runai_obj_uri("./relative/path"))
+        self.assertFalse(is_runai_obj_uri("meta-llama/Llama-3.2-1B"))
+
+    def test_is_runai_obj_uri_other_schemes(self):
+        self.assertFalse(is_runai_obj_uri("http://example.com/model"))
+        self.assertFalse(is_runai_obj_uri("https://example.com/model"))
+        self.assertFalse(is_runai_obj_uri("ftp://example.com/model"))
+
+    def test_is_runai_obj_uri_pathlib(self):
+        self.assertFalse(is_runai_obj_uri(Path("/local/model")))
+
+    def test_get_path_deterministic(self):
+        path1 = ObjectStorageModel.get_path("s3://bucket/model/")
+        path2 = ObjectStorageModel.get_path("s3://bucket/model/")
+        self.assertEqual(path1, path2)
+
+    def test_get_path_different_uris(self):
+        path1 = ObjectStorageModel.get_path("s3://bucket/model-a/")
+        path2 = ObjectStorageModel.get_path("s3://bucket/model-b/")
+        self.assertNotEqual(path1, path2)
+
+    def test_get_path_contains_model_streamer(self):
+        path = ObjectStorageModel.get_path("s3://bucket/model/")
+        self.assertIn("model_streamer", path)
+
+    def test_load_format_enum(self):
+        self.assertEqual(LoadFormat.RUNAI_STREAMER.value, "runai_streamer")
+
+
+if __name__ == "__main__":
+    unittest.main()