test to detect (platform specific) plugged-in scheduler

.buildkite/test-pipeline.yaml

@@ -521,6 +521,9 @@ steps:
   - pip install -e ./plugins/vllm_add_dummy_platform
   - pytest -v -s plugins_tests/test_platform_plugins.py
   - pip uninstall vllm_add_dummy_platform -y
+  - pip install -e ./plugins/vllm_add_dummy_scheduler
+  - pytest -v -s plugins_tests/test_scheduler_plugins.py
+  - pip uninstall vllm_add_dummy_scheduler -y
   # end platform plugin tests
   # other tests continue here:
   - pip install -e ./plugins/vllm_add_dummy_model

tests/plugins/vllm_add_dummy_scheduler/setup.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from setuptools import setup
+
+setup(
+    name='vllm_add_dummy_scheduler',
+    version='0.1',
+    packages=['vllm_add_dummy_scheduler'],
+    entry_points={
+        'vllm.platform_plugins': [
+            "dummy_scheduler_plugin = vllm_add_dummy_scheduler:dummy_scheduler_plugin"  # noqa
+        ]
+    })

tests/plugins/vllm_add_dummy_scheduler/vllm_add_dummy_scheduler/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+
+def dummy_scheduler_plugin() -> Optional[str]:
+    return "vllm_add_dummy_scheduler.dummy_platform.DummyPlatform"

tests/plugins/vllm_add_dummy_scheduler/vllm_add_dummy_scheduler/dummy_platform.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.platforms.cuda import CudaPlatform
+
+
+class DummyPlatform(CudaPlatform):
+    device_name = "DummyDevice"
+
+    # function copied from vllm-project/vllm/platforms/cuda.py
+    # last line added to change scheduler_cls to platform specific scheduler
+    @classmethod
+    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        parallel_config = vllm_config.parallel_config
+        scheduler_config = vllm_config.scheduler_config
+
+        if parallel_config.worker_cls == "auto":
+            if scheduler_config.is_multi_step:
+                if envs.VLLM_USE_V1:
+                    raise NotImplementedError(
+                        "Multi-step scheduling is not supported (and not "
+                        "needed) on VLLM V1. Please launch without "
+                        "--num-scheduler-steps.")
+                else:
+                    parallel_config.worker_cls = \
+                        "vllm.worker.multi_step_worker.MultiStepWorker"
+            elif vllm_config.speculative_config:
+                if envs.VLLM_USE_V1:
+                    raise NotImplementedError(
+                        "Speculative decoding is not yet supported on VLLM V1."
+                    )
+                else:
+                    parallel_config.worker_cls = \
+                        "vllm.spec_decode.spec_decode_worker.create_spec_worker"
+                    parallel_config.sd_worker_cls = \
+                        "vllm.worker.worker.Worker"
+            else:
+                if envs.VLLM_USE_V1:
+                    parallel_config.worker_cls = \
+                            "vllm.v1.worker.gpu_worker.Worker"
+                else:
+                    parallel_config.worker_cls = "vllm.worker.worker.Worker"
+
+        cache_config = vllm_config.cache_config
+        if cache_config and cache_config.block_size is None:
+            cache_config.block_size = 16
+
+        # additional line to load plugin specific scheduler instead of default
+        parallel_config.scheduler_cls = \
+            "vllm_add_dummy_scheduler.dummy_scheduler.DummyScheduler"

tests/plugins/vllm_add_dummy_scheduler/vllm_add_dummy_scheduler/dummy_scheduler.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from vllm.core.scheduler import Scheduler
+
+
+class CustomException(Exception):
+    pass
+
+
+class DummyScheduler(Scheduler):
+
+    def schedule(self):
+        raise CustomException("Exception raised by DummyScheduler")

tests/plugins_tests/test_scheduler_plugins.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+def test_scheduler_plugins():
+    # simulate workload by running an example
+    import runpy
+
+    from vllm_add_dummy_scheduler.dummy_scheduler import CustomException
+    current_file = __file__
+    import os
+    example_file = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.dirname(current_file))),
+        "examples", "offline_inference/basic.py")
+    try:
+        runpy.run_path(example_file)
+    except CustomException as e:
+        assert str(e) == "Exception raised by DummyScheduler"