|
| 1 | +""" |
| 2 | +Inference-only entrypoint. |
| 3 | +
|
| 4 | +Spins up the SkyRL HTTP inference servers (vLLM server groups + ``vllm-router``) |
| 5 | +from a provided inference configuration and keeps them alive so that client code |
| 6 | +and generation configurations can be iterated against a fixed deployment without |
| 7 | +re-launching the engines on every run. |
| 8 | +
|
| 9 | +Only the new inference codepath (``_SKYRL_USE_NEW_INFERENCE=1``) is supported. |
| 10 | +
|
| 11 | +Example:: |
| 12 | +
|
| 13 | + uv run --isolated --extra fsdp -m skyrl.train.entrypoints.serve \ |
| 14 | + trainer.policy.model.path=Qwen/Qwen2.5-1.5B-Instruct \ |
| 15 | + trainer.placement.colocate_all=false \ |
| 16 | + generator.inference_engine.num_engines=4 \ |
| 17 | + generator.inference_engine.tensor_parallel_size=2 |
| 18 | +
|
| 19 | +All ``generator.inference_engine.*`` knobs (plus ``trainer.policy.model.path``) |
| 20 | +are accepted as overrides. |
| 21 | +""" |
| 22 | + |
| 23 | +import asyncio |
| 24 | +import sys |
| 25 | +import time |
| 26 | + |
| 27 | +import ray |
| 28 | +from loguru import logger |
| 29 | + |
| 30 | +from skyrl.env_vars import _SKYRL_USE_NEW_INFERENCE |
| 31 | +from skyrl.train.config import SkyRLTrainConfig |
| 32 | +from skyrl.train.entrypoints.main_base import BasePPOExp |
| 33 | +from skyrl.train.utils.utils import initialize_ray, validate_inference_engine_cfg |
| 34 | + |
| 35 | + |
| 36 | +class InferenceOnlyEntrypoint(BasePPOExp): |
| 37 | + """Sets up and serves the inference servers without a trainer. |
| 38 | +
|
| 39 | + Reuses :class:`BasePPOExp` for tokenizer construction and inference-client |
| 40 | + setup (``get_inference_client`` -> ``_get_new_inference_client``), but skips |
| 41 | + all training-only state (datasets, trainer, tracker). |
| 42 | + """ |
| 43 | + |
| 44 | + def get_train_dataset(self): |
| 45 | + """No training dataset is needed for inference-only serving.""" |
| 46 | + return None |
| 47 | + |
| 48 | + def get_eval_dataset(self): |
| 49 | + """No eval dataset is needed for inference-only serving.""" |
| 50 | + return None |
| 51 | + |
| 52 | + def _teardown(self) -> None: |
| 53 | + """Tear down the router and server groups started during setup.""" |
| 54 | + logger.info("Tearing down inference servers...") |
| 55 | + if self._inference_router is not None: |
| 56 | + self._inference_router.shutdown() |
| 57 | + for group in ( |
| 58 | + (self._server_groups or []) + (self._prefill_server_groups or []) + (self._decode_server_groups or []) |
| 59 | + ): |
| 60 | + group.shutdown() |
| 61 | + |
| 62 | + def run(self) -> None: |
| 63 | + # Builds the vLLM server groups + router and returns an HTTP client. |
| 64 | + # The server group actor handles / router process are stored on ``self`` |
| 65 | + # by ``_get_new_inference_client`` and kept alive while this task runs. |
| 66 | + client = self.get_inference_client() |
| 67 | + |
| 68 | + logger.info( |
| 69 | + "Inference servers are up.\n" |
| 70 | + f" proxy_url (data plane): {client.proxy_url}\n" |
| 71 | + f" server_urls (control plane): {client.server_urls}\n" |
| 72 | + "Point your client at proxy_url for OpenAI-compatible requests. " |
| 73 | + "Press Ctrl+C to shut down." |
| 74 | + ) |
| 75 | + |
| 76 | + try: |
| 77 | + while True: |
| 78 | + time.sleep(3600) |
| 79 | + # TODO: Currently this is run inside a ray task so a KeyboardInterrupt on |
| 80 | + # driver never reaches here - Ray sends a SIGTERM. We should propagate the |
| 81 | + # interrupt for better shutdown |
| 82 | + except KeyboardInterrupt: |
| 83 | + logger.info("Received interrupt, shutting down inference servers...") |
| 84 | + finally: |
| 85 | + try: |
| 86 | + asyncio.run(client.teardown()) |
| 87 | + finally: |
| 88 | + self._teardown() |
| 89 | + |
| 90 | + |
| 91 | +@ray.remote(num_cpus=1) |
| 92 | +def skyrl_entrypoint(cfg: SkyRLTrainConfig): |
| 93 | + # make sure that the serving loop is not run on the head node. |
| 94 | + exp = InferenceOnlyEntrypoint(cfg) |
| 95 | + exp.run() |
| 96 | + |
| 97 | + |
| 98 | +def _validate_serve_cfg(cfg: SkyRLTrainConfig) -> None: |
| 99 | + """Validate the config for the inference-only serving path.""" |
| 100 | + if not _SKYRL_USE_NEW_INFERENCE: |
| 101 | + raise ValueError( |
| 102 | + "The serve entrypoint only supports the new inference codepath. " |
| 103 | + "Unset `_SKYRL_USE_NEW_INFERENCE=0` (the default is enabled)." |
| 104 | + ) |
| 105 | + |
| 106 | + if cfg.trainer.policy.model.path is None: |
| 107 | + raise ValueError("trainer.policy.model.path must be set to the model to serve.") |
| 108 | + |
| 109 | + ie_cfg = cfg.generator.inference_engine |
| 110 | + |
| 111 | + if not ie_cfg.run_engines_locally: |
| 112 | + raise ValueError( |
| 113 | + "The serve entrypoint launches engines locally; set " |
| 114 | + "generator.inference_engine.run_engines_locally=true (the default)." |
| 115 | + ) |
| 116 | + |
| 117 | + if ie_cfg.external_proxy_url is not None or ie_cfg.external_server_urls is not None: |
| 118 | + raise ValueError( |
| 119 | + "The serve entrypoint starts its own servers; " |
| 120 | + "generator.inference_engine.external_proxy_url / external_server_urls " |
| 121 | + "must not be set." |
| 122 | + ) |
| 123 | + |
| 124 | + if cfg.trainer.placement.colocate_all: |
| 125 | + raise ValueError( |
| 126 | + "trainer.placement.colocate_all must be false for inference-only " |
| 127 | + "serving. Colocation puts the engines to sleep after startup so they " |
| 128 | + "can share GPUs with training workers, which do not exist here." |
| 129 | + ) |
| 130 | + |
| 131 | + # Shared inference-engine validation (PD, parallelism, executor backend, |
| 132 | + # new inference layer). Resolves `override_existing_update_group="auto"`. |
| 133 | + validate_inference_engine_cfg(cfg) |
| 134 | + |
| 135 | + |
| 136 | +def main() -> None: |
| 137 | + # Parse CLI args and build typed config |
| 138 | + cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:]) |
| 139 | + |
| 140 | + # validate inference cfg |
| 141 | + _validate_serve_cfg(cfg) |
| 142 | + |
| 143 | + # init and run |
| 144 | + initialize_ray(cfg) |
| 145 | + ray.get(skyrl_entrypoint.remote(cfg)) |
| 146 | + |
| 147 | + |
| 148 | +if __name__ == "__main__": |
| 149 | + main() |
0 commit comments