chore: import upstream snapshot with attribution

python/ray/llm/tests/common/utils/test_callback_base.py

@@ -0,0 +1,112 @@
+import asyncio
+
+import pytest
+
+from ray.llm._internal.common.callbacks.base import (
+    CallbackBase,
+)
+from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
+from ray.llm._internal.serve.core.configs.llm_config import (
+    LLMConfig,
+    ModelLoadingConfig,
+)
+
+
+class TestingCallback(CallbackBase):
+    def __init__(self, llm_config, raise_error_on_callback: bool = True, **kwargs):
+        super().__init__(llm_config, raise_error_on_callback, **kwargs)
+        self.before_init_called = False
+        self.after_init_called = False
+        self.before_init_ctx = None
+        self.after_init_ctx = None
+        assert kwargs["kwargs_test_key"] == "kwargs_test_value"
+
+    async def on_before_node_init(self) -> None:
+        assert (
+            self.ctx.worker_node_download_model
+            == NodeModelDownloadable.MODEL_AND_TOKENIZER
+        )
+        self.ctx.worker_node_download_model = NodeModelDownloadable.NONE
+
+        self.ctx.custom_data["ctx_test_key"] = "ctx_test_value"
+        self.before_init_called = True
+        self.ctx.run_init_node = False
+
+    async def on_after_node_init(self) -> None:
+        assert self.ctx.worker_node_download_model == NodeModelDownloadable.NONE
+
+        self.after_init_called = True
+        assert self.ctx.custom_data["ctx_test_key"] == "ctx_test_value"
+
+
+class TestCallbackBase:
+    @pytest.fixture
+    def llm_config(self):
+        config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test-model"),
+            llm_engine="vLLM",
+            callback_config={
+                "callback_class": TestingCallback,
+                "callback_kwargs": {"kwargs_test_key": "kwargs_test_value"},
+            },
+        )
+        return config
+
+    def test_callback_methods_called(self, llm_config):
+        """Test that callback methods are called during initialization."""
+
+        # Run initialization
+        async def run_initialization():
+            callback = llm_config.get_or_create_callback()
+            await callback.run_callback("on_before_node_init")
+            if callback.ctx.run_init_node:
+                raise Exception("run_init_node is True")
+            await callback.run_callback("on_after_node_init")
+
+        asyncio.run(run_initialization())
+        # Verify callback was created and methods were called
+
+        callback = llm_config.get_or_create_callback()
+        assert callback is not None
+        assert isinstance(callback, TestingCallback)
+        assert callback.before_init_called is True
+        assert callback.after_init_called is True
+
+    def test_callback_singleton_behavior(self, llm_config):
+        """Test that callback instance is cached (singleton pattern)."""
+        # Get callback multiple times
+        callback1 = llm_config.get_or_create_callback()
+        callback2 = llm_config.get_or_create_callback()
+
+        # Should be the same instance
+        assert callback1 is callback2
+
+    def test_callback_must_inherit_from_callback_class(self):
+        """Test that callback_class must be a subclass of Callback, not just implement the same methods."""
+
+        class FakeCallback:
+            """A class that implements the same methods as Callback but doesn't inherit from it."""
+
+            def __init__(self, **kwargs):
+                pass
+
+            async def on_before_node_init(self):
+                pass
+
+            async def on_after_node_init(self):
+                pass
+
+        # Should raise an error when trying to create callback
+        with pytest.raises(Exception, match="is-subclass"):
+            LLMConfig(
+                model_loading_config=ModelLoadingConfig(model_id="test-model"),
+                llm_engine="vLLM",
+                callback_config={
+                    "callback_class": FakeCallback,
+                    "callback_kwargs": {},
+                },
+            )
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

python/ray/llm/tests/common/utils/test_upload_utils.py

@@ -0,0 +1,101 @@
+import os
+import sys
+import tempfile
+from pathlib import Path
+from unittest.mock import ANY, call, patch
+
+import pytest
+
+from ray.llm._internal.common.utils.upload_utils import upload_model_files
+
+
+@patch("pyarrow.fs.copy_files")
+def test_upload_undownloaded_model(mock_copy_files):
+    model_id = "fake-model-id"
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        model_dir = os.path.join(tempdir, model_id)
+        # mock snapshot_download to create a subdir in the model_dir
+        def create_subdir_side_effect(*args, **kwargs):
+            os.makedirs(model_dir, exist_ok=True)
+
+        with patch(
+            "ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
+            return_value=model_dir,
+        ), patch(
+            "huggingface_hub.snapshot_download", side_effect=create_subdir_side_effect
+        ):
+            # upload the model
+            upload_model_files(model_id, "gs://bucket/model-id")
+
+        # check that the model was uploaded1
+        mock_copy_files.assert_called_once_with(
+            source=Path(model_dir),
+            destination="bucket/model-id",
+            source_filesystem=ANY,
+            destination_filesystem=ANY,
+        )
+
+
+@patch("pyarrow.fs.copy_files")
+def test_upload_downloaded_hf_model(mock_copy_files):
+    model_id = "fake-model-id"
+    hash = "0123456789abcdef"
+    with tempfile.TemporaryDirectory() as tempdir:
+        model_dir = os.path.join(tempdir, model_id)
+        os.makedirs(model_dir, exist_ok=True)
+        os.makedirs(os.path.join(model_dir, "snapshots", hash), exist_ok=True)
+        model_rev_path = os.path.join(model_dir, "refs")
+        os.makedirs(model_rev_path, exist_ok=True)
+        with open(os.path.join(model_rev_path, "main"), "w") as f:
+            f.write(hash)
+
+        with patch(
+            "ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
+            return_value=model_dir,
+        ):
+            upload_model_files(model_id, "pyarrow-s3://bucket/model-id")
+
+        assert mock_copy_files.call_count == 2
+        mock_copy_files.assert_has_calls(
+            [
+                call(
+                    source=os.path.join(model_dir, "snapshots", hash),
+                    destination="bucket/model-id",
+                    source_filesystem=ANY,
+                    destination_filesystem=ANY,
+                ),
+                call(
+                    source=os.path.join(model_rev_path, "main"),
+                    destination="bucket/model-id/hash",
+                    source_filesystem=ANY,
+                    destination_filesystem=ANY,
+                ),
+            ],
+            any_order=True,
+        )
+
+
+@patch("pyarrow.fs.copy_files")
+def test_upload_custom_model(mock_copy_files):
+    model_id = "fake-model-id"
+    with tempfile.TemporaryDirectory() as tempdir:
+        model_dir = os.path.join(tempdir, model_id)
+        os.makedirs(model_dir, exist_ok=True)
+
+        with patch(
+            "ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
+            return_value=model_dir,
+        ):
+            upload_model_files(model_id, "pyarrow-s3://bucket/model-id")
+
+        mock_copy_files.assert_called_once_with(
+            source=Path(model_dir),
+            destination="bucket/model-id",
+            source_filesystem=ANY,
+            destination_filesystem=ANY,
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))