diff --git a/README.md b/README.md index caeb152..3a9691a 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,8 @@ training = service.create_lora_training_client( ## Shared deployment quick start +Shared deployments use Python recipes inheriting from `BaseConfig`. See [Python deployment configs](docs/deployment-configs.md). Keep the active Python config list in [scripts/deploy_models.sh](scripts/deploy_models.sh); run it to deploy the complete list. + Install Lilo into your own Python project, deploy it once to Modal, then call its API from your training scripts. The commands below work in Bash or Zsh. @@ -60,7 +62,7 @@ clients do not need Modal deployment credentials or sampler proxy tokens. With [uv](https://docs.astral.sh/uv/) installed: ```bash -uv init my-lilo-project +uv init --python 3.12 my-lilo-project cd my-lilo-project uv add 'lilo @ git+https://github.com/modal-projects/lilo.git' ``` @@ -127,15 +129,17 @@ uv run modal secret create lilo-proxy \ ### 3. Deploy the installed package -Deploying the entire Tinker server can be done with a single modal deploy command: +Create a configuration from a preset, validate it, and deploy it with Python 3.12: ```bash -uv run modal deploy -m lilo.providers.modal.app +uv run lilo config init --preset qwen35-9b-lora-16k > deployment.py +uv run lilo config validate deployment.py +uv run lilo deploy deployment.py ``` -This deploys the control plane and bundled model definitions, then prints the -`server` URL to use in step 4. Reuse the deployment across training runs and -redeploy after updating Lilo. +This deploys the shared app and prints its `server` URL. Add more Python config files to the same command to serve more recipes. Always supply the complete current set. The Miles commit is pinned in `miles_image.py`; see [Python deployment configs](docs/deployment-configs.md). + +From a repository checkout, maintain the list in `scripts/deploy_models.sh` and run that script. `lilo deploy` supplies the current configs and frontend platform settings to Modal. Deploying the server doesn't allocate any GPUs; rather, this allocation for both the training and sampling sides are done on demand. See [cold starts and capacity configuration](docs/full-fine-tunes.md#performance-and-behavior-considerations) before running a larger workload. @@ -150,8 +154,8 @@ has finished in the Modal dashboard or list apps with: uv run modal app list ``` -To tear down the deployment, stop its `lilo-fft-...` sampler apps, then `lilo`, -using `uv run modal app stop `. Stopping `lilo` does not stop sampler apps. +To tear down the deployment, stop its `lilo-fft-...` sampler apps, then the frontend selected with `--app` (`lilo` by default), +using `uv run modal app stop `. Stopping the frontend does not stop sampler apps. ## Next steps diff --git a/docs/deployment-configs.md b/docs/deployment-configs.md new file mode 100644 index 0000000..d73bbe4 --- /dev/null +++ b/docs/deployment-configs.md @@ -0,0 +1,136 @@ +# Python deployment configs + +A recipe subclasses `BaseConfig` and exports `config = Config()`. Settings are flat, untyped Python attributes. Backend options are ordinary dictionaries. + +```python +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "my-9b" + model = "Qwen/Qwen3.5-9B-Base" + max_context_length = 16384 + backend = "miles" + trainer_gpu = "H100" + trainer_gpus_per_node = 4 + trainer_cpu = 16 + trainer_memory_mib = 65536 + trainer_max_clients_per_instance = 6 + inference_gpu = "H200" + inference_max_replicas = 8 + miles_cfg = { + "model_type": "qwen3.5-9B", + "tensor_model_parallel_size": 4, + "max_lora_slots": 6, + "max_lora_rank": 32, + "target_modules": [ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", "lm_head", + ], + } + sglang_cfg = {"max_running_requests": 32} + + +config = Config() +``` + +See the [9B LoRA recipe](../src/lilo/configs/qwen35_9b_lora_16k.py) and [4B FFT recipe](../src/lilo/configs/qwen35_4b_fft_64k.py) for complete examples. Weights are downloaded to `/assets/`. + +## Variants + +Override ordinary attributes directly. Use dotted `overrides` to change individual backend options: + +```python +from lilo.configs.qwen35_9b_lora_16k import Config as Parent + + +class Config(Parent): + name = "my-9b-more-memory" + trainer_memory_mib = 98304 + overrides = {"miles_cfg.max_tokens_per_gpu": 8192} + + +config = Config() +``` + +Each parent's settings and overrides apply before its child's. Constructor fields and overrides apply last: `Config(trainer_gpu="H200", overrides={"sglang_cfg.max_running_requests": 16})`. Assigning a dictionary or list replaces that value; `trainer_env = {}` clears inherited environment settings. Instances own independent copies of mutable values and can also be edited directly. + +## Backend options + +| Setting | Consumed by | +| --- | --- | +| `trainer_*`, `inference_*` | Modal GPU/CPU/memory allocation, scaling, timeouts and Lilo admission limits | +| `megatron_cfg` | Existing Megatron `EngineModelConfig`; provider, optimizer and distributed options use its native dictionaries | +| `miles_cfg` | Existing `MilesBackendConfig`; `cli_options` supplies additional Miles arguments | +| `sglang_cfg` | SGLang `ServerArgs` | + +Modal and backend libraries validate their own options. Lilo checks integration requirements such as trainer slot capacity, supported training modes, and parallelism agreeing with allocated GPUs. It supplies managed model paths, context length and adapter settings; conflicting backend overrides are rejected. + +`BaseConfig` does not enforce field types or reject arbitrary attributes. Extra backend options belong in the corresponding backend dictionary. A misspelled top-level attribute is ordinary Python data and may be unused. + +Miles argument conversion lives in [miles_arguments.py](../src/lilo/backends/miles_arguments.py). SGLang receives `ServerArgs(**settings)` in its [worker entrypoint](../src/lilo/inference/sglang.py). Backend libraries validate native options when workers start; frontend config imports remain CPU-only. + +## Prepare and launch + +~~~text +load(config.py) → recipe: BaseConfig + → DeploymentConfig.create(recipe) + → deploy named trainer and inference apps if they are missing or changed + → deploy the frontend with the current config set +~~~ + +The launcher consumes the saved settings. It does not reparse backend configuration or add another set of backend defaults. JSON decoding in a trainer reconstructs the saved config, without importing the author's config file. + +| File | Responsibility | +| --- | --- | +| [configuration.py](../src/lilo/configuration.py) | BaseConfig defaults, inheritance and overrides | +| [deployments.py](../src/lilo/deployments.py) | Python recipe loader and DeploymentConfig | +| [backends/deployment.py](../src/lilo/backends/deployment.py) | Backend settings used when creating a DeploymentConfig | +| [megatron_runtime/common/settings.py](../src/lilo/backends/megatron_runtime/common/settings.py) | Shared Megatron ownership rules and constructor dictionaries | +| [deployment_apps.py](../src/lilo/providers/modal/deployment_apps.py) | Trainer, inference, and pool app builders | +| [deployment_configs.py](../src/lilo/providers/modal/deployment_configs.py) | Read saved configs and frontend platform settings | +| [deployment_cli.py](../src/lilo/deployment_cli.py) | Deploy trainer and inference apps by calling the builders, then deploy the frontend | + +## Multi-node Miles + +[qwen38_27b_lora_256k.py](../src/lilo/configs/qwen38_27b_lora_256k.py) configures two nodes with eight H200s per node, TP2 and CP8. The trainer app uses Modal's clustered launcher and RDMA. Every rank mounts the same volumes; rank 0 starts the engine, and the other ranks join Ray using the launcher merged in #39. The driver receives the Ray address. GPU counts cannot be independently overridden through Miles options. + +This path has CPU construction/topology tests. This refactor has not been redeployed or tested on multiple GPU nodes. + +## Deploy and update + +~~~bash +lilo config init --preset qwen35-9b-lora-16k > my_model.py +lilo config validate my_model.py +lilo deploy my_model.py +~~~ + +Validation loads the recipes and checks that they can share one frontend. Backend integration settings are checked when creating a DeploymentConfig at deploy time. Native backend options are checked at engine startup. + +The checked-in [deploy_models.sh](../scripts/deploy_models.sh) lists the complete active config set. Add a config path there, then run it. The deployment command owns frontend selection and trainer/inference updates: + +~~~bash +./scripts/deploy_models.sh --app my-lilo --env dev --region us-west +./scripts/deploy_models.sh --refresh-trainer qwen35-9b-lora-16k +./scripts/deploy_models.sh --refresh-inference qwen35-9b-lora-16k +~~~ + +These settings are not model-config fields. Secret/volume names come from provider defaults. Credentials remain in Modal secrets. + +One frontend serves all models through Tinker’s `base_model`. When recipes share a model, the first matching recipe in the `lilo deploy` argument list is used. Training also matches the requested LoRA/FFT mode; base sampling uses the first recipe regardless of training mode. To select a specific recipe, pass its name from `/api/v1/lilo/deployments` as `base_model`. The frontend lists the current config set only. + +## Update isolation + +Trainer apps are named `lilo-trainer-{name}` and inference apps `lilo-inference-{name}`. An inference-only change reuses the trainer app. A trainer-only change reuses the inference app. Source-only upgrades use `--refresh-trainer` / `--refresh-inference`. Updating inference settings also redeploys existing pools for that recipe before publishing the new frontend configuration. An interrupted pool update can be retried with the same command. Changes to region, storage, or secrets refresh both worker apps. + +Recipe names must be unique across frontends in the same Modal environment, since worker and pool names include the recipe name. Drain jobs before removing a recipe or making incompatible model, context, or adapter changes. + +Trainers receive the frontend App ID at launch and use its shared registry for engine state and reconciliation. A startup failure pauses new starts for five minutes; the failure then expires, allowing subsequent requests and reconciliation to retry. `lilo deployment retry` clears it immediately. + +The frontend carries the current configs and exposes them through `deployed_configs`. The CLI reads that function to compare settings, and checks trainer and inference app existence directly with Modal. There is no retained history of earlier configs. Run deploy commands sequentially. + +`--app`, `--env`, `--region`, secrets, and storage names are frontend/platform settings. They are not copied onto each recipe. + +LoRA pools are named `lilo-lora-{name}`. FFT pools retain the existing `lilo-fft-{session_digest}-latest` or `lilo-fft-{session_digest}-v{version}` names to fit Modal’s 63-character limit. The digest identifies the recipe/client pair; it is not a config revision. + +See [validation results](deployment-validation.md) for CPU coverage, live GPU checks, and their limits. diff --git a/docs/deployment-validation.md b/docs/deployment-validation.md new file mode 100644 index 0000000..cccf927 --- /dev/null +++ b/docs/deployment-validation.md @@ -0,0 +1,48 @@ +# Deployment configuration validation + +PR #55 was checked against `main` on 2026-09-28. These results cover the flat Python recipes and independently deployed trainer/inference apps. + +## CPU and packaging + +- **701 tests passed, 1 skipped**, using Python 3.12 and the CI version of PyTorch (2.10.0 CPU). The skipped test imports SGLang, which is installed in the GPU image rather than the CPU test environment. +- All **15 packaged presets** resolve backend settings and round-trip through saved JSON from the built wheel. +- Ruff import/error checks, formatting, the public PyPI lockfile check, and `git diff --check` passed. Changed Python files contain no function-local imports. + +Regression coverage includes inherited overrides, independent mutable settings, backend constructor forwarding, separate trainer/frontend registries, startup-failure expiry, existing-pool updates, platform-setting changes, and retrying a partially failed deployment. Megatron provider replacement preserves pretrained-weight loading hooks. LoRA presets select explicit language-model targets and retain the corresponding Tinker training flags. FFT session pool names fit Modal's app-name limit. + +## Live GPU checks + +The isolated frontend is `lilo-pr55-validation-20260928` in `modal-labs/lilo-deploy`, region `us-west`. It uses uniquely named variants of the 9B LoRA/16K and 4B FFT/64K recipes. Each trainer requests 4 H100s. Both test inference pools use one H100 replica kept warm during validation. LoRA inference was changed from the preset’s H200 request after Modal reported insufficient H200 capacity in `us-west`; the production preset still requests H200. + +| Check | Result | +| --- | --- | +| 9B LoRA client A | 3 complete training/publication/sampling steps passed | +| 9B LoRA client B, sharing client A’s trainer | 2 initial steps plus 1 after the inference update passed | +| 4B full-parameter training | 3 complete training/publication/sampling steps passed | +| Inference-only update | SGLang `max_running_requests` changed from 32 to 24; the running worker and frontend lookup both reported 24 | +| Trainer continuity | Client B retained the same trainer instance and boot ID through the inference updates | + +Every step returned finite training and sampling log probabilities. The frontend lookup initially returned the prior configuration during the deployment transition; it was checked again, along with the worker’s actual settings, before releasing the continuation. All completed clients unloaded without reported cleanup errors. The isolated validation apps and pools were stopped afterward. + +These checks use 512 training tokens and a 16-token generation cap. They exercise the configured processes and interfaces, not maximum-context capacity, convergence, autoscaling under load, or throughput. Other model presets and multi-node topology have CPU coverage only. + +## Reproduce + +Use Python 3.12, authenticated Modal credentials, and the deployment's `TINKER_API_KEY`. Create isolated recipe files with distinct names so they do not share worker apps with other frontends: + +```bash +lilo config init --preset qwen35-9b-lora-16k > lora.py +lilo config init --preset qwen35-4b-fft-64k > fft.py +# Give the two recipes different, unique names before deploying. +lilo deploy lora.py fft.py --app YOUR_TEST_FRONTEND --env YOUR_ENVIRONMENT +MODAL_ENVIRONMENT=YOUR_ENVIRONMENT python scripts/deployment_smoke.py \ + --frontend YOUR_TEST_FRONTEND --name YOUR_LORA_RECIPE \ + --steps 3 --output /tmp/lora-result.json \ + --continue-file /tmp/continue-after-redeploy +``` + +Run another instance for the FFT recipe. Each step checks finite training log probabilities, applies an optimizer update, publishes weights, and checks a generated sample with finite log probabilities. + +The continuation option keeps the client alive after the initial steps. Change one recipe's inference settings, deploy the complete recipe set again, then create the continuation file. The client performs another complete step and asserts that its trainer boot ID is unchanged. Inspect pool startup settings to confirm that the inference change took effect. + +The script unloads its model on exit. Stop the dedicated test frontend, worker apps, and pools after validation. diff --git a/docs/design.md b/docs/design.md index 6070b53..c3cb213 100644 --- a/docs/design.md +++ b/docs/design.md @@ -122,51 +122,34 @@ sampling scales according to rollout traffic. We best-effort sticky-route groups - [`providers/modal/fft_pool.py`](../src/lilo/providers/modal/fft_pool.py): creates, finds, wakes, and stops each FFT model's sampling service using Modal flash proxy -## Adding a new model deployment +## Adding a new model deployment -"Model deployment" in this context refers to a particular training configuration for a base model, defined by its parameterization (full parameter, LoRA, etc.), desired context/sampling length, parallelism, quantization, and so forth. To add a new deployment that can be spun up by the control plane: +A Python deployment dataclass specifies the base model, training mode, context length, GPUs, parallelism, and inference settings. Shared deployments are defined only through these files. -1. Existing model definitions are in [`providers/modal/definitions`](../src/lilo/providers/modal/definitions) (one per file). Create a new file with the desired configuration details (model, checkpoint, context-length, GPU, and parallelism settings). -2. Keep the module filename, `DEFINITION_ID`, and engine function name the same. -3. Import the module in - [`providers/modal/app.py`](../src/lilo/providers/modal/app.py) and append it - to `DEFINITIONS`. +1. Create a Python `Config` subclass under `src/lilo/configs/`, optionally inheriting from a packaged config. +2. Add its path to the list in [`scripts/deploy_models.sh`](../scripts/deploy_models.sh). +3. Run the script to apply the complete list to the shared frontend. -Every definition exports: +No central model catalog registration is needed. The generic builders in [`deployment_apps.py`](../src/lilo/providers/modal/deployment_apps.py) construct trainer functions and inference apps from the saved config. See [Python deployment configs](deployment-configs.md) for the configuration schema and app structure. -- `DEFINITION_ID`, `MODEL_NAME`, `PARAMETERIZATION`, and `CATALOG_VISIBLE`; -- `TRAINER_MODELS_PER_INSTANCE`; -- the model asset Volume and backend configuration; -- a Modal `app` and engine function that calls - [`run_engine_with_backend`](../src/lilo/providers/modal/serve.py); and -- `ENGINE_FUNCTION`, referencing that engine function. +Set each configuration's trainer limit with `trainer.max_instances` and its inference limits with `inference.min_replicas` and `inference.max_replicas`. Trainer limits are read from the config; `LILO_TRAINER_MAX_CONTAINERS` is no longer used. -The existing model definitions show the complete template for parameterization-specific settings. FOr example, FFT definitions include rollout resources + Stitch bulletin volume, whereas LoRA definitions include adapter rank, slot capacity, and adapter storage. Our FFT engines are currently only capable of hosting one model (but if multiple FFT experiments are submitted to the control plane, it will spin up as many engine replicas as necessary to support these concurrently). LoRA engines are multi-lora and so can host multiple adapters. - -Adding the module to `DEFINITIONS` automatically includes its Modal -sub-application and makes it available to model lookup, trainer provisioning, -parameterization lookup, and the public catalog. The registry tests also cover -the new definition automatically. To run the tests before deploying: +To check training, publication, and sampling against a deployed configuration: ```bash -uv run pytest tests/providers/test_definition_registry.py -uv run modal deploy -m lilo.providers.modal.app +uv run python scripts/deployment_smoke.py \ + --frontend lilo-yaml --name qwen35-9b-lora-16k \ + --output /tmp/lilo-smoke.json ``` -Trainer container limits are deployment-specific. Set -`LILO_TRAINER_MAX_CONTAINERS` to a positive integer when deploying to apply the -same limit to every definition. Leaving it unset makes trainer containers -unlimited. - -The following script tests e2e deployment for one model definition, launching the Modal app, creating the particular model, running forward_backward + optim_step operations, and sampler weight publication/rollouts: +For the longer FFT checkpoint and sampler-recovery checks: ```bash uv run python scripts/e2e_engine_definition.py \ - --definition-id + --frontend lilo-yaml --name qwen35-4b-fft-64k --checkpoint-only ``` -This validates all steps of the training cycle for the particular model definition. For an FFT checkpoint round trip, add -`--checkpoint-only`. +Both scripts read the deployed configuration rather than importing model-specific Python definitions. ## Adding a new backend diff --git a/docs/observability.md b/docs/observability.md index 62b9646..0196f66 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -45,7 +45,7 @@ Deploy Lilo after updating the secret: ```bash MODAL_PROFILE=your-workspace MODAL_ENVIRONMENT=your-environment \ - uv run modal deploy -m lilo.providers.modal.app + uv run lilo deploy deployment.py ``` The configuration applies to the control plane, sampling workers, and new diff --git a/docs/profiling.md b/docs/profiling.md index 5136d9d..dfac5ed 100644 --- a/docs/profiling.md +++ b/docs/profiling.md @@ -14,7 +14,7 @@ Pick the step you want to trace and set one environment variable in the shell you deploy from. The trainer inherits it: ```bash -LILO_TORCH_PROFILE_STEP=2 uv run modal deploy -m lilo.providers.modal.app +LILO_TORCH_PROFILE_STEP=2 uv run lilo deploy deployment.py ``` Then run your training loop as usual. Steps are counted from 0, so `2` traces diff --git a/examples/codeforces-codegolf/uv.lock b/examples/codeforces-codegolf/uv.lock index 8cca0e9..66a9bd1 100644 --- a/examples/codeforces-codegolf/uv.lock +++ b/examples/codeforces-codegolf/uv.lock @@ -497,11 +497,13 @@ source = { editable = "../../" } dependencies = [ { name = "fastapi" }, { name = "httpx" }, + { name = "huggingface-hub" }, { name = "modal" }, { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, { name = "protobuf" }, { name = "pydantic" }, + { name = "pyyaml" }, { name = "stitch" }, { name = "tinker" }, { name = "uvicorn" }, @@ -513,11 +515,13 @@ dependencies = [ requires-dist = [ { name = "fastapi", specifier = ">=0.141.1" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "huggingface-hub", specifier = ">=0.34" }, { name = "modal", specifier = ">=1.5.3" }, { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.39,<2" }, { name = "opentelemetry-sdk", specifier = ">=1.39,<2" }, { name = "protobuf", specifier = ">=5.29" }, { name = "pydantic", specifier = ">=2.13.4" }, + { name = "pyyaml", specifier = ">=6.0.2" }, { name = "stitch", git = "https://github.com/modal-projects/stitch.git?rev=375a9396a7b05770dc4ed9cc5fe34fc4d5a472d5" }, { name = "tinker", specifier = ">=0.24.1,<0.25" }, { name = "uvicorn", specifier = ">=0.52.0" }, diff --git a/pyproject.toml b/pyproject.toml index d39e2a8..26e74d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "opentelemetry-sdk>=1.39,<2", "fastapi>=0.141.1", "httpx>=0.28.1", + "huggingface-hub>=0.34", "modal>=1.5.3", "protobuf>=5.29", "pydantic>=2.13.4", @@ -31,3 +32,6 @@ testpaths = ["tests"] dev = [ "pytest>=9.1.1", ] + +[project.scripts] +lilo = "lilo.deployment_cli:main" diff --git a/scripts/definition_smoke.py b/scripts/definition_smoke.py index 97d4c42..5dd5caf 100644 --- a/scripts/definition_smoke.py +++ b/scripts/definition_smoke.py @@ -15,11 +15,8 @@ --definition-id qwen3_5_9b_miles_lora_16k \ --definition-id qwen3_5_4b_full_64k -Any id from ``lilo.providers.modal.definitions`` works; the client type -(full or LoRA) and the LoRA target flags are read from the definition module so -the request matches what the deployment accepts. The definition id is passed as -``base_model`` so non-cataloged definitions can be targeted directly. -Definitions run sequentially unless ``--parallel`` is set. +IDs come from the deployed records selected by --app / --env. Use --list to +see active configurations. No local model catalog or GPU worker imports are needed. """ from __future__ import annotations @@ -40,9 +37,9 @@ from lilo.backends.miles_config import lora_target_flags from lilo.client import create_full_training_client -from lilo.providers.modal.app import DEFINITIONS, module_for +from lilo.deployments import DeploymentConfig +from lilo.providers.modal.deployment_configs import deployed_configs -DEFAULT_DEFINITION = "qwen3_8_27b_miles_lora_64k" TIMEOUT = 60 * 60 PROMPT = "Question: What is two plus two?\nAnswer:" @@ -168,13 +165,15 @@ def _unload(base_url: str, api_key: str, model_id: str) -> None: def _create_training_client( service: tinker.ServiceClient, definition: Any, rank: int | None ) -> tuple[tinker.TrainingClient, dict[str, Any]]: - definition_id = definition.DEFINITION_ID - if definition.PARAMETERIZATION == "full": + definition_id = definition.definition_id + if definition.parameterization == "full": training = create_full_training_client(service, definition_id) return training, {"parameterization": "full"} - train_attn, train_mlp, train_unembed = lora_target_flags(definition.TARGET_MODULES) + train_attn, train_mlp, train_unembed = lora_target_flags( + definition.trainer_settings["miles"]["target_modules"] + ) if rank is None: - rank = min(16, definition.MAX_LORA_RANK) + rank = min(16, definition.trainer_settings["miles"]["max_lora_rank"]) training = service.create_lora_training_client( base_model=definition_id, rank=rank, @@ -192,13 +191,14 @@ def _create_training_client( def _run_definition( - definition_id: str, + definition: Any, *, base_url: str, api_key: str, rank: int | None, max_tokens: int, ) -> dict: + definition_id = definition.definition_id report: dict[str, Any] = { "definition_id": definition_id, "status": "running", @@ -207,7 +207,6 @@ def _run_definition( } training = None try: - definition = module_for(definition_id) service = tinker.ServiceClient(base_url=base_url, api_key=api_key) started = time.perf_counter() training, spec = _create_training_client(service, definition, rank) @@ -272,14 +271,12 @@ def _run_definition( def main() -> None: parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) - known = [definition.DEFINITION_ID for definition in DEFINITIONS] parser.add_argument( "--definition-id", action="append", dest="definition_ids", - choices=known, metavar="ID", - help=f"engine definition id (repeatable); default {DEFAULT_DEFINITION}", + help="engine definition id (repeatable); defaults to the first active deployment", ) parser.add_argument( "--list", action="store_true", help="print known definitions and exit" @@ -297,14 +294,25 @@ def main() -> None: type=Path, default=Path("scripts/results/definition_smoke.json"), ) + parser.add_argument("--app", default="lilo") + parser.add_argument("--env") args = parser.parse_args() + rows = deployed_configs(args.app, args.env) + definitions = { + row.definition_id: row + for data in rows + if (row := DeploymentConfig.model_validate(data)) + } + if not definitions: + parser.error("no deployments in the saved configs") + if set(args.definition_ids or []) - definitions.keys(): + parser.error("unknown deployment ID; use --list to see deployed configurations") if args.list: - for definition in DEFINITIONS: + for definition in definitions.values(): print( - f"{definition.DEFINITION_ID:40} {definition.PARAMETERIZATION:5} " - f"{definition.GPUS}x{definition.GPU_TYPE} " - f"ctx={definition.MAX_CONTEXT_LENGTH}" - + ("" if definition.CATALOG_VISIBLE else " (not cataloged)") + f"{definition.definition_id:40} {definition.parameterization:5} " + f"{definition.recipe.trainer_nodes} nodes x {definition.recipe.trainer_gpu}:{definition.recipe.trainer_gpus_per_node} " + f"ctx={definition.max_context_length}" ) return if not args.base_url: @@ -312,11 +320,11 @@ def main() -> None: api_key = os.environ.get("TINKER_API_KEY") if not api_key: parser.error("TINKER_API_KEY is required") - definition_ids = args.definition_ids or [DEFAULT_DEFINITION] + definition_ids = args.definition_ids or [next(iter(definitions))] def run(definition_id: str) -> dict: return _run_definition( - definition_id, + definitions[definition_id], base_url=args.base_url, api_key=api_key, rank=args.rank, diff --git a/scripts/deploy_models.sh b/scripts/deploy_models.sh new file mode 100755 index 0000000..6d06c17 --- /dev/null +++ b/scripts/deploy_models.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +set -e + +# Run from the repository root so the Python config paths below work from any directory. +cd "$(dirname "$0")/.." + +# Add a model by creating its Python config in src/lilo/configs/ and adding it here. +# Put the preferred recipe first when multiple configs share a model. +deployment_files=( + src/lilo/configs/qwen35_9b_lora_16k.py + src/lilo/configs/qwen35_9b_lora_64k.py + src/lilo/configs/qwen35_4b_fft_64k.py +) + +lilo deploy "${deployment_files[@]}" "$@" diff --git a/scripts/deployment_smoke.py b/scripts/deployment_smoke.py new file mode 100644 index 0000000..3aecc1e --- /dev/null +++ b/scripts/deployment_smoke.py @@ -0,0 +1,194 @@ +"""Small real-GPU training/publication/sampling check for a Python-configured deployment. + +Run from an authenticated operator environment. Results are written to --output; +credentials are read from TINKER_API_KEY and never included in the report. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import time +from pathlib import Path + +import httpx +import modal +import tinker +from tinker import types + +from lilo.client import create_full_training_client +from lilo.providers.modal.deployment_configs import deployed_configs +from lilo.providers.modal.kv import app_store_name + + +def write(path, report): + temp = path.with_suffix(".tmp") + temp.write_text(json.dumps(report, indent=2) + "\n") + temp.replace(path) + + +def trainer_record(app_id, model_id): + placement = modal.Dict.from_name(app_store_name("lilo-models", app_id)).get( + f"placement:{model_id}" + ) + if not placement: + raise RuntimeError(f"No trainer placement for {model_id}") + row = modal.Dict.from_name(app_store_name("lilo-engines", app_id)).get( + f"engine_instance:{placement['engine_instance_id']}" + ) + return { + key: row.get(key) + for key in ( + "instance_id", + "definition_id", + "state", + "boot_id", + "call_id", + "revision", + ) + } + + +def step(training, tokens): + datum = types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + "target_tokens": tokens[1:], + "weights": [1.0] * (len(tokens) - 1), + }, + ) + start = time.monotonic() + result = training.forward_backward([datum], "cross_entropy").result(timeout=3600) + if len(result.loss_fn_outputs[0]["logprobs"].data) != len(tokens) - 1: + raise AssertionError("wrong forward/backward output length") + if not all( + math.isfinite(float(x)) for x in result.loss_fn_outputs[0]["logprobs"].data + ): + raise AssertionError("non-finite trainer logprobs") + train_seconds = time.monotonic() - start + optim = training.optim_step(types.AdamParams(learning_rate=1e-4)).result( + timeout=3600 + ) + sample = training.save_weights_and_get_sampling_client() + published = time.monotonic() + output = sample.sample( + prompt=types.ModelInput.from_ints(tokens[:16]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=16, temperature=0.0, seed=42), + ).result(timeout=3600) + sequence = output.sequences[0] + assert len(sequence.tokens) > 0 + assert sequence.logprobs is not None and len(sequence.logprobs) == len( + sequence.tokens + ) + assert all(math.isfinite(float(v)) for v in sequence.logprobs) + return { + "train_seconds": train_seconds, + "step_seconds": time.monotonic() - start, + "sample_seconds": time.monotonic() - published, + "train_metrics": result.metrics, + "optim_metrics": optim.metrics, + "generated_tokens": len(sequence.tokens), + "sample_tokens": sequence.tokens, + "sample_logprobs": sequence.logprobs, + } + + +def run(args): + app = modal.App.lookup(args.frontend) + server = modal.Function.from_name(args.frontend, "server") + url = server.get_web_url() + headers = {"X-API-Key": os.environ["TINKER_API_KEY"]} + rows = deployed_configs(args.frontend) + row = next(r for r in rows if r["recipe"]["name"] == args.name) + definition_id = args.name + report = { + "frontend": args.frontend, + "app_id": app.app_id, + "name": args.name, + "definition_id": definition_id, + "recipe": row["recipe"], + "status": "creating", + "steps": [], + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + write(output, report) + service = tinker.ServiceClient(base_url=url, api_key=os.environ["TINKER_API_KEY"]) + training = None + try: + start = time.monotonic() + training = ( + service.create_lora_training_client(base_model=definition_id, rank=32) + if row["recipe"]["parameterization"] == "lora" + else create_full_training_client(service, definition_id) + ) + report.update( + model_id=training.model_id, + create_seconds=time.monotonic() - start, + trainer_before=trainer_record(app.app_id, training.model_id), + status="training", + ) + write(output, report) + tokenizer = training.get_tokenizer() + text = "The sum of one and one is two. The sum of two and two is four.\n" * 40 + tokens = tokenizer.encode(text, add_special_tokens=True)[:513] + for index in range(args.steps): + report["steps"].append(step(training, tokens)) + write(output, report) + print(f"{args.name}: completed step {index + 1}", flush=True) + if args.continue_file: + report["status"] = "waiting_for_redeploy" + write(output, report) + deadline = time.monotonic() + 3600 + while not Path(args.continue_file).exists(): + if time.monotonic() > deadline: + raise TimeoutError("redeploy gate did not open") + time.sleep(5) + report["after_redeploy_step"] = step(training, tokens) + report["trainer_after"] = trainer_record(app.app_id, training.model_id) + assert ( + report["trainer_before"]["boot_id"] + == report["trainer_after"]["boot_id"] + ), "unchanged trainer restarted" + report["status"] = "passed" + except BaseException as exc: + report.update(status="failed", error=f"{type(exc).__name__}: {exc}") + raise + finally: + write(output, report) + if training is not None: + try: + with httpx.Client(base_url=url, headers=headers, timeout=60) as client: + response = client.post( + "/api/v1/unload_model", json={"model_id": training.model_id} + ) + response.raise_for_status() + request = response.json()["request_id"] + deadline = time.monotonic() + 180 + while time.monotonic() < deadline: + response = client.post( + "/api/v1/retrieve_future", json={"request_id": request} + ) + if response.status_code != 408: + response.raise_for_status() + break + except Exception as exc: + report["cleanup_error"] = str(exc) + write(output, report) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--frontend", required=True) + parser.add_argument("--name", required=True) + parser.add_argument("--steps", type=int, default=3) + parser.add_argument("--output", required=True) + parser.add_argument("--continue-file") + run(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/scripts/e2e_engine_definition.py b/scripts/e2e_engine_definition.py index 39a22ff..45780b8 100644 --- a/scripts/e2e_engine_definition.py +++ b/scripts/e2e_engine_definition.py @@ -7,8 +7,8 @@ import random import time import uuid -from contextlib import nullcontext from pathlib import Path +from types import SimpleNamespace from typing import Any import httpx @@ -16,17 +16,40 @@ import tinker from tinker import types -from lilo.providers.modal.app import app, cleaner, module_for, server +from lilo.backends.deployment import backend_config +from lilo.client import create_full_training_client +from lilo.deployments import DeploymentConfig +from lilo.providers.modal.deployment_configs import deployed_configs -DEFAULT_DEFINITION = "qwen3_5_4b_full_64k" TIMEOUT = 3 * 60 * 60 -def _definition(definition_id: str) -> tuple[Any, str]: - module = module_for(definition_id) - if not module.CATALOG_VISIBLE: - raise ValueError(f"definition is not cataloged: {definition_id}") - return module, module.PARAMETERIZATION +def _definition(frontend: str, name: str) -> tuple[Any, str]: + rows = deployed_configs(frontend) + matches = [ + DeploymentConfig.model_validate(row) + for row in rows + if row["recipe"]["name"] == name + ] + if len(matches) != 1: + raise ValueError(f"expected one configuration named {name} in {frontend}") + resolved = matches[0] + spec = resolved.recipe + settings = backend_config(spec)[spec.backend] + definition = SimpleNamespace( + DEFINITION_ID=resolved.definition_id, + MODEL_NAME=spec.model, + PARAMETERIZATION=spec.parameterization, + MAX_CONTEXT_LENGTH=spec.max_context_length, + MAX_TOKENS_PER_MICROBATCH=settings.get( + "max_tokens_per_microbatch", settings.get("max_tokens_per_gpu") + ), + MICRO_BATCH_SIZE=settings.get("micro_batch_size", 1), + GPU_TYPE=spec.trainer_gpu, + GPUS=spec.trainer_gpus_per_node, + LORA_RANK=settings.get("max_lora_rank"), + ) + return definition, definition.PARAMETERIZATION def _timestamped(path: Path) -> Path: @@ -205,9 +228,7 @@ def _checkpoint_roundtrip( load_finished = time.perf_counter() restore_error = _max_error(restored, before) if restore_error > 1e-5: - raise RuntimeError( - f"checkpoint restore max logprob error: {restore_error}" - ) + raise RuntimeError(f"checkpoint restore max logprob error: {restore_error}") resumed_step = _forward_step(resumed, [datum], [length], trained_tokens) continued = _forward_logprobs(resumed, datum, trained_tokens) @@ -415,11 +436,9 @@ def _sampler_recovery_roundtrip(training, tokenizer) -> dict: def _create_training(service, module, parameterization: str): if parameterization == "full": - from lilo.client import create_full_training_client - - return create_full_training_client(service, module.MODEL_NAME) + return create_full_training_client(service, module.DEFINITION_ID) return service.create_lora_training_client( - base_model=module.MODEL_NAME, + base_model=module.DEFINITION_ID, rank=module.LORA_RANK, ) @@ -595,14 +614,14 @@ def _correctness( def _run(args: argparse.Namespace, base_url: str, output: Path) -> dict: - module, parameterization = _definition(args.definition_id) + module, parameterization = _definition(args.frontend, args.name) api_key = os.environ["TINKER_API_KEY"] context_length = int(module.MAX_CONTEXT_LENGTH) packed_capacity = int(module.MAX_TOKENS_PER_MICROBATCH) report: dict[str, Any] = { "status": "running", "definition": { - "definition_id": args.definition_id, + "definition_id": module.DEFINITION_ID, "base_model": module.MODEL_NAME, "parameterization": parameterization, "context_length": context_length, @@ -668,7 +687,7 @@ def _run(args: argparse.Namespace, base_url: str, output: Path) -> dict: sampling, warmup = _warm( training, tokenizer, - args.definition_id, + module.DEFINITION_ID, ) report["phases"]["warmup"] = warmup _write(output, report) @@ -705,7 +724,8 @@ def _run(args: argparse.Namespace, base_url: str, output: Path) -> dict: def main() -> None: parser = argparse.ArgumentParser() - parser.add_argument("--definition-id", default=DEFAULT_DEFINITION) + parser.add_argument("--frontend", required=True) + parser.add_argument("--name", required=True) parser.add_argument("--base-url") parser.add_argument("--skip-max-context", action="store_true") parser.add_argument("--checkpoint-only", action="store_true") @@ -717,30 +737,30 @@ def main() -> None: default=Path("scripts/results/engine_definition_e2e.json"), ) args = parser.parse_args() - if sum( - ( - args.checkpoint_only, - args.hf_roundtrip_only, - args.sampler_recovery_only, + if ( + sum( + ( + args.checkpoint_only, + args.hf_roundtrip_only, + args.sampler_recovery_only, + ) ) - ) > 1: + > 1 + ): parser.error( "--checkpoint-only, --hf-roundtrip-only, and " "--sampler-recovery-only are exclusive" ) output = _timestamped(args.output) - context = nullcontext(args.base_url) - if args.base_url is None: - context = app.run(name=f"tinker-e2e-{uuid.uuid4().hex[:12]}") try: - with modal.enable_output(), context: - base_url = args.base_url or server.get_web_url() - if not base_url: - raise RuntimeError("Modal did not provide a control-plane URL") - report = _run(args, base_url, output) - if args.base_url is None: - cleaner.remote() + base_url = ( + args.base_url + or modal.Function.from_name(args.frontend, "server").get_web_url() + ) + if not base_url: + raise RuntimeError("Modal did not provide a control-plane URL") + report = _run(args, base_url, output) finally: print(output) summary = { diff --git a/src/lilo/backends/deployment.py b/src/lilo/backends/deployment.py new file mode 100644 index 0000000..3af8479 --- /dev/null +++ b/src/lilo/backends/deployment.py @@ -0,0 +1,158 @@ +"""Resolve backend settings before launch; reserve fields owned by Lilo.""" + +from dataclasses import asdict + +from lilo.backends.megatron_config import parse_backend_config +from lilo.backends.miles_config import MilesBackendConfig +from lilo.config_validation import reject_managed_options + +MILES_MANAGED = { + "context_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "lora_alpha", + "lora_dropout", + "lora_rank", + "max_tokens_per_gpu", + "multi_lora_n_adapters", + "target_modules", + "tensor_model_parallel_size", + "hf_checkpoint", + "load", + "pretrained_checkpoint", + "train_backend", + "actor_num_nodes", + "actor_num_gpus_per_node", + "rollout_num_gpus", + "debug_train_only", + "megatron_to_hf_mode", + "seq_length", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "colocate", + "custom_actor", + "sglang_model_path", + "use_dynamic_global_batch_size", + "delay_split_train_data_by_dp", + "use_dynamic_batch_size", + "optimizer", + "gradient_accumulation_fusion", + "save", + "save_interval", + "ckpt_step", + "rollout_num_gpus_per_engine", + "num_gpus_per_node", +} + + +SGLANG_MANAGED = { + "model_path", + "model", + "host", + "port", + "context_length", + "enable_lora", + "max_lora_rank", + "lora_target_modules", + "enable_cpu_weight_cache", + "api_key", + "pp_size", + "lora_paths", + "dist_init_addr", + "nnodes", + "node_rank", + "tokenizer_path", + "tokenizer_revision", + "revision", + "grpc_mode", + "smg_grpc_mode", + "encoder_only", + "use_ray", + "disaggregation_mode", + "skip_tokenizer_init", +} + + +def backend_config(spec, asset_path="/assets/pending"): + if spec.backend == "megatron": + if spec.trainer_nodes != 1: + raise ValueError("multi-node training currently requires Miles") + if spec.parameterization != "full": + raise ValueError("Megatron requires full parameterization") + if spec.trainer_max_clients_per_instance != 1: + raise ValueError("FFT trainers admit one client per instance") + if spec.sampler_persistence_concurrency != 1: + raise ValueError("Megatron requires sampler_persistence_concurrency: 1") + settings = spec.megatron_cfg + reject_managed_options(settings, {"hf_checkpoint", "seq_length"}) + config, _ = parse_backend_config( + { + "megatron": { + **settings, + "hf_checkpoint": asset_path, + "seq_length": spec.max_context_length, + } + } + ) + if config.optimizer.optimizer != "adam": + raise ValueError("Tinker optim_step requires an Adam optimizer") + config.validate(spec.trainer_gpus_per_node) + return {"megatron": asdict(config), "checkpoint_dir": "/checkpoints"} + if spec.backend != "miles": + raise ValueError(f"unknown backend: {spec.backend}") + if spec.parameterization != "lora": + raise ValueError("Miles requires lora parameterization") + settings = spec.miles_cfg + reject_managed_options( + settings, + {"hf_checkpoint", "actor_num_gpus_per_node", "actor_num_nodes", "extra_args"}, + ) + reject_managed_options(settings.get("cli_options", {}), MILES_MANAGED) + config = MilesBackendConfig( + hf_checkpoint=asset_path, + actor_num_gpus_per_node=spec.trainer_gpus_per_node, + actor_num_nodes=spec.trainer_nodes, + extra_args=("--seq-length", str(spec.max_context_length)), + **settings, + ) + config.validate() + if config.world_size % ( + config.expert_model_parallel_size * config.expert_tensor_parallel_size + ): + raise ValueError("expert parallel sizes must divide the trainer GPU allocation") + if spec.trainer_max_clients_per_instance > config.max_lora_slots: + raise ValueError("max_clients_per_instance exceeds max_lora_slots") + return {"miles": asdict(config), "checkpoint_dir": "/checkpoints"} + + +def serving_options(spec): + options = dict(spec.sglang_cfg) + reject_managed_options(options, SGLANG_MANAGED) + tp = options.get("tp_size", spec.inference_gpus_per_node) + if tp != spec.inference_gpus_per_node: + raise ValueError("sglang.tp_size must equal the replica GPU allocation") + return options + + +def resolve_backend_settings(spec, asset_path): + trainer = backend_config(spec, asset_path) + inference = { + "context_length": spec.max_context_length, + "tp_size": spec.inference_gpus_per_node, + "mem_fraction_static": 0.8, + "max_running_requests": 32, + "weight_loader_disable_mmap": True, + **serving_options(spec), + } + if spec.parameterization == "lora": + miles = MilesBackendConfig(**trainer["miles"]) + inference.update( + enable_lora=True, + max_lora_rank=miles.max_lora_rank, + lora_target_modules=list(miles.peft_target_modules), + ) + inference.setdefault("max_loaded_loras", 64) + inference.setdefault("max_loras_per_batch", 8) + else: + inference["enable_cpu_weight_cache"] = True + return trainer, inference diff --git a/src/lilo/backends/megatron_runtime/common/config.py b/src/lilo/backends/megatron_runtime/common/config.py index 785852a..494083a 100644 --- a/src/lilo/backends/megatron_runtime/common/config.py +++ b/src/lilo/backends/megatron_runtime/common/config.py @@ -3,6 +3,8 @@ import math from dataclasses import dataclass, field +from .settings import distributed_settings, optimizer_settings, provider_settings + @dataclass(frozen=True, slots=True) class OptimizerConfig: @@ -61,6 +63,8 @@ class EngineModelConfig: defer_fp32_logits: bool = False fp32_lm_head: bool = False provider_overrides: dict[str, object] = field(default_factory=dict) + optimizer_overrides: dict[str, object] = field(default_factory=dict) + distributed_overrides: dict[str, object] = field(default_factory=dict) overlap_grad_reduce: bool = False align_grad_reduce: bool = True @@ -107,6 +111,9 @@ def packed_token_capacity(self) -> int: ) def validate(self, world_size: int) -> None: + provider_settings(self, None) + optimizer_settings(self, None, self.use_distributed_optimizer) + distributed_settings(self, self.use_distributed_optimizer) parallel_sizes = { "tensor_model_parallel_size": self.tensor_model_parallel_size, "pipeline_model_parallel_size": self.pipeline_model_parallel_size, diff --git a/src/lilo/backends/megatron_runtime/common/modeling.py b/src/lilo/backends/megatron_runtime/common/modeling.py index 97cdbf7..8bc4b00 100644 --- a/src/lilo/backends/megatron_runtime/common/modeling.py +++ b/src/lilo/backends/megatron_runtime/common/modeling.py @@ -1,5 +1,7 @@ from __future__ import annotations +from dataclasses import replace + import torch from megatron.bridge import AutoBridge from megatron.core.distributed import DistributedDataParallelConfig @@ -7,6 +9,7 @@ from megatron.core.transformer.enums import AttnBackend from .config import EngineModelConfig +from .settings import distributed_settings, optimizer_settings, provider_settings def model_provider(config: EngineModelConfig): @@ -16,27 +19,20 @@ def model_provider(config: EngineModelConfig): trust_remote_code=True, ) provider = bridge.to_megatron_provider() - provider.tensor_model_parallel_size = config.tensor_model_parallel_size - provider.pipeline_model_parallel_size = config.pipeline_model_parallel_size - provider.virtual_pipeline_model_parallel_size = ( - config.virtual_pipeline_model_parallel_size + settings = provider_settings(config, dtype) + if provider.moe_token_dispatcher_type == "allgather": + settings.setdefault("moe_token_dispatcher_type", "alltoall") + configured = replace( + provider, + **{**settings, "attention_backend": AttnBackend[config.attention_backend]}, ) - provider.context_parallel_size = config.context_parallel_size - provider.expert_model_parallel_size = config.expert_model_parallel_size - provider.expert_tensor_parallel_size = config.expert_tensor_parallel_size - provider.sequence_parallel = config.sequence_parallel - provider.variable_seq_lengths = True - if getattr(provider, "moe_token_dispatcher_type", None) == "allgather": - provider.moe_token_dispatcher_type = "alltoall" - provider.calculate_per_token_loss = config.calculate_per_token_loss - provider.attention_backend = AttnBackend[config.attention_backend] - provider.cross_entropy_loss_fusion = config.cross_entropy_loss_fusion - provider.params_dtype = dtype - for name, value in config.provider_overrides.items(): - if not hasattr(provider, name): - raise ValueError(f"unknown Megatron provider override: {name}") - setattr(provider, name, value) - return bridge, provider, dtype + # Bridge attaches weight loading as a hook outside the dataclass fields. + # Replacing the provider must preserve it before constructing the model. + if provider.pre_wrap_hook is not None: + configured.register_pre_wrap_hook(provider.pre_wrap_hook) + if provider.post_wrap_hook is not None: + configured.register_post_wrap_hook(provider.post_wrap_hook) + return bridge, configured, dtype def parameter_dtype(config: EngineModelConfig): @@ -57,10 +53,7 @@ def distributed_model( ): return provider.provide_distributed_model( ddp_config=DistributedDataParallelConfig( - use_distributed_optimizer=distributed_optimizer, - overlap_grad_reduce=config.overlap_grad_reduce, - overlap_param_gather=config.overlap_param_gather, - align_param_gather=config.align_param_gather, + **distributed_settings(config, distributed_optimizer) ), bf16=config.bf16, fp16=config.fp16, @@ -74,17 +67,5 @@ def optimizer_config( distributed_optimizer: bool, ): return MCoreOptimizerConfig( - optimizer=config.optimizer.optimizer, - lr=config.optimizer.lr, - weight_decay=config.optimizer.weight_decay, - adam_beta1=config.optimizer.adam_beta1, - adam_beta2=config.optimizer.adam_beta2, - adam_eps=config.optimizer.adam_eps, - clip_grad=config.optimizer.clip_grad, - loss_scale=config.optimizer.loss_scale, - bf16=config.bf16, - fp16=config.fp16, - params_dtype=dtype, - use_distributed_optimizer=distributed_optimizer, - overlap_param_gather=config.overlap_param_gather, + **optimizer_settings(config, dtype, distributed_optimizer) ) diff --git a/src/lilo/backends/megatron_runtime/common/settings.py b/src/lilo/backends/megatron_runtime/common/settings.py new file mode 100644 index 0000000..6504a8a --- /dev/null +++ b/src/lilo/backends/megatron_runtime/common/settings.py @@ -0,0 +1,57 @@ +"""One ownership rule: backend extras may add fields, never replace Lilo settings.""" + +from dataclasses import asdict + +from lilo.config_validation import reject_managed_options + + +def provider_settings(config, dtype): + owned = { + "tensor_model_parallel_size": config.tensor_model_parallel_size, + "pipeline_model_parallel_size": config.pipeline_model_parallel_size, + "virtual_pipeline_model_parallel_size": config.virtual_pipeline_model_parallel_size, + "context_parallel_size": config.context_parallel_size, + "expert_model_parallel_size": config.expert_model_parallel_size, + "expert_tensor_parallel_size": config.expert_tensor_parallel_size, + "sequence_parallel": config.sequence_parallel, + "variable_seq_lengths": True, + "calculate_per_token_loss": config.calculate_per_token_loss, + "attention_backend": config.attention_backend, + "cross_entropy_loss_fusion": config.cross_entropy_loss_fusion, + "params_dtype": dtype, + } + reject_managed_options(config.provider_overrides, owned.keys() | {"seq_length"}) + return {**owned, **config.provider_overrides} + + +def optimizer_settings(config, dtype, distributed_optimizer): + owned = { + "optimizer": config.optimizer.optimizer, + "lr": config.optimizer.lr, + "weight_decay": config.optimizer.weight_decay, + "adam_beta1": config.optimizer.adam_beta1, + "adam_beta2": config.optimizer.adam_beta2, + "adam_eps": config.optimizer.adam_eps, + "clip_grad": config.optimizer.clip_grad, + "loss_scale": config.optimizer.loss_scale, + "bf16": config.bf16, + "fp16": config.fp16, + "params_dtype": dtype, + "use_distributed_optimizer": distributed_optimizer, + "overlap_param_gather": config.overlap_param_gather, + } + reject_managed_options( + config.optimizer_overrides, owned.keys() | asdict(config.optimizer).keys() + ) + return {**owned, **config.optimizer_overrides} + + +def distributed_settings(config, distributed_optimizer): + owned = { + "use_distributed_optimizer": distributed_optimizer, + "overlap_grad_reduce": config.overlap_grad_reduce, + "overlap_param_gather": config.overlap_param_gather, + "align_param_gather": config.align_param_gather, + } + reject_managed_options(config.distributed_overrides, owned.keys()) + return {**owned, **config.distributed_overrides} diff --git a/src/lilo/backends/megatron_runtime/fft/checkpoint.py b/src/lilo/backends/megatron_runtime/fft/checkpoint.py index 4aa0f14..eadf0b9 100644 --- a/src/lilo/backends/megatron_runtime/fft/checkpoint.py +++ b/src/lilo/backends/megatron_runtime/fft/checkpoint.py @@ -7,7 +7,7 @@ import os import shutil from copy import deepcopy -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any @@ -56,9 +56,18 @@ class FFTCheckpointMetadata: optimizer_config: dict[str, Any] has_optimizer: bool optimizer_state_format: str | None + base_model_revision: str | None = None + native_optimizer_config: dict[str, Any] = field(default_factory=dict) + native_distributed_config: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: - return asdict(self) + value = asdict(self) + if self.base_model_revision is None: + value.pop("base_model_revision") + for name in ("native_optimizer_config", "native_distributed_config"): + if not value[name]: + value.pop(name) + return value def identity(self) -> str: encoded = json.dumps( @@ -95,11 +104,15 @@ def create_fft_checkpoint_metadata( precision="bf16" if config.bf16 else "fp16" if config.fp16 else "fp32", use_distributed_optimizer=config.use_distributed_optimizer, optimizer_config=asdict(config.optimizer), + native_optimizer_config=dict(config.optimizer_overrides), + native_distributed_config=dict(config.distributed_overrides), has_optimizer=include_optimizer, optimizer_state_format=( DISTRIBUTED_OPTIMIZER_STATE_FORMAT if include_optimizer and config.use_distributed_optimizer - else REGULAR_OPTIMIZER_STATE_FORMAT if include_optimizer else None + else REGULAR_OPTIMIZER_STATE_FORMAT + if include_optimizer + else None ), ) @@ -176,8 +189,7 @@ def restore_fft_optimizer_state( ) if not isinstance(state, dict) or state.get("format") != expected_format: raise ValueError( - "checkpoint optimizer state format mismatch: " - f"expected {expected_format!r}" + f"checkpoint optimizer state format mismatch: expected {expected_format!r}" ) state_dict = state.get("state_dict") if not isinstance(state_dict, dict | list): @@ -291,6 +303,8 @@ def load_fft_training_checkpoint( "precision", "use_distributed_optimizer", "optimizer_config", + "native_optimizer_config", + "native_distributed_config", "optimizer_state_format", ) for name in comparable_fields: @@ -363,7 +377,9 @@ def synchronize_checkpoint_preflight( def _materialize_local_sharded_state(value): """Detach Megatron's local torch_dist representation from live optimizer buffers.""" if isinstance(value, ShardedTensorFactory): - raise TypeError("distributed optimizer state contains an unresolved tensor factory") + raise TypeError( + "distributed optimizer state contains an unresolved tensor factory" + ) if isinstance(value, ShardedTensor): if value.data is None: raise ValueError("distributed optimizer sharded tensor has no local data") @@ -374,8 +390,7 @@ def _materialize_local_sharded_state(value): return _materialize_local_sharded_state(value.unwrap()) if isinstance(value, dict): return { - key: _materialize_local_sharded_state(item) - for key, item in value.items() + key: _materialize_local_sharded_state(item) for key, item in value.items() } if isinstance(value, list): return [_materialize_local_sharded_state(item) for item in value] @@ -388,7 +403,9 @@ def _validate_distributed_optimizer_state_dict(state_dict: dict | list) -> None: """Require every DistributedOptimizer leaf to contain its local tensor shards.""" if isinstance(state_dict, list): if not state_dict: - raise ValueError("distributed optimizer checkpoint has no optimizer entries") + raise ValueError( + "distributed optimizer checkpoint has no optimizer entries" + ) for item in state_dict: if not isinstance(item, dict | list): raise TypeError("distributed optimizer entry must be a dict or list") @@ -397,17 +414,22 @@ def _validate_distributed_optimizer_state_dict(state_dict: dict | list) -> None: if state_dict.get("param_state_sharding_type") == "dp_reshardable": if not isinstance(state_dict.get("optimizer"), dict): - raise ValueError("distributed optimizer checkpoint has no optimizer metadata") - if not isinstance(state_dict.get("param_state"), dict) or not state_dict[ - "param_state" - ]: + raise ValueError( + "distributed optimizer checkpoint has no optimizer metadata" + ) + if ( + not isinstance(state_dict.get("param_state"), dict) + or not state_dict["param_state"] + ): raise ValueError("distributed optimizer checkpoint has no parameter state") return if state_dict and all(isinstance(key, int) for key in state_dict): for item in state_dict.values(): if not isinstance(item, dict | list): - raise TypeError("chained distributed optimizer entry must be a dict or list") + raise TypeError( + "chained distributed optimizer entry must be a dict or list" + ) _validate_distributed_optimizer_state_dict(item) return diff --git a/src/lilo/backends/miles_arguments.py b/src/lilo/backends/miles_arguments.py new file mode 100644 index 0000000..1a8af52 --- /dev/null +++ b/src/lilo/backends/miles_arguments.py @@ -0,0 +1,51 @@ +"""Translate config dictionaries into arguments for Miles. + +Miles exposes an argparse parser. Append configured values after their +preset arguments so argparse itself handles types, choices and required options. +Boolean flags need defaults because a store_true flag cannot express False. +""" + +import argparse + + +def apply_config_overrides(parser, options, argv): + """Mutate argv/defaults before the backend calls parse_args.""" + boolean_actions = ( + argparse._StoreTrueAction, + argparse._StoreFalseAction, + argparse.BooleanOptionalAction, + ) + actions_by_name = {} + for action in parser._actions: + actions_by_name.setdefault(action.dest, []).append(action) + + for name, value in options.items(): + actions = actions_by_name.get(name, []) + if not actions or not all(action.option_strings for action in actions): + raise ValueError(f"unknown backend option: {name}") + if value is None: + raise ValueError(f"backend option {name} cannot be null") + if all(isinstance(action, boolean_actions) for action in actions): + if not isinstance(value, bool): + raise ValueError(f"backend option {name} requires a boolean") + flags = {flag for action in actions for flag in action.option_strings} + argv[:] = [arg for arg in argv if arg.split("=", 1)[0] not in flags] + for action in actions: + action.required = False + parser.set_defaults(**{name: value}) + continue + if len(actions) != 1 or type(actions[0]) is not argparse._StoreAction: + raise ValueError( + f"backend option {name} uses an unsupported argparse action" + ) + action = actions[0] + multiple = action.nargs in ("+", "*") or isinstance(action.nargs, int) + flag = action.option_strings[0] + if multiple: + if not isinstance(value, list): + raise ValueError(f"backend option {name} requires a list") + argv.extend([flag, *(str(item) for item in value)]) + else: + if isinstance(value, (dict, list, bool)): + raise ValueError(f"backend option {name} requires a scalar") + argv.append(f"{flag}={value}") diff --git a/src/lilo/backends/miles_config.py b/src/lilo/backends/miles_config.py index cf108af..b83029a 100644 --- a/src/lilo/backends/miles_config.py +++ b/src/lilo/backends/miles_config.py @@ -1,11 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any -MILES_REF = "main" - _PEFT_TARGETS = { "linear_qkv": ("q_proj", "k_proj", "v_proj"), "linear_q": ("q_proj",), @@ -19,12 +17,30 @@ "output_layer": ("lm_head",), } _ATTN_LEAVES = frozenset( - {"linear_qkv", "linear_q", "linear_k", "linear_v", "linear_proj"} + { + "linear_qkv", + "linear_q", + "linear_k", + "linear_v", + "linear_proj", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + } ) _MLP_LEAVES = frozenset( - {"linear_fc1", "linear_fc1_gate", "linear_fc1_up", "linear_fc2"} + { + "linear_fc1", + "linear_fc1_gate", + "linear_fc1_up", + "linear_fc2", + "gate_proj", + "up_proj", + "down_proj", + } ) -_UNEMBED_LEAVES = frozenset({"output_layer"}) +_UNEMBED_LEAVES = frozenset({"output_layer", "lm_head"}) def lora_target_flags(target_modules: tuple[str, ...]) -> tuple[bool, bool, bool]: @@ -63,6 +79,7 @@ class MilesBackendConfig: max_tokens_per_gpu: int = 8192 align_sequences_to_parallel_layout: bool = False extra_args: tuple[str, ...] = () + cli_options: dict[str, Any] = field(default_factory=dict) @property def world_size(self) -> int: @@ -103,14 +120,20 @@ def validate(self) -> None: "max_tokens_per_gpu": self.max_tokens_per_gpu, } for name, value in positive.items(): - if value < 1: - raise ValueError(f"{name} must be at least 1") + if type(value) is not int or value < 1: + raise ValueError(f"{name} must be a positive integer") if not self.hf_checkpoint: raise ValueError("hf_checkpoint is required") - if not self.model_type: - raise ValueError("model_type is required") - if not self.target_modules: - raise ValueError("target_modules must not be empty") + if not self.model_type and not self.cli_options: + raise ValueError( + "model_type or explicit native architecture options are required" + ) + if ( + not isinstance(self.target_modules, (list, tuple)) + or not self.target_modules + or not all(isinstance(name, str) and name for name in self.target_modules) + ): + raise ValueError("target_modules must be a nonempty list of module names") if ( self.default_lora_alpha <= 0 or not float(self.default_lora_alpha).is_integer() diff --git a/src/lilo/backends/miles_runtime/runtime.py b/src/lilo/backends/miles_runtime/runtime.py index 48e8bc5..b6580c9 100644 --- a/src/lilo/backends/miles_runtime/runtime.py +++ b/src/lilo/backends/miles_runtime/runtime.py @@ -13,6 +13,7 @@ from pathlib import Path from typing import Any +from lilo.backends.miles_arguments import apply_config_overrides from lilo.backends.miles_config import MilesBackendConfig from lilo.errors import BackendFailed @@ -215,10 +216,23 @@ async def _start(self) -> None: install_bridge_replay(tinker_runtime) _configure_actor_spec(train_specs) _allow_context_parallel_multi_lora(lora_arguments) - architecture = shlex.split(load_model_args(self.config.model_type)) + architecture = ( + shlex.split(load_model_args(self.config.model_type)) + if self.config.model_type + else [] + ) with _temporary_argv([*architecture, *self.config.miles_arguments()]): - args = parse_args(entry="serve") + if self.config.cli_options: + + def configure(parser): + apply_config_overrides(parser, self.config.cli_options, sys.argv) + return parser + + args = parse_args(add_custom_arguments=configure, entry="serve") + else: + args = parse_args(entry="serve") self._args = args + args.use_dynamic_global_batch_size = True args.delay_split_train_data_by_dp = True configure_logger(args, source=MainProcessIdentity()) diff --git a/src/lilo/config_validation.py b/src/lilo/config_validation.py new file mode 100644 index 0000000..a9a10ad --- /dev/null +++ b/src/lilo/config_validation.py @@ -0,0 +1,9 @@ +"""Checks for backend options that would override Lilo-managed settings.""" + + +def reject_managed_options(options, managed): + if not isinstance(options, dict): + raise ValueError("backend configuration must be a mapping") + conflicts = options.keys() & managed + if conflicts: + raise ValueError(f"options managed by Lilo: {sorted(conflicts)}") diff --git a/src/lilo/configs/__init__.py b/src/lilo/configs/__init__.py new file mode 100644 index 0000000..1fe58e9 --- /dev/null +++ b/src/lilo/configs/__init__.py @@ -0,0 +1 @@ +"""Example BaseConfig recipes; compose variants with ordinary Python inheritance.""" diff --git a/src/lilo/configs/qwen35_35b_a3b_fft_64k.py b/src/lilo/configs/qwen35_35b_a3b_fft_64k.py new file mode 100644 index 0000000..d659372 --- /dev/null +++ b/src/lilo/configs/qwen35_35b_a3b_fft_64k.py @@ -0,0 +1,57 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen35-35b-a3b-fft-64k" + parameterization = "full" + model = "Qwen/Qwen3.5-35B-A3B" + max_context_length = 65536 + trainer_gpu = "H200" + trainer_gpus_per_node = 8 + backend = "megatron" + megatron_cfg = { + "tensor_model_parallel_size": 4, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 2, + "expert_model_parallel_size": 8, + "expert_tensor_parallel_size": 1, + "sequence_parallel": True, + "micro_batch_size": 1, + "max_tokens_per_microbatch": 65536, + "bf16": True, + "fp16": False, + "gpu_memory_fraction": 0.9, + "use_distributed_optimizer": True, + "provider_overrides": { + "mtp_num_layers": 0, + "recompute_granularity": "selective", + "moe_layer_recompute": True, + "moe_token_dispatcher_type": "alltoall", + "moe_router_fusion": True, + "moe_permute_fusion": True, + "moe_grouped_gemm": True, + "moe_shared_expert_overlap": False, + "moe_aux_loss_coeff": 0.0, + }, + "optimizer": {"optimizer": "adam", "lr": 0.0001, "min_lr": 0.0001}, + } + trainer_env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + sampler_persistence_concurrency = 1 + inference_gpu = "H200" + inference_gpus_per_node = 4 + sglang_cfg = { + "tp_size": 4, + "ep_size": 4, + "mem_fraction_static": 0.9, + "max_running_requests": 32, + "max_queued_requests": 4, + "cpu_weight_cache_max_compile_group_gb": 32, + "dp_size": 4, + "enable_dp_attention": True, + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_4b_fft_64k.py b/src/lilo/configs/qwen35_4b_fft_64k.py new file mode 100644 index 0000000..aaab108 --- /dev/null +++ b/src/lilo/configs/qwen35_4b_fft_64k.py @@ -0,0 +1,40 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen35-4b-fft-64k" + parameterization = "full" + model = "Qwen/Qwen3.5-4B" + max_context_length = 65536 + trainer_gpu = "H100" + trainer_gpus_per_node = 4 + backend = "megatron" + megatron_cfg = { + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "sequence_parallel": True, + "micro_batch_size": 1, + "max_tokens_per_microbatch": 65536, + "defer_fp32_logits": True, + "fp32_lm_head": True, + "use_distributed_optimizer": True, + "provider_overrides": { + "mtp_num_layers": 0, + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + "optimizer": {"lr": 0.0001, "min_lr": 0.0001, "loss_scale": 1.0}, + } + sampler_persistence_concurrency = 1 + inference_gpu = "H100" + sglang_cfg = { + "tp_size": 1, + "mem_fraction_static": 0.85, + "max_running_requests": 32, + "max_queued_requests": 4, + "cpu_weight_cache_max_compile_group_gb": 16, + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_fft_64k.py b/src/lilo/configs/qwen35_9b_fft_64k.py new file mode 100644 index 0000000..b8682ed --- /dev/null +++ b/src/lilo/configs/qwen35_9b_fft_64k.py @@ -0,0 +1,18 @@ +from lilo.configs.qwen35_4b_fft_64k import Config as Parent + + +class Config(Parent): + name = "qwen35-9b-fft-64k" + model = "Qwen/Qwen3.5-9B" + overrides = { + "trainer_gpu": "H200", + "trainer_env": { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + }, + "inference_gpu": "H200", + "sglang_cfg.ep_size": 1, + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_instruct_lora_16k.py b/src/lilo/configs/qwen35_9b_instruct_lora_16k.py new file mode 100644 index 0000000..dc52519 --- /dev/null +++ b/src/lilo/configs/qwen35_9b_instruct_lora_16k.py @@ -0,0 +1,50 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen35-9b-instruct-lora-16k" + model = "Qwen/Qwen3.5-9B" + max_context_length = 16384 + trainer_gpu = "H100" + trainer_gpus_per_node = 8 + miles_cfg = { + "model_type": "qwen3.5-9B", + "tensor_model_parallel_size": 8, + "target_modules": [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + "max_tokens_per_gpu": 16384, + "max_lora_slots": 6, + "max_lora_rank": 32, + "default_lora_alpha": 32, + "cli_options": { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + } + trainer_env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + trainer_max_clients_per_instance = 6 + inference_gpu = "H200" + sglang_cfg = { + "tp_size": 1, + "ep_size": 1, + "mem_fraction_static": 0.8, + "max_running_requests": 32, + "max_queued_requests": 8, + "max_loaded_loras": 256, + "max_loras_per_batch": 8, + "schedule_policy": "lpm", + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_instruct_lora_16k_dp2.py b/src/lilo/configs/qwen35_9b_instruct_lora_16k_dp2.py new file mode 100644 index 0000000..5f7fbfe --- /dev/null +++ b/src/lilo/configs/qwen35_9b_instruct_lora_16k_dp2.py @@ -0,0 +1,9 @@ +from lilo.configs.qwen35_9b_instruct_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen35-9b-instruct-lora-16k-dp2" + overrides = {"miles_cfg.tensor_model_parallel_size": 4} + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_lora_16k.py b/src/lilo/configs/qwen35_9b_lora_16k.py new file mode 100644 index 0000000..916c649 --- /dev/null +++ b/src/lilo/configs/qwen35_9b_lora_16k.py @@ -0,0 +1,51 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen35-9b-lora-16k" + model = "Qwen/Qwen3.5-9B-Base" + max_context_length = 16384 + trainer_gpu = "H100" + trainer_gpus_per_node = 4 + trainer_cpu = 16 + trainer_memory_mib = 65536 + miles_cfg = { + "model_type": "qwen3.5-9B", + "tensor_model_parallel_size": 4, + "max_lora_slots": 6, + "max_lora_rank": 32, + "default_lora_alpha": 32, + "target_modules": [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "lm_head", + ], + "max_tokens_per_gpu": 16384, + "cli_options": { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + } + trainer_env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + trainer_max_clients_per_instance = 6 + inference_gpu = "H200" + sglang_cfg = { + "tp_size": 1, + "mem_fraction_static": 0.8, + "max_running_requests": 32, + "max_queued_requests": 8, + "max_loaded_loras": 64, + "max_loras_per_batch": 8, + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_lora_16k_single.py b/src/lilo/configs/qwen35_9b_lora_16k_single.py new file mode 100644 index 0000000..7a596e9 --- /dev/null +++ b/src/lilo/configs/qwen35_9b_lora_16k_single.py @@ -0,0 +1,9 @@ +from lilo.configs.qwen35_9b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen35-9b-lora-16k-single" + overrides = {"trainer_max_clients_per_instance": 1} + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_lora_2k.py b/src/lilo/configs/qwen35_9b_lora_2k.py new file mode 100644 index 0000000..d900ca4 --- /dev/null +++ b/src/lilo/configs/qwen35_9b_lora_2k.py @@ -0,0 +1,20 @@ +from lilo.configs.qwen35_9b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen35-9b-lora-2k" + max_context_length = 2048 + overrides = { + "trainer_gpu": "H200", + "trainer_cpu": 8, + "trainer_memory_mib": 32768, + "trainer_max_clients_per_instance": 4, + "miles_cfg.max_tokens_per_gpu": 2048, + "miles_cfg.max_lora_slots": 4, + "sglang_cfg.ep_size": 1, + "sglang_cfg.max_loaded_loras": 32, + "sglang_cfg.schedule_policy": "lpm", + } + + +config = Config() diff --git a/src/lilo/configs/qwen35_9b_lora_64k.py b/src/lilo/configs/qwen35_9b_lora_64k.py new file mode 100644 index 0000000..a0354d3 --- /dev/null +++ b/src/lilo/configs/qwen35_9b_lora_64k.py @@ -0,0 +1,15 @@ +from lilo.configs.qwen35_9b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen35-9b-lora-64k" + max_context_length = 65536 + overrides = { + "trainer_gpu": "H200", + "trainer_gpus_per_node": 8, + "miles_cfg.tensor_model_parallel_size": 8, + "miles_cfg.max_tokens_per_gpu": 65536, + } + + +config = Config() diff --git a/src/lilo/configs/qwen36_27b_fft_64k.py b/src/lilo/configs/qwen36_27b_fft_64k.py new file mode 100644 index 0000000..0af1b67 --- /dev/null +++ b/src/lilo/configs/qwen36_27b_fft_64k.py @@ -0,0 +1,48 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen36-27b-fft-64k" + parameterization = "full" + model = "Qwen/Qwen3.6-27B" + max_context_length = 65536 + trainer_gpu = "H200" + trainer_gpus_per_node = 8 + backend = "megatron" + megatron_cfg = { + "tensor_model_parallel_size": 4, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 2, + "sequence_parallel": True, + "micro_batch_size": 1, + "max_tokens_per_microbatch": 65536, + "bf16": True, + "fp16": False, + "gpu_memory_fraction": 0.9, + "use_distributed_optimizer": True, + "provider_overrides": { + "mtp_num_layers": 0, + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + "optimizer": {"optimizer": "adam", "lr": 0.0001, "min_lr": 0.0001}, + } + trainer_env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + sampler_persistence_concurrency = 1 + inference_gpu = "H200" + inference_gpus_per_node = 4 + sglang_cfg = { + "tp_size": 4, + "ep_size": 1, + "mem_fraction_static": 0.9, + "max_running_requests": 32, + "max_queued_requests": 4, + "cpu_weight_cache_max_compile_group_gb": 32, + } + + +config = Config() diff --git a/src/lilo/configs/qwen36_35b_a3b_fft_64k.py b/src/lilo/configs/qwen36_35b_a3b_fft_64k.py new file mode 100644 index 0000000..25a854a --- /dev/null +++ b/src/lilo/configs/qwen36_35b_a3b_fft_64k.py @@ -0,0 +1,10 @@ +from lilo.configs.qwen35_35b_a3b_fft_64k import Config as Parent + + +class Config(Parent): + name = "qwen36-35b-a3b-fft-64k" + model = "Qwen/Qwen3.6-35B-A3B" + overrides = {"sglang_cfg.dp_size": 1, "sglang_cfg.enable_dp_attention": False} + + +config = Config() diff --git a/src/lilo/configs/qwen38_27b_lora_128k.py b/src/lilo/configs/qwen38_27b_lora_128k.py new file mode 100644 index 0000000..b18b718 --- /dev/null +++ b/src/lilo/configs/qwen38_27b_lora_128k.py @@ -0,0 +1,19 @@ +from lilo.configs.qwen38_27b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen38-27b-lora-128k" + max_context_length = 131072 + overrides = { + "miles_cfg.tensor_model_parallel_size": 2, + "miles_cfg.max_tokens_per_gpu": 32768, + "miles_cfg.context_parallel_size": 4, + "inference_gpus_per_node": 2, + "inference_max_replicas": 4, + "inference_target_concurrency": 4, + "sglang_cfg.tp_size": 2, + "sglang_cfg.max_running_requests": 8, + } + + +config = Config() diff --git a/src/lilo/configs/qwen38_27b_lora_16k.py b/src/lilo/configs/qwen38_27b_lora_16k.py new file mode 100644 index 0000000..02a9426 --- /dev/null +++ b/src/lilo/configs/qwen38_27b_lora_16k.py @@ -0,0 +1,50 @@ +from lilo.configuration import BaseConfig + + +class Config(BaseConfig): + name = "qwen38-27b-lora-16k" + model = "Qwen/Qwen3.8-27B" + max_context_length = 16384 + trainer_gpu = "H200" + trainer_gpus_per_node = 8 + miles_cfg = { + "model_type": "qwen3.8-27B", + "tensor_model_parallel_size": 4, + "target_modules": [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + "max_tokens_per_gpu": 16384, + "max_lora_slots": 6, + "max_lora_rank": 32, + "default_lora_alpha": 32, + "cli_options": { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + } + trainer_env = { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + trainer_max_clients_per_instance = 6 + inference_gpu = "H200" + sglang_cfg = { + "tp_size": 1, + "ep_size": 1, + "mem_fraction_static": 0.8, + "max_running_requests": 32, + "max_queued_requests": 8, + "max_loaded_loras": 256, + "max_loras_per_batch": 8, + "schedule_policy": "lpm", + } + + +config = Config() diff --git a/src/lilo/configs/qwen38_27b_lora_256k.py b/src/lilo/configs/qwen38_27b_lora_256k.py new file mode 100644 index 0000000..0ab78e9 --- /dev/null +++ b/src/lilo/configs/qwen38_27b_lora_256k.py @@ -0,0 +1,24 @@ +from lilo.configs.qwen38_27b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen38-27b-lora-256k" + max_context_length = 262144 + overrides = { + "trainer_nodes": 2, + "miles_cfg.tensor_model_parallel_size": 2, + "miles_cfg.context_parallel_size": 8, + "miles_cfg.max_tokens_per_gpu": 32768, + "miles_cfg.cli_options.distributed_timeout_minutes": 120, + "inference_gpus_per_node": 4, + "inference_min_replicas": 2, + "inference_max_replicas": 2, + "inference_target_concurrency": 2, + "sglang_cfg.tp_size": 4, + "sglang_cfg.max_running_requests": 4, + "sglang_cfg.max_queued_requests": 8, + "sglang_cfg.max_loaded_loras": 256, + } + + +config = Config() diff --git a/src/lilo/configs/qwen38_27b_lora_64k.py b/src/lilo/configs/qwen38_27b_lora_64k.py new file mode 100644 index 0000000..d864a3e --- /dev/null +++ b/src/lilo/configs/qwen38_27b_lora_64k.py @@ -0,0 +1,15 @@ +from lilo.configs.qwen38_27b_lora_16k import Config as Parent + + +class Config(Parent): + name = "qwen38-27b-lora-64k" + max_context_length = 65536 + overrides = { + "miles_cfg.max_tokens_per_gpu": 32768, + "miles_cfg.context_parallel_size": 2, + "inference_target_concurrency": 8, + "sglang_cfg.max_running_requests": 16, + } + + +config = Config() diff --git a/src/lilo/configuration.py b/src/lilo/configuration.py new file mode 100644 index 0000000..32fb097 --- /dev/null +++ b/src/lilo/configuration.py @@ -0,0 +1,70 @@ +"""Flat Python recipes. Modal and the backend validate their own options.""" + +from copy import deepcopy + + +def _apply_overrides(values, overrides): + for path, value in overrides.items(): + parts = path.split(".") + target = values + for part in parts[:-1]: + target = target.setdefault(part, {}) + target[parts[-1]] = deepcopy(value) + + +class BaseConfig: + name = "" + model = "" + max_context_length = 16384 + parameterization = "lora" + backend = "miles" + + trainer_gpu = "H100" + trainer_gpus_per_node = 1 + trainer_nodes = 1 + trainer_cpu = 8 + trainer_memory_mib = 32768 + trainer_max_instances = 1 + trainer_max_clients_per_instance = 1 + trainer_timeout_s = 86400 + trainer_env = {} + sampler_persistence_concurrency = 8 + + inference_gpu = "H100" + inference_gpus_per_node = 1 + inference_cpu = 8 + inference_memory_mib = 32768 + inference_min_replicas = 0 + inference_max_replicas = 8 + inference_target_concurrency = 16 + inference_scaledown_window_s = 300 + inference_startup_timeout_s = 1200 + inference_env = {} + + miles_cfg = {} + megatron_cfg = {} + sglang_cfg = {} + session_idle_timeout_s = 300 + pool_idle_timeout_s = 300 + sweep_interval_s = 300 + + def __init__(self, **kwargs): + # Copy each ancestor's values before applying its overrides. + values = {} + for cls in reversed(type(self).__mro__): + values.update( + deepcopy( + { + key: value + for key, value in vars(cls).items() + if not key.startswith("_") + and key != "overrides" + and not callable(value) + } + ) + ) + _apply_overrides(values, vars(cls).get("overrides", {})) + overrides = kwargs.pop("overrides", {}) + values.update(deepcopy(kwargs)) + _apply_overrides(values, overrides) + self.__dict__.update(values) diff --git a/src/lilo/control_plane/deployments.py b/src/lilo/control_plane/deployments.py new file mode 100644 index 0000000..cd3bb33 --- /dev/null +++ b/src/lilo/control_plane/deployments.py @@ -0,0 +1,35 @@ +"""Select recipes in deployment order, or explicitly by definition ID.""" + + +class DeploymentRoutes: + def __init__(self, definitions): + self.definitions = tuple(definitions) + + def select(self, model, mode=None): + candidates = [ + d for d in self.definitions if mode is None or d.parameterization == mode + ] + for definition in candidates: + if definition.definition_id == model: + return definition + for definition in candidates: + if definition.model == model: + return definition + return None + + def capabilities(self): + selected = {} + for definition in self.definitions: + selected.setdefault( + (definition.model, definition.parameterization), definition + ) + contexts = {} + for (model, _), definition in selected.items(): + contexts[model] = min( + contexts.get(model, definition.max_context_length), + definition.max_context_length, + ) + return [ + {"model_name": model, "max_context_length": context} + for model, context in contexts.items() + ] diff --git a/src/lilo/control_plane/http.py b/src/lilo/control_plane/http.py index 9fbba5b..7d70d65 100644 --- a/src/lilo/control_plane/http.py +++ b/src/lilo/control_plane/http.py @@ -26,7 +26,9 @@ from lilo.providers.contracts import Parameterization from lilo.request_timing import enabled as request_timing_enabled from lilo.request_timing import mark +from lilo.telemetry.trainer import CommandMiddleware +from .deployments import DeploymentRoutes from .service import ControlPlane, FutureResolutionStatus ERROR_STATUSES: tuple[tuple[type[Exception], int, str], ...] = ( @@ -160,39 +162,17 @@ def create_control_plane_app( retrieve_window: float = 30.0, checkpoint_volume: str = "lilo-checkpoints", ) -> FastAPI: - all_definitions = tuple(definitions) - definitions = tuple( - definition for definition in all_definitions if definition.CATALOG_VISIBLE - ) + definitions = tuple(definitions) - def definition_for( - model_name: str, - parameterization: Parameterization, - ) -> str | None: - explicit = [ - d - for d in all_definitions - if d.DEFINITION_ID == model_name and d.PARAMETERIZATION == parameterization - ] - if explicit: - return explicit[0].DEFINITION_ID - matches = [ - definition.DEFINITION_ID - for definition in definitions - if definition.MODEL_NAME == model_name - and definition.PARAMETERIZATION == parameterization - ] - if len(matches) > 1: - raise ValueError( - f"duplicate {parameterization} definition for {model_name}" - ) - return matches[0] if matches else None + routes = DeploymentRoutes(definitions) + + def definition_for(model_name, parameterization): + selected = routes.select(model_name, parameterization) + return selected.definition_id if selected else None def supports_model(model_name: str) -> bool: - return any( - definition.MODEL_NAME == model_name for definition in definitions - ) or any( - definition.DEFINITION_ID == model_name for definition in all_definitions + return any(definition.model == model_name for definition in definitions) or any( + definition.definition_id == model_name for definition in definitions ) async def authorize(request: Request) -> None: @@ -200,7 +180,6 @@ async def authorize(request: Request) -> None: raise HTTPException(status_code=401, detail="unauthorized") app = FastAPI(dependencies=[Depends(authorize)]) - from lilo.telemetry.trainer import CommandMiddleware app.add_middleware(CommandMiddleware) @@ -283,20 +262,21 @@ async def healthz() -> dict[str, str]: @app.get("/api/v1/get_server_capabilities") async def get_server_capabilities() -> dict[str, object]: + return {"supported_models": routes.capabilities()} + + @app.get("/api/v1/lilo/deployments") + async def list_deployments(): return { - "supported_models": [ + "deployments": [ { - "model_name": name, - "max_context_length": min( - definition.MAX_CONTEXT_LENGTH - for definition in definitions - if definition.MODEL_NAME == name - ), + "name": d.name, + "definition_id": d.definition_id, + "base_model": d.model, + "parameterization": d.parameterization, + "max_context_length": d.max_context_length, } - for name in dict.fromkeys( - definition.MODEL_NAME for definition in definitions - ) - ], + for d in definitions + ] } @app.post("/api/v1/client/config") @@ -371,9 +351,7 @@ async def create_model(body: CreateModelBody) -> dict[str, object]: definition_id=definition_id, spec={ "base_model": next( - d.MODEL_NAME - for d in all_definitions - if d.DEFINITION_ID == definition_id + d.model for d in definitions if d.definition_id == definition_id ), "lora_config": body.lora_config, "parameterization": {"type": parameterization}, @@ -395,12 +373,8 @@ async def create_sampling_session( status_code=400, detail="base_model or model_path is required", ) - definition_id = ( - definition_for(body.base_model, "full") - or definition_for(body.base_model, "lora") - if body.base_model - else None - ) + selected = routes.select(body.base_model) if body.base_model else None + definition_id = selected.definition_id if selected else None if body.model_path is None and definition_id is None: raise HTTPException( status_code=400, @@ -410,11 +384,7 @@ async def create_sampling_session( session_id=body.session_id, sampling_session_seq_id=body.sampling_session_seq_id, base_model=( - next( - d.MODEL_NAME - for d in all_definitions - if d.DEFINITION_ID == definition_id - ) + next(d.model for d in definitions if d.definition_id == definition_id) if definition_id else body.base_model ), @@ -510,7 +480,7 @@ async def load_weights(body: LoadWeightsBody) -> dict[str, str]: base_model=body.base_model, user_metadata=body.user_metadata, optimizer=body.optimizer, - definition_ids={d.DEFINITION_ID for d in all_definitions}, + definition_ids={d.definition_id for d in definitions}, ) return { "request_id": creation.request_id, diff --git a/src/lilo/control_plane/service.py b/src/lilo/control_plane/service.py index d3b895f..f6b4f91 100644 --- a/src/lilo/control_plane/service.py +++ b/src/lilo/control_plane/service.py @@ -29,13 +29,13 @@ SamplingTaskStatus, SessionKeyValueStores, ) +from lilo.telemetry.metadata import experiment_tags from .keys import ( model_creation_key, model_key, placement_claim_key, placement_key, - trainer_demand_key, sample_task_key, sampler_artifact_key, sampler_export_result_key, @@ -45,6 +45,7 @@ session_closed_key, session_key, session_last_seen_key, + trainer_demand_key, ) from .records import ( ModelCreationRecord, @@ -137,6 +138,7 @@ def __init__( list_checkpoints: CheckpointListing | None = None, delete_checkpoint: Callable[[str], Awaitable[None]] | None = None, checkpoint_root: str = "/checkpoints", + creation_error: Callable[[str], Awaitable[str | None]] | None = None, ) -> None: self.kv = kv self.engines = engines @@ -153,6 +155,7 @@ def __init__( self.list_checkpoints = list_checkpoints self.delete_checkpoint = delete_checkpoint self.checkpoint_root = checkpoint_root + self.creation_error = creation_error async def create_session( self, @@ -992,7 +995,6 @@ async def _finalize_sampler_export( model.model_id, publish_version, ) - from lilo.telemetry.metadata import experiment_tags telemetry_tags = experiment_tags((model.spec or {}).get("user_metadata")) latest = SamplerArtifactRecord( @@ -1140,6 +1142,10 @@ async def _retrieve_creation( ) -> FutureResolution: try: placement = await self._place(model) + if placement is None and self.creation_error is not None: + error = await self.creation_error(model.engine_definition_id) + if error: + raise ValueError(error) except ModelLost: return FutureResolution( request_id, @@ -1162,7 +1168,7 @@ async def _retrieve_creation( if placement is None: return FutureResolution(request_id, FutureResolutionStatus.PENDING) try: - instance = await self._live_instance(placement) + await self._live_instance(placement) except ModelLost: return FutureResolution( request_id, diff --git a/src/lilo/deployment_cli.py b/src/lilo/deployment_cli.py new file mode 100644 index 0000000..c8e38f2 --- /dev/null +++ b/src/lilo/deployment_cli.py @@ -0,0 +1,203 @@ +"""Operator CLI for Python deployments. Only `deploy` provisions Modal resources.""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys + +import modal + +from lilo.deployments import ( + PLATFORM_DEFAULTS, + DeploymentConfig, + config_path, + load, + platform_defaults, + validate_frontend, +) +from lilo.providers.modal.deployment_apps import ( + build_inference_app, + build_rollout_app, + build_trainer_app, +) +from lilo.providers.modal.deployment_configs import ( + CONFIGS_ENV, + PLATFORM_ENV, + deployed_configs, + deployed_platform, + deployed_pools, +) +from lilo.providers.modal.fft_pool import FFTPoolSpec +from lilo.providers.modal.lora_pool import LoraPoolSpec + + +def compile_configs(paths): + recipes = [load(path) for path in paths] + validate_frontend(recipes) + return [DeploymentConfig.create(recipe) for recipe in recipes] + + +def app_is_missing(app_name, environment): + try: + modal.App.lookup(app_name, environment_name=environment) + except modal.exception.NotFoundError: + return True + return False + + +def apps_to_publish(configs, current_configs, environment, trainers, inference): + """Add names whose apps are missing or whose trainer/inference settings changed.""" + current_by_name = {deployment.name: deployment for deployment in current_configs} + for deployment in configs: + current = current_by_name.get(deployment.name) + if app_is_missing(deployment.trainer_app_name, environment) or ( + current is not None and not deployment.same_trainer(current) + ): + trainers.add(deployment.name) + if app_is_missing(deployment.inference_app_name, environment) or ( + current is not None and not deployment.same_inference(current) + ): + inference.add(deployment.name) + return trainers, inference + + +def deploy(configs, platform=None, *, refresh_trainers=(), refresh_inference=()): + """Deploy missing or changed trainer and inference apps, then the frontend.""" + if sys.version_info[:2] != (3, 12): + raise ValueError( + "Python deployment requires Python 3.12 to match the serialized GPU runtime images" + ) + trainers = set(refresh_trainers) + inference = set(refresh_inference) + names = {deployment.name for deployment in configs} + unknown = (trainers | inference) - names + if unknown: + raise ValueError(f"unknown configs to refresh: {sorted(unknown)}") + if platform is None: + platform = platform_defaults() + environment = platform["modal"]["environment"] + current_configs = [ + DeploymentConfig.model_validate(row) + for row in deployed_configs(platform["frontend"], environment) + ] + current_platform = deployed_platform(platform["frontend"], environment) + if current_platform is not None and current_platform != platform: + trainers.update(names) + inference.update(names) + trainers, inference = apps_to_publish( + configs, current_configs, environment, trainers, inference + ) + pools = deployed_pools(platform["frontend"], environment) if inference else [] + + with modal.enable_output(): + for deployment in configs: + if deployment.name in trainers: + build_trainer_app(deployment, platform)[0].deploy( + environment_name=environment + ) + if deployment.name in inference: + build_inference_app(deployment, platform)[0].deploy( + environment_name=environment + ) + for record in pools: + if record["definition_id"] != deployment.definition_id: + continue + pool = ( + LoraPoolSpec.from_dict(record) + if deployment.parameterization == "lora" + else FFTPoolSpec.from_dict(record) + ) + if not app_is_missing(pool.app_name, environment): + build_rollout_app(deployment, pool, platform)[0].deploy( + environment_name=environment + ) + command = [sys.executable, "-m", "modal", "deploy"] + if environment: + command += ["--env", environment] + subprocess.run( + [*command, "-m", "lilo.providers.modal.app"], + check=True, + env={ + **os.environ, + CONFIGS_ENV: json.dumps( + [deployment.model_dump(mode="json") for deployment in configs] + ), + PLATFORM_ENV: json.dumps(platform), + "LILO_APP_NAME": platform["frontend"], + }, + ) + + +def parser(): + result = argparse.ArgumentParser(prog="lilo") + commands = result.add_subparsers(dest="command", required=True) + config = commands.add_parser("config").add_subparsers(dest="action", required=True) + init = config.add_parser("init") + init.add_argument("--preset", required=True) + validate = config.add_parser("validate") + validate.add_argument("files", nargs="+") + apply = commands.add_parser( + "deploy", + help="Deploy the complete active Python config set behind one frontend", + ) + apply.add_argument("files", nargs="+") + apply.add_argument("--app", default=PLATFORM_DEFAULTS["frontend"]) + apply.add_argument("--env") + apply.add_argument("--region", default=PLATFORM_DEFAULTS["modal"]["region"]) + apply.add_argument( + "--refresh-trainer", action="append", default=[], metavar="CONFIG_NAME" + ) + apply.add_argument( + "--refresh-inference", action="append", default=[], metavar="CONFIG_NAME" + ) + management = commands.add_parser("deployment").add_subparsers( + dest="action", required=True + ) + retry = management.add_parser("retry") + retry.add_argument("--frontend", required=True) + retry.add_argument("--env") + retry.add_argument("definition_id") + return result + + +def main(argv=None): + cli = parser() + args = cli.parse_args(argv) + try: + if args.command == "config": + if args.action == "init": + module = config_path(args.preset).stem + if not config_path(args.preset).is_file(): + raise ValueError(f"unknown example config: {args.preset}") + print( + f'from lilo.configs.{module} import Config as Parent\n\n\nclass Config(Parent):\n name = "my-model"\n\n\nconfig = Config()' + ) + else: + recipes = [load(path) for path in args.files] + validate_frontend(recipes) + print( + f"Validated {len(recipes)} deployment(s). Backend integration settings are checked when preparing trainers and pools; native options are checked at engine startup." + ) + elif args.command == "deploy": + platform = platform_defaults() + platform["frontend"] = args.app + platform["modal"].update(environment=args.env, region=args.region) + deploy( + compile_configs(args.files), + platform, + refresh_trainers=args.refresh_trainer, + refresh_inference=args.refresh_inference, + ) + else: + modal.Function.from_name( + args.frontend, "clear_deployment_failure", environment_name=args.env + ).remote(args.definition_id) + except (ValueError, OSError, subprocess.CalledProcessError) as exc: + cli.exit(1, f"{exc}\n") + + +if __name__ == "__main__": + main() diff --git a/src/lilo/deployments.py b/src/lilo/deployments.py new file mode 100644 index 0000000..e28f102 --- /dev/null +++ b/src/lilo/deployments.py @@ -0,0 +1,233 @@ +"""Python deployment configs: the recipe plus backend settings used to deploy it.""" + +from __future__ import annotations + +import json +import re +import runpy +import sys +from copy import deepcopy +from importlib.resources import files +from pathlib import Path + +from pydantic import BaseModel, ConfigDict, field_serializer, field_validator + +from lilo.backends.deployment import resolve_backend_settings +from lilo.configuration import BaseConfig + +PLATFORM_DEFAULTS = { + "frontend": "lilo", + "modal": {"environment": None, "region": "us-west"}, + "secrets": { + "api": "lilo-api", + "sampler_proxy": "lilo-proxy", + "huggingface": "huggingface-secret", + }, + "storage": { + "assets": "lilo-model-assets", + "checkpoints": "lilo-checkpoints", + "bulletin": "lilo-snapshot-bulletin", + }, +} + + +def platform_defaults(): + return deepcopy(PLATFORM_DEFAULTS) + + +def _jsonable(value): + return json.loads(json.dumps(value)) + + +class DeploymentConfig(BaseModel): + """Recipe plus the backend settings used to deploy one model.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + recipe: BaseConfig + trainer_settings: dict + inference_settings: dict + + @field_validator("recipe", mode="before") + @classmethod + def restore_recipe(cls, value): + return BaseConfig(**value) if isinstance(value, dict) else value + + @field_serializer("recipe") + def serialize_recipe(self, recipe): + return vars(recipe) + + @classmethod + def create(cls, recipe: BaseConfig) -> DeploymentConfig: + """Copy the recipe and attach backend settings.""" + pinned = BaseConfig(**_jsonable(vars(recipe))) + trainer_settings, inference_settings = resolve_backend_settings( + pinned, f"/assets/{pinned.model}" + ) + return cls( + recipe=pinned, + trainer_settings=_jsonable(trainer_settings), + inference_settings=_jsonable(inference_settings), + ) + + @property + def name(self) -> str: + return self.recipe.name + + @property + def model(self) -> str: + return self.recipe.model + + @property + def parameterization(self) -> str: + return self.recipe.parameterization + + @property + def max_context_length(self) -> int: + return self.recipe.max_context_length + + @property + def definition_id(self) -> str: + return self.recipe.name + + @property + def trainer_app_name(self) -> str: + return f"lilo-trainer-{self.recipe.name}" + + @property + def inference_app_name(self) -> str: + return f"lilo-inference-{self.recipe.name}" + + @property + def asset_path(self) -> str: + return f"/assets/{self.recipe.model}" + + @property + def rollout_tensor_parallel_size(self) -> int: + serving = self.inference_settings + return serving.get("tp_size", self.recipe.inference_gpus_per_node) // ( + serving.get("dp_size", 1) if serving.get("enable_dp_attention") else 1 + ) + + def same_trainer(self, other) -> bool: + recipe = self.recipe + current = other.recipe + return ( + recipe.name, + recipe.model, + recipe.parameterization, + recipe.max_context_length, + recipe.backend, + recipe.miles_cfg, + recipe.megatron_cfg, + recipe.sampler_persistence_concurrency, + recipe.trainer_gpu, + recipe.trainer_gpus_per_node, + recipe.trainer_nodes, + recipe.trainer_cpu, + recipe.trainer_memory_mib, + recipe.trainer_max_instances, + recipe.trainer_max_clients_per_instance, + recipe.trainer_timeout_s, + recipe.trainer_env, + self.trainer_settings, + ) == ( + current.name, + current.model, + current.parameterization, + current.max_context_length, + current.backend, + current.miles_cfg, + current.megatron_cfg, + current.sampler_persistence_concurrency, + current.trainer_gpu, + current.trainer_gpus_per_node, + current.trainer_nodes, + current.trainer_cpu, + current.trainer_memory_mib, + current.trainer_max_instances, + current.trainer_max_clients_per_instance, + current.trainer_timeout_s, + current.trainer_env, + other.trainer_settings, + ) + + def same_inference(self, other) -> bool: + recipe = self.recipe + current = other.recipe + return ( + recipe.name, + recipe.model, + recipe.parameterization, + recipe.max_context_length, + recipe.sglang_cfg, + recipe.inference_gpu, + recipe.inference_gpus_per_node, + recipe.inference_cpu, + recipe.inference_memory_mib, + recipe.inference_min_replicas, + recipe.inference_max_replicas, + recipe.inference_target_concurrency, + recipe.inference_scaledown_window_s, + recipe.inference_startup_timeout_s, + recipe.inference_env, + self.inference_settings, + ) == ( + current.name, + current.model, + current.parameterization, + current.max_context_length, + current.sglang_cfg, + current.inference_gpu, + current.inference_gpus_per_node, + current.inference_cpu, + current.inference_memory_mib, + current.inference_min_replicas, + current.inference_max_replicas, + current.inference_target_concurrency, + current.inference_scaledown_window_s, + current.inference_startup_timeout_s, + current.inference_env, + other.inference_settings, + ) + + +def config_path(name: str) -> Path: + """Locate an installed example config without maintaining a model catalog.""" + if not re.fullmatch(r"[a-zA-Z0-9_-]+", name): + raise ValueError("invalid config name") + return Path(str(files("lilo").joinpath("configs", name.replace("-", "_") + ".py"))) + + +def load(path: str | Path) -> BaseConfig: + """Execute a Python config file and read its exported config object.""" + path = Path(path).resolve() + if path.suffix != ".py": + raise ValueError("deployment configs must be Python .py files") + original_path = sys.path[:] + sys.path.insert(0, str(path.parent)) + try: + namespace = runpy.run_path(str(path)) + config = namespace.get("config") + if not isinstance(config, BaseConfig): + raise ValueError(f"{path} must export a BaseConfig instance named config") + return config + finally: + sys.path[:] = original_path + + +def validate_frontend(recipes: list[BaseConfig]) -> None: + if not recipes: + raise ValueError("at least one deployment is required") + if len({recipe.name for recipe in recipes}) != len(recipes): + raise ValueError("duplicate deployment name") + first = recipes[0] + for recipe in recipes: + if ( + recipe.session_idle_timeout_s != first.session_idle_timeout_s + or recipe.pool_idle_timeout_s != first.pool_idle_timeout_s + or recipe.sweep_interval_s != first.sweep_interval_s + ): + raise ValueError( + "deployments on one frontend must share lifecycle settings" + ) diff --git a/src/lilo/inference/sglang.py b/src/lilo/inference/sglang.py new file mode 100644 index 0000000..04a5416 --- /dev/null +++ b/src/lilo/inference/sglang.py @@ -0,0 +1,30 @@ +"""SGLang worker entrypoint. The frontend never imports this GPU-only module.""" + +import json +import logging +import os +import sys + +from sglang.launch_server import run_server +from sglang.srt.plugins import load_plugins +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree + + +def main(): + load_plugins() + args = ServerArgs( + model_path=sys.argv[1], + host="127.0.0.1", + port=8001, + **json.loads(sys.argv[2]), + ) + logging.basicConfig(level=args.log_level.upper()) + try: + run_server(args) + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +if __name__ == "__main__": + main() diff --git a/src/lilo/providers/modal/app.py b/src/lilo/providers/modal/app.py index 2ef67bf..381bc84 100644 --- a/src/lilo/providers/modal/app.py +++ b/src/lilo/providers/modal/app.py @@ -7,38 +7,30 @@ from dataclasses import asdict import modal +from huggingface_hub import snapshot_download from stitch.pools.modal_flash import ModalFlashPool +from lilo.control_plane import ControlPlane, create_control_plane_app +from lilo.control_plane.keys import model_key, placement_key, trainer_demand_key +from lilo.control_plane.records import ModelRecord +from lilo.deployments import validate_frontend +from lilo.inference.sampling import sample_task from lilo.providers.contracts import ( Parameterization, SamplingTask, ) +from lilo.telemetry.otlp import sample_trace from .checkpoint_storage import ( CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, ModalCheckpointStorage, - checkpoint_volume, ) -from .definitions import ( - qwen3_5_4b_full_64k, - qwen3_5_9b_base_miles_lora_2k, - qwen3_5_9b_base_miles_lora_16k, - qwen3_5_9b_base_miles_lora_16k_single, - qwen3_5_9b_full_64k, - qwen3_5_9b_miles_lora_16k, - qwen3_5_9b_miles_lora_16k_dp2, - qwen3_5_35b_a3b_full_64k, - qwen3_6_27b_full_64k, - qwen3_6_35b_a3b_full_64k, - qwen3_8_27b_miles_lora_16k, - qwen3_8_27b_miles_lora_64k, - qwen3_8_27b_miles_lora_128k, - qwen3_8_27b_miles_lora_256k, -) -from .deployment import ( - trainer_deployment_env, - trainer_max_containers, +from .deployment import trainer_deployment_env +from .deployment_configs import ( + CONFIGS_ENV, + PLATFORM_ENV, + configs_from_env, + platform_from_env, ) from .engines import ModalEnginePlatform from .fft_pool import ( @@ -48,9 +40,15 @@ proxy_auth_headers, stop_pool, ) -from .image_dependencies import STITCH_PACKAGE +from .image_dependencies import ( + CORE_PACKAGES, + STITCH_PACKAGE, + TINKER_PACKAGE, + ignore_config_source, +) from .kv import ( ModalSessionKeyValueStores, + current_app_id, fft_pool_kv, shared_kv, ) @@ -67,42 +65,41 @@ stop_pool as stop_lora_pool, ) from .sampling import ModalSamplingTaskPlatform +from .trainer_reconciler import ( + complete_reconcile, + pending_reconciliations, + reconcile_trainers, + release_reconcile_call, + request_reconcile, +) -APP_NAME = os.environ.get("LILO_APP_NAME", "lilo") -ROUTING_REGION = "us-west" +DEFINITIONS = tuple(configs_from_env()) +validate_frontend([definition.recipe for definition in DEFINITIONS]) +SETTINGS = DEFINITIONS[0] +PLATFORM = platform_from_env() +APP_NAME = PLATFORM["frontend"] +ROUTING_REGION = PLATFORM["modal"]["region"] MODEL_ASSET_ROOT = "/assets" -SESSION_IDLE_TIMEOUT = 300.0 -FFT_POOL_IDLE_TIMEOUT = 300.0 -LORA_POOL_IDLE_TIMEOUT = 300.0 +SESSION_IDLE_TIMEOUT = SETTINGS.recipe.session_idle_timeout_s +FFT_POOL_IDLE_TIMEOUT = LORA_POOL_IDLE_TIMEOUT = SETTINGS.recipe.pool_idle_timeout_s FFT_POOL_TOUCH_INTERVAL = 60.0 LORA_POOL_CHECK_INTERVAL = 60.0 -SWEEP_PERIOD = modal.Period(minutes=5) +SWEEP_PERIOD = modal.Period(seconds=SETTINGS.recipe.sweep_interval_s) CHECKPOINT_READ_LOCK = asyncio.Lock() _pool_touches: dict[str, float] = {} _lora_pool_gateways: dict[str, tuple[float, str]] = {} _lora_pool_checks: dict[str, asyncio.Lock] = {} - -DEFINITIONS = ( - qwen3_5_4b_full_64k, - qwen3_5_9b_full_64k, - qwen3_5_9b_base_miles_lora_2k, - qwen3_5_9b_base_miles_lora_16k, - qwen3_5_9b_base_miles_lora_16k_single, - qwen3_5_9b_miles_lora_16k, - qwen3_5_9b_miles_lora_16k_dp2, - qwen3_5_35b_a3b_full_64k, - qwen3_6_27b_full_64k, - qwen3_6_35b_a3b_full_64k, - qwen3_8_27b_miles_lora_16k, - qwen3_8_27b_miles_lora_64k, - qwen3_8_27b_miles_lora_128k, - qwen3_8_27b_miles_lora_256k, +CHECKPOINT_VOLUME_NAME = PLATFORM["storage"]["checkpoints"] +checkpoint_volume = modal.Volume.from_name( + CHECKPOINT_VOLUME_NAME, create_if_missing=True, version=2 ) -TRAINER_MAX_CONTAINERS = trainer_max_containers() -TRAINER_DEPLOYMENT_ENV = trainer_deployment_env() +TRAINER_DEPLOYMENT_ENV = { + **trainer_deployment_env(), + CONFIGS_ENV: os.environ[CONFIGS_ENV], + PLATFORM_ENV: os.environ.get(PLATFORM_ENV, ""), + "LILO_APP_NAME": APP_NAME, +} app = modal.App(APP_NAME) -for definition in DEFINITIONS: - app.include(definition.app) async def _read_checkpoint_metadata(uri: str) -> dict[str, object]: @@ -124,15 +121,45 @@ async def _delete_checkpoint(uri: str) -> None: image = ( - modal.Image.debian_slim(python_version="3.11") + modal.Image.debian_slim(python_version="3.12") .apt_install("git") - .pip_install_from_pyproject("pyproject.toml") + .pip_install(*CORE_PACKAGES, TINKER_PACKAGE) .pip_install(STITCH_PACKAGE, "huggingface-hub") - .add_local_python_source("lilo") + .env(TRAINER_DEPLOYMENT_ENV) + .add_local_python_source("lilo", ignore=ignore_config_source) +) + + +@app.function(image=image) +def deployed_configs(): + return [row.model_dump(mode="json") for row in configs_from_env()] + + +@app.function(image=image) +def deployed_platform(): + return PLATFORM + + +@app.function(image=image) +async def deployed_pools(): + return [ + value + for registry, prefix in ( + (shared_kv(), "lora_pool:"), + (fft_pool_kv(), "fft_pool:"), + ) + for _, value in await registry.list_items(prefix) + ] + + +model_assets = modal.Volume.from_name( + PLATFORM["storage"]["assets"], + create_if_missing=True, ) -model_assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) +API_SECRET_NAME = PLATFORM["secrets"]["api"] +HF_SECRET_NAME = PLATFORM["secrets"]["huggingface"] proxy_secret = modal.Secret.from_name( - "lilo-proxy", + PLATFORM["secrets"]["sampler_proxy"], required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], ) @@ -140,22 +167,23 @@ async def _delete_checkpoint(uri: str) -> None: @app.function( image=image, volumes={MODEL_ASSET_ROOT: model_assets}, - secrets=[modal.Secret.from_name("huggingface-secret")], + secrets=[modal.Secret.from_name(HF_SECRET_NAME)] if HF_SECRET_NAME else [], timeout=4 * 60 * 60, max_containers=1, retries=2, ) def prepare_model_assets(definition_id: str) -> None: - from huggingface_hub import snapshot_download - definition = module_for(definition_id) - checkpoint = os.path.normpath(definition.HF_CHECKPOINT) + checkpoint = os.path.normpath(definition.asset_path) if ( os.path.commonpath((MODEL_ASSET_ROOT, checkpoint)) != MODEL_ASSET_ROOT or checkpoint == MODEL_ASSET_ROOT ): raise ValueError(f"invalid model asset path: {checkpoint}") - snapshot_download(repo_id=definition.MODEL_NAME, local_dir=checkpoint) + snapshot_download( + repo_id=definition.model, + local_dir=checkpoint, + ) model_assets.commit() @@ -237,27 +265,23 @@ async def _ready_lora_pool(spec: LoraPoolSpec) -> str: min_containers=0, timeout=60 * 60, retries=2, - secrets=[proxy_secret, modal.Secret.from_name("lilo-api")], + secrets=[proxy_secret, modal.Secret.from_name(API_SECRET_NAME)], ) @modal.concurrent(max_inputs=128) async def execute_sample(task: dict) -> dict: - from lilo.telemetry.otlp import sample_trace - stats: dict = {} with sample_trace(task, stats): return await _execute_sample(task, stats) async def _execute_sample(task: dict, stats: dict) -> dict: - from lilo.inference.sampling import sample_task - definition_id = str(task["engine_definition_id"]) parameterization = parameterization_for(definition_id) if parameterization not in {"full", "lora"}: raise ValueError(f"unsupported sampling definition: {definition_id}") definition = module_for(definition_id) - rollout_world_size = definition.ROLLOUT_GPUS - rollout_tensor_parallel_size = definition.ROLLOUT_TENSOR_PARALLEL_SIZE + rollout_world_size = definition.recipe.inference_gpus_per_node + rollout_tensor_parallel_size = definition.rollout_tensor_parallel_size if rollout_world_size % rollout_tensor_parallel_size: raise ValueError("rollout GPU count must be divisible by tensor parallel size") rollout_data_parallel_size = rollout_world_size // rollout_tensor_parallel_size @@ -275,7 +299,7 @@ async def keep_pool_ready() -> None: gateway, data_parallel_size=rollout_data_parallel_size, headers=proxy_auth_headers(), - context_length=definition.MAX_CONTEXT_LENGTH, + context_length=definition.max_context_length, on_wait=keep_pool_ready, stats=stats, ) @@ -296,7 +320,7 @@ async def keep_pool_ready() -> None: data_parallel_size=rollout_data_parallel_size, headers=proxy_auth_headers(), on_wait=lambda: _touch_fft_pool(spec), - context_length=definition.MAX_CONTEXT_LENGTH, + context_length=definition.max_context_length, stats=stats, ) @@ -312,22 +336,19 @@ def _latest_pool(model) -> FFTPoolSpec: async def _model_record(kv, model_id: str): - from lilo.control_plane.keys import model_key - from lilo.control_plane.records import ModelRecord - return ModelRecord.model_validate(await kv.get(model_key(model_id))) def module_for(definition_id: str): for definition in DEFINITIONS: - if definition.DEFINITION_ID == definition_id: + if definition.definition_id == definition_id: return definition raise KeyError(definition_id) def parameterization_for(definition_id: str) -> Parameterization | None: try: - return module_for(definition_id).PARAMETERIZATION + return module_for(definition_id).parameterization except KeyError: return None @@ -344,34 +365,25 @@ def trainer_autoscaling(definition_id: str) -> bool: retries=3, ) async def trainer_reconciler(delay_seconds: float = 0.0) -> None: - from .trainer_reconciler import ( - complete_reconcile, - pending_reconciliations, - reconcile_trainers, - release_reconcile_call, - ) - if delay_seconds > 0: await asyncio.sleep(delay_seconds) call_id = modal.current_function_call_id() async def run(definition_id: str, token: str) -> None: parameterization = parameterization_for(definition_id) - if parameterization is None: + if parameterization is None or await deployment_error(definition_id): await complete_reconcile(definition_id, token) return module = module_for(definition_id) - maximum_instances = TRAINER_MAX_CONTAINERS + maximum_instances = module.recipe.trainer_max_instances try: await reconcile_trainers( shared_kv(), ModalEnginePlatform(shared_kv(), _spawn_engine), definition_id, revision=None, - maximum_instances=( - int(maximum_instances) if maximum_instances is not None else None - ), - models_per_instance=module.TRAINER_MODELS_PER_INSTANCE, + maximum_instances=maximum_instances, + models_per_instance=module.recipe.trainer_max_clients_per_instance, scale_up=trainer_autoscaling(definition_id), ) except Exception: @@ -397,7 +409,6 @@ async def run(definition_id: str, token: str) -> None: async def kick_trainer_reconciler(definition_id: str) -> None: if parameterization_for(definition_id) is None: return - from .trainer_reconciler import request_reconcile async def spawn(delay_seconds: float) -> str: call = await trainer_reconciler.spawn.aio(delay_seconds) @@ -407,14 +418,40 @@ async def spawn(delay_seconds: float) -> str: async def _spawn_engine(definition_id: str, instance_id: str) -> str: - engine = module_for(definition_id).ENGINE_FUNCTION - call = await engine.spawn.aio(instance_id) + if error := await deployment_error(definition_id): + raise ValueError(error) + definition = module_for(definition_id) + engine = modal.Function.from_name( + definition.trainer_app_name, + "trainer", + environment_name=PLATFORM["modal"]["environment"], + ) + call = await engine.spawn.aio( + instance_id, definition.model_dump_json(), current_app_id() + ) return call.object_id -def _plane(): - from lilo.control_plane import ControlPlane +async def deployment_error(definition_id: str) -> str | None: + registry = shared_kv() + key = f"deployment_failure:{definition_id}" + record = await registry.get(key) + if record: + if time.time() - record.get("failed_at", 0) >= 300: + await registry.delete(key) + return None + return f"Trainer startup failed for {definition_id}: {record['error']}. See Modal call logs for instance {record['instance_id']}; automatic retries resume after five minutes, or run lilo deployment retry." + return None + + +@app.function(image=image) +async def clear_deployment_failure(definition_id: str) -> None: + module_for(definition_id) + await shared_kv().delete(f"deployment_failure:{definition_id}") + await kick_trainer_reconciler(definition_id) + +def _plane(): kv = shared_kv() task_stores = ModalSessionKeyValueStores() engines = ModalEnginePlatform(kv, _spawn_engine) @@ -439,6 +476,8 @@ async def ensure_pool(session) -> None: definition_id = session.engine_definition_id parameterization = parameterization_for(definition_id) if parameterization == "lora": + if session.model_id is None: + await prepare_model_assets.remote.aio(definition_id) await _ready_lora_pool(LoraPoolSpec(definition_id)) return if parameterization != "full": @@ -470,9 +509,7 @@ async def kick_trainers(definition_id: str) -> bool: await kick_trainer_reconciler(definition_id) if not trainer_autoscaling(definition_id): return False - maximum = TRAINER_MAX_CONTAINERS - if maximum is None: - return True + maximum = module_for(definition_id).recipe.trainer_max_instances instances = [ instance for instance in await engines.list_instances() @@ -489,6 +526,7 @@ async def kick_trainers(definition_id: str) -> bool: session_idle_timeout=SESSION_IDLE_TIMEOUT, ensure_sampling_pool=ensure_pool, prepare_model=prepare_model, + creation_error=deployment_error, sampling_task_stores=task_stores, read_checkpoint_metadata=_read_checkpoint_metadata, list_checkpoints=_list_checkpoints, @@ -506,13 +544,11 @@ async def kick_trainers(definition_id: str) -> bool: routing_region=ROUTING_REGION, timeout=20 * 60, volumes={CHECKPOINT_ROOT: checkpoint_volume}, - secrets=[modal.Secret.from_name("lilo-api", required_keys=["TINKER_API_KEY"])], + secrets=[modal.Secret.from_name(API_SECRET_NAME, required_keys=["TINKER_API_KEY"])], ) @modal.concurrent(max_inputs=128) @modal.asgi_app(requires_proxy_auth=False) def server(): - from lilo.control_plane import create_control_plane_app - return create_control_plane_app( _plane(), DEFINITIONS, @@ -522,9 +558,6 @@ def server(): async def _lose_undefined_models() -> tuple[str, ...]: - from lilo.control_plane.keys import placement_key, trainer_demand_key - from lilo.control_plane.records import ModelRecord - kv = shared_kv() lost = [] for _, value in await kv.list_items("model:"): @@ -538,8 +571,6 @@ async def _lose_undefined_models() -> tuple[str, ...]: async def _cleanup_fft_pools() -> tuple[str, ...]: - from lilo.control_plane.records import ModelRecord - active_latest = { FFTPoolSpec( model.engine_definition_id, @@ -588,8 +619,6 @@ async def _cleanup_fft_pools() -> tuple[str, ...]: async def _cleanup_lora_pools() -> tuple[str, ...]: - from lilo.control_plane.records import ModelRecord - registry = shared_kv() active = { LoraPoolSpec(model.engine_definition_id).app_name @@ -618,8 +647,6 @@ async def _cleanup_lora_pools() -> tuple[str, ...]: @app.function(image=image, env=TRAINER_DEPLOYMENT_ENV, schedule=SWEEP_PERIOD) def cleaner(): - import asyncio - async def run() -> None: plane = _plane() await plane.sweep_idle_sessions(SESSION_IDLE_TIMEOUT) diff --git a/src/lilo/providers/modal/definitions/__init__.py b/src/lilo/providers/modal/definitions/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/lilo/providers/modal/definitions/qwen3_5_35b_a3b_full_64k.py b/src/lilo/providers/modal/definitions/qwen3_5_35b_a3b_full_64k.py deleted file mode 100644 index 2a1c1e4..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_35b_a3b_full_64k.py +++ /dev/null @@ -1,155 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-35B-A3B" -HF_CHECKPOINT = "/assets/Qwen3.5-35B-A3B" -MICRO_BATCH_SIZE = 1 -MAX_CONTEXT_LENGTH = 65_536 -MAX_TOKENS_PER_MICROBATCH = MAX_CONTEXT_LENGTH -DEFINITION_ID = "qwen3_5_35b_a3b_full_64k" -PARAMETERIZATION = "full" -CATALOG_VISIBLE = True -TRAINER_MODELS_PER_INSTANCE = 1 -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -PIPELINE_MODEL_PARALLEL_SIZE = 1 -CONTEXT_PARALLEL_SIZE = 2 -DATA_PARALLEL_SIZE = 1 -EXPERT_MODEL_PARALLEL_SIZE = 8 -EXPERT_TENSOR_PARALLEL_SIZE = 1 -SEQUENCE_PARALLEL = True -GPU_MEMORY_FRACTION = 0.90 -MTP_NUM_LAYERS = 0 -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 4 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 4 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.90 -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = 32 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 4 -ROLLOUT_TARGET_CONCURRENCY = 16 -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..megatron_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_5_35b_a3b_full_64k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "megatron": { - "hf_checkpoint": HF_CHECKPOINT, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "pipeline_model_parallel_size": PIPELINE_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "expert_model_parallel_size": EXPERT_MODEL_PARALLEL_SIZE, - "expert_tensor_parallel_size": EXPERT_TENSOR_PARALLEL_SIZE, - "sequence_parallel": SEQUENCE_PARALLEL, - "micro_batch_size": MICRO_BATCH_SIZE, - "max_tokens_per_microbatch": MAX_TOKENS_PER_MICROBATCH, - "seq_length": MAX_CONTEXT_LENGTH, - "bf16": True, - "fp16": False, - "gpu_memory_fraction": GPU_MEMORY_FRACTION, - "use_distributed_optimizer": True, - "provider_overrides": { - "mtp_num_layers": MTP_NUM_LAYERS, - "recompute_granularity": "selective", - "moe_layer_recompute": True, - "moe_token_dispatcher_type": "alltoall", - "moe_router_fusion": True, - "moe_permute_fusion": True, - "moe_grouped_gemm": True, - "moe_shared_expert_overlap": False, - "moe_aux_loss_coeff": 0.0, - }, - "optimizer": { - "optimizer": "adam", - "lr": 1e-4, - "min_lr": 1e-4, - }, - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.megatron_fft:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - }, - nproc=GPUS, - max_models=1, - ) - - -ENGINE_FUNCTION = qwen3_5_35b_a3b_full_64k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_4b_full_64k.py b/src/lilo/providers/modal/definitions/qwen3_5_4b_full_64k.py deleted file mode 100644 index 7396c1f..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_4b_full_64k.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-4B" -HF_CHECKPOINT = "/assets/Qwen3.5-4B" -MICRO_BATCH_SIZE = 1 -MAX_CONTEXT_LENGTH = 65_536 -MAX_TOKENS_PER_MICROBATCH = MAX_CONTEXT_LENGTH -DEFINITION_ID = "qwen3_5_4b_full_64k" -PARAMETERIZATION = "full" -CATALOG_VISIBLE = True -TRAINER_MODELS_PER_INSTANCE = 1 -GPU_TYPE = "H100" -GPUS = 4 -TENSOR_MODEL_PARALLEL_SIZE = 2 -CONTEXT_PARALLEL_SIZE = 2 -DATA_PARALLEL_SIZE = GPUS // (TENSOR_MODEL_PARALLEL_SIZE * CONTEXT_PARALLEL_SIZE) -MTP_NUM_LAYERS = 0 -RECOMPUTE_GRANULARITY = "full" -RECOMPUTE_METHOD = "uniform" -RECOMPUTE_NUM_LAYERS = 1 -LOSS_SCALE = 1.0 -DEFER_FP32_LOGITS = True -FP32_LM_HEAD = True -ROLLOUT_GPU_TYPE = "H100" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.85 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 4 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = 16 -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..megatron_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_5_4b_full_64k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "megatron": { - "hf_checkpoint": HF_CHECKPOINT, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "sequence_parallel": True, - "micro_batch_size": MICRO_BATCH_SIZE, - "max_tokens_per_microbatch": MAX_TOKENS_PER_MICROBATCH, - "seq_length": MAX_CONTEXT_LENGTH, - "defer_fp32_logits": DEFER_FP32_LOGITS, - "fp32_lm_head": FP32_LM_HEAD, - "use_distributed_optimizer": True, - "provider_overrides": { - "mtp_num_layers": MTP_NUM_LAYERS, - "recompute_granularity": RECOMPUTE_GRANULARITY, - "recompute_method": RECOMPUTE_METHOD, - "recompute_num_layers": RECOMPUTE_NUM_LAYERS, - }, - "optimizer": { - "lr": 1e-4, - "min_lr": 1e-4, - "loss_scale": LOSS_SCALE, - }, - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.megatron_fft:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=GPUS, - max_models=1, - ) - - -ENGINE_FUNCTION = qwen3_5_4b_full_64k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k.py deleted file mode 100644 index b64c080..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k.py +++ /dev/null @@ -1,172 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-9B-Base" -HF_CHECKPOINT = "/assets/Qwen3.5-9B-Base" -DEFINITION_ID = "qwen3_5_9b_base_miles_lora_16k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = True -MAX_CONTEXT_LENGTH = 16_384 - -GPU_TYPE = "H100" -GPUS = 4 -TENSOR_MODEL_PARALLEL_SIZE = 4 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", - "output_layer", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 16 -# Bound retained publication versions per replica; older versions reload on demand. -ROLLOUT_MAX_LOADED_LORAS = 64 -ROLLOUT_MIN_CONTAINERS = 8 -ROLLOUT_MAX_CONTAINERS = 8 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - "lm_head", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - env=trainer_deployment_env(), - single_use_containers=True, -) -def qwen3_5_9b_base_miles_lora_16k(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.5-9B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_5_9b_base_miles_lora_16k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k_single.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k_single.py deleted file mode 100644 index bf3c9ed..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_16k_single.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Isolated single-tenant option; identical hardware and backend to the shared trainer.""" -# ruff: noqa: F401 -import modal -from ..deployment import trainer_deployment_env, trainer_max_containers -from .qwen3_5_9b_base_miles_lora_16k import ( - MODEL_NAME, - HF_CHECKPOINT, - PARAMETERIZATION, - MAX_CONTEXT_LENGTH, - GPU_TYPE, - GPUS, - MAX_LORA_SLOTS, - MAX_LORA_RANK, - DEFAULT_LORA_ALPHA, - TARGET_MODULES, - TENSOR_MODEL_PARALLEL_SIZE, - BULLETIN_ROOT, - BULLETIN_VOLUME_NAME, - TRAINER_VOLUMES, - assets, - bulletin, - image, - api_secret, - proxy_secret, - run_trainer, - ROLLOUT_GPU_TYPE, - ROLLOUT_GPUS, - ROLLOUT_TENSOR_PARALLEL_SIZE, - ROLLOUT_EXPERT_PARALLEL_SIZE, - ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE, - ROLLOUT_MEMORY_FRACTION, - ROLLOUT_MAX_RUNNING_REQUESTS, - ROLLOUT_MAX_QUEUED_REQUESTS, - ROLLOUT_TARGET_CONCURRENCY, - ROLLOUT_MAX_LOADED_LORAS, - ROLLOUT_MIN_CONTAINERS, - ROLLOUT_MAX_CONTAINERS, - ROLLOUT_MAX_LORAS_PER_BATCH, - ROLLOUT_LORA_TARGET_MODULES, -) - -DEFINITION_ID = "qwen3_5_9b_base_miles_lora_16k_single" -CATALOG_VISIBLE = False -TRAINER_MODELS_PER_INSTANCE = 1 -app = modal.App(f"lilo-{DEFINITION_ID}") - -@app.function(image=image, gpu=f"{GPU_TYPE}:{GPUS}", volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], timeout=86_400, - max_containers=trainer_max_containers(), env=trainer_deployment_env(), - single_use_containers=True) -def qwen3_5_9b_base_miles_lora_16k_single(instance_id: str) -> None: - run_trainer(instance_id, definition_id=DEFINITION_ID, max_models=1) - -ENGINE_FUNCTION = qwen3_5_9b_base_miles_lora_16k_single diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_2k.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_2k.py deleted file mode 100644 index 41b0d29..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_base_miles_lora_2k.py +++ /dev/null @@ -1,158 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-9B-Base" -HF_CHECKPOINT = "/assets/Qwen3.5-9B-Base" -DEFINITION_ID = "qwen3_5_9b_base_miles_lora_2k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = False -MAX_CONTEXT_LENGTH = 2048 - -GPU_TYPE = "H200" -GPUS = 4 -TENSOR_MODEL_PARALLEL_SIZE = 4 -MAX_LORA_SLOTS = 4 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", - "output_layer", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_MAX_LOADED_LORAS = 32 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - "lm_head", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - env=trainer_deployment_env(), - single_use_containers=True, -) -def qwen3_5_9b_base_miles_lora_2k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - from modal.config import config - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.5-9B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH, - "extra_args": ( - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=MAX_LORA_SLOTS, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_5_9b_base_miles_lora_2k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_full_64k.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_full_64k.py deleted file mode 100644 index 3383d7a..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_full_64k.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-9B" -HF_CHECKPOINT = "/assets/Qwen3.5-9B" -MICRO_BATCH_SIZE = 1 -MAX_CONTEXT_LENGTH = 65_536 -MAX_TOKENS_PER_MICROBATCH = MAX_CONTEXT_LENGTH -DEFINITION_ID = "qwen3_5_9b_full_64k" -PARAMETERIZATION = "full" -CATALOG_VISIBLE = True -TRAINER_MODELS_PER_INSTANCE = 1 -GPU_TYPE = "H200" -GPUS = 4 -TENSOR_MODEL_PARALLEL_SIZE = 2 -CONTEXT_PARALLEL_SIZE = 2 -DATA_PARALLEL_SIZE = GPUS // (TENSOR_MODEL_PARALLEL_SIZE * CONTEXT_PARALLEL_SIZE) -MTP_NUM_LAYERS = 0 -RECOMPUTE_GRANULARITY = "full" -RECOMPUTE_METHOD = "uniform" -RECOMPUTE_NUM_LAYERS = 1 -LOSS_SCALE = 1.0 -DEFER_FP32_LOGITS = True -FP32_LM_HEAD = True -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.85 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 4 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = 16 -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..megatron_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_5_9b_full_64k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "megatron": { - "hf_checkpoint": HF_CHECKPOINT, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "sequence_parallel": True, - "micro_batch_size": MICRO_BATCH_SIZE, - "max_tokens_per_microbatch": MAX_TOKENS_PER_MICROBATCH, - "seq_length": MAX_CONTEXT_LENGTH, - "defer_fp32_logits": DEFER_FP32_LOGITS, - "fp32_lm_head": FP32_LM_HEAD, - "use_distributed_optimizer": True, - "provider_overrides": { - "mtp_num_layers": MTP_NUM_LAYERS, - "recompute_granularity": RECOMPUTE_GRANULARITY, - "recompute_method": RECOMPUTE_METHOD, - "recompute_num_layers": RECOMPUTE_NUM_LAYERS, - }, - "optimizer": { - "lr": 1e-4, - "min_lr": 1e-4, - "loss_scale": LOSS_SCALE, - }, - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.megatron_fft:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=GPUS, - max_models=1, - ) - - -ENGINE_FUNCTION = qwen3_5_9b_full_64k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k.py deleted file mode 100644 index e48008c..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-9B" -HF_CHECKPOINT = "/assets/Qwen3.5-9B" -DEFINITION_ID = "qwen3_5_9b_miles_lora_16k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = True -MAX_CONTEXT_LENGTH = 16_384 - -GPU_TYPE = "H100" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 8 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 8 -ROLLOUT_MAX_CONTAINERS = 8 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_5_9b_miles_lora_16k(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.5-9B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - if deterministic_training: - backend_config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_5_9b_miles_lora_16k diff --git a/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k_dp2.py b/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k_dp2.py deleted file mode 100644 index f95f190..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_5_9b_miles_lora_16k_dp2.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.5-9B" -HF_CHECKPOINT = "/assets/Qwen3.5-9B" -DEFINITION_ID = "qwen3_5_9b_miles_lora_16k_dp2" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = False -MAX_CONTEXT_LENGTH = 16_384 - -GPU_TYPE = "H100" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 8 -ROLLOUT_MAX_CONTAINERS = 8 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_5_9b_miles_lora_16k_dp2(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.5-9B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - if deterministic_training: - backend_config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_5_9b_miles_lora_16k_dp2 diff --git a/src/lilo/providers/modal/definitions/qwen3_6_27b_full_64k.py b/src/lilo/providers/modal/definitions/qwen3_6_27b_full_64k.py deleted file mode 100644 index 6a95ac4..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_6_27b_full_64k.py +++ /dev/null @@ -1,146 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.6-27B" -HF_CHECKPOINT = "/assets/Qwen3.6-27B" -MICRO_BATCH_SIZE = 1 -MAX_CONTEXT_LENGTH = 65_536 -MAX_TOKENS_PER_MICROBATCH = MAX_CONTEXT_LENGTH -DEFINITION_ID = "qwen3_6_27b_full_64k" -PARAMETERIZATION = "full" -CATALOG_VISIBLE = True -TRAINER_MODELS_PER_INSTANCE = 1 -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -PIPELINE_MODEL_PARALLEL_SIZE = 1 -CONTEXT_PARALLEL_SIZE = 2 -DATA_PARALLEL_SIZE = 1 -SEQUENCE_PARALLEL = True -GPU_MEMORY_FRACTION = 0.90 -MTP_NUM_LAYERS = 0 -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 4 -ROLLOUT_TENSOR_PARALLEL_SIZE = 4 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.90 -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = 32 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 4 -ROLLOUT_TARGET_CONCURRENCY = 16 -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..megatron_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_6_27b_full_64k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "megatron": { - "hf_checkpoint": HF_CHECKPOINT, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "pipeline_model_parallel_size": PIPELINE_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "sequence_parallel": SEQUENCE_PARALLEL, - "micro_batch_size": MICRO_BATCH_SIZE, - "max_tokens_per_microbatch": MAX_TOKENS_PER_MICROBATCH, - "seq_length": MAX_CONTEXT_LENGTH, - "bf16": True, - "fp16": False, - "gpu_memory_fraction": GPU_MEMORY_FRACTION, - "use_distributed_optimizer": True, - "provider_overrides": { - "mtp_num_layers": MTP_NUM_LAYERS, - "recompute_granularity": "full", - "recompute_method": "uniform", - "recompute_num_layers": 1, - }, - "optimizer": { - "optimizer": "adam", - "lr": 1e-4, - "min_lr": 1e-4, - }, - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.megatron_fft:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - }, - nproc=GPUS, - max_models=1, - ) - - -ENGINE_FUNCTION = qwen3_6_27b_full_64k diff --git a/src/lilo/providers/modal/definitions/qwen3_6_35b_a3b_full_64k.py b/src/lilo/providers/modal/definitions/qwen3_6_35b_a3b_full_64k.py deleted file mode 100644 index 3a3c08a..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_6_35b_a3b_full_64k.py +++ /dev/null @@ -1,155 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.6-35B-A3B" -HF_CHECKPOINT = "/assets/Qwen3.6-35B-A3B" -MICRO_BATCH_SIZE = 1 -MAX_CONTEXT_LENGTH = 65_536 -MAX_TOKENS_PER_MICROBATCH = MAX_CONTEXT_LENGTH -DEFINITION_ID = "qwen3_6_35b_a3b_full_64k" -PARAMETERIZATION = "full" -CATALOG_VISIBLE = True -TRAINER_MODELS_PER_INSTANCE = 1 -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -PIPELINE_MODEL_PARALLEL_SIZE = 1 -CONTEXT_PARALLEL_SIZE = 2 -DATA_PARALLEL_SIZE = 1 -EXPERT_MODEL_PARALLEL_SIZE = 8 -EXPERT_TENSOR_PARALLEL_SIZE = 1 -SEQUENCE_PARALLEL = True -GPU_MEMORY_FRACTION = 0.90 -MTP_NUM_LAYERS = 0 -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 4 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 4 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.90 -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = 32 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 4 -ROLLOUT_TARGET_CONCURRENCY = 16 -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..megatron_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret], - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_6_35b_a3b_full_64k(instance_id: str) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "megatron": { - "hf_checkpoint": HF_CHECKPOINT, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "pipeline_model_parallel_size": PIPELINE_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "expert_model_parallel_size": EXPERT_MODEL_PARALLEL_SIZE, - "expert_tensor_parallel_size": EXPERT_TENSOR_PARALLEL_SIZE, - "sequence_parallel": SEQUENCE_PARALLEL, - "micro_batch_size": MICRO_BATCH_SIZE, - "max_tokens_per_microbatch": MAX_TOKENS_PER_MICROBATCH, - "seq_length": MAX_CONTEXT_LENGTH, - "bf16": True, - "fp16": False, - "gpu_memory_fraction": GPU_MEMORY_FRACTION, - "use_distributed_optimizer": True, - "provider_overrides": { - "mtp_num_layers": MTP_NUM_LAYERS, - "recompute_granularity": "selective", - "moe_layer_recompute": True, - "moe_token_dispatcher_type": "alltoall", - "moe_router_fusion": True, - "moe_permute_fusion": True, - "moe_grouped_gemm": True, - "moe_shared_expert_overlap": False, - "moe_aux_loss_coeff": 0.0, - }, - "optimizer": { - "optimizer": "adam", - "lr": 1e-4, - "min_lr": 1e-4, - }, - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - run_engine_with_backend( - shared_kv(), - "lilo.backends.megatron_fft:build_executor", - definition_id=DEFINITION_ID, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": DEFINITION_ID, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - }, - nproc=GPUS, - max_models=1, - ) - - -ENGINE_FUNCTION = qwen3_6_35b_a3b_full_64k diff --git a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_128k.py b/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_128k.py deleted file mode 100644 index 5bc7983..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_128k.py +++ /dev/null @@ -1,177 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.8-27B" -HF_CHECKPOINT = "/assets/Qwen3.8-27B" -DEFINITION_ID = "qwen3_8_27b_miles_lora_128k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = False -MAX_CONTEXT_LENGTH = 131_072 - -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 2 -CONTEXT_PARALLEL_SIZE = 4 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 2 -ROLLOUT_TENSOR_PARALLEL_SIZE = 2 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 8 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 4 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 4 -ROLLOUT_MAX_CONTAINERS = 4 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) -huggingface_secret = modal.Secret.from_name("huggingface-secret") - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret, huggingface_secret], - env=trainer_deployment_env(), - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_8_27b_miles_lora_128k(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.8-27B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH // CONTEXT_PARALLEL_SIZE, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - if deterministic_training: - backend_config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_8_27b_miles_lora_128k diff --git a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_16k.py b/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_16k.py deleted file mode 100644 index 823676e..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_16k.py +++ /dev/null @@ -1,175 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.8-27B" -HF_CHECKPOINT = "/assets/Qwen3.8-27B" -DEFINITION_ID = "qwen3_8_27b_miles_lora_16k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = True -MAX_CONTEXT_LENGTH = 16_384 - -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 32 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 16 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 8 -ROLLOUT_MAX_CONTAINERS = 8 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) -huggingface_secret = modal.Secret.from_name("huggingface-secret") - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret, huggingface_secret], - env=trainer_deployment_env(), - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_8_27b_miles_lora_16k(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.8-27B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - if deterministic_training: - backend_config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_8_27b_miles_lora_16k diff --git a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_256k.py b/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_256k.py deleted file mode 100644 index 6ceaff2..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_256k.py +++ /dev/null @@ -1,224 +0,0 @@ -from __future__ import annotations - -import json -import os - -import modal -import modal.experimental -from modal.config import config - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume -from ..kv import shared_kv -from ..serve import run_engine_with_backend - -MODEL_NAME = "Qwen/Qwen3.8-27B" -HF_CHECKPOINT = "/assets/Qwen3.8-27B" -DEFINITION_ID = "qwen3_8_27b_miles_lora_256k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = False -MAX_CONTEXT_LENGTH = 262_144 - -GPU_TYPE = "H200" -GPUS = 8 -TRAINER_NODES = 2 -# 16 = TP2 x CP8 x DP1, the topology raw Miles runs 256k on. TP2 keeps the 27B -# base weights at ~27 GB/GPU of the 141 GB H200, and spending the rest of the -# world size on CP is what shrinks the activation working set: 32k tokens per -# rank instead of 87k under TP8 x CP3. -TENSOR_MODEL_PARALLEL_SIZE = 2 -CONTEXT_PARALLEL_SIZE = 8 -# Divisible by 2 * cp (zigzag chunks) and by tp (sequence parallelism). -_SEQ_ALIGNMENT = 2 * CONTEXT_PARALLEL_SIZE * TENSOR_MODEL_PARALLEL_SIZE -SEQ_LENGTH = -(-MAX_CONTEXT_LENGTH // _SEQ_ALIGNMENT) * _SEQ_ALIGNMENT -MAX_TOKENS_PER_GPU = SEQ_LENGTH // CONTEXT_PARALLEL_SIZE -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS -# Megatron's default 10-minute collective timeout is the NCCL watchdog budget -# for a single collective. At 250k tokens the inter-node context-parallel -# all-to-all runs behind a straggler's recompute, so a rank can sit in one -# collective far longer than ten minutes without anything being wrong. -DISTRIBUTED_TIMEOUT_MINUTES = 120 - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 4 -ROLLOUT_TENSOR_PARALLEL_SIZE = 4 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 4 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 2 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 2 -ROLLOUT_MAX_CONTAINERS = 2 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -with image.imports(): - from huggingface_hub import snapshot_download - - from ..ray_cluster import start_trainer_cluster - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) -huggingface_secret = modal.Secret.from_name("huggingface-secret") - - -def ensure_assets() -> None: - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret, huggingface_secret], - env=trainer_deployment_env(), - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, - experimental_options={"efa_enabled": True}, -) -@modal.experimental.clustered(TRAINER_NODES, rdma=True) -def qwen3_8_27b_miles_lora_256k(instance_id: str) -> None: - ray_address = start_trainer_cluster( - TRAINER_NODES, - before_head=ensure_assets, - before_worker_join=assets.reload, - ) - if ray_address is None: - return - run_trainer(instance_id, ray_address=ray_address) - - -def backend_config( - instance_id: str, - *, - deterministic_training: bool = False, -) -> dict: - config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.8-27B", - "actor_num_gpus_per_node": GPUS, - "actor_num_nodes": TRAINER_NODES, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_TOKENS_PER_GPU, - "extra_args": ( - "--seq-length", - str(SEQ_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - "--distributed-timeout-minutes", - str(DISTRIBUTED_TIMEOUT_MINUTES), - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - "capture_dir": f"{CHECKPOINT_ROOT}/.captures/{instance_id}", - } - if deterministic_training: - config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - return config - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, - ray_address: str | None = None, -) -> None: - ensure_assets() - config_payload = backend_config( - instance_id, - deterministic_training=deterministic_training, - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(config_payload), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - **({"LILO_RAY_ADDRESS": ray_address} if ray_address else {}), - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_8_27b_miles_lora_256k diff --git a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_64k.py b/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_64k.py deleted file mode 100644 index a927159..0000000 --- a/src/lilo/providers/modal/definitions/qwen3_8_27b_miles_lora_64k.py +++ /dev/null @@ -1,177 +0,0 @@ -from __future__ import annotations - -import os - -import modal - -from ..checkpoint_storage import ( - CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, -) -from ..deployment import trainer_deployment_env, trainer_max_containers -from ..kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - -MODEL_NAME = "Qwen/Qwen3.8-27B" -HF_CHECKPOINT = "/assets/Qwen3.8-27B" -DEFINITION_ID = "qwen3_8_27b_miles_lora_64k" -PARAMETERIZATION = "lora" -CATALOG_VISIBLE = False -MAX_CONTEXT_LENGTH = 65_536 - -GPU_TYPE = "H200" -GPUS = 8 -TENSOR_MODEL_PARALLEL_SIZE = 4 -CONTEXT_PARALLEL_SIZE = 2 -MAX_LORA_SLOTS = 6 -MAX_LORA_RANK = 32 -DEFAULT_LORA_ALPHA = 32 -TARGET_MODULES = ( - "linear_qkv", - "linear_proj", - "linear_fc1", - "linear_fc2", -) -TRAINER_MODELS_PER_INSTANCE = MAX_LORA_SLOTS - -ROLLOUT_GPU_TYPE = "H200" -ROLLOUT_GPUS = 1 -ROLLOUT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_PARALLEL_SIZE = 1 -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = 1 -ROLLOUT_MEMORY_FRACTION = 0.8 -ROLLOUT_MAX_RUNNING_REQUESTS = 16 -ROLLOUT_MAX_QUEUED_REQUESTS = 8 -ROLLOUT_TARGET_CONCURRENCY = 8 -ROLLOUT_MAX_LOADED_LORAS = 256 -ROLLOUT_MIN_CONTAINERS = 8 -ROLLOUT_MAX_CONTAINERS = 8 -ROLLOUT_MAX_LORAS_PER_BATCH = 8 -ROLLOUT_LORA_TARGET_MODULES = ( - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", -) - -BULLETIN_ROOT = "/bulletin" -BULLETIN_VOLUME_NAME = "lilo-snapshot-bulletin" -app = modal.App(f"lilo-{DEFINITION_ID}") - -if modal.is_local(): - from ..miles_image import image -else: - image = modal.Image.debian_slim() - -assets = modal.Volume.from_name("lilo-model-assets", create_if_missing=True) -bulletin = modal.Volume.from_name( - BULLETIN_VOLUME_NAME, - create_if_missing=True, - version=2, -) -TRAINER_VOLUMES = { - "/assets": assets, - BULLETIN_ROOT: bulletin, - CHECKPOINT_ROOT: checkpoint_volume, - KERNEL_CACHE_ROOT: kernel_cache_volume, -} -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -proxy_secret = modal.Secret.from_name( - "lilo-proxy", - required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], -) -huggingface_secret = modal.Secret.from_name("huggingface-secret") - - -@app.function( - image=image, - gpu=f"{GPU_TYPE}:{GPUS}", - volumes=TRAINER_VOLUMES, - secrets=[api_secret, proxy_secret, huggingface_secret], - env=trainer_deployment_env(), - timeout=86_400, - max_containers=trainer_max_containers(), - single_use_containers=True, -) -def qwen3_8_27b_miles_lora_64k(instance_id: str) -> None: - run_trainer(instance_id) - - -def run_trainer( - instance_id: str, - *, - definition_id: str = DEFINITION_ID, - max_models: int = MAX_LORA_SLOTS, - deterministic_training: bool = False, -) -> None: - import json - - from huggingface_hub import snapshot_download - from modal.config import config - - from lilo.providers.modal.kv import shared_kv - from lilo.providers.modal.serve import run_engine_with_backend - - if not os.path.exists(HF_CHECKPOINT): - snapshot_download(repo_id=MODEL_NAME, local_dir=HF_CHECKPOINT) - assets.commit() - backend_config = { - "miles": { - "hf_checkpoint": HF_CHECKPOINT, - "model_type": "qwen3.8-27B", - "actor_num_gpus_per_node": GPUS, - "tensor_model_parallel_size": TENSOR_MODEL_PARALLEL_SIZE, - "context_parallel_size": CONTEXT_PARALLEL_SIZE, - "max_lora_slots": MAX_LORA_SLOTS, - "max_lora_rank": MAX_LORA_RANK, - "default_lora_alpha": DEFAULT_LORA_ALPHA, - "target_modules": TARGET_MODULES, - "max_tokens_per_gpu": MAX_CONTEXT_LENGTH // CONTEXT_PARALLEL_SIZE, - "extra_args": ( - "--seq-length", - str(MAX_CONTEXT_LENGTH), - "--recompute-granularity", - "full", - "--recompute-method", - "uniform", - "--recompute-num-layers", - "1", - ), - }, - "checkpoint_dir": CHECKPOINT_ROOT, - } - if deterministic_training: - backend_config["miles"].update( - tp_reduce_precision="float64", deterministic_attention=True - ) - run_engine_with_backend( - shared_kv(), - "lilo.backends.miles_lora:build_executor", - definition_id=definition_id, - revision=config["image_id"], - instance_id=instance_id, - backend_env={ - **KERNEL_CACHE_ENV, - "LILO_BACKEND_CONFIG": json.dumps(backend_config), - "LILO_BASE_MODEL": MODEL_NAME, - "LILO_DEFINITION_ID": definition_id, - "LILO_CHECKPOINT_VOLUME": CHECKPOINT_VOLUME_NAME, - "LILO_BULLETIN_ROOT": BULLETIN_ROOT, - "LILO_BULLETIN_VOLUME": BULLETIN_VOLUME_NAME, - "LILO_DEFINITION_REVISION": config["image_id"], - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "TORCHINDUCTOR_COMPILE_THREADS": "1", - }, - nproc=1, - max_models=max_models, - sampler_persistence_concurrency=8, - ) - - -ENGINE_FUNCTION = qwen3_8_27b_miles_lora_64k diff --git a/src/lilo/providers/modal/deployment.py b/src/lilo/providers/modal/deployment.py index a91ab02..b95b4ef 100644 --- a/src/lilo/providers/modal/deployment.py +++ b/src/lilo/providers/modal/deployment.py @@ -1,26 +1,9 @@ import os -TRAINER_MAX_CONTAINERS_ENV = "LILO_TRAINER_MAX_CONTAINERS" APP_NAME_ENV = "LILO_APP_NAME" -def trainer_max_containers() -> int | None: - value = os.environ.get(TRAINER_MAX_CONTAINERS_ENV) - if value is None: - return None - try: - limit = int(value) - except ValueError as exc: - raise ValueError( - f"{TRAINER_MAX_CONTAINERS_ENV} must be a positive integer" - ) from exc - if limit < 1: - raise ValueError(f"{TRAINER_MAX_CONTAINERS_ENV} must be a positive integer") - return limit - - FORWARDED_DEPLOYMENT_ENVS = ( - TRAINER_MAX_CONTAINERS_ENV, APP_NAME_ENV, "LILO_TORCH_PROFILE_STEP", "LILO_TORCH_PROFILE_DIR", diff --git a/src/lilo/providers/modal/deployment_apps.py b/src/lilo/providers/modal/deployment_apps.py new file mode 100644 index 0000000..769799c --- /dev/null +++ b/src/lilo/providers/modal/deployment_apps.py @@ -0,0 +1,320 @@ +"""Modal app builders shared by all Python-configured model deployments. + +Trainers and inference provisioners are independently deployed apps. Pools are +created on demand with the existing LoRA/FFT pool lifecycle. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time + +import modal +import modal.experimental +from modal.config import config + +from lilo.deployments import DeploymentConfig, platform_defaults +from lilo.inference.serving import ( + start_fft_sidecar, + start_lora_sidecar, + supervise_children, + terminate, + wait_http, +) + +from .deployment import trainer_deployment_env +from .deployment_configs import platform_from_env +from .fft_pool import FFTPoolSpec +from .fft_pool import deploy_pool as deploy_fft +from .image_dependencies import ( + CORE_PACKAGES, + STITCH_PACKAGE, + TINKER_PACKAGE, + ignore_config_source, +) +from .kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume +from .kv import shared_kv +from .lora_pool import LoraPoolSpec +from .lora_pool import deploy_pool as deploy_lora +from .megatron_image import image as megatron_image +from .miles_image import image as miles_image +from .ray_cluster import start_trainer_cluster +from .rollout_image import image as rollout_image +from .serve import run_engine_with_backend + + +def image_for(backend): + if not modal.is_local(): + return modal.Image.debian_slim() + return {"miles": miles_image, "megatron": megatron_image, "sglang": rollout_image}[ + backend + ] + + +def volumes_for(platform): + storage = platform["storage"] + return { + KERNEL_CACHE_ROOT: kernel_cache_volume, + "/assets": modal.Volume.from_name(storage["assets"], create_if_missing=True), + "/checkpoints": modal.Volume.from_name( + storage["checkpoints"], create_if_missing=True, version=2 + ), + "/bulletin": modal.Volume.from_name( + storage["bulletin"], create_if_missing=True, version=2 + ), + } + + +def secrets_for(platform, *, training=False): + names = platform["secrets"] + result = [modal.Secret.from_name(names["api"], required_keys=["TINKER_API_KEY"])] + if training: + result.append( + modal.Secret.from_name( + names["sampler_proxy"], + required_keys=["MODAL_PROXY_TOKEN_ID", "MODAL_PROXY_TOKEN_SECRET"], + ) + ) + if names["huggingface"]: + result.append(modal.Secret.from_name(names["huggingface"])) + return result + + +def deployment_env(values): + """Keep user environment overrides separate from Lilo's deployment wiring.""" + if any(key.startswith("LILO_") for key in values): + raise ValueError("LILO_ environment variables are managed by Lilo") + return dict(values) + + +def build_trainer_app(deployment: DeploymentConfig, platform=None, *, image=None): + if platform is None: + platform = platform_defaults() + recipe = deployment.recipe + trainer_name = recipe.name + app = modal.App(deployment.trainer_app_name) + + env = { + **trainer_deployment_env(), + **deployment_env(recipe.trainer_env), + "LILO_APP_NAME": platform["frontend"], + } + + def trainer(instance_id: str, config_json: str, frontend_app_id: str): + saved = DeploymentConfig.model_validate_json(config_json) + if saved.name != trainer_name: + raise ValueError("trainer settings do not match the deployed app") + os.environ["LILO_FRONTEND_APP_ID"] = frontend_app_id + run_trainer(saved, instance_id, platform) + + if recipe.trainer_nodes > 1: + trainer = modal.experimental.clustered(recipe.trainer_nodes, rdma=True)(trainer) + trainer = app.function( + name="trainer", + serialized=True, + image=image if image is not None else image_for(recipe.backend), + gpu=f"{recipe.trainer_gpu}:{recipe.trainer_gpus_per_node}", + region=platform["modal"]["region"], + cpu=recipe.trainer_cpu, + memory=recipe.trainer_memory_mib, + timeout=recipe.trainer_timeout_s, + max_containers=None, + min_containers=0, + single_use_containers=True, + volumes=volumes_for(platform), + secrets=secrets_for(platform, training=True), + env=env, + experimental_options={"efa_enabled": True} if recipe.trainer_nodes > 1 else {}, + )(trainer) + + return app, trainer + + +def run_trainer(deployment, instance_id, platform=None): + if platform is None: + platform = platform_defaults() + recipe = deployment.recipe + settings = deployment.trainer_settings + assets = volumes_for(platform)["/assets"] + if recipe.trainer_nodes > 1: + ray_address = start_trainer_cluster( + recipe.trainer_nodes, + before_head=assets.reload, + before_worker_join=assets.reload, + ) + if ray_address is None: + return + else: + assets.reload() + ray_address = None + env = { + **KERNEL_CACHE_ENV, + **deployment_env(recipe.trainer_env), + "LILO_APP_NAME": platform["frontend"], + "LILO_BACKEND_CONFIG": json.dumps(settings), + "LILO_BASE_MODEL": recipe.model, + "LILO_DEFINITION_ID": deployment.definition_id, + "LILO_CHECKPOINT_VOLUME": platform["storage"]["checkpoints"], + "LILO_BULLETIN_ROOT": "/bulletin", + "LILO_BULLETIN_VOLUME": platform["storage"]["bulletin"], + "LILO_DEFINITION_REVISION": deployment.definition_id, + } + if ray_address is not None: + env["LILO_RAY_ADDRESS"] = ray_address + executor = ( + "lilo.backends.miles_lora:build_executor" + if recipe.backend == "miles" + else "lilo.backends.megatron_fft:build_executor" + ) + + async def failed(error): + await shared_kv().put( + f"deployment_failure:{deployment.definition_id}", + {"error": str(error), "instance_id": instance_id, "failed_at": time.time()}, + ) + + run_engine_with_backend( + shared_kv(), + executor, + definition_id=deployment.definition_id, + revision=config["image_id"], + instance_id=instance_id, + backend_env=env, + nproc=1 if recipe.backend == "miles" else recipe.trainer_gpus_per_node, + max_models=recipe.trainer_max_clients_per_instance, + sampler_persistence_concurrency=recipe.sampler_persistence_concurrency, + on_startup_error=failed, + ) + + +def build_rollout_app(deployment, pool, platform=None, *, image=None): + """Create one frozen-base LoRA pool or one FFT latest/pinned/base pool.""" + if platform is None: + platform = platform_from_env() + recipe = deployment.recipe + lora = recipe.parameterization == "lora" + if pool.definition_id != deployment.definition_id: + raise ValueError("pool does not match deployment") + app = modal.App(pool.app_name) + options = deployment.inference_settings + minimum = recipe.inference_min_replicas + maximum = recipe.inference_max_replicas + window = recipe.inference_scaledown_window_s + if isinstance(pool, FFTPoolSpec): + minimum = minimum if pool.min_containers is None else pool.min_containers + maximum = maximum if pool.max_containers is None else pool.max_containers + window = window if pool.scaledown_window is None else pool.scaledown_window + + @app.server( + name="Server", + serialized=True, + image=image if image is not None else image_for("sglang"), + gpu=f"{recipe.inference_gpu}:{recipe.inference_gpus_per_node}", + cpu=recipe.inference_cpu, + memory=recipe.inference_memory_mib, + volumes=volumes_for(platform), + secrets=secrets_for(platform), + env=deployment_env(recipe.inference_env), + min_containers=minimum, + max_containers=maximum, + target_concurrency=recipe.inference_target_concurrency, + scaledown_window=window, + startup_timeout=recipe.inference_startup_timeout_s, + exit_grace_period=300, + port=8000, + routing_region=platform["modal"]["region"], + compute_region=platform["modal"]["region"], + ) + class Server: + @modal.enter() + def start(self): + self.sidecar = None + self.sglang = None + + self.sglang = subprocess.Popen( + [ + sys.executable, + "-m", + "lilo.inference.sglang", + deployment.asset_path, + json.dumps(options), + ], + start_new_session=True, + ) + try: + wait_http( + "http://127.0.0.1:8001/health", + self.sglang, + recipe.inference_startup_timeout_s, + ) + kwargs = dict( + port=8000, + sglang_port=8001, + bulletin_root="/bulletin", + bulletin_volume=platform["storage"]["bulletin"], + ) + self.sidecar = ( + start_lora_sidecar(**kwargs) + if lora + else start_fft_sidecar( + **kwargs, + model_path=deployment.asset_path, + run_id=pool.model_id, + pinned_version=None if pool.latest else pool.version, + ) + ) + self.supervisor = supervise_children(self.sglang, self.sidecar) + wait_http( + "http://127.0.0.1:8000/health", + self.sidecar, + recipe.inference_startup_timeout_s, + ) + except BaseException: + terminate(self.sidecar) + terminate(self.sglang) + raise + + @modal.exit() + def stop(self): + terminate(self.sidecar) + terminate(self.sglang) + + return app, Server + + +def build_inference_app(deployment, platform=None, *, image=None): + """Deploy the code that provisions this configuration's inference pools.""" + if platform is None: + platform = platform_defaults() + + if image is None: + image = ( + modal.Image.debian_slim(python_version="3.12") + .apt_install("git") + .pip_install( + *CORE_PACKAGES, STITCH_PACKAGE, TINKER_PACKAGE, "huggingface-hub" + ) + .add_local_python_source("lilo", copy=True, ignore=ignore_config_source) + ) + app = modal.App(deployment.inference_app_name) + inference_name = deployment.name + + @app.function(name="provision", image=image, serialized=True, timeout=1800) + def provision(config_json: str, pool_data: dict): + saved = DeploymentConfig.model_validate_json(config_json) + if saved.name != inference_name: + raise ValueError("inference settings do not match the deployed app") + if pool_data["definition_id"] != saved.definition_id: + raise ValueError("pool definition does not match deployment") + if saved.parameterization == "lora": + return deploy_lora( + LoraPoolSpec.from_dict(pool_data), config=saved, platform=platform + ) + return deploy_fft( + FFTPoolSpec.from_dict(pool_data), config=saved, platform=platform + ) + + return app, provision diff --git a/src/lilo/providers/modal/deployment_configs.py b/src/lilo/providers/modal/deployment_configs.py new file mode 100644 index 0000000..c30e435 --- /dev/null +++ b/src/lilo/providers/modal/deployment_configs.py @@ -0,0 +1,77 @@ +"""Deployment configs and frontend-only platform settings.""" + +import json +import os + +import modal + +from lilo.deployments import DeploymentConfig, platform_defaults + +CONFIGS_ENV = "LILO_DEPLOYMENT_CONFIGS" +PLATFORM_ENV = "LILO_PLATFORM" +POOL_CONFIG_ENV = "LILO_POOL_DEPLOYMENT" + + +def configs_from_env(): + data = os.environ.get(CONFIGS_ENV) + if not data: + raise ValueError( + "Missing deployment configs. Use lilo deploy with your Python config files." + ) + rows = json.loads(data) + if not isinstance(rows, list) or not rows: + raise ValueError("Deployment configs must be a nonempty list") + return [DeploymentConfig.model_validate(row) for row in rows] + + +def platform_from_env(): + data = os.environ.get(PLATFORM_ENV) + if not data: + return platform_defaults() + return json.loads(data) + + +def deployed_configs(frontend, environment=None): + """Read the configuration carried by the currently deployed frontend.""" + try: + return modal.Function.from_name( + frontend, "deployed_configs", environment_name=environment + ).remote() + except modal.exception.NotFoundError: + return [] + + +def deployed_pools(frontend, environment=None): + """Read pools owned by this frontend before updating its inference settings.""" + try: + return modal.Function.from_name( + frontend, "deployed_pools", environment_name=environment + ).remote() + except modal.exception.NotFoundError: + return [] + + +def deployed_platform(frontend, environment=None): + try: + return modal.Function.from_name( + frontend, "deployed_platform", environment_name=environment + ).remote() + except modal.exception.NotFoundError: + return None + + +def pool_config(definition_id): + for deployment in configs_from_env(): + if deployment.definition_id == definition_id: + return deployment + return None + + +def provision_pool(deployment, pool, platform): + """Ask the inference app to create a pool using its original code.""" + provision = modal.Function.from_name( + deployment.inference_app_name, + "provision", + environment_name=platform["modal"]["environment"], + ) + return provision.remote(deployment.model_dump_json(), pool.as_dict()) diff --git a/src/lilo/providers/modal/deployment_pool_app.py b/src/lilo/providers/modal/deployment_pool_app.py new file mode 100644 index 0000000..787661e --- /dev/null +++ b/src/lilo/providers/modal/deployment_pool_app.py @@ -0,0 +1,27 @@ +"""Generic rollout app constructed in the pool deployment subprocess.""" + +import os + +from lilo.deployments import DeploymentConfig + +from .deployment_apps import build_rollout_app +from .deployment_configs import POOL_CONFIG_ENV, platform_from_env +from .fft_pool import FFTPoolSpec +from .lora_pool import LoraPoolSpec + +deployment = DeploymentConfig.model_validate_json(os.environ[POOL_CONFIG_ENV]) +if deployment.parameterization == "lora": + pool = LoraPoolSpec(deployment.definition_id) +else: + pool = FFTPoolSpec( + definition_id=deployment.definition_id, + model_id=os.environ["LILO_FFT_POOL_MODEL_ID"], + latest=os.environ["LILO_FFT_POOL_LATEST"] == "1", + version=int(os.environ["LILO_FFT_POOL_VERSION"]), + **{ + key: int(os.environ[f"LILO_FFT_POOL_{key.upper()}"]) + for key in ("min_containers", "max_containers", "scaledown_window") + if f"LILO_FFT_POOL_{key.upper()}" in os.environ + }, + ) +app, Server = build_rollout_app(deployment, pool, platform_from_env()) diff --git a/src/lilo/providers/modal/fft_pool.py b/src/lilo/providers/modal/fft_pool.py index 8155fb7..36683b7 100644 --- a/src/lilo/providers/modal/fft_pool.py +++ b/src/lilo/providers/modal/fft_pool.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import json import logging import os import shutil @@ -10,9 +11,18 @@ from typing import Any import httpx +import modal from stitch.pools.modal_flash import ModalFlashPool, list_flash_containers from stitch.types import VersionRef +from .deployment_configs import ( + PLATFORM_ENV, + POOL_CONFIG_ENV, + platform_from_env, + pool_config, + provision_pool, +) + logger = logging.getLogger(__name__) @@ -73,8 +83,6 @@ def __init__(self, definition_id: str, model_id: str) -> None: ) def discover_replicas(self) -> list[str]: - import modal - try: return super().discover_replicas() except modal.exception.NotFoundError: @@ -119,28 +127,37 @@ async def pool_gateway(spec: FFTPoolSpec) -> str: return await ModalFlashPool(spec.app_name, "Server").gateway_url_async() -def deploy_pool(spec: FFTPoolSpec) -> str: +def deploy_pool(spec: FFTPoolSpec, *, config=None, platform=None) -> str: pool = ModalFlashPool(spec.app_name, "Server") try: return pool.gateway_url() except Exception as exc: - import modal - if not isinstance(exc, modal.exception.NotFoundError): raise + if config is None: + saved = pool_config(spec.definition_id) + if saved is None: + raise ValueError(f"missing deployment config: {spec.definition_id}") + return provision_pool(saved, spec, platform or platform_from_env()) modal_cli = shutil.which("modal") if modal_cli is None: raise RuntimeError("modal CLI is unavailable") - env = {**os.environ, **spec.env()} + + platform = platform or platform_from_env() + recipe_env = { + POOL_CONFIG_ENV: config.model_dump_json(), + PLATFORM_ENV: json.dumps(platform), + } + env = {**os.environ, **spec.env(), **recipe_env} command = [ modal_cli, "deploy", "-m", - "lilo.providers.modal.fft_pool_app", + "lilo.providers.modal.deployment_pool_app", "--name", spec.app_name, ] - environment = os.environ.get("MODAL_ENVIRONMENT") + environment = platform["modal"]["environment"] if environment: command.extend(["--env", environment]) subprocess.run(command, env=env, check=True) diff --git a/src/lilo/providers/modal/fft_pool_app.py b/src/lilo/providers/modal/fft_pool_app.py deleted file mode 100644 index 036b5ad..0000000 --- a/src/lilo/providers/modal/fft_pool_app.py +++ /dev/null @@ -1,144 +0,0 @@ -from __future__ import annotations - -import importlib -import os - -import modal - -from lilo.inference.serving import ( - start_fft_sidecar, - start_sglang, - supervise_children, - terminate, - wait_http, -) - -from .rollout_image import image - -APP_NAME = os.environ["LILO_FFT_POOL_APP_NAME"] -DEFINITION_ID = os.environ["LILO_FFT_POOL_DEFINITION_ID"] -MODEL_ID = os.environ["LILO_FFT_POOL_MODEL_ID"] -LATEST = os.environ["LILO_FFT_POOL_LATEST"] == "1" -VERSION = int(os.environ["LILO_FFT_POOL_VERSION"]) -definition = importlib.import_module( - f"lilo.providers.modal.definitions.{DEFINITION_ID}" -) -ROLLOUT_GPU_TYPE = definition.ROLLOUT_GPU_TYPE -ROLLOUT_GPUS = definition.ROLLOUT_GPUS -ROLLOUT_TENSOR_PARALLEL_SIZE = definition.ROLLOUT_TENSOR_PARALLEL_SIZE -ROLLOUT_EXPERT_PARALLEL_SIZE = definition.ROLLOUT_EXPERT_PARALLEL_SIZE -ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE = definition.ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE -ROLLOUT_MEMORY_FRACTION = definition.ROLLOUT_MEMORY_FRACTION -ROLLOUT_MAX_RUNNING_REQUESTS = definition.ROLLOUT_MAX_RUNNING_REQUESTS -ROLLOUT_MAX_QUEUED_REQUESTS = definition.ROLLOUT_MAX_QUEUED_REQUESTS -ROLLOUT_TARGET_CONCURRENCY = definition.ROLLOUT_TARGET_CONCURRENCY - - -def _pool_setting(name: str, default: int | None) -> int | None: - value = os.environ.get(f"LILO_FFT_POOL_{name}") - return default if value is None else int(value) - - -ROLLOUT_MIN_CONTAINERS = _pool_setting( - "MIN_CONTAINERS", - getattr(definition, "ROLLOUT_MIN_CONTAINERS", None), -) -ROLLOUT_MAX_CONTAINERS = _pool_setting( - "MAX_CONTAINERS", - getattr(definition, "ROLLOUT_MAX_CONTAINERS", None), -) -ROLLOUT_SCALEDOWN_WINDOW = _pool_setting( - "SCALEDOWN_WINDOW", - getattr(definition, "ROLLOUT_SCALEDOWN_WINDOW", 5 * 60), -) -ROLLOUT_EXIT_GRACE_PERIOD = getattr( - definition, - "ROLLOUT_EXIT_GRACE_PERIOD", - 5 * 60, -) -ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB = ( - definition.ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB -) -SGLANG_PORT = 8001 -SIDECAR_PORT = 8000 - -pool_secret = modal.Secret.from_dict( - { - key: value - for key, value in os.environ.items() - if key.startswith("LILO_FFT_POOL_") - } -) -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -app = modal.App(APP_NAME) - - -@app.server( - image=image, - gpu=f"{ROLLOUT_GPU_TYPE}:{ROLLOUT_GPUS}", - volumes={ - "/assets": definition.assets, - definition.BULLETIN_ROOT: definition.bulletin, - }, - secrets=[api_secret, pool_secret], - target_concurrency=ROLLOUT_TARGET_CONCURRENCY, - min_containers=ROLLOUT_MIN_CONTAINERS, - max_containers=ROLLOUT_MAX_CONTAINERS, - scaledown_window=ROLLOUT_SCALEDOWN_WINDOW, - startup_timeout=20 * 60, - exit_grace_period=ROLLOUT_EXIT_GRACE_PERIOD, - port=SIDECAR_PORT, - routing_region="us-west", -) -class Server: - @modal.enter() - def start(self) -> None: - self.sglang = start_sglang( - definition.HF_CHECKPOINT, - port=SGLANG_PORT, - context_length=definition.MAX_CONTEXT_LENGTH, - max_loras_per_batch=1, - max_loaded_loras=1, - max_lora_rank=1, - max_running_requests=ROLLOUT_MAX_RUNNING_REQUESTS, - max_queued_requests=ROLLOUT_MAX_QUEUED_REQUESTS, - tensor_parallel_size=ROLLOUT_TENSOR_PARALLEL_SIZE, - expert_parallel_size=ROLLOUT_EXPERT_PARALLEL_SIZE, - expert_tensor_parallel_size=ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE, - parallel_world_size=ROLLOUT_GPUS, - enable_lora=False, - enable_cpu_weight_cache=True, - cpu_weight_cache_max_compile_group_gb=( - ROLLOUT_CPU_WEIGHT_CACHE_MAX_COMPILE_GROUP_GB - ), - memory_fraction=ROLLOUT_MEMORY_FRACTION, - schedule_policy="lpm", - ) - wait_http( - f"http://127.0.0.1:{SGLANG_PORT}/health", - self.sglang, - 20 * 60, - ) - self.sidecar = start_fft_sidecar( - port=SIDECAR_PORT, - sglang_port=SGLANG_PORT, - model_path=definition.HF_CHECKPOINT, - bulletin_root=definition.BULLETIN_ROOT, - bulletin_volume=definition.BULLETIN_VOLUME_NAME, - run_id=MODEL_ID, - pinned_version=None if LATEST else VERSION, - ) - self.supervisor = supervise_children(self.sglang, self.sidecar) - wait_http( - f"http://127.0.0.1:{SIDECAR_PORT}/health", - self.sidecar, - 20 * 60, - ) - - @modal.exit() - def stop(self) -> None: - terminate(getattr(self, "sidecar", None)) - terminate(getattr(self, "sglang", None)) diff --git a/src/lilo/providers/modal/image_dependencies.py b/src/lilo/providers/modal/image_dependencies.py index ec39f81..6d8c071 100644 --- a/src/lilo/providers/modal/image_dependencies.py +++ b/src/lilo/providers/modal/image_dependencies.py @@ -1,3 +1,8 @@ +def ignore_config_source(path): + """Workers receive computed configs as JSON; authoring files stay local.""" + return path.suffix != ".py" or path.parts[0] == "configs" + + CORE_PACKAGES = ( "fastapi>=0.141.1", "httpx>=0.28.1", @@ -5,6 +10,7 @@ "opentelemetry-exporter-otlp-proto-http==1.43.0", "modal>=1.5.3", "protobuf>=5.29", + "pyyaml>=6.0.2", "pydantic>=2.13.4", "uvicorn>=0.52.0", "xxhash>=3.8.1", diff --git a/src/lilo/providers/modal/kv.py b/src/lilo/providers/modal/kv.py index 16af6c2..89cb572 100644 --- a/src/lilo/providers/modal/kv.py +++ b/src/lilo/providers/modal/kv.py @@ -2,11 +2,11 @@ import asyncio import hashlib +import os from collections.abc import Awaitable, Callable -from grpclib.exceptions import StreamTerminatedError - import modal +from grpclib.exceptions import StreamTerminatedError from ..contracts import InsertResult @@ -42,6 +42,7 @@ "trainer_demand": "models", "engine_instance": "engines", "lora_pool": "engines", + "deployment_failure": "engines", "engine_call": "engines", "engine_current": "engines", "trainer_plan": "engines", @@ -157,11 +158,16 @@ def _store(self, key: str) -> ModalKeyValueStore: except KeyError: raise ValueError(f"unknown key family: {family}") from None - def _matching_stores(self, prefixes: tuple[str, ...]) -> tuple[ModalKeyValueStore, ...]: + def _matching_stores( + self, prefixes: tuple[str, ...] + ) -> tuple[ModalKeyValueStore, ...]: domains = { domain for family, domain in KEY_STORES.items() - if any(f"{family}:".startswith(prefix) or prefix.startswith(f"{family}:") for prefix in prefixes) + if any( + f"{family}:".startswith(prefix) or prefix.startswith(f"{family}:") + for prefix in prefixes + ) } if not domains: raise ValueError(f"unknown key prefixes: {prefixes}") @@ -193,7 +199,7 @@ async def list_items( def shared_kv(app_id: str | None = None) -> RoutedKeyValueStore: - app_id = app_id or current_app_id() + app_id = app_id or os.environ.get("LILO_FRONTEND_APP_ID") or current_app_id() return RoutedKeyValueStore( { domain: ModalKeyValueStore( diff --git a/src/lilo/providers/modal/lora_pool.py b/src/lilo/providers/modal/lora_pool.py index 6c58b8c..44151b3 100644 --- a/src/lilo/providers/modal/lora_pool.py +++ b/src/lilo/providers/modal/lora_pool.py @@ -1,44 +1,39 @@ from __future__ import annotations -import ast -import hashlib +import json import os import shutil import subprocess from dataclasses import asdict, dataclass from pathlib import Path +import modal from stitch.pools.modal_flash import ModalFlashPool +from .deployment_configs import ( + PLATFORM_ENV, + POOL_CONFIG_ENV, + platform_from_env, + pool_config, + provision_pool, +) + @dataclass(frozen=True) class LoraPoolSpec: definition_id: str - revision: str = "" def __post_init__(self) -> None: if Path(self.definition_id).name != self.definition_id: raise ValueError(f"invalid definition id: {self.definition_id!r}") - if not self.revision: - object.__setattr__( - self, - "revision", - _implementation_revision(self.definition_id), - ) @classmethod def from_dict(cls, value: dict) -> LoraPoolSpec: - return cls( - definition_id=str(value["definition_id"]), - revision=str(value.get("revision", "")), - ) + return cls(definition_id=str(value["definition_id"])) @property def app_name(self) -> str: - digest = hashlib.sha256( - f"{self.definition_id}\0{self.revision}".encode() - ).hexdigest()[:16] - return f"lilo-lora-{digest}" + return f"lilo-lora-{self.definition_id}" def as_dict(self) -> dict[str, str]: return asdict(self) @@ -54,30 +49,39 @@ async def pool_gateway(spec: LoraPoolSpec) -> str: return await ModalFlashPool(spec.app_name, "Server").gateway_url_async() -def deploy_pool(spec: LoraPoolSpec) -> str: +def deploy_pool(spec: LoraPoolSpec, *, config=None, platform=None) -> str: pool = ModalFlashPool(spec.app_name, "Server") try: return pool.gateway_url() except Exception as exc: - import modal - if not isinstance(exc, modal.exception.NotFoundError): raise + if config is None: + saved = pool_config(spec.definition_id) + if saved is None: + raise ValueError(f"missing deployment config: {spec.definition_id}") + return provision_pool(saved, spec, platform or platform_from_env()) modal_cli = shutil.which("modal") if modal_cli is None: raise RuntimeError("modal CLI is unavailable") + + platform = platform or platform_from_env() + recipe_env = { + POOL_CONFIG_ENV: config.model_dump_json(), + PLATFORM_ENV: json.dumps(platform), + } command = [ modal_cli, "deploy", "-m", - "lilo.providers.modal.lora_pool_app", + "lilo.providers.modal.deployment_pool_app", "--name", spec.app_name, ] - environment = os.environ.get("MODAL_ENVIRONMENT") + environment = platform["modal"]["environment"] if environment: command.extend(["--env", environment]) - subprocess.run(command, env={**os.environ, **spec.env()}, check=True) + subprocess.run(command, env={**os.environ, **spec.env(), **recipe_env}, check=True) return pool.gateway_url() @@ -96,43 +100,3 @@ def stop_pool(spec: LoraPoolSpec) -> None: ) if not already_stopped: result.check_returncode() - - -def _implementation_revision(definition_id: str) -> str: - here = Path(__file__) - files = ( - here, - here.with_name("lora_pool_app.py"), - here.with_name("rollout_image.py"), - here.with_name("image_dependencies.py"), - *_definition_sources(here.with_name("definitions") / f"{definition_id}.py"), - here.parents[2] / "inference" / "bulletin.py", - here.parents[2] / "inference" / "lora_sidecar.py", - here.parents[2] / "inference" / "serving.py", - ) - digest = hashlib.sha256() - for path in files: - digest.update(path.name.encode()) - digest.update(path.read_bytes()) - return digest.hexdigest() - - -def _definition_sources(path: Path): - """Include inherited sibling definitions without importing deployment code.""" - pending, seen = [path], set() - while pending: - source = pending.pop() - if source in seen: - continue - seen.add(source) - yield source - for node in ast.walk(ast.parse(source.read_text())): - if not isinstance(node, ast.ImportFrom) or node.level != 1: - continue - modules = ( - [node.module] if node.module else [alias.name for alias in node.names] - ) - for module in modules: - sibling = source.parent / (module.replace(".", "/") + ".py") - if sibling.is_file(): - pending.append(sibling) diff --git a/src/lilo/providers/modal/lora_pool_app.py b/src/lilo/providers/modal/lora_pool_app.py deleted file mode 100644 index 427e34f..0000000 --- a/src/lilo/providers/modal/lora_pool_app.py +++ /dev/null @@ -1,112 +0,0 @@ -from __future__ import annotations - -import importlib -import os - -import modal - -from lilo.inference.serving import ( - start_lora_sidecar, - start_sglang, - supervise_children, - terminate, - wait_http, -) - -from .rollout_image import image - -APP_NAME = os.environ["LILO_LORA_POOL_APP_NAME"] -DEFINITION_ID = os.environ["LILO_LORA_POOL_DEFINITION_ID"] -definition = importlib.import_module( - f"lilo.providers.modal.definitions.{DEFINITION_ID}" -) -SGLANG_PORT = 8001 -SIDECAR_PORT = 8000 - -api_secret = modal.Secret.from_name( - "lilo-api", - required_keys=["TINKER_API_KEY"], -) -pool_secret = modal.Secret.from_dict( - { - key: value - for key, value in os.environ.items() - if key.startswith("LILO_LORA_POOL_") - } -) -app = modal.App(APP_NAME) - - -@app.server( - image=image, - gpu=f"{definition.ROLLOUT_GPU_TYPE}:{definition.ROLLOUT_GPUS}", - volumes={ - "/assets": definition.assets, - definition.BULLETIN_ROOT: definition.bulletin, - }, - secrets=[api_secret, pool_secret], - target_concurrency=definition.ROLLOUT_TARGET_CONCURRENCY, - min_containers=getattr(definition, "ROLLOUT_MIN_CONTAINERS", 0), - max_containers=getattr(definition, "ROLLOUT_MAX_CONTAINERS", None), - scaledown_window=getattr(definition, "ROLLOUT_SCALEDOWN_WINDOW", 5 * 60), - startup_timeout=20 * 60, - exit_grace_period=getattr( - definition, - "ROLLOUT_EXIT_GRACE_PERIOD", - 5 * 60, - ), - port=SIDECAR_PORT, - routing_region="us-west", -) -class Server: - @modal.enter() - def start(self) -> None: - self.sglang = start_sglang( - definition.HF_CHECKPOINT, - port=SGLANG_PORT, - context_length=definition.MAX_CONTEXT_LENGTH, - max_loras_per_batch=getattr( - definition, - "ROLLOUT_MAX_LORAS_PER_BATCH", - 8, - ), - max_loaded_loras=definition.ROLLOUT_MAX_LOADED_LORAS, - max_lora_rank=definition.MAX_LORA_RANK, - max_running_requests=definition.ROLLOUT_MAX_RUNNING_REQUESTS, - max_queued_requests=definition.ROLLOUT_MAX_QUEUED_REQUESTS, - tensor_parallel_size=definition.ROLLOUT_TENSOR_PARALLEL_SIZE, - expert_parallel_size=definition.ROLLOUT_EXPERT_PARALLEL_SIZE, - expert_tensor_parallel_size=( - definition.ROLLOUT_EXPERT_TENSOR_PARALLEL_SIZE - ), - parallel_world_size=definition.ROLLOUT_GPUS, - lora_target_modules=definition.ROLLOUT_LORA_TARGET_MODULES, - enable_lora=True, - enable_return_routed_experts=getattr( - definition, "ROLLOUT_RETURN_ROUTED_EXPERTS", False - ), - memory_fraction=definition.ROLLOUT_MEMORY_FRACTION, - schedule_policy="lpm", - ) - wait_http( - f"http://127.0.0.1:{SGLANG_PORT}/health", - self.sglang, - 20 * 60, - ) - self.sidecar = start_lora_sidecar( - port=SIDECAR_PORT, - sglang_port=SGLANG_PORT, - bulletin_root=definition.BULLETIN_ROOT, - bulletin_volume=definition.BULLETIN_VOLUME_NAME, - ) - self.supervisor = supervise_children(self.sglang, self.sidecar) - wait_http( - f"http://127.0.0.1:{SIDECAR_PORT}/health", - self.sidecar, - 20 * 60, - ) - - @modal.exit() - def stop(self) -> None: - terminate(getattr(self, "sidecar", None)) - terminate(getattr(self, "sglang", None)) diff --git a/src/lilo/providers/modal/megatron_image.py b/src/lilo/providers/modal/megatron_image.py index 6105296..2ef03c2 100644 --- a/src/lilo/providers/modal/megatron_image.py +++ b/src/lilo/providers/modal/megatron_image.py @@ -5,6 +5,7 @@ MEGATRON_RUNTIME_CHECK, MEGATRON_RUNTIME_PACKAGES, STITCH_PACKAGE, + ignore_config_source, ) BASE_IMAGE = "radixark/miles:v0.1.0" @@ -37,12 +38,12 @@ "pip uninstall -y megatron-bridge megatron_bridge || true", 'python3 -c "import importlib.util as u, os, shutil; ' "s = u.find_spec('megatron.bridge'); " - "p = os.path.dirname(s.origin) if s and getattr(s, 'origin', None) " + "p = os.path.dirname(s.origin) if s and s.origin " 'else None; shutil.rmtree(p) if p else None" || true', f"pip install --no-build-isolation --no-deps -e {MEGATRON_BRIDGE_PATH}", ) .pip_install(*CORE_PACKAGES, STITCH_PACKAGE) .pip_install(*MEGATRON_RUNTIME_PACKAGES) .run_commands(MEGATRON_RUNTIME_CHECK) - .add_local_python_source("lilo") + .add_local_python_source("lilo", ignore=ignore_config_source) ) diff --git a/src/lilo/providers/modal/miles_image.py b/src/lilo/providers/modal/miles_image.py index be46765..e64ea93 100644 --- a/src/lilo/providers/modal/miles_image.py +++ b/src/lilo/providers/modal/miles_image.py @@ -1,16 +1,17 @@ import modal -from .miles_revision import MILES_REPOSITORY, resolve_miles_commit - from .image_dependencies import ( CORE_PACKAGES, MEGATRON_RUNTIME_CHECK, MEGATRON_RUNTIME_PACKAGES, STITCH_PACKAGE, + ignore_config_source, ) BASE_IMAGE = "radixark/miles:v0.1.0" -MILES_COMMIT = resolve_miles_commit() +MILES_REPOSITORY = "https://github.com/radixark/miles.git" +# Update this when Miles main should be picked up, then refresh the trainer app. +MILES_COMMIT = "5510af675238be8271c0a24740f70d5116f6d32b" MILES_PATH = "/root/miles" MEGATRON_REPOSITORY = "https://github.com/radixark/Megatron-LM.git" MEGATRON_REVISION = "8c1e05747eb612b382df2632783df5c83a853646" @@ -60,10 +61,10 @@ .run_commands( "pip install --no-deps 'peft>=0.18.1'", MEGATRON_RUNTIME_CHECK, - "python -c \"from miles.ray.train.group import TrainerController as T; " + 'python -c "from miles.ray.train.group import TrainerController as T; ' "assert all(hasattr(T, name) for name in " "('load_slot', 'unload_slot', 'forward_backward', " "'forward_only', 'optim_step', 'save_slot', 'export_slot'))\"", ) - .add_local_python_source("lilo") + .add_local_python_source("lilo", ignore=ignore_config_source) ) diff --git a/src/lilo/providers/modal/miles_revision.py b/src/lilo/providers/modal/miles_revision.py deleted file mode 100644 index 7b24e91..0000000 --- a/src/lilo/providers/modal/miles_revision.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Resolve a moving Miles ref before constructing the cached Modal image.""" - -import os -import re -import subprocess -from functools import lru_cache - -from lilo.backends.miles_config import MILES_REF - -MILES_REPOSITORY = "https://github.com/radixark/miles.git" - - -def validate_commit(value: str) -> str: - if re.fullmatch(r"[0-9a-f]{40}", value) is None: - raise ValueError("LILO_MILES_COMMIT must be a full lowercase Git commit SHA") - return value - - -@lru_cache(maxsize=1) -def resolve_miles_commit() -> str: - """Resolve once per deployment process; allow an exact reproducibility override.""" - override = os.environ.get("LILO_MILES_COMMIT") - if override is not None: - return validate_commit(override) - ref = f"refs/heads/{MILES_REF}" - result = subprocess.run( - ["git", "ls-remote", "--exit-code", MILES_REPOSITORY, ref], - check=True, capture_output=True, text=True, timeout=30, - ) - entries = [line.split() for line in result.stdout.splitlines()] - commits = [sha for sha, name in entries if name == ref] - if len(commits) != 1: - raise RuntimeError(f"Expected exactly one Miles {ref} commit") - return validate_commit(commits[0]) diff --git a/src/lilo/providers/modal/rollout_image.py b/src/lilo/providers/modal/rollout_image.py index 0aaa9ff..16b2e4a 100644 --- a/src/lilo/providers/modal/rollout_image.py +++ b/src/lilo/providers/modal/rollout_image.py @@ -1,6 +1,11 @@ import modal -from .image_dependencies import CORE_PACKAGES, STITCH_PACKAGE, TINKER_PACKAGE +from .image_dependencies import ( + CORE_PACKAGES, + STITCH_PACKAGE, + TINKER_PACKAGE, + ignore_config_source, +) SGLANG_IMAGE = "lmsysorg/sglang:v0.5.17" SGLANG_REPOSITORY = "https://github.com/modal-projects/sglang.git" @@ -240,5 +245,5 @@ "SGLANG_DISABLE_CUDNN_CHECK": "1", } ) - .add_local_python_source("lilo") + .add_local_python_source("lilo", ignore=ignore_config_source) ) diff --git a/src/lilo/providers/modal/scoped.py b/src/lilo/providers/modal/scoped.py index 01d468a..00fcd0c 100644 --- a/src/lilo/providers/modal/scoped.py +++ b/src/lilo/providers/modal/scoped.py @@ -4,19 +4,48 @@ import asyncio import json +import logging import os import time +import uuid from dataclasses import asdict, replace from types import SimpleNamespace import modal +from fastapi import HTTPException +from huggingface_hub import snapshot_download +from modal.config import config +from lilo.control_plane import create_control_plane_app from lilo.engines import Engine, gpu_count +from lilo.inference.sampling import sample_task +from lilo.inference.serving import ( + start_fft_sidecar, + start_sglang, + supervise_children, + terminate, + wait_http, +) +from lilo.providers.modal.scoped_pool import set_minimum +from lilo.telemetry.otlp import sample_trace + +from .checkpoint_storage import ModalCheckpointStorage +from .engines import ModalEnginePlatform +from .fft_pool import proxy_auth_headers +from .image_dependencies import CORE_PACKAGES, STITCH_PACKAGE, TINKER_PACKAGE +from .kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume +from .kv import ModalSessionKeyValueStores, shared_kv +from .megatron_image import image as default_trainer_image +from .rollout_image import image as default_sampler_image +from .sampling import ModalSamplingTaskPlatform +from .scoped_assignment import claim_model +from .scoped_control import ScopedControlPlane +from .scoped_pins import forget_pin_route, pin_demand, publish_pin, touch_pin +from .scoped_pool import ScopedFlashPool +from .serve import run_engine_with_backend def control_image(*extra_packages): - from .image_dependencies import CORE_PACKAGES, STITCH_PACKAGE, TINKER_PACKAGE - return ( modal.Image.debian_slim(python_version="3.12") .apt_install("git") @@ -46,14 +75,6 @@ def register_sampler( proxy_secret, name, ): - from lilo.inference.serving import ( - start_fft_sidecar, - start_sglang, - supervise_children, - terminate, - wait_http, - ) - model_path = engine.training.hf_checkpoint # Latest min is activated only after its model has been assigned. @@ -75,6 +96,9 @@ def register_sampler( routing_region="us-west", ) class Sampler: + sidecar = None + sglang = None + @modal.enter() def start(self): run_id = model_id @@ -110,8 +134,8 @@ def start(self): @modal.exit() def stop(self): - terminate(getattr(self, "sidecar", None)) - terminate(getattr(self, "sglang", None)) + terminate(self.sidecar) + terminate(self.sglang) return Sampler @@ -129,20 +153,6 @@ def build_app( *, telemetry_secret=None, ): - from lilo.control_plane import create_control_plane_app - from lilo.inference.sampling import sample_task - - from .engines import ModalEnginePlatform - from .checkpoint_storage import ModalCheckpointStorage - from .fft_pool import proxy_auth_headers - from .kernel_cache import KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, kernel_cache_volume - from .kv import ModalSessionKeyValueStores, shared_kv - from .megatron_image import image as default_trainer_image - from .rollout_image import image as default_sampler_image - from .sampling import ModalSamplingTaskPlatform - from .scoped_control import ScopedControlPlane - from .scoped_pool import ScopedFlashPool - engine = replace( engine, training=replace( @@ -173,8 +183,6 @@ def build_app( timeout=3600, ) def prepare_assets(): - from huggingface_hub import snapshot_download - # Let HF validate/resume the snapshot. config.json alone can survive an # interrupted download without the model's weight shards or tokenizer. snapshot_download( @@ -206,10 +214,6 @@ def prepare_assets(): secrets=[*telemetry_secrets, api_secret, proxy_secret], ) def trainer(instance_id): - from modal.config import config - - from .serve import run_engine_with_backend - backend_config = { "megatron": asdict(engine.training), "checkpoint_dir": "/checkpoints", @@ -292,8 +296,6 @@ async def manage(action, model_id=None, version=None, lease=None, route=None): await registry.put.aio("closing", True) return if action == "forget_pin_route": - from .scoped_pins import forget_pin_route - return await forget_pin_route(registry, model_id, route) if await registry.get.aio("closing"): raise RuntimeError("deployment is closing") @@ -310,17 +312,12 @@ async def manage(action, model_id=None, version=None, lease=None, route=None): await asyncio.sleep(2) raise TimeoutError("trainer warmup exceeded deadline") if action == "claim": - from .scoped_assignment import claim_model - route = await claim_model( registry, shared_kv(), engines, engine.name, model_id ) if latest.min_containers: - from lilo.providers.modal.scoped_pool import set_minimum - await set_minimum.aio(route["function_id"], latest.min_containers) return route - from .scoped_pins import pin_demand, publish_pin, touch_pin if action == "pin_demand": return await pin_demand(registry, now=time.time(), routes=route) @@ -349,8 +346,6 @@ async def route_for(model_id, version, is_latest): return (await registry.get.aio("routes"))[0] if is_latest: if await registry.get.aio("slot:0") != model_id: - from fastapi import HTTPException - raise HTTPException( 410, "sampling model was replaced; use a new sampling client" ) @@ -370,8 +365,6 @@ async def route_for(model_id, version, is_latest): ) @modal.concurrent(max_inputs=128) async def execute_sample(task): - import uuid - pinned_request = task["model_id"] is not None and not task.get("latest") lease = uuid.uuid4().hex if pinned_request else None if pinned_request: @@ -389,8 +382,6 @@ async def gateway(): ) gateway = ScopedFlashPool(route).gateway_url() try: - from lilo.telemetry.otlp import sample_trace - stats = {} with sample_trace(task, stats): return await sample_task( @@ -414,7 +405,6 @@ async def gateway(): except Exception: # A cleanup outage must not discard a completed sample. # The bounded lease still permits eventual reclamation. - import logging logging.getLogger(__name__).exception( "pinned lease release failed; lease will expire" @@ -462,11 +452,11 @@ async def spawn_sampling(task): checkpoint_root=storage.root, ) definition = SimpleNamespace( - CATALOG_VISIBLE=True, - DEFINITION_ID=engine.name, - MODEL_NAME=engine.model, - PARAMETERIZATION="full", - MAX_CONTEXT_LENGTH=engine.training.seq_length, + definition_id=engine.name, + name=engine.name, + model=engine.model, + parameterization="full", + max_context_length=engine.training.seq_length, ) return create_control_plane_app( plane, diff --git a/src/lilo/providers/modal/serve.py b/src/lilo/providers/modal/serve.py index 18f23ad..8f0d082 100644 --- a/src/lilo/providers/modal/serve.py +++ b/src/lilo/providers/modal/serve.py @@ -13,13 +13,16 @@ import httpx import modal +import uvicorn from lilo.engine import Engine from lilo.engine.backend_http import HttpBackendClient from lilo.engine.http import create_engine_app +from lilo.telemetry.trainer import TrainerTelemetry from .engines import EngineInstanceRecord, instance_key from .kv import ModalKeyValueStore +from .trainer_reconciler import request_reconcile ENGINE_PORT = 8000 BACKEND_SHUTDOWN_TIMEOUT = 120.0 @@ -29,8 +32,6 @@ async def _kick_trainer_reconciler(definition_id: str) -> None: - from .trainer_reconciler import request_reconcile - async def spawn(delay_seconds: float) -> str: function = modal.Function.from_name( os.environ.get("LILO_APP_NAME", "lilo"), @@ -51,8 +52,6 @@ async def serve_engine( instance_id: str, notify_reconciler: bool = True, ) -> None: - import uvicorn - record = EngineInstanceRecord( instance_id=instance_id, definition_id=definition_id, @@ -66,7 +65,6 @@ async def serve_engine( trainer_telemetry = None try: engine = await make_server() - from lilo.telemetry.trainer import TrainerTelemetry trainer_telemetry = TrainerTelemetry( instance_id, @@ -120,6 +118,7 @@ def run_engine_with_backend( startup_timeout: float = BACKEND_STARTUP_TIMEOUT, operation_timeout: float = BACKEND_OPERATION_TIMEOUT, notify_reconciler: bool = True, + on_startup_error: Callable[[Exception], Awaitable[None]] | None = None, ) -> None: if sampler_persistence_concurrency > 1 and nproc != 1: raise ValueError( @@ -169,7 +168,10 @@ def signal_backend(sig: signal.Signals) -> None: on_transport_error=lambda: signal_backend(signal.SIGKILL), ) + ready = False + async def make_server() -> Engine: + nonlocal ready try: async with asyncio.timeout(startup_timeout): while True: @@ -179,6 +181,7 @@ async def make_server() -> Engine: ) try: if (await executor.http.get("/healthz")).is_success: + ready = True return Engine( executor, max_models=max_models, @@ -212,6 +215,10 @@ async def serve_until_backend_exits() -> None: await serving return raise RuntimeError(f"backend exited with code {backend.returncode}") + except Exception as exc: + if not ready and on_startup_error is not None: + await on_startup_error(exc) + raise finally: serving.cancel() await asyncio.gather(serving, return_exceptions=True) diff --git a/tests/backends/test_megatron_constructor_settings.py b/tests/backends/test_megatron_constructor_settings.py new file mode 100644 index 0000000..7c4c2af --- /dev/null +++ b/tests/backends/test_megatron_constructor_settings.py @@ -0,0 +1,76 @@ +"""CPU checks for native config forwarding into Megatron constructors.""" + +from dataclasses import make_dataclass +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from runtime_stubs import backend_runtime_imports + +from lilo.backends.deployment import backend_config +from lilo.backends.megatron_config import parse_backend_config +from lilo.deployments import config_path, load + +with backend_runtime_imports(): + from lilo.backends.megatron_runtime.common import modeling + + +def test_config_overrides_reach_megatron(monkeypatch): + spec = load(config_path("qwen35-4b-fft-64k")) + spec.megatron_cfg["optimizer_overrides"] = {"native_optimizer_setting": False} + spec.megatron_cfg["distributed_overrides"] = {"native_ddp_setting": 123} + spec.megatron_cfg["provider_overrides"]["native_provider_setting"] = [1, 2] + config, _ = parse_backend_config(backend_config(spec)) + # These stand in for an installed upstream version with extra fields. The + # deployment reader must not need its own list of those fields. + # Model providers in Megatron Bridge are dataclasses. + fields = { + **modeling.provider_settings(config, "bf16"), + "moe_token_dispatcher_type": "allgather", + "provide_distributed_model": Mock(return_value="model"), + } + + class Hooks: + def register_pre_wrap_hook(self, hook): + self.pre_wrap_hook = hook + + def register_post_wrap_hook(self, hook): + self.post_wrap_hook = hook + + Provider = make_dataclass( + "Provider", [(name, object) for name in fields], bases=(Hooks,) + ) + provider = Provider(**fields) + provider.pre_wrap_hook = Mock(name="load_pretrained_weights") + provider.post_wrap_hook = Mock(name="post_wrap") + bridge = SimpleNamespace(to_megatron_provider=lambda: provider) + monkeypatch.setattr( + modeling, + "AutoBridge", + SimpleNamespace(from_hf_pretrained=lambda *a, **k: bridge), + ) + monkeypatch.setattr(modeling, "parameter_dtype", lambda c: "bf16") + optimizer_constructor = Mock(side_effect=lambda **kwargs: kwargs) + ddp_constructor = Mock(side_effect=lambda **kwargs: kwargs) + monkeypatch.setattr(modeling, "MCoreOptimizerConfig", optimizer_constructor) + monkeypatch.setattr(modeling, "DistributedDataParallelConfig", ddp_constructor) + _, actual_provider, _ = modeling.model_provider(config) + assert actual_provider.native_provider_setting == [1, 2] + assert actual_provider is not provider + assert actual_provider.pre_wrap_hook is provider.pre_wrap_hook + assert actual_provider.post_wrap_hook is provider.post_wrap_hook + assert actual_provider.tensor_model_parallel_size == 2 + assert actual_provider.moe_token_dispatcher_type == "alltoall" + optimizer = modeling.optimizer_config(config, "bf16", distributed_optimizer=True) + assert optimizer["native_optimizer_setting"] is False + assert optimizer["lr"] == 0.0001 + assert ( + modeling.distributed_model(provider, config, distributed_optimizer=True) + == "model" + ) + assert ddp_constructor.call_args.kwargs["native_ddp_setting"] == 123 + assert ddp_constructor.call_args.kwargs["use_distributed_optimizer"] is True + # An unsupported native field is the installed backend's error at startup. + config.provider_overrides["unknown_field"] = True + with pytest.raises(TypeError, match="unknown_field"): + modeling.model_provider(config) diff --git a/tests/backends/test_megatron_fft.py b/tests/backends/test_megatron_fft.py index d54c37c..a1f6434 100644 --- a/tests/backends/test_megatron_fft.py +++ b/tests/backends/test_megatron_fft.py @@ -9,6 +9,7 @@ from stitch.types import VersionRef from lilo.backends import ModelSpec, SamplerPublication +from lilo.backends.megatron_runtime.fft.model import apply_fp32_lm_head with backend_runtime_imports(): from lilo.backends import megatron_fft as fft_backend @@ -416,9 +417,7 @@ def sharded_state_dict(self, **kwargs): "metadata": {"distrib_optim_sharding_type": "dp_reshardable"}, } ] - assert captured["format"] == ( - fft_checkpoint.DISTRIBUTED_OPTIMIZER_STATE_FORMAT - ) + assert captured["format"] == (fft_checkpoint.DISTRIBUTED_OPTIMIZER_STATE_FORMAT) tensors = captured["state_dict"]["param_state"][0]["float32"][0][0] assert tensors["param"].tolist() == [1.25] assert tensors["exp_avg"].tolist() == [2.5] @@ -619,6 +618,19 @@ def test_fft_backend_hf_load_never_reads_native(monkeypatch) -> None: ), "optimizer_config", ), + ( + lambda metadata: replace( + metadata, + native_optimizer_config={"use_precision_aware_optimizer": True}, + ), + "native_optimizer_config", + ), + ( + lambda metadata: replace( + metadata, native_distributed_config={"grad_reduce_in_fp32": True} + ), + "native_distributed_config", + ), ], ) def test_native_metadata_mismatch_prevents_model_mutation( @@ -868,8 +880,6 @@ def fail(*args, **kwargs): def test_fp32_lm_head_upcasts_output_projection() -> None: torch = pytest.importorskip("torch") - from lilo.backends.megatron_runtime.fft.model import apply_fp32_lm_head - calls: list[dict] = [] class Layer(torch.nn.Module): @@ -903,3 +913,16 @@ def _impl(self, **kwargs): assert calls[0]["sequence_parallel"] is True assert model.output_layer.weight.grad is not None assert not hasattr(model.decoder, "_forward_impl") + + +def test_checkpoint_metadata_omits_base_revision(): + config = EngineModelConfig(hf_checkpoint="/model") + metadata = fft_checkpoint.create_fft_checkpoint_metadata( + config, + checkpoint_id="snapshot", + base_model=BASE_MODEL, + include_optimizer=True, + world_size=1, + ) + assert metadata.base_model_revision is None + assert "base_model_revision" not in metadata.to_dict() diff --git a/tests/backends/test_miles.py b/tests/backends/test_miles.py index f1aff73..4658839 100644 --- a/tests/backends/test_miles.py +++ b/tests/backends/test_miles.py @@ -11,7 +11,11 @@ from tinker import AdamParams, Datum, LoraConfig, ModelInput, TensorData from lilo.backends import ForwardBatch, ForwardItem, ModelSpec, miles_lora -from lilo.backends.miles_config import MilesBackendConfig, parse_backend_config +from lilo.backends.miles_config import ( + MilesBackendConfig, + lora_target_flags, + parse_backend_config, +) from lilo.backends.miles_lora import ( MilesCommandBackend, _adam_parameters, @@ -844,3 +848,28 @@ def test_a_capture_short_of_a_nodes_shards_is_refused(monkeypatch, tmp_path): source, tmp_path / "000000" / "model-a", overwrite=False, world_size=2 ) assert copies == [] + + +@pytest.mark.parametrize( + "targets,expected", + [ + ( + ( + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "lm_head", + ), + (True, True, True), + ), + (("model.layers.0.self_attn.q_proj",), (True, False, False)), + (("gate_proj", "up_proj", "down_proj"), (False, True, False)), + (("lm_head",), (False, False, True)), + ], +) +def test_hf_targets_match_tinker_training_flags(targets, expected): + assert lora_target_flags(targets) == expected diff --git a/tests/conftest.py b/tests/conftest.py index 877cfd8..54a9f63 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,5 @@ -"""Keep image-definition imports offline during CPU tests.""" +"""The Miles runtime reads this commit from the environment.""" + import os -# No image is built by this suite. Resolver tests explicitly clear this override. os.environ.setdefault("LILO_MILES_COMMIT", "a" * 40) diff --git a/tests/control_plane/test_http.py b/tests/control_plane/test_http.py index 22df0eb..8b360b4 100644 --- a/tests/control_plane/test_http.py +++ b/tests/control_plane/test_http.py @@ -16,18 +16,18 @@ DEFINITION = "qwen3_8b" DEFINITIONS = ( SimpleNamespace( - DEFINITION_ID=DEFINITION, - MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="lora", - CATALOG_VISIBLE=True, - MAX_CONTEXT_LENGTH=16_384, + definition_id=DEFINITION, + name=DEFINITION, + model=BASE_MODEL, + parameterization="lora", + max_context_length=16_384, ), SimpleNamespace( - DEFINITION_ID=f"{DEFINITION}_full", - MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="full", - CATALOG_VISIBLE=True, - MAX_CONTEXT_LENGTH=65_536, + definition_id=f"{DEFINITION}_full", + name=f"{DEFINITION}_full", + model=BASE_MODEL, + parameterization="full", + max_context_length=65_536, ), ) @@ -283,13 +283,14 @@ async def run() -> None: asyncio.run(run()) -def test_base_sampling_session_prefers_full_definition() -> None: +def test_base_sampling_session_uses_first_deployment() -> None: async def run() -> None: plane = ControlPlane( InMemoryKeyValueStore(), LocalEnginePlatform(DEFINITION, EchoExecutor), ) - app = create_control_plane_app(plane, DEFINITIONS, api_key=None) + definitions = list(reversed(DEFINITIONS)) + app = create_control_plane_app(plane, definitions, api_key=None) client = httpx.AsyncClient( base_url="http://control-plane", transport=httpx.ASGITransport(app=app), @@ -303,9 +304,7 @@ async def run() -> None: "base_model": BASE_MODEL, }, ) - stored = await plane.get_sampling_session( - created.json()["sampling_session_id"] - ) + stored = await plane.get_sampling_session(created.json()["sampling_session_id"]) assert stored.engine_definition_id == f"{DEFINITION}_full" await client.aclose() @@ -466,27 +465,70 @@ def create(seq: int, **body): asyncio.run(run()) -def test_explicit_hidden_deployment_keeps_canonical_model_name() -> None: +def test_explicit_deployment_keeps_canonical_model_name() -> None: async def run(): - hidden = SimpleNamespace(DEFINITION_ID="isolated", MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="lora", CATALOG_VISIBLE=False, MAX_CONTEXT_LENGTH=16384) - plane = ControlPlane(InMemoryKeyValueStore(), LocalEnginePlatform("isolated", EchoExecutor)) - app = create_control_plane_app(plane, (*DEFINITIONS, hidden), retrieve_window=1.0) - async with httpx.AsyncClient(base_url="http://test", transport=httpx.ASGITransport(app=app)) as client: - session = (await client.post("/api/v1/create_session", json={"tags": [], "sdk_version": "0.5.0"})).json()["session_id"] - response = await client.post("/api/v1/create_model", json={"session_id": session, - "model_seq_id": 0, "base_model": "isolated", "lora_config": {"rank": 16}}) + explicit = SimpleNamespace( + definition_id="isolated", + name="isolated", + model=BASE_MODEL, + parameterization="lora", + max_context_length=16384, + ) + plane = ControlPlane( + InMemoryKeyValueStore(), LocalEnginePlatform("isolated", EchoExecutor) + ) + app = create_control_plane_app( + plane, (*DEFINITIONS, explicit), retrieve_window=1.0 + ) + async with httpx.AsyncClient( + base_url="http://test", transport=httpx.ASGITransport(app=app) + ) as client: + listed = (await client.get("/api/v1/lilo/deployments")).json()[ + "deployments" + ] + assert [row["definition_id"] for row in listed] == [ + d.definition_id for d in (*DEFINITIONS, explicit) + ] + session = ( + await client.post( + "/api/v1/create_session", json={"tags": [], "sdk_version": "0.5.0"} + ) + ).json()["session_id"] + response = await client.post( + "/api/v1/create_model", + json={ + "session_id": session, + "model_seq_id": 0, + "base_model": "isolated", + "lora_config": {"rank": 16}, + }, + ) assert response.status_code == 200, response.text created = response.json() - result = await client.post("/api/v1/retrieve_future", json={"request_id": created["request_id"]}) + result = await client.post( + "/api/v1/retrieve_future", json={"request_id": created["request_id"]} + ) assert result.json()["model_id"] == created["model_id"] - info = await client.post("/api/v1/get_info", json={"model_id": created["model_id"]}) + info = await client.post( + "/api/v1/get_info", json={"model_id": created["model_id"]} + ) assert info.json()["model_data"]["model_name"] == BASE_MODEL - assert (await plane.get_model(created["model_id"])).engine_definition_id == "isolated" - sampling = await client.post("/api/v1/create_sampling_session", json={ - "session_id": session, "sampling_session_seq_id": 0, "base_model": "isolated"}) + assert ( + await plane.get_model(created["model_id"]) + ).engine_definition_id == "isolated" + sampling = await client.post( + "/api/v1/create_sampling_session", + json={ + "session_id": session, + "sampling_session_seq_id": 0, + "base_model": "isolated", + }, + ) assert sampling.status_code == 200, sampling.text - record = await plane.get_sampling_session(sampling.json()["sampling_session_id"]) + record = await plane.get_sampling_session( + sampling.json()["sampling_session_id"] + ) assert record.base_model == BASE_MODEL assert record.engine_definition_id == "isolated" + asyncio.run(run()) diff --git a/tests/control_plane/test_sdk_e2e.py b/tests/control_plane/test_sdk_e2e.py index 2ec1b37..7b01de1 100644 --- a/tests/control_plane/test_sdk_e2e.py +++ b/tests/control_plane/test_sdk_e2e.py @@ -29,20 +29,20 @@ MAX_CONTEXT_LENGTH = 32_768 DEFINITIONS = ( SimpleNamespace( - DEFINITION_ID=DEFINITION, - MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="lora", - CATALOG_VISIBLE=True, - MAX_CONTEXT_LENGTH=MAX_CONTEXT_LENGTH, + definition_id=DEFINITION, + name=DEFINITION, + model=BASE_MODEL, + parameterization="lora", + max_context_length=MAX_CONTEXT_LENGTH, ), ) FULL_DEFINITIONS = ( SimpleNamespace( - DEFINITION_ID=FULL_DEFINITION, - MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="full", - CATALOG_VISIBLE=True, - MAX_CONTEXT_LENGTH=MAX_CONTEXT_LENGTH, + definition_id=FULL_DEFINITION, + name=FULL_DEFINITION, + model=BASE_MODEL, + parameterization="full", + max_context_length=MAX_CONTEXT_LENGTH, ), ) @@ -143,9 +143,11 @@ def test_real_sdk_recreates_lora_client_from_state(base_url: str) -> None: training = service.create_lora_training_client(base_model=BASE_MODEL, rank=32) saved = training.save_state("lora-resume").result(timeout=30) assert saved.path == f"tinker://{training.model_id}/weights/snapshot" - info = service.create_rest_client().get_weights_info_by_tinker_path( - saved.path - ).result(timeout=30) + info = ( + service.create_rest_client() + .get_weights_info_by_tinker_path(saved.path) + .result(timeout=30) + ) assert info.base_model == BASE_MODEL assert info.is_lora assert info.lora_rank == 32 @@ -164,9 +166,11 @@ def test_real_sdk_recreates_full_client_from_state(full_base_url: str) -> None: service = tinker.ServiceClient(base_url=full_base_url, api_key=API_KEY) training = create_full_training_client(service, BASE_MODEL) saved = training.save_state("full-resume").result(timeout=30) - info = service.create_rest_client().get_weights_info_by_tinker_path( - saved.path - ).result(timeout=30) + info = ( + service.create_rest_client() + .get_weights_info_by_tinker_path(saved.path) + .result(timeout=30) + ) assert info.base_model == BASE_MODEL assert not info.is_lora assert info.lora_rank is None @@ -412,9 +416,9 @@ def test_real_sdk_lists_and_deletes_checkpoints(tmp_path, monkeypatch) -> None: assert loaded[-1] == str(second) with pytest.raises(tinker.APIStatusError) as archive: - rest.get_checkpoint_archive_url_from_tinker_path( - newest.tinker_path - ).result(timeout=30) + rest.get_checkpoint_archive_url_from_tinker_path(newest.tinker_path).result( + timeout=30 + ) assert archive.value.status_code == 405 assert archive.value.body["error"] == "unsupported" assert ( diff --git a/tests/providers/conftest.py b/tests/providers/conftest.py new file mode 100644 index 0000000..2e43282 --- /dev/null +++ b/tests/providers/conftest.py @@ -0,0 +1,16 @@ +"""Shared provider tests construct the app from explicit offline configs.""" + +import json +import os + +from lilo.deployments import DeploymentConfig, config_path, load + +os.environ.setdefault( + "LILO_DEPLOYMENT_CONFIGS", + json.dumps( + [ + DeploymentConfig.create(load(config_path(name))).model_dump(mode="json") + for name in ("qwen35-9b-fft-64k", "qwen35-9b-lora-16k") + ] + ), +) diff --git a/tests/providers/test_checkpoint_storage.py b/tests/providers/test_checkpoint_storage.py index c9f4da2..8ccb3ad 100644 --- a/tests/providers/test_checkpoint_storage.py +++ b/tests/providers/test_checkpoint_storage.py @@ -1,34 +1,24 @@ -import ast +import asyncio +import importlib import runpy from pathlib import Path from unittest.mock import patch, sentinel import modal -from lilo.providers.modal.app import DEFINITIONS +from lilo.backends.deployment import backend_config +from lilo.providers.modal.app import DEFINITIONS, PLATFORM from lilo.providers.modal.checkpoint_storage import ( CHECKPOINT_ROOT, - CHECKPOINT_VOLUME_NAME, - checkpoint_volume, + _scan_checkpoints, ) +from lilo.providers.modal.deployment_apps import volumes_for FULL_DEFINITIONS = tuple( - definition for definition in DEFINITIONS if definition.PARAMETERIZATION == "full" + definition for definition in DEFINITIONS if definition.parameterization == "full" ) -def source_tree(definition) -> ast.Module: - return ast.parse(Path(definition.__file__).read_text()) - - -def string_dict_entries(node: ast.Dict) -> dict[str, ast.expr]: - return { - key.value: value - for key, value in zip(node.keys, node.values, strict=True) - if isinstance(key, ast.Constant) and isinstance(key.value, str) - } - - def test_checkpoint_storage_creates_one_v2_volume_without_live_lookup() -> None: storage_path = Path(__file__).parents[2] / ( "src/lilo/providers/modal/checkpoint_storage.py" @@ -48,66 +38,26 @@ def test_checkpoint_storage_creates_one_v2_volume_without_live_lookup() -> None: assert storage["CHECKPOINT_ROOT"] == "/checkpoints" -def test_all_definitions_share_checkpoint_storage() -> None: - assert len(FULL_DEFINITIONS) == 5 - assert CHECKPOINT_VOLUME_NAME == "lilo-checkpoints" - assert CHECKPOINT_ROOT == "/checkpoints" - +def test_deployments_use_configured_checkpoint_storage(): for definition in DEFINITIONS: - assert definition.TRAINER_VOLUMES[CHECKPOINT_ROOT] is checkpoint_volume - assert ( - definition.TRAINER_VOLUMES[definition.BULLETIN_ROOT] is definition.bulletin + with patch.object( + modal.Volume, "from_name", side_effect=lambda name, **kwargs: (name, kwargs) + ): + volumes = volumes_for(PLATFORM) + assert volumes[CHECKPOINT_ROOT] == ( + PLATFORM["storage"]["checkpoints"], + {"create_if_missing": True, "version": 2}, ) - assert definition.bulletin is not checkpoint_volume - - -def test_full_definitions_configure_checkpoint_dir_and_environment() -> None: - for definition in FULL_DEFINITIONS: - tree = source_tree(definition) - config_assignments = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Assign) - and any( - isinstance(target, ast.Name) and target.id == "backend_config" - for target in node.targets - ) - and isinstance(node.value, ast.Dict) - ] - assert len(config_assignments) == 1 - config = string_dict_entries(config_assignments[0].value) - checkpoint_dir = config["checkpoint_dir"] - assert isinstance(checkpoint_dir, ast.Name) - assert checkpoint_dir.id == "CHECKPOINT_ROOT" - - engine_calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "run_engine_with_backend" - ] - assert len(engine_calls) == 1 - backend_env = next( - keyword.value - for keyword in engine_calls[0].keywords - if keyword.arg == "backend_env" - ) - assert isinstance(backend_env, ast.Dict) - env = string_dict_entries(backend_env) - volume_name = env["LILO_CHECKPOINT_VOLUME"] - assert isinstance(volume_name, ast.Name) - assert volume_name.id == "CHECKPOINT_VOLUME_NAME" + assert volumes["/bulletin"][0] == PLATFORM["storage"]["bulletin"] + assert backend_config(definition.recipe)["checkpoint_dir"] == CHECKPOINT_ROOT def test_same_checkpoint_name_isolated_by_model(tmp_path, monkeypatch) -> None: - import asyncio - import importlib - app = importlib.import_module("lilo.providers.modal.app") - from lilo.providers.modal.checkpoint_storage import _scan_checkpoints + def scan(model_id): return _scan_checkpoints(str(tmp_path), model_id) + monkeypatch.setattr(app, "CHECKPOINT_ROOT", str(tmp_path)) for relative in ("final/run-a", "final/run-b"): checkpoint = tmp_path / relative @@ -117,7 +67,8 @@ def scan(model_id): (tmp_path / "notes.txt").write_text("notes") entries = scan(None) assert {(e["model_id"], e["name"]) for e in entries} == { - ("run-a", "final"), ("run-b", "final") + ("run-a", "final"), + ("run-b", "final"), } assert len(entries) == 2 assert scan("run-a")[0]["path"] == str(tmp_path / "final/run-a") diff --git a/tests/providers/test_definition_registry.py b/tests/providers/test_definition_registry.py index 74e829e..cf48119 100644 --- a/tests/providers/test_definition_registry.py +++ b/tests/providers/test_definition_registry.py @@ -7,17 +7,8 @@ def test_definition_registry_resolves_every_definition() -> None: for definition in DEFINITIONS: - assert module_for(definition.DEFINITION_ID) is definition - assert parameterization_for(definition.DEFINITION_ID) == ( - definition.PARAMETERIZATION + assert module_for(definition.definition_id) is definition + assert parameterization_for(definition.definition_id) == ( + definition.parameterization ) - assert definition.ENGINE_FUNCTION is not None - - -def test_public_model_parameterizations_are_unique() -> None: - visible = [ - (definition.MODEL_NAME, definition.PARAMETERIZATION) - for definition in DEFINITIONS - if definition.CATALOG_VISIBLE - ] - assert len(visible) == len(set(visible)) + assert definition.trainer_app_name diff --git a/tests/providers/test_deployment_apps.py b/tests/providers/test_deployment_apps.py new file mode 100644 index 0000000..a3932ae --- /dev/null +++ b/tests/providers/test_deployment_apps.py @@ -0,0 +1,477 @@ +import asyncio +import importlib +import json +import os +import subprocess +import sys +from copy import deepcopy +from types import SimpleNamespace + +import modal +import pytest +from modal._serialization import serialize + +import lilo.backends.deployment as backend +from lilo.deployments import DeploymentConfig, config_path, load, platform_defaults +from lilo.providers.local import InMemoryKeyValueStore +from lilo.providers.modal import ( + deployment_apps, + deployment_configs, + fft_pool, + lora_pool, +) +from lilo.providers.modal.fft_pool import FFTPoolSpec +from lilo.providers.modal.lora_pool import LoraPoolSpec + + +def deployment(preset="qwen35-9b-lora-16k"): + return DeploymentConfig.create(load(config_path(preset))) + + +class App: + def __init__(self, name): + self.name = name + self.functions = {} + self.servers = {} + + def function(self, **settings): + def decorate(fn): + self.functions[settings["name"]] = (settings, fn) + return fn + + return decorate + + def server(self, **settings): + def decorate(cls): + self.servers[settings["name"]] = (settings, cls) + return cls + + return decorate + + +@pytest.fixture +def builders(monkeypatch): + monkeypatch.setattr(modal, "App", App) + monkeypatch.setattr(modal, "enter", lambda: lambda fn: fn) + monkeypatch.setattr(modal, "exit", lambda: lambda fn: fn) + + +@pytest.mark.parametrize( + "preset,backend,clients,nproc", + [ + ("qwen35-9b-lora-16k", "miles_lora", 6, 1), + ("qwen35-4b-fft-64k", "megatron_fft", 1, 4), + ], +) +def test_trainer_declaration_and_executor_configuration( + builders, monkeypatch, preset, backend, clients, nproc +): + row = deployment(preset) + platform = platform_defaults() + platform["storage"]["checkpoints"] = "test-custom-checkpoints" + image = object() + app, trainer = deployment_apps.build_trainer_app(row, platform, image=image) + declaration, _ = app.functions["trainer"] + assert declaration["gpu"] == "H100:4" + assert declaration["region"] == "us-west" + assert declaration["max_containers"] is None + assert declaration["single_use_containers"] is True + assert declaration["image"] is image + calls = [] + reloaded = [] + + monkeypatch.setattr(deployment_apps, "shared_kv", lambda: "store") + monkeypatch.setattr( + deployment_apps, + "volumes_for", + lambda spec: {"/assets": SimpleNamespace(reload=lambda: reloaded.append(True))}, + ) + monkeypatch.setattr( + deployment_apps, + "run_engine_with_backend", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + monkeypatch.setenv("LILO_FRONTEND_APP_ID", "previous") + trainer("instance-a", row.model_dump_json(), "ap-frontend") + assert os.environ["LILO_FRONTEND_APP_ID"] == "ap-frontend" + args, kwargs = calls[0] + assert args == ("store", f"lilo.backends.{backend}:build_executor") + assert kwargs["max_models"] == clients + assert kwargs["nproc"] == nproc + assert kwargs["backend_env"]["LILO_CHECKPOINT_VOLUME"] == "test-custom-checkpoints" + assert kwargs["backend_env"]["LILO_BASE_MODEL"] == row.model + config = json.loads(kwargs["backend_env"]["LILO_BACKEND_CONFIG"]) + assert config[row.recipe.backend]["hf_checkpoint"] == row.asset_path + assert config["checkpoint_dir"] == "/checkpoints" + assert reloaded == [True] + + +@pytest.mark.parametrize("kind", ["lora", "base", "latest", "pinned"]) +def test_pool_starts_native_server_and_correct_sidecar(builders, monkeypatch, kind): + row = deployment("qwen35-9b-lora-16k" if kind == "lora" else "qwen35-4b-fft-64k") + pool = ( + LoraPoolSpec(row.definition_id) + if kind == "lora" + else ( + FFTPoolSpec.base(row.definition_id) + if kind == "base" + else FFTPoolSpec(row.definition_id, "job", kind == "latest", 7) + ) + ) + app, server = deployment_apps.build_rollout_app(row, pool, image="test-image") + settings, _ = app.servers["Server"] + assert app.name == pool.app_name + assert ( + settings["gpu"] + == f"{row.recipe.inference_gpu}:{row.recipe.inference_gpus_per_node}" + ) + assert settings["min_containers"] == 0 + assert settings["target_concurrency"] == 16 + assert settings["compute_region"] == "us-west" + + calls, commands, stops = [], [], [] + process = object() + monkeypatch.setattr( + subprocess, "Popen", lambda argv, **kw: (commands.append(argv) or process) + ) + monkeypatch.setattr(deployment_apps, "wait_http", lambda *args: None) + monkeypatch.setattr(deployment_apps, "supervise_children", lambda *args: None) + monkeypatch.setattr( + deployment_apps, + "start_lora_sidecar", + lambda **kw: (calls.append(("lora", kw)) or process), + ) + monkeypatch.setattr( + deployment_apps, + "start_fft_sidecar", + lambda **kw: (calls.append(("fft", kw)) or process), + ) + monkeypatch.setattr(deployment_apps, "terminate", stops.append) + replica = server() + replica.start() + assert commands[0][2] == "lilo.inference.sglang" + assert commands[0][3] == row.asset_path + native = json.loads(commands[0][4]) + assert native["context_length"] == row.max_context_length + if kind == "lora": + assert native["enable_lora"] is True + assert native["max_lora_rank"] == 32 + assert native["lora_target_modules"] == [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "lm_head", + ] + else: + assert native["enable_cpu_weight_cache"] is True + assert calls[0][1]["pinned_version"] == ( + None if kind == "latest" else 0 if kind == "base" else 7 + ) + replica.stop() + assert stops == [process, process] + + +def test_pool_lookup_uses_definition_name(monkeypatch): + row = deployment() + monkeypatch.setenv(deployment_configs.CONFIGS_ENV, json.dumps([row.model_dump()])) + saved = deployment_configs.pool_config(row.definition_id) + assert saved.definition_id == row.definition_id + assert deployment_configs.pool_config("missing") is None + + +def test_startup_failure_is_visible_and_blocks_new_spawns(monkeypatch): + app = importlib.import_module("lilo.providers.modal.app") + store = InMemoryKeyValueStore() + row = deployment() + monkeypatch.setattr(app, "shared_kv", lambda: store) + + async def run(): + await store.put( + f"deployment_failure:{row.definition_id}", + { + "error": "backend exited with code 1", + "instance_id": "failed-instance", + "failed_at": app.time.time(), + }, + ) + with pytest.raises(ValueError, match="Trainer startup failed.*failed-instance"): + await app._spawn_engine(row.definition_id, "new-instance") + assert await app.deployment_error("legacy") is None + + asyncio.run(run()) + + +def test_real_modal_app_constructs_from_configs_without_legacy_catalog(monkeypatch): + row = deployment() + env = { + **os.environ, + deployment_configs.CONFIGS_ENV: json.dumps([row.model_dump()]), + } + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import importlib, sys, modal +from lilo.providers.modal import deployment_apps, deployment_configs +deployment_apps.image_for = lambda backend: modal.Image.debian_slim() +app = importlib.import_module('lilo.providers.modal.app') +assert len(app.DEFINITIONS) == 1 +assert app.APP_NAME == 'lilo' +assert app.deployed_configs.local() == [d.model_dump(mode='json') for d in app.configs_from_env()] +assert app.DEFINITIONS[0].trainer_app_name +assert not any(name.startswith('lilo.providers.modal.definitions.') for name in sys.modules) +print('constructed') +""", + ], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert "constructed" in result.stdout + + +def test_admission_changes_preserve_serialized_trainer(builders): + first = deployment() + old_bytes = serialize(deployment_apps.build_trainer_app(first, image="test")[1]) + changed = first.model_copy(deep=True) + changed.recipe.trainer_timeout_s = 1 + new_bytes = serialize(deployment_apps.build_trainer_app(changed, image="test")[1]) + assert new_bytes == old_bytes + changed.recipe.trainer_gpu = "H200" + assert ( + serialize(deployment_apps.build_trainer_app(changed, image="test")[1]) + == old_bytes + ) + changed.recipe.name = "other-model" + assert ( + serialize(deployment_apps.build_trainer_app(changed, image="test")[1]) + != old_bytes + ) + + +@pytest.mark.parametrize("kind", ["lora", "full"]) +def test_pool_launch_uses_only_generic_deployment_app(monkeypatch, kind): + row = deployment("qwen35-9b-lora-16k" if kind == "lora" else "qwen35-4b-fft-64k") + monkeypatch.setenv(deployment_configs.CONFIGS_ENV, json.dumps([row.model_dump()])) + module = lora_pool if kind == "lora" else fft_pool + spec = ( + LoraPoolSpec(row.definition_id) + if kind == "lora" + else FFTPoolSpec(row.definition_id, "model", True, 0) + ) + calls = [] + + class Pool: + def __init__(self, *args): + self.lookups = 0 + + def gateway_url(self): + self.lookups += 1 + if self.lookups == 1: + raise modal.exception.NotFoundError("not deployed") + return "https://pool" + + monkeypatch.setattr(module, "ModalFlashPool", Pool) + monkeypatch.setattr(module.shutil, "which", lambda _: "/bin/modal") + monkeypatch.setattr( + module.subprocess, + "run", + lambda command, **kwargs: calls.append((command, kwargs)), + ) + assert module.deploy_pool(spec, config=row) == "https://pool" + command, kwargs = calls[0] + assert ( + command[command.index("-m") + 1] == "lilo.providers.modal.deployment_pool_app" + ) + assert ( + json.loads(kwargs["env"][deployment_configs.POOL_CONFIG_ENV])["recipe"]["name"] + == row.name + ) + + +@pytest.mark.parametrize("kind", ["lora", "full"]) +def test_missing_pool_uses_saved_provisioner(monkeypatch, kind): + row = deployment("qwen35-9b-lora-16k" if kind == "lora" else "qwen35-4b-fft-64k") + monkeypatch.setenv(deployment_configs.CONFIGS_ENV, json.dumps([row.model_dump()])) + module = lora_pool if kind == "lora" else fft_pool + spec = ( + LoraPoolSpec(row.definition_id) + if kind == "lora" + else FFTPoolSpec.base(row.definition_id) + ) + + class MissingPool: + def __init__(self, *args): + pass + + def gateway_url(self): + raise modal.exception.NotFoundError("not deployed") + + monkeypatch.setattr(module, "ModalFlashPool", MissingPool) + calls = [] + monkeypatch.setattr( + modal.Function, + "from_name", + lambda name, function, **kw: ( + calls.append((name, function)) + or SimpleNamespace(remote=lambda record, pool: "https://saved-runtime") + ), + ) + monkeypatch.setattr( + module.subprocess, "run", lambda *a, **k: pytest.fail("frontend rebuilt pool") + ) + assert module.deploy_pool(spec) == "https://saved-runtime" + assert calls == [(row.inference_app_name, "provision")] + + +def test_provisioner_rejects_wrong_settings_and_uses_saved_config( + builders, monkeypatch +): + row = deployment() + app, provision = deployment_apps.build_inference_app(row, image="test") + assert app.name == row.inference_app_name + calls = [] + monkeypatch.setattr( + deployment_apps, + "deploy_lora", + lambda pool, *, config, platform=None: calls.append((pool, config)) + or "https://pool", + ) + pool = LoraPoolSpec(row.definition_id) + assert provision(row.model_dump_json(), pool.as_dict()) == "https://pool" + assert calls[0][1].name == row.name + changed = row.model_copy(deep=True) + changed.recipe.name = "other-model" + with pytest.raises(ValueError, match="inference settings"): + provision(changed.model_dump_json(), pool.as_dict()) + + +@pytest.mark.parametrize("role", ["trainer", "inference"]) +def test_trainer_and_inference_apps_construct_offline(builders, role): + row = deployment() + if role == "trainer": + app, _ = deployment_apps.build_trainer_app(row, image="test") + else: + app, _ = deployment_apps.build_inference_app(row, image="test") + assert app.name == ( + row.trainer_app_name if role == "trainer" else row.inference_app_name + ) + + +def test_spawn_passes_job_configuration_to_saved_trainer(monkeypatch): + app = importlib.import_module("lilo.providers.modal.app") + row = deployment() + calls = [] + + async def spawn(instance_id, config_json, frontend_app_id): + calls.append((instance_id, config_json, frontend_app_id)) + return SimpleNamespace(object_id="call-id") + + async def no_error(definition_id): + return None + + monkeypatch.setattr(app, "deployment_error", no_error) + monkeypatch.setattr(app, "current_app_id", lambda: "ap-frontend") + monkeypatch.setattr(app, "module_for", lambda _: row) + monkeypatch.setattr( + modal.Function, + "from_name", + lambda name, function, **kwargs: SimpleNamespace( + spawn=SimpleNamespace(aio=spawn) + ), + ) + assert asyncio.run(app._spawn_engine(row.definition_id, "instance")) == "call-id" + assert calls == [("instance", row.model_dump_json(), "ap-frontend")] + + +def test_declared_compute_settings_reach_modal(builders): + spec = deepcopy(deployment().recipe) + spec.trainer_timeout_s = 90 + spec.trainer_cpu = 12 + spec.trainer_memory_mib = 123456 + spec.inference_startup_timeout_s = 90 + spec.inference_min_replicas = 1 + spec.inference_max_replicas = 3 + spec.inference_cpu = 6 + spec.inference_memory_mib = 45000 + row = DeploymentConfig.create(spec) + trainer_app, _ = deployment_apps.build_trainer_app(row, image="test") + trainer, _ = trainer_app.functions["trainer"] + assert (trainer["cpu"], trainer["memory"], trainer["timeout"]) == (12, 123456, 90) + pool_app, _ = deployment_apps.build_rollout_app( + row, LoraPoolSpec(row.definition_id), image="test" + ) + server, _ = pool_app.servers["Server"] + assert (server["cpu"], server["memory"], server["startup_timeout"]) == ( + 6, + 45000, + 90, + ) + assert (server["min_containers"], server["max_containers"]) == (1, 3) + + +def test_multinode_trainer_uses_cluster_launcher(builders, monkeypatch): + row = deployment("qwen38-27b-lora-256k") + clusters = [] + + def clustered(nodes, *, rdma): + clusters.append((nodes, rdma)) + + def decorate(fn): + return fn + + return decorate + + monkeypatch.setattr(modal.experimental, "clustered", clustered) + app, _ = deployment_apps.build_trainer_app(row, image="test") + settings, _ = app.functions["trainer"] + assert clusters == [(2, True)] + assert settings["gpu"] == "H200:8" + assert settings["experimental_options"] == {"efa_enabled": True} + calls = [] + monkeypatch.setattr(deployment_apps, "shared_kv", lambda: "store") + monkeypatch.setattr( + deployment_apps, + "volumes_for", + lambda _: {"/assets": SimpleNamespace(reload=lambda: None)}, + ) + monkeypatch.setattr( + deployment_apps, + "start_trainer_cluster", + lambda nodes, **kwargs: "10.0.0.1:6379", + ) + monkeypatch.setattr( + deployment_apps, "run_engine_with_backend", lambda *a, **kw: calls.append(kw) + ) + deployment_apps.run_trainer(row, "instance") + assert calls[0]["backend_env"]["LILO_RAY_ADDRESS"] == "10.0.0.1:6379" + assert ( + json.loads(calls[0]["backend_env"]["LILO_BACKEND_CONFIG"])["miles"][ + "actor_num_nodes" + ] + == 2 + ) + monkeypatch.setattr(deployment_apps, "start_trainer_cluster", lambda *a, **k: None) + deployment_apps.run_trainer(row, "worker") + assert len(calls) == 1 + + +def test_launchers_do_not_reparse_backend_config(builders, monkeypatch): + row = deployment() + + def unexpected(*args, **kwargs): + pytest.fail("launcher must use the saved resolved settings") + + monkeypatch.setattr(backend, "backend_config", unexpected) + monkeypatch.setattr(backend, "serving_options", unexpected) + deployment_apps.build_trainer_app(row, image="test") + deployment_apps.build_rollout_app( + row, LoraPoolSpec(row.definition_id), image="test" + ) diff --git a/tests/providers/test_deployment_e2e_helper.py b/tests/providers/test_deployment_e2e_helper.py new file mode 100644 index 0000000..cf4efab --- /dev/null +++ b/tests/providers/test_deployment_e2e_helper.py @@ -0,0 +1,71 @@ +import runpy +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lilo.deployments import DeploymentConfig, config_path, load +from lilo.providers.modal import deployment_configs + + +@pytest.mark.parametrize("preset", ["qwen35-9b-lora-16k", "qwen35-4b-fft-64k"]) +def test_e2e_helper_reads_active_deployed_configuration(monkeypatch, preset): + row = DeploymentConfig.create(load(config_path(preset))) + monkeypatch.setattr( + deployment_configs, + "deployed_configs", + lambda frontend: [row.model_dump()], + ) + helper = runpy.run_path( + str(Path(__file__).parents[2] / "scripts/e2e_engine_definition.py") + ) + definition, mode = helper["_definition"]("test-frontend", row.name) + assert definition.DEFINITION_ID == row.definition_id + assert definition.MAX_CONTEXT_LENGTH == row.max_context_length + assert definition.GPUS == row.recipe.trainer_gpus_per_node + assert mode == row.parameterization + assert definition.MAX_TOKENS_PER_MICROBATCH > 0 + with pytest.raises(ValueError, match="one configuration"): + helper["_definition"]("test-frontend", "missing") + + +@pytest.mark.parametrize("preset", ["qwen35-9b-lora-16k", "qwen35-4b-fft-64k"]) +def test_gpu_smoke_selects_training_mode_from_current_recipe( + monkeypatch, tmp_path, preset +): + row = DeploymentConfig.create(load(config_path(preset))) + helper = runpy.run_path( + str(Path(__file__).parents[2] / "scripts/deployment_smoke.py") + ) + namespace = helper["run"].__globals__ + monkeypatch.setenv("TINKER_API_KEY", "test-key") + monkeypatch.setitem( + namespace, "deployed_configs", lambda frontend: [row.model_dump()] + ) + monkeypatch.setattr( + namespace["modal"].App, + "lookup", + lambda name: SimpleNamespace(app_id="test-app"), + ) + monkeypatch.setattr( + namespace["modal"].Function, + "from_name", + lambda *args: SimpleNamespace(get_web_url=lambda: "https://test.invalid"), + ) + create = Mock(side_effect=RuntimeError("selected training mode")) + service = SimpleNamespace(create_lora_training_client=create) + monkeypatch.setattr(namespace["tinker"], "ServiceClient", lambda **kwargs: service) + full = Mock(side_effect=RuntimeError("selected training mode")) + monkeypatch.setitem(namespace, "create_full_training_client", full) + args = SimpleNamespace( + frontend="test", name=row.name, output=str(tmp_path / "result.json") + ) + with pytest.raises(RuntimeError, match="selected training mode"): + helper["run"](args) + if row.parameterization == "lora": + create.assert_called_once_with(base_model=row.name, rank=32) + full.assert_not_called() + else: + full.assert_called_once_with(service, row.name) + create.assert_not_called() diff --git a/tests/providers/test_deployment_presets.py b/tests/providers/test_deployment_presets.py new file mode 100644 index 0000000..e89e8ee --- /dev/null +++ b/tests/providers/test_deployment_presets.py @@ -0,0 +1,73 @@ +import pytest + +from lilo.backends.deployment import backend_config, serving_options +from lilo.deployments import DeploymentConfig, config_path, load +from lilo.providers.modal.deployment_configs import configs_from_env + + +@pytest.mark.parametrize( + "path", + sorted(config_path("qwen35-9b-lora-16k").parent.glob("qwen*.py")), + ids=lambda p: p.stem, +) +def test_all_packaged_recipes_validate_offline(path): + spec = load(path) + config = backend_config(spec) + assert config[spec.backend]["hf_checkpoint"] == "/assets/pending" + serving_options(spec) + if spec.parameterization == "lora": + resolved = DeploymentConfig.create(spec) + # Legacy Megatron selectors also match vision layers in current Bridge. + assert set(spec.miles_cfg["target_modules"]) == set( + resolved.inference_settings["lora_target_modules"] + ) + + +@pytest.mark.parametrize("preset", ["qwen35-35b-a3b-fft-64k", "qwen36-35b-a3b-fft-64k"]) +def test_moe_recipes_preserve_trainer_expert_parallelism(preset): + config = backend_config(load(config_path(preset)))["megatron"] + assert config["tensor_model_parallel_size"] == 4 + assert config["context_parallel_size"] == 2 + assert config["expert_model_parallel_size"] == 8 + assert config["provider_overrides"]["moe_token_dispatcher_type"] == "alltoall" + + +def test_moe_rollout_preserves_attention_data_parallelism(): + spec = load(config_path("qwen35-35b-a3b-fft-64k")) + options = serving_options(spec) + assert options["tp_size"] == options["dp_size"] == options["ep_size"] == 4 + assert options["enable_dp_attention"] is True + definition = DeploymentConfig.create(spec) + assert definition.recipe.inference_gpus_per_node == 4 + assert definition.rollout_tensor_parallel_size == 1 + + +@pytest.mark.parametrize("context,cp", [(16384, 1), (65536, 2), (131072, 4)]) +def test_qwen38_context_parallel_token_budget(context, cp): + spec = load(config_path(f"qwen38-27b-lora-{context // 1024}k")) + config = backend_config(spec)["miles"] + assert config["context_parallel_size"] == cp + assert config["max_tokens_per_gpu"] == context // cp + assert config["actor_num_gpus_per_node"] == 8 + + +def test_single_client_recipe_keeps_shared_backend_capacity(): + shared = load(config_path("qwen35-9b-lora-16k")) + single = load(config_path("qwen35-9b-lora-16k-single")) + assert single.trainer_max_clients_per_instance == 1 + assert (single.trainer_gpu, single.trainer_gpus_per_node, single.trainer_nodes) == ( + shared.trainer_gpu, + shared.trainer_gpus_per_node, + shared.trainer_nodes, + ) + assert backend_config(single) == backend_config(shared) + + +@pytest.mark.parametrize("value", [None, "", "[]", "{}"]) +def test_missing_configs_have_no_python_catalog_fallback(monkeypatch, value): + if value is None: + monkeypatch.delenv("LILO_DEPLOYMENT_CONFIGS", raising=False) + else: + monkeypatch.setenv("LILO_DEPLOYMENT_CONFIGS", value) + with pytest.raises(ValueError, match="configs"): + configs_from_env() diff --git a/tests/providers/test_fft_pool.py b/tests/providers/test_fft_pool.py index 411d759..2c2617d 100644 --- a/tests/providers/test_fft_pool.py +++ b/tests/providers/test_fft_pool.py @@ -1,5 +1,7 @@ +import subprocess from types import SimpleNamespace +import pytest from stitch.types import VersionRef from lilo.providers.modal import fft_pool @@ -47,9 +49,7 @@ def post(self, url, *, headers): pool = fft_pool.FFTLatestPool("definition", "model") pool.wake(["replica-a", "replica-b"], VersionRef("model", 3)) - assert { - headers["modal-flash-upstream"]: url for url, headers in calls - } == { + assert {headers["modal-flash-upstream"]: url for url, headers in calls} == { "replica-a.modal.host:443": "https://rollout.modal.direct/wake", "replica-b.modal.host:8443": "https://rollout.modal.direct/wake", } @@ -79,15 +79,13 @@ def test_pool_spec_round_trips_sizing_through_dict_and_env() -> None: def test_stop_pool_accepts_already_stopped_but_propagates_other_errors(monkeypatch): - import subprocess - - import pytest - spec = fft_pool.FFTPoolSpec("definition", "model", True, 0) monkeypatch.setattr(fft_pool.shutil, "which", lambda _: "/bin/modal") monkeypatch.setenv("MODAL_ENVIRONMENT", "test-env") calls = [] - result = subprocess.CompletedProcess([], 1, "", "App is already stopped. (Stopped yesterday).\n") + result = subprocess.CompletedProcess( + [], 1, "", "App is already stopped. (Stopped yesterday).\n" + ) def run(command, **kwargs): calls.append(command) @@ -96,9 +94,25 @@ def run(command, **kwargs): monkeypatch.setattr(fft_pool.subprocess, "run", run) fft_pool.stop_pool(spec) - assert calls == [["/bin/modal", "app", "stop", "-y", spec.app_name, "--env", "test-env"]] + assert calls == [ + ["/bin/modal", "app", "stop", "-y", spec.app_name, "--env", "test-env"] + ] result = subprocess.CompletedProcess([], 0, "Stopped", "") fft_pool.stop_pool(spec) result = subprocess.CompletedProcess([], 1, "", "Permission denied") with pytest.raises(subprocess.CalledProcessError): fft_pool.stop_pool(spec) + + +def test_fft_pool_names_fit_modal_limit_and_distinguish_sessions(): + recipe = "qwen35-35b-a3b-fft-64k" + model_id = "a" * 32 + pools = [ + fft_pool.FFTPoolSpec(recipe, model_id, True, 0), + fft_pool.FFTPoolSpec(recipe, model_id, False, 1000000), + fft_pool.FFTPoolSpec(recipe, "b" * 32, True, 0), + fft_pool.FFTPoolSpec.base(recipe), + fft_pool.FFTPoolSpec.base(recipe + "-other"), + ] + assert all(len(pool.app_name) < 64 for pool in pools) + assert len({pool.app_name for pool in pools}) == len(pools) diff --git a/tests/providers/test_kernel_cache.py b/tests/providers/test_kernel_cache.py index f6d23c5..4f8d9a7 100644 --- a/tests/providers/test_kernel_cache.py +++ b/tests/providers/test_kernel_cache.py @@ -1,11 +1,11 @@ -import ast import runpy from pathlib import Path from unittest.mock import patch, sentinel import modal -from lilo.providers.modal.app import DEFINITIONS +from lilo.deployments import DeploymentConfig, config_path, load, platform_defaults +from lilo.providers.modal import deployment_apps from lilo.providers.modal.kernel_cache import ( KERNEL_CACHE_ENV, KERNEL_CACHE_ROOT, @@ -35,38 +35,26 @@ def test_kernel_cache_creates_one_v2_volume_without_live_lookup() -> None: } -def test_all_definitions_mount_the_shared_kernel_cache() -> None: +def test_configured_trainers_mount_the_shared_kernel_cache() -> None: assert KERNEL_CACHE_VOLUME_NAME == "lilo-kernel-cache" - for definition in DEFINITIONS: - assert definition.TRAINER_VOLUMES[KERNEL_CACHE_ROOT] is kernel_cache_volume + volumes = deployment_apps.volumes_for(platform_defaults()) + assert volumes[KERNEL_CACHE_ROOT] is kernel_cache_volume assert KERNEL_CACHE_ENV["TRITON_CACHE_DIR"].startswith(KERNEL_CACHE_ROOT + "/") assert KERNEL_CACHE_ENV["TORCHINDUCTOR_CACHE_DIR"].startswith( KERNEL_CACHE_ROOT + "/" ) -def test_all_definitions_point_compilers_at_the_kernel_cache() -> None: - """Every trainer's backend_env starts from KERNEL_CACHE_ENV so definition-specific - entries can still override the cache locations.""" - for definition in DEFINITIONS: - tree = ast.parse(Path(definition.__file__).read_text()) - engine_calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "run_engine_with_backend" - ] - if not engine_calls: - continue - assert len(engine_calls) == 1, definition.DEFINITION_ID - backend_env = next( - keyword.value - for keyword in engine_calls[0].keywords - if keyword.arg == "backend_env" - ) - assert isinstance(backend_env, ast.Dict), definition.DEFINITION_ID - first_key, first_value = backend_env.keys[0], backend_env.values[0] - assert first_key is None, definition.DEFINITION_ID - assert isinstance(first_value, ast.Name), definition.DEFINITION_ID - assert first_value.id == "KERNEL_CACHE_ENV", definition.DEFINITION_ID +def test_configured_trainers_point_compilers_at_the_kernel_cache() -> None: + for preset in ("qwen35-9b-lora-16k", "qwen35-4b-fft-64k"): + deployment = DeploymentConfig.create(load(config_path(preset))) + with ( + patch.object(deployment_apps, "volumes_for") as volumes, + patch.object(deployment_apps, "shared_kv"), + patch.object(deployment_apps, "run_engine_with_backend") as run, + ): + deployment_apps.run_trainer(deployment, "test-instance") + volumes.return_value.__getitem__.return_value.reload.assert_called_once() + env = run.call_args.kwargs["backend_env"] + for key, value in KERNEL_CACHE_ENV.items(): + assert env[key] == value diff --git a/tests/providers/test_lora_pool.py b/tests/providers/test_lora_pool.py index 902b4e6..53a0c27 100644 --- a/tests/providers/test_lora_pool.py +++ b/tests/providers/test_lora_pool.py @@ -1,14 +1,18 @@ +import subprocess + +import pytest + +from lilo.providers.modal import lora_pool from lilo.providers.modal.lora_pool import LoraPoolSpec def test_lora_pool_is_shared_by_every_adapter_for_definition() -> None: - first = LoraPoolSpec("qwen3_5_9b_base_miles_lora_2k") + first = LoraPoolSpec("qwen35-9b-lora-16k") second = LoraPoolSpec.from_dict(first.as_dict()) assert first == second assert first.app_name == second.app_name - assert first.app_name != LoraPoolSpec(first.definition_id, "old").app_name - assert first.app_name.startswith("lilo-lora-") + assert first.app_name == "lilo-lora-qwen35-9b-lora-16k" assert first.env() == { "LILO_LORA_POOL_APP_NAME": first.app_name, "LILO_LORA_POOL_DEFINITION_ID": first.definition_id, @@ -18,46 +22,13 @@ def test_lora_pool_is_shared_by_every_adapter_for_definition() -> None: def test_stop_already_stopped_lora_pool_succeeds_but_real_failure_propagates( monkeypatch, ): - import subprocess - import pytest - from lilo.providers.modal import lora_pool - monkeypatch.setattr(lora_pool.shutil, "which", lambda _: "/bin/modal") result = subprocess.CompletedProcess( [], 1, "", "App is already stopped. (Stopped yesterday).\n" ) monkeypatch.setattr(lora_pool.subprocess, "run", lambda *args, **kwargs: result) - spec = LoraPoolSpec("definition", "revision") + spec = LoraPoolSpec("qwen35-9b-lora-16k") lora_pool.stop_pool(spec) result.stderr = "authentication failed" with pytest.raises(subprocess.CalledProcessError): lora_pool.stop_pool(spec) - - -def test_pool_revision_tracks_inherited_settings_and_bulletin(monkeypatch): - from pathlib import Path - - original = Path.read_bytes - changed = None - - def read(path): - value = original(path) - return value + b"\n# changed\n" if path.name == changed else value - - monkeypatch.setattr(Path, "read_bytes", read) - name = "qwen3_5_9b_base_miles_lora_16k_single" - original_pool = LoraPoolSpec(name).app_name - for changed in ("qwen3_5_9b_base_miles_lora_16k.py", "bulletin.py"): - assert LoraPoolSpec(name).app_name != original_pool - changed = "qwen3_5_4b_full_64k.py" - assert LoraPoolSpec(name).app_name == original_pool - - -def test_definition_dependency_scan_handles_cycles_and_relative_imports(tmp_path): - from lilo.providers.modal.lora_pool import _definition_sources - - (tmp_path / "child.py").write_text("from .parent import CONFIG\n") - (tmp_path / "parent.py").write_text("from . import shared\n") - (tmp_path / "shared.py").write_text("from .child import CONFIG\n") - sources = list(_definition_sources(tmp_path / "child.py")) - assert [path.name for path in sources] == ["child.py", "parent.py", "shared.py"] diff --git a/tests/providers/test_miles_definition.py b/tests/providers/test_miles_definition.py deleted file mode 100644 index a4a3342..0000000 --- a/tests/providers/test_miles_definition.py +++ /dev/null @@ -1,56 +0,0 @@ -import ast -from pathlib import Path - -from lilo.backends.miles_config import MILES_REF -from lilo.providers.modal import miles_image -from lilo.providers.modal.definitions import qwen3_5_9b_base_miles_lora_2k as definition - - -def test_miles_definition_uses_one_lilo_driver_for_all_ray_workers() -> None: - assert definition.PARAMETERIZATION == "lora" - assert definition.TRAINER_MODELS_PER_INSTANCE == definition.MAX_LORA_SLOTS - assert definition.GPUS == 4 - assert MILES_REF == "main" - assert len(miles_image.MILES_COMMIT) == 40 - - tree = ast.parse(Path(definition.__file__).read_text()) - calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "run_engine_with_backend" - ] - assert len(calls) == 1 - call = calls[0] - assert isinstance(call.args[1], ast.Constant) - assert call.args[1].value == "lilo.backends.miles_lora:build_executor" - keywords = {keyword.arg: keyword.value for keyword in call.keywords} - assert isinstance(keywords["nproc"], ast.Constant) - assert keywords["nproc"].value == 1 - assert isinstance(keywords["max_models"], ast.Name) - assert keywords["max_models"].id == "MAX_LORA_SLOTS" - - -def test_long_context_miles_definition_bounds_retained_adapter_versions(): - from lilo.providers.modal.definitions import qwen3_5_9b_base_miles_lora_16k as long - - assert long.MAX_CONTEXT_LENGTH >= 8192 + 2048 - assert long.MAX_LORA_SLOTS >= 6 - assert long.ROLLOUT_MAX_LOADED_LORAS == 64 - assert long.ROLLOUT_MAX_LOADED_LORAS >= long.ROLLOUT_MAX_LORAS_PER_BATCH - assert long.ROLLOUT_MAX_CONTAINERS == 8 - assert long.CATALOG_VISIBLE - assert not definition.CATALOG_VISIBLE - - - -def test_single_tenant_definition_preserves_shared_hardware_and_backend(): - from lilo.providers.modal.definitions import qwen3_5_9b_base_miles_lora_16k as shared - from lilo.providers.modal.definitions import qwen3_5_9b_base_miles_lora_16k_single as single - - assert single.TRAINER_MODELS_PER_INSTANCE == 1 - assert not single.CATALOG_VISIBLE - assert single.run_trainer is shared.run_trainer - for name in ("MODEL_NAME", "GPUS", "GPU_TYPE", "MAX_CONTEXT_LENGTH", - "TENSOR_MODEL_PARALLEL_SIZE", "MAX_LORA_SLOTS", "MAX_LORA_RANK", - "ROLLOUT_MIN_CONTAINERS", "ROLLOUT_MAX_CONTAINERS", "ROLLOUT_GPU_TYPE"): - assert getattr(single, name) == getattr(shared, name) diff --git a/tests/providers/test_miles_revision.py b/tests/providers/test_miles_revision.py deleted file mode 100644 index 9e71524..0000000 --- a/tests/providers/test_miles_revision.py +++ /dev/null @@ -1,55 +0,0 @@ -import subprocess -from types import SimpleNamespace - -import pytest - -from lilo.providers.modal.miles_revision import resolve_miles_commit - - -@pytest.fixture(autouse=True) -def fresh_resolution(monkeypatch): - monkeypatch.delenv("LILO_MILES_COMMIT", raising=False) - resolve_miles_commit.cache_clear() - yield - resolve_miles_commit.cache_clear() - - -def test_main_resolves_once_and_next_deployment_can_advance(monkeypatch): - calls = [] - head = ["b" * 40] - - def run(argv, **kwargs): - calls.append(argv) - assert argv[-1] == "refs/heads/main" - assert kwargs["check"] and kwargs["timeout"] == 30 - return SimpleNamespace(stdout=f"{head[0]}\trefs/heads/main\n") - - monkeypatch.setattr(subprocess, "run", run) - assert resolve_miles_commit() == "b" * 40 - head[0] = "c" * 40 - assert resolve_miles_commit() == "b" * 40 - assert len(calls) == 1 - resolve_miles_commit.cache_clear() - assert resolve_miles_commit() == "c" * 40 - - -def test_exact_override_does_not_lookup_main(monkeypatch): - monkeypatch.setenv("LILO_MILES_COMMIT", "d" * 40) - def unexpected(*args, **kwargs): - pytest.fail("Pinned reproduction should not query main") - monkeypatch.setattr(subprocess, "run", unexpected) - assert resolve_miles_commit() == "d" * 40 - resolve_miles_commit.cache_clear() - monkeypatch.setenv("LILO_MILES_COMMIT", "main") - with pytest.raises(ValueError, match="full lowercase Git commit"): - resolve_miles_commit() - - -def test_lookup_failure_is_not_cached(monkeypatch): - def fail(*args, **kwargs): - raise subprocess.CalledProcessError(2, args[0]) - monkeypatch.setattr(subprocess, "run", fail) - with pytest.raises(subprocess.CalledProcessError): - resolve_miles_commit() - monkeypatch.setattr(subprocess, "run", lambda *a, **k: SimpleNamespace(stdout=f'{"e" * 40}\trefs/heads/main\n')) - assert resolve_miles_commit() == "e" * 40 diff --git a/tests/providers/test_modal_app.py b/tests/providers/test_modal_app.py index 055cc64..3765c21 100644 --- a/tests/providers/test_modal_app.py +++ b/tests/providers/test_modal_app.py @@ -1,17 +1,28 @@ import asyncio import importlib import json +import subprocess +from contextlib import contextmanager from types import SimpleNamespace import pytest +from lilo.control_plane.keys import model_key, placement_key, trainer_demand_key +from lilo.control_plane.records import ModelRecord +from lilo.deployments import DeploymentConfig, config_path, load from lilo.errors import RecordNotFound from lilo.providers.local import InMemoryKeyValueStore +from lilo.providers.modal import fft_pool from lilo.providers.modal.fft_pool import FFTPoolSpec from lilo.providers.modal.lora_pool import LoraPoolSpec -FULL_DEFINITION = "qwen3_5_9b_full_64k" -LORA_DEFINITION = "qwen3_5_9b_base_miles_lora_2k" + +def definition_id(preset): + return DeploymentConfig.create(load(config_path(preset))).definition_id + + +FULL_DEFINITION = definition_id("qwen35-9b-fft-64k") +LORA_DEFINITION = definition_id("qwen35-9b-lora-16k") @pytest.fixture(autouse=True) @@ -21,10 +32,11 @@ def reset_lora_pool_cache(monkeypatch): monkeypatch.setattr(modal_app, "_lora_pool_checks", {}) -def test_definitions_exclude_stale_128k_definition() -> None: +def test_definitions_come_only_from_the_configured_records() -> None: modal_app = importlib.import_module("lilo.providers.modal.app") - assert "qwen3_5_9b_full_128k" not in { - definition.DEFINITION_ID for definition in modal_app.DEFINITIONS + assert {definition.definition_id for definition in modal_app.DEFINITIONS} == { + FULL_DEFINITION, + LORA_DEFINITION, } @@ -60,7 +72,7 @@ async def kick(_definition_id: str) -> None: monkeypatch.setattr(modal_app, "ModalSessionKeyValueStores", SimpleNamespace) monkeypatch.setattr(modal_app, "ModalEnginePlatform", lambda *args: engines) monkeypatch.setattr(modal_app, "kick_trainer_reconciler", kick) - monkeypatch.setattr(modal_app, "TRAINER_MAX_CONTAINERS", "1") + modal_app.module_for(LORA_DEFINITION).recipe.trainer_max_instances = 1 plane = modal_app._plane() assert asyncio.run(plane.reconcile_trainers(LORA_DEFINITION)) is available @@ -177,7 +189,7 @@ def test_prepare_model_assets_validates_snapshot_before_commit(monkeypatch) -> N def download(*, repo_id: str, local_dir: str) -> None: events.append(("download", repo_id, local_dir)) - monkeypatch.setattr("huggingface_hub.snapshot_download", download) + monkeypatch.setattr("lilo.providers.modal.app.snapshot_download", download) monkeypatch.setattr( modal_app, "model_assets", @@ -188,15 +200,12 @@ def download(*, repo_id: str, local_dir: str) -> None: definition = modal_app.module_for(FULL_DEFINITION) assert events == [ - ("download", definition.MODEL_NAME, definition.HF_CHECKPOINT), + ("download", definition.model, definition.asset_path), ("commit",), ] def test_ensure_pool_sizes_latest_pool_from_model_rollout_config(monkeypatch) -> None: - from lilo.control_plane.keys import model_key - from lilo.control_plane.records import ModelRecord - modal_app = importlib.import_module("lilo.providers.modal.app") kv = InMemoryKeyValueStore() registry = InMemoryKeyValueStore() @@ -251,7 +260,6 @@ async def run() -> None: def test_execute_sample_routes_base_session_without_version(monkeypatch) -> None: modal_app = importlib.import_module("lilo.providers.modal.app") - sampling = importlib.import_module("lilo.inference.sampling") specs = [] async def gateway(spec: FFTPoolSpec) -> str: @@ -265,7 +273,7 @@ async def sample(task, gateway, *, data_parallel_size, **_): monkeypatch.setenv("MODAL_PROXY_TOKEN_SECRET", "ws-a") monkeypatch.setattr(modal_app, "pool_gateway", gateway) monkeypatch.setattr(modal_app, "fft_pool_kv", InMemoryKeyValueStore) - monkeypatch.setattr(sampling, "sample_task", sample) + monkeypatch.setattr(modal_app, "sample_task", sample) task = { "engine_definition_id": FULL_DEFINITION, "model_id": None, @@ -281,7 +289,6 @@ async def sample(task, gateway, *, data_parallel_size, **_): def test_execute_sample_routes_lora_models_to_shared_pool(monkeypatch) -> None: modal_app = importlib.import_module("lilo.providers.modal.app") - sampling = importlib.import_module("lilo.inference.sampling") ensured = [] specs = [] @@ -295,15 +302,12 @@ async def gateway(spec: LoraPoolSpec) -> str: exported_stats = {} - from contextlib import contextmanager - from lilo.telemetry import otlp - @contextmanager def trace(task, stats): yield exported_stats.update(stats) - monkeypatch.setattr(otlp, "sample_trace", trace) + monkeypatch.setattr(modal_app, "sample_trace", trace) async def sample(task, gateway, *, data_parallel_size, stats, **_): stats.update( @@ -323,7 +327,7 @@ async def sample(task, gateway, *, data_parallel_size, stats, **_): ) monkeypatch.setattr(modal_app, "lora_pool_gateway", gateway) monkeypatch.setattr(modal_app, "shared_kv", InMemoryKeyValueStore) - monkeypatch.setattr(sampling, "sample_task", sample) + monkeypatch.setattr(modal_app, "sample_task", sample) task = { "engine_definition_id": LORA_DEFINITION, "model_id": "model-a", @@ -476,13 +480,10 @@ async def run() -> tuple[str, ...]: def test_cleanup_stops_superseded_lora_pool(monkeypatch) -> None: - from lilo.control_plane.keys import model_key - from lilo.control_plane.records import ModelRecord - modal_app = importlib.import_module("lilo.providers.modal.app") registry = InMemoryKeyValueStore() current = LoraPoolSpec(LORA_DEFINITION) - superseded = LoraPoolSpec(LORA_DEFINITION, revision="superseded") + superseded = LoraPoolSpec("retired-lora") stopped = [] model = ModelRecord( @@ -514,9 +515,6 @@ async def run() -> tuple[str, ...]: def test_cleaner_loses_models_on_removed_definitions(monkeypatch) -> None: - from lilo.control_plane.keys import model_key, placement_key, trainer_demand_key - from lilo.control_plane.records import ModelRecord - modal_app = importlib.import_module("lilo.providers.modal.app") kv = InMemoryKeyValueStore() registry = InMemoryKeyValueStore() @@ -663,10 +661,6 @@ async def run(): def test_cleanup_removes_already_stopped_pool_from_registry(monkeypatch): - import subprocess - - from lilo.providers.modal import fft_pool - modal_app = importlib.import_module("lilo.providers.modal.app") registry = InMemoryKeyValueStore() spec = FFTPoolSpec("definition", "stopped", True, 0) @@ -834,7 +828,9 @@ async def deploy(record): "ensure_lora_pool", SimpleNamespace(remote=SimpleNamespace(aio=deploy)), ) - session = SimpleNamespace(engine_definition_id=LORA_DEFINITION) + session = SimpleNamespace( + engine_definition_id=LORA_DEFINITION, model_id="existing-model" + ) async def run(): plane = modal_app._plane() @@ -845,7 +841,7 @@ async def run(): assert lookups == [LoraPoolSpec(LORA_DEFINITION)] -def test_lora_readiness_cache_is_revision_specific(monkeypatch): +def test_lora_readiness_cache_is_per_definition(monkeypatch): modal_app = importlib.import_module("lilo.providers.modal.app") lookups = [] @@ -857,27 +853,25 @@ async def lookup(spec): monkeypatch.setattr(modal_app, "shared_kv", InMemoryKeyValueStore) async def run(): - first = LoraPoolSpec(LORA_DEFINITION, "old") - second = LoraPoolSpec(LORA_DEFINITION, "new") + first = LoraPoolSpec(LORA_DEFINITION) + second = LoraPoolSpec("other-lora") assert await modal_app._ready_lora_pool( first ) != await modal_app._ready_lora_pool(second) await modal_app._ready_lora_pool(first) asyncio.run(run()) - assert [spec.revision for spec in lookups] == ["old", "new"] + assert [spec.definition_id for spec in lookups] == [LORA_DEFINITION, "other-lora"] def test_lora_cleanup_continues_after_failure_and_retries(monkeypatch, caplog): modal_app = importlib.import_module("lilo.providers.modal.app") registry = InMemoryKeyValueStore() - specs = [ - LoraPoolSpec(LORA_DEFINITION, revision) for revision in ("failed", "healthy") - ] + specs = [LoraPoolSpec("failed-lora"), LoraPoolSpec("healthy-lora")] fail = True def stop(spec): - if spec.revision == "failed" and fail: + if spec.definition_id == "failed-lora" and fail: raise RuntimeError("stop unavailable") monkeypatch.setattr(modal_app, "shared_kv", lambda: registry) @@ -896,3 +890,29 @@ async def run(): asyncio.run(run()) assert "stop unavailable" in caplog.text + + +def test_startup_failure_expires_so_transient_errors_do_not_disable_training( + monkeypatch, +): + app = importlib.import_module("lilo.providers.modal.app") + registry = InMemoryKeyValueStore() + monkeypatch.setattr(app, "shared_kv", lambda: registry) + monkeypatch.setattr(app.time, "time", lambda: 1000.0) + + async def check(): + key = f"deployment_failure:{LORA_DEFINITION}" + await registry.put( + key, + { + "error": "temporary host failure", + "instance_id": "engine-a", + "failed_at": 999.0, + }, + ) + assert "automatic retries" in await app.deployment_error(LORA_DEFINITION) + monkeypatch.setattr(app.time, "time", lambda: 1300.0) + assert await app.deployment_error(LORA_DEFINITION) is None + assert await registry.get(key) is None + + asyncio.run(check()) diff --git a/tests/providers/test_modal_deployment.py b/tests/providers/test_modal_deployment.py index 164378f..8562f95 100644 --- a/tests/providers/test_modal_deployment.py +++ b/tests/providers/test_modal_deployment.py @@ -1,37 +1,6 @@ -import pytest +from lilo.providers.modal.deployment import trainer_deployment_env -from lilo.providers.modal.deployment import ( - TRAINER_MAX_CONTAINERS_ENV, - trainer_deployment_env, - trainer_max_containers, -) - -def test_trainer_max_containers_is_configured_for_deployment(monkeypatch) -> None: - monkeypatch.setenv(TRAINER_MAX_CONTAINERS_ENV, "3") - - assert trainer_max_containers() == 3 - assert trainer_deployment_env() == {TRAINER_MAX_CONTAINERS_ENV: "3"} - - -def test_trainer_max_containers_is_unlimited_when_unset(monkeypatch) -> None: - monkeypatch.delenv(TRAINER_MAX_CONTAINERS_ENV, raising=False) - - assert trainer_max_containers() is None - assert trainer_deployment_env() == {} - - -@pytest.mark.parametrize( - "value", - [ - "invalid", - "0", - "-1", - "1.5", - ], -) -def test_trainer_max_containers_rejects_invalid_config(monkeypatch, value) -> None: - monkeypatch.setenv(TRAINER_MAX_CONTAINERS_ENV, value) - - with pytest.raises(ValueError, match=TRAINER_MAX_CONTAINERS_ENV): - trainer_max_containers() +def test_capacity_is_not_forwarded_from_legacy_environment(monkeypatch): + monkeypatch.setenv("LILO_TRAINER_MAX_CONTAINERS", "3") + assert "LILO_TRAINER_MAX_CONTAINERS" not in trainer_deployment_env() diff --git a/tests/providers/test_modal_kv.py b/tests/providers/test_modal_kv.py index 030f6da..5ffe90a 100644 --- a/tests/providers/test_modal_kv.py +++ b/tests/providers/test_modal_kv.py @@ -1,11 +1,29 @@ import asyncio -import pytest from types import SimpleNamespace from unittest.mock import AsyncMock +import pytest + +from lilo.providers.local.kv import InMemoryKeyValueStore from lilo.providers.modal import kv +@pytest.mark.parametrize( + "explicit,expected", [(None, "ap-frontend"), ("ap-explicit", "ap-explicit")] +) +def test_independent_trainer_uses_frontend_registry(monkeypatch, explicit, expected): + names = [] + monkeypatch.setenv("LILO_FRONTEND_APP_ID", "ap-frontend") + monkeypatch.setattr(kv, "current_app_id", lambda: "ap-trainer") + monkeypatch.setattr( + kv.modal.Dict, "from_name", lambda name, **kwargs: names.append(name) + ) + kv.shared_kv(explicit) + assert set(names) == { + kv.app_store_name(name, expected) for name in kv.STORE_NAMES.values() + } + + class RetryableStreamError(Exception): pass @@ -62,7 +80,9 @@ async def run() -> None: monkeypatch.setattr(kv, "StreamTerminatedError", RetryableStreamError) monkeypatch.setattr(kv.asyncio, "sleep", sleep) - assert await store.list_items("placement:") == (("placement:model-a", {"instance": "engine-a"}),) + assert await store.list_items("placement:") == ( + ("placement:model-a", {"instance": "engine-a"}), + ) assert items.attempts == 3 assert [call.args[0] for call in sleep.await_args_list] == [0.2, 0.4] @@ -131,8 +151,6 @@ async def run() -> None: def test_lora_pool_registry_supports_publication_and_cleanup(): async def run(): - from lilo.providers.local.kv import InMemoryKeyValueStore - stores = {name: InMemoryKeyValueStore() for name in kv.STORE_NAMES} routed = kv.RoutedKeyValueStore(stores) key = "lora_pool:lilo-test" diff --git a/tests/providers/test_modal_serve.py b/tests/providers/test_modal_serve.py index 3f3ff92..18d7cc9 100644 --- a/tests/providers/test_modal_serve.py +++ b/tests/providers/test_modal_serve.py @@ -43,7 +43,9 @@ async def transport(request): return httpx.Response(200, json={"result": None}) def executor(url, **kwargs): - return HttpBackendClient(url, transport=httpx.MockTransport(transport), **kwargs) + return HttpBackendClient( + url, transport=httpx.MockTransport(transport), **kwargs + ) async def serving(kv, make_server, **kwargs): engine = await make_server() @@ -88,3 +90,31 @@ async def serving(kv, make_server, **kwargs): assert signals and all(sig == signal.SIGKILL for sig in signals) assert posted.count("/execute_forward_backward_batch") == 1 assert closed == [True] + + +def test_backend_startup_failure_is_reported_before_health_ready(monkeypatch): + class Backend: + pid = 12345 + returncode = 1 + + def poll(self): + return self.returncode + + backend = Backend() + errors = [] + monkeypatch.setattr(serve.subprocess, "Popen", lambda *args, **kwargs: backend) + monkeypatch.setattr(serve.os, "killpg", lambda *args: None) + + async def record(error): + errors.append(str(error)) + + with pytest.raises(RuntimeError, match="backend exited with code 1"): + serve.run_engine_with_backend( + None, + "unused:executor", + definition_id="test", + revision="test", + instance_id="test", + on_startup_error=record, + ) + assert errors == ["backend exited with code 1"] diff --git a/tests/providers/test_qwen36_definitions.py b/tests/providers/test_qwen36_definitions.py deleted file mode 100644 index ee33a1d..0000000 --- a/tests/providers/test_qwen36_definitions.py +++ /dev/null @@ -1,33 +0,0 @@ -from lilo.providers.modal.app import ( - DEFINITIONS, - module_for, - parameterization_for, -) -from lilo.providers.modal.definitions import ( - qwen3_6_27b_full_64k, - qwen3_6_35b_a3b_full_64k, -) - - -def test_qwen36_full_definitions_are_registered() -> None: - for definition in ( - qwen3_6_27b_full_64k, - qwen3_6_35b_a3b_full_64k, - ): - assert definition in DEFINITIONS - assert module_for(definition.DEFINITION_ID) is definition - assert parameterization_for(definition.DEFINITION_ID) == "full" - - -def test_qwen36_35b_reuses_qwen35_moe_topology() -> None: - assert qwen3_6_35b_a3b_full_64k.TENSOR_MODEL_PARALLEL_SIZE == 4 - assert qwen3_6_35b_a3b_full_64k.CONTEXT_PARALLEL_SIZE == 2 - assert qwen3_6_35b_a3b_full_64k.EXPERT_MODEL_PARALLEL_SIZE == 8 - assert qwen3_6_35b_a3b_full_64k.ROLLOUT_EXPERT_PARALLEL_SIZE == 4 - - -def test_qwen36_27b_uses_dense_tp4_topology() -> None: - assert qwen3_6_27b_full_64k.TENSOR_MODEL_PARALLEL_SIZE == 4 - assert qwen3_6_27b_full_64k.CONTEXT_PARALLEL_SIZE == 2 - assert qwen3_6_27b_full_64k.ROLLOUT_TENSOR_PARALLEL_SIZE == 4 - assert qwen3_6_27b_full_64k.ROLLOUT_EXPERT_PARALLEL_SIZE == 1 diff --git a/tests/providers/test_sglang_entrypoint.py b/tests/providers/test_sglang_entrypoint.py new file mode 100644 index 0000000..ef83c72 --- /dev/null +++ b/tests/providers/test_sglang_entrypoint.py @@ -0,0 +1,46 @@ +"""The SGLang worker receives its resolved constructor options without argparse.""" + +import json +import runpy +import sys +from types import ModuleType + +from lilo.deployments import DeploymentConfig, config_path, load + + +def test_sglang_constructor_receives_recorded_settings(monkeypatch): + row = DeploymentConfig.create(load(config_path("qwen35-9b-lora-16k"))) + calls = [] + + class ServerArgs: + def __init__(self, **kwargs): + calls.append(kwargs) + self.log_level = "info" + + modules = { + "sglang": {"__path__": []}, + "sglang.srt": {"__path__": []}, + "sglang.launch_server": {"run_server": lambda args: calls.append("run")}, + "sglang.srt.plugins": {"load_plugins": lambda: calls.append("plugins")}, + "sglang.srt.server_args": {"ServerArgs": ServerArgs}, + "sglang.srt.utils": {"kill_process_tree": lambda *a, **k: calls.append("stop")}, + } + for name, attrs in modules.items(): + module = ModuleType(name) + module.__dict__.update(attrs) + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr( + sys, "argv", ["sglang", row.asset_path, json.dumps(row.inference_settings)] + ) + runpy.run_module("lilo.inference.sglang", run_name="__main__") + assert calls == [ + "plugins", + { + "model_path": row.asset_path, + "host": "127.0.0.1", + "port": 8001, + **row.inference_settings, + }, + "run", + "stop", + ] diff --git a/tests/replay/test_sdk.py b/tests/replay/test_sdk.py index d09f174..2ec1d98 100644 --- a/tests/replay/test_sdk.py +++ b/tests/replay/test_sdk.py @@ -52,11 +52,11 @@ async def sampler(task): } definition = SimpleNamespace( - DEFINITION_ID="replay", - MODEL_NAME="Qwen/Qwen3-8B", - PARAMETERIZATION="lora", - CATALOG_VISIBLE=True, - MAX_CONTEXT_LENGTH=32768, + definition_id="replay", + name="replay", + model="Qwen/Qwen3-8B", + parameterization="lora", + max_context_length=32768, ) plane = ControlPlane( InMemoryKeyValueStore(), @@ -68,7 +68,7 @@ async def sampler(task): ) with serve(app) as url: service = tinker.ServiceClient(base_url=url, api_key="tml-test-key") - sampling = service.create_sampling_client(base_model=definition.MODEL_NAME) + sampling = service.create_sampling_client(base_model=definition.model) sequence = ( sample_with_replay( sampling, @@ -84,7 +84,7 @@ async def sampler(task): assert sequence.logprobs == [-2.0, -3.0] assert sequence.replay.sampling_logprobs == [-0.3, 0.0] training = service.create_lora_training_client( - base_model=definition.MODEL_NAME, rank=16 + base_model=definition.model, rank=16 ) datum = types.Datum( model_input=types.ModelInput.from_ints([1, 2, 3]), diff --git a/tests/scoped/test_assets.py b/tests/scoped/test_assets.py index 21f260e..6cc3ce2 100644 --- a/tests/scoped/test_assets.py +++ b/tests/scoped/test_assets.py @@ -13,14 +13,23 @@ def test_config_only_asset_cache_is_resumed_before_commit(tmp_path): """A failed download can leave config.json without weights or a tokenizer.""" (tmp_path / "config.json").write_text("{}") engine = qwen3_5_4b_full_64k() - engine = replace(engine, training=replace(engine.training, hf_checkpoint=str(tmp_path))) + engine = replace( + engine, training=replace(engine.training, hf_checkpoint=str(tmp_path)) + ) resources = build_app( - engine, "asset-cache-test", "asset-cache-test", "test-key", 1, - Pool(), Pool(), "test-checkpoints", modal.Secret.from_dict({}), + engine, + "asset-cache-test", + "asset-cache-test", + "test-key", + 1, + Pool(), + Pool(), + "test-checkpoints", + modal.Secret.from_dict({}), ) prepare = resources[4] with ( - patch("huggingface_hub.snapshot_download") as download, + patch("lilo.providers.modal.scoped.snapshot_download") as download, patch.object(modal.Volume, "commit") as commit, ): download.side_effect = OSError("interrupted download") diff --git a/tests/test_deployment_cli.py b/tests/test_deployment_cli.py new file mode 100644 index 0000000..ca7d17e --- /dev/null +++ b/tests/test_deployment_cli.py @@ -0,0 +1,284 @@ +import json +import subprocess +from contextlib import nullcontext +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import modal +import pytest + +from lilo import deployment_cli as cli +from lilo.deployments import DeploymentConfig, config_path, load, platform_defaults +from lilo.providers.modal.deployment_configs import ( + CONFIGS_ENV, + PLATFORM_ENV, + deployed_configs, +) +from lilo.providers.modal.fft_pool import FFTPoolSpec +from lilo.providers.modal.image_dependencies import ignore_config_source +from lilo.providers.modal.lora_pool import LoraPoolSpec + + +def deployment(): + return DeploymentConfig.create(load(config_path("qwen35-9b-lora-16k"))) + + +@pytest.fixture +def deployed(monkeypatch): + monkeypatch.setattr(cli.sys, "version_info", (3, 12, 0)) + state = SimpleNamespace( + apps=set(), configs=[], platform=None, pools=[], calls=[], fail=None + ) + + class FakeApp: + def __init__(self, name, role): + self.name = name + self.role = role + + def deploy(self, environment_name=None): + state.calls.append(self.role) + if state.fail == self.role: + raise RuntimeError(f"{self.role} deploy failed") + state.apps.add(self.name) + + def read_configs(frontend, environment): + return state.configs + + def lookup(name, **kwargs): + if name not in state.apps: + raise modal.exception.NotFoundError("not deployed") + + def run(command, *, env, check): + state.calls.append("frontend") + if state.fail == "frontend": + raise subprocess.CalledProcessError(1, command) + state.configs = json.loads(env[CONFIGS_ENV]) + state.platform = json.loads(env[PLATFORM_ENV]) + + monkeypatch.setattr(cli, "deployed_configs", read_configs) + monkeypatch.setattr(cli, "deployed_pools", lambda *args: state.pools) + monkeypatch.setattr(cli, "deployed_platform", lambda *args: state.platform) + monkeypatch.setattr(modal.App, "lookup", lookup) + monkeypatch.setattr( + cli, + "build_trainer_app", + lambda row, platform=None, **k: ( + FakeApp(row.trainer_app_name, "trainer"), + None, + ), + ) + monkeypatch.setattr( + cli, + "build_inference_app", + lambda row, platform=None, **k: ( + FakeApp(row.inference_app_name, "inference"), + None, + ), + ) + monkeypatch.setattr(cli.modal, "enable_output", lambda: nullcontext()) + monkeypatch.setattr( + cli, + "build_rollout_app", + lambda row, pool, platform: (FakeApp(pool.app_name, "pool"), None), + ) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr( + modal.Dict, + "from_name", + lambda *a, **k: pytest.fail("deployment must not use a Dict registry"), + ) + return state + + +def test_deploy_reuses_unchanged_apps(deployed): + row = deployment() + cli.deploy([row]) + assert deployed.calls == ["trainer", "inference", "frontend"] + deployed.calls.clear() + cli.deploy([row]) + assert deployed.calls == ["frontend"] + + spec = deepcopy(row.recipe) + spec.inference_max_replicas = 6 + changed = DeploymentConfig.create(spec) + deployed.calls.clear() + cli.deploy([changed]) + assert deployed.calls == ["inference", "frontend"] + + assert [r["recipe"]["name"] for r in deployed.configs] == [changed.name] + + # Modal's actual state wins over a frontend config that mentions an old app. + deployed.apps.remove(changed.inference_app_name) + deployed.calls.clear() + cli.deploy([changed]) + assert deployed.calls == ["inference", "frontend"] + + +@pytest.mark.parametrize("mode", ["lora", "full"]) +def test_inference_changes_refresh_existing_pools_and_retry_before_frontend( + deployed, mode +): + row = DeploymentConfig.create( + load( + config_path("qwen35-9b-lora-16k" if mode == "lora" else "qwen35-4b-fft-64k") + ) + ) + cli.deploy([row]) + pool = ( + LoraPoolSpec(row.name) + if mode == "lora" + else FFTPoolSpec(row.name, "job", True, 0) + ) + deployed.apps.add(pool.app_name) + deployed.pools = [pool.as_dict(), {"definition_id": "unrelated"}] + recipe = deepcopy(row.recipe) + recipe.inference_max_replicas = 3 + changed = DeploymentConfig.create(recipe) + deployed.calls.clear() + deployed.fail = "pool" + with pytest.raises(RuntimeError, match="pool deploy failed"): + cli.deploy([changed]) + assert deployed.calls == ["inference", "pool"] + assert deployed.configs[0]["recipe"]["inference_max_replicas"] == 8 + deployed.fail = None + deployed.calls.clear() + cli.deploy([changed]) + assert deployed.calls == ["inference", "pool", "frontend"] + + +@pytest.mark.parametrize( + "failure,error", + [("inference", RuntimeError), ("frontend", subprocess.CalledProcessError)], +) +def test_retry_discovers_completed_apps_without_pending_configs( + deployed, failure, error +): + row = deployment() + deployed.fail = failure + with pytest.raises(error): + cli.deploy([row]) + assert deployed.configs == [] + assert row.trainer_app_name in deployed.apps + deployed.calls.clear() + deployed.fail = None + cli.deploy([row]) + assert deployed.calls == ( + ["inference", "frontend"] if failure == "inference" else ["frontend"] + ) + + +def test_refresh_redeploys_only_that_trainer(deployed): + miles = deployment() + fft = DeploymentConfig.create(load(config_path("qwen35-4b-fft-64k"))) + cli.deploy([miles, fft]) + deployed.calls.clear() + cli.deploy([miles, fft], refresh_trainers=[miles.name]) + assert deployed.calls == ["trainer", "frontend"] + deployed.calls.clear() + cli.deploy([miles, fft]) + assert deployed.calls == ["frontend"] + + +def test_region_change_updates_worker_placement(deployed): + row = deployment() + cli.deploy([row]) + platform = platform_defaults() + platform["modal"]["region"] = "us-east" + deployed.calls.clear() + cli.deploy([row], platform) + assert deployed.calls == ["trainer", "inference", "frontend"] + assert deployed.platform["modal"]["region"] == "us-east" + + +def test_existing_frontend_without_deployment_metadata_can_be_updated(deployed): + row = deployment() + deployed.apps.add("lilo") + cli.deploy([row]) + assert deployed.configs[0]["recipe"]["name"] == row.name + + +def test_read_configs_from_deployed_function(monkeypatch): + function = SimpleNamespace(remote=Mock(return_value=[{"configuration": "saved"}])) + lookup = Mock(return_value=function) + monkeypatch.setattr(modal.Function, "from_name", lookup) + assert deployed_configs("my-app", "dev") == [{"configuration": "saved"}] + lookup.assert_called_once_with("my-app", "deployed_configs", environment_name="dev") + function.remote.side_effect = modal.exception.NotFoundError("no function") + assert deployed_configs("my-app", "dev") == [] + function.remote.side_effect = RuntimeError("frontend failed") + with pytest.raises(RuntimeError, match="frontend failed"): + deployed_configs("my-app", "dev") + + +def test_validate_never_resolves_or_deploys(monkeypatch, capsys): + monkeypatch.setattr( + cli, + "compile_configs", + lambda *args, **kwargs: pytest.fail("unexpected compile"), + ) + monkeypatch.setattr( + cli, "deploy", lambda *args: pytest.fail("unexpected deployment") + ) + cli.main(["config", "validate", str(config_path("qwen35-9b-lora-16k"))]) + assert "Validated 1 deployment" in capsys.readouterr().out + + +def test_deploy_rejects_python_mismatch_before_remote_changes(monkeypatch): + monkeypatch.setattr(cli.sys, "version_info", (3, 11, 0)) + monkeypatch.setattr( + modal.Dict, + "from_name", + lambda *a, **k: pytest.fail("must reject before touching Modal"), + ) + with pytest.raises(ValueError, match="requires Python 3.12"): + cli.deploy([deployment()]) + + +def test_source_mount_excludes_authoring_configs(): + assert ignore_config_source(Path("configs/example.py")) + assert ignore_config_source(Path("configs/__init__.py")) + assert ignore_config_source(Path("data.json")) + assert not ignore_config_source(Path("deployments.py")) + assert not ignore_config_source(Path("backends/miles_config.py")) + + +def test_deploy_command_owns_platform_settings(monkeypatch): + seen = {} + + def compile(paths): + seen["compiled"] = paths + return ["config"] + + def deploy(configs, platform, **kwargs): + seen["configs"] = configs + seen["platform"] = platform + seen.update(kwargs) + + monkeypatch.setattr(cli, "compile_configs", compile) + monkeypatch.setattr(cli, "deploy", deploy) + cli.main( + [ + "deploy", + "model.py", + "--app", + "my-lilo", + "--env", + "dev", + "--region", + "us-east", + "--refresh-trainer", + "my-model", + ] + ) + assert seen["platform"]["frontend"] == "my-lilo" + assert seen["platform"]["modal"] == {"environment": "dev", "region": "us-east"} + assert seen["refresh_trainers"] == ["my-model"] + assert seen["configs"] == ["config"] + + +def test_builtin_config_keeps_the_model_path(): + (row,) = cli.compile_configs([config_path("qwen35-9b-lora-16k")]) + assert row.asset_path == "/assets/Qwen/Qwen3.5-9B-Base" + assert not hasattr(load(config_path("qwen35-9b-lora-16k")), "revision") diff --git a/tests/test_deployments.py b/tests/test_deployments.py new file mode 100644 index 0000000..fbd1ade --- /dev/null +++ b/tests/test_deployments.py @@ -0,0 +1,514 @@ +import argparse +import asyncio +import json +import sys +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import httpx +import pytest + +from lilo.backends.deployment import backend_config, serving_options +from lilo.backends.megatron_config import parse_backend_config +from lilo.backends.miles_arguments import apply_config_overrides +from lilo.configs.qwen35_9b_lora_16k import Config as Parent +from lilo.configuration import BaseConfig +from lilo.control_plane import ControlPlane, create_control_plane_app +from lilo.control_plane.deployments import DeploymentRoutes +from lilo.control_plane.keys import model_key +from lilo.deployments import ( + DeploymentConfig, + config_path, + load, + validate_frontend, +) +from lilo.providers.local import InMemoryKeyValueStore +from lilo.providers.modal.deployment_apps import deployment_env + + +def recipe(preset="qwen35-9b-lora-16k", **changes): + data = vars(load(config_path(preset))) + for path, value in changes.items(): + keys = path.split("__") + target = data + for key in keys[:-1]: + target = target[key] + target[keys[-1]] = value + return BaseConfig(**data) + + +def resolved(spec=None, **changes): + return DeploymentConfig.create(spec or recipe(**changes)) + + +def test_presets_context_topology_and_backend_options(): + spec = recipe() + config = backend_config(spec)["miles"] + assert ( + config["actor_num_gpus_per_node"], + config["tensor_model_parallel_size"], + config["max_lora_slots"], + ) == (4, 4, 6) + assert config["cli_options"]["recompute_num_layers"] == 1 + assert config["extra_args"] == ("--seq-length", "16384") + large = recipe("qwen35-9b-lora-64k") + assert large.max_context_length == 65536 + assert backend_config(large)["miles"]["actor_num_gpus_per_node"] == 8 + fft = backend_config(recipe("qwen35-4b-fft-64k"))["megatron"] + assert (fft["tensor_model_parallel_size"], fft["context_parallel_size"]) == (2, 2) + assert fft["provider_overrides"]["recompute_granularity"] == "full" + + +def test_no_model_catalog_required(): + spec = recipe( + model="my-org/new-model", + miles_cfg__model_type="", + miles_cfg__cli_options={ + "num_layers": 12, + "hidden_size": 768, + "num_attention_heads": 12, + }, + ) + assert resolved(spec).model == "my-org/new-model" + assert backend_config(spec)["miles"]["model_type"] == "" + + +@pytest.mark.parametrize( + "changes,match", + [ + ({"miles_cfg__cli_options": {"hf_checkpoint": "other"}}, "managed"), + ( + {"miles_cfg__cli_options": {"pipeline_model_parallel_size": 2}}, + "managed", + ), + ({"sglang_cfg": {"model_path": "other"}}, "managed"), + ({"sglang_cfg": {"tp_size": 2}}, "replica GPU"), + ({"trainer_max_clients_per_instance": 7}, "max_lora_slots"), + ], +) +def test_invalid_integrations_fail_when_building_backend_settings(changes, match): + spec = recipe(**changes) + with pytest.raises(ValueError, match=match): + backend_config(spec) + serving_options(spec) + + +def test_asset_paths_follow_the_model(): + a = resolved() + assert a.asset_path == f"/assets/{a.recipe.model}" + b = resolved(recipe(model="other/Qwen3.5-9B-Base")) + assert a.asset_path != b.asset_path + + +def test_routing_uses_deployment_order(): + small = resolved() + large = resolved(recipe("qwen35-9b-lora-64k")) + routes = DeploymentRoutes([small, large]) + assert routes.select(small.model, "lora").definition_id == small.definition_id + assert routes.capabilities()[0]["max_context_length"] == 16384 + validate_frontend([small.recipe, large.recipe]) + + routes = DeploymentRoutes([large]) + assert routes.select(small.model, "lora").definition_id == large.definition_id + assert routes.capabilities()[0]["max_context_length"] == 65536 + + other = resolved(recipe(name="other", model="org/other")) + routes = DeploymentRoutes([other]) + assert routes.select(other.model, "lora").definition_id == other.definition_id + assert {row["model_name"] for row in routes.capabilities()} == {other.model} + + +def test_sampling_uses_order_and_training_filters_parameterization(): + lora = resolved() + fft = resolved(recipe("qwen35-4b-fft-64k", model=lora.model)) + for first, second in ((lora, fft), (fft, lora)): + routes = DeploymentRoutes([first, second]) + assert routes.select(lora.model).definition_id == first.definition_id + assert routes.select(lora.model, "lora").definition_id == lora.definition_id + assert routes.select(lora.model, "full").definition_id == fft.definition_id + assert routes.select(second.definition_id).definition_id == second.definition_id + assert routes.select("missing") is None + assert routes.select(lora.definition_id, "full") is None + + +def test_native_false_list_aliases_and_scalar_overrides(): + parser = argparse.ArgumentParser() + parser.add_argument("--use-feature", action="store_true") + parser.add_argument("--layers", nargs="+", type=int) + parser.add_argument("--tp", "--tensor-parallel-size", dest="tp_size", type=int) + parser.add_argument("--unchanged") + argv = ["--use-feature", "--layers", "1", "2", "--tp=4", "--unchanged", "keep"] + apply_config_overrides( + parser, {"use_feature": False, "layers": [3], "tp_size": 8}, argv + ) + parsed = parser.parse_args(argv) + assert vars(parsed) == { + "use_feature": False, + "layers": [3], + "tp_size": 8, + "unchanged": "keep", + } + with pytest.raises(ValueError, match="unknown backend option"): + apply_config_overrides(parser, {"typo": 1}, []) + with pytest.raises(ValueError, match="boolean"): + apply_config_overrides(parser, {"use_feature": "false"}, []) + with pytest.raises(ValueError, match="list"): + apply_config_overrides(parser, {"layers": "1,2"}, []) + + +def test_multiple_models_same_http_service_and_old_binding_survives_switch(): + async def run(): + store = InMemoryKeyValueStore() + plane = ControlPlane(store, SimpleNamespace()) + first = resolved() + other = resolved(recipe(name="other", model="org/other-model")) + app = create_control_plane_app(plane, [first, other], api_key="test") + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://one-frontend", + headers={"x-api-key": "test"}, + ) as client: + session = (await client.post("/api/v1/create_session", json={})).json()[ + "session_id" + ] + for seq, row in enumerate([first, other]): + response = await client.post( + "/api/v1/create_model", + json={ + "session_id": session, + "model_seq_id": seq, + "base_model": row.model, + "lora_config": {"rank": 32}, + }, + ) + assert response.status_code == 200, response.text + record = await store.get(model_key(response.json()["model_id"])) + assert record["engine_definition_id"] == row.definition_id + configs = await client.get("/api/v1/lilo/deployments") + assert len(configs.json()["deployments"]) == 2 + client.headers.clear() + assert (await client.get("/api/v1/lilo/deployments")).status_code == 401 + + asyncio.run(run()) + + +def test_native_boolean_opposite_flags_and_optional_value(): + parser = argparse.ArgumentParser() + parser.add_argument("--bias", dest="bias", action="store_true") + parser.add_argument("--no-bias", dest="bias", action="store_false") + parser.add_argument("--optional", nargs="?") + parser.add_argument("--keep", action="store_true") + argv = ["--bias", "--optional", "--keep"] + apply_config_overrides(parser, {"bias": False, "optional": "supplied"}, argv) + assert vars(parser.parse_args(argv)) == { + "bias": False, + "optional": "supplied", + "keep": True, + } + parser.add_argument("--custom", action="append") + with pytest.raises(ValueError, match="unsupported argparse action"): + apply_config_overrides(parser, {"custom": [1]}, []) + + +def test_native_type_callbacks_receive_text(): + def readable_int(value): + return int(value.strip().removesuffix("k")) * ( + 1000 if value.endswith("k") else 1 + ) + + parser = argparse.ArgumentParser() + parser.add_argument("--context-length", type=readable_int) + parser.add_argument("--sizes", nargs="+", type=readable_int) + argv = [] + apply_config_overrides(parser, {"context_length": 65536, "sizes": [32, "2k"]}, argv) + args = parser.parse_args(argv) + assert args.context_length == 65536 + assert args.sizes == [32, 2000] + + +def test_native_sections_survive_serialization_without_allowlist(): + spec = recipe("qwen35-4b-fft-64k") + data = vars(spec) + data["megatron_cfg"]["provider_overrides"]["future_provider_option"] = { + "layers": [1, 4], + "enabled": False, + } + data["megatron_cfg"]["optimizer_overrides"] = {"future_optimizer_option": 0.125} + data["megatron_cfg"]["distributed_overrides"] = {"future_ddp_option": False} + spec = BaseConfig(**data) + settings = backend_config(spec, "/assets/pinned") + config, _ = parse_backend_config(json.loads(json.dumps(settings))) + assert config.hf_checkpoint == "/assets/pinned" + assert config.seq_length == spec.max_context_length + assert config.provider_overrides["future_provider_option"] == { + "layers": [1, 4], + "enabled": False, + } + assert config.optimizer_overrides == {"future_optimizer_option": 0.125} + assert config.distributed_overrides == {"future_ddp_option": False} + assert config.optimizer.lr == 0.0001 + assert vars(spec) == data # Building does not consume or mutate the config. + + +@pytest.mark.parametrize( + "section,options,match", + [ + ("provider_overrides", {"context_parallel_size": 4}, "managed"), + ("optimizer_overrides", {"bf16": False}, "managed"), + ("distributed_overrides", {"use_distributed_optimizer": False}, "managed"), + ("provider_overrides", [], "mapping"), + ("optimizer", {"optimizer": "sgd"}, "Adam"), + ], +) +def test_megatron_cli_options_preserve_integration_contract(section, options, match): + spec = recipe("qwen35-4b-fft-64k", **{f"megatron_cfg__{section}": options}) + with pytest.raises(ValueError, match=match): + backend_config(spec) + + +def test_backend_dispatch_rejects_unknown_backend(): + with pytest.raises(ValueError, match="backend"): + backend_config(recipe(backend="missing")) + + +def test_new_miles_and_sglang_options_need_no_deployment_schema_change(): + spec = recipe( + miles_cfg__cli_options__future_miles_option=[1, 2], + sglang_cfg__future_sglang_option=False, + ) + assert backend_config(spec)["miles"]["cli_options"]["future_miles_option"] == [ + 1, + 2, + ] + assert serving_options(spec)["future_sglang_option"] is False + + +def test_fft_capacity_is_checked_by_backend_setup(): + with pytest.raises(ValueError, match="FFT trainers admit one client"): + backend_config(recipe("qwen35-4b-fft-64k", trainer_max_clients_per_instance=2)) + + +def test_reserved_environment_is_checked_by_modal_setup(): + with pytest.raises(ValueError, match="managed"): + deployment_env({"LILO_BACKEND_CONFIG": "oops"}) + assert deployment_env({"MY_SETTING": "value"}) == {"MY_SETTING": "value"} + + +def test_record_creation_copies_without_reparsing(): + spec = recipe() + original = vars(spec) + row = DeploymentConfig.create(spec) + assert vars(spec) == original + assert row.asset_path == f"/assets/{spec.model}" + row.recipe.miles_cfg["max_lora_rank"] = 64 + assert spec.miles_cfg["max_lora_rank"] == 32 + saved = row.model_dump_json() + assert DeploymentConfig.model_validate_json(saved).model_dump( + mode="json" + ) == row.model_dump(mode="json") + + +def test_python_config_composition(tmp_path): + path = tmp_path / "model.py" + path.write_text( + "from lilo.configs.qwen35_9b_lora_16k import Config as Parent\n" + "class Config(Parent):\n" + " name = 'custom'\n" + " overrides = {'trainer_memory_mib': 123456}\n" + "config = Config()\n" + ) + custom = load(path) + original = recipe() + assert custom.name == "custom" + assert custom.trainer_memory_mib == 123456 + assert custom.miles_cfg == original.miles_cfg + assert original.trainer_memory_mib == 65536 + custom.trainer_max_instances = 9 + assert custom.trainer_max_instances == 9 + assert original.trainer_max_instances == 1 + + +def test_worker_record_contains_resolved_settings(monkeypatch): + record = resolved() + assert record.trainer_settings["miles"]["actor_num_gpus_per_node"] == 4 + assert record.inference_settings["max_lora_rank"] == 32 + assert DeploymentConfig.model_validate_json(record.model_dump_json()).model_dump( + mode="json" + ) == record.model_dump(mode="json") + + +@pytest.mark.parametrize("source", ["Config = {}", "class Config: pass", "value = 1"]) +def test_config_file_must_export_config_subclass(tmp_path, source): + path = tmp_path / "model.py" + path.write_text(source) + with pytest.raises(ValueError, match="BaseConfig instance"): + load(path) + + +def test_config_import_error_preserves_traceback_and_restores_path(tmp_path): + path = tmp_path / "model.py" + path.write_text("raise RuntimeError('bad user config')") + before = list(sys.path) + with pytest.raises(RuntimeError, match="bad user config"): + load(path) + assert sys.path == before + + +def test_no_yaml_config_ingestion(tmp_path): + with pytest.raises(ValueError, match="Python .py"): + load(tmp_path / "old.yaml") + + +def test_recipe_app_names_follow_the_recipe_name(): + base = resolved() + assert base.trainer_app_name == f"lilo-trainer-{base.name}" + assert base.inference_app_name == f"lilo-inference-{base.name}" + + inference = resolved(recipe(sglang_cfg__max_running_requests=24)) + assert inference.trainer_app_name == base.trainer_app_name + assert inference.same_trainer(base) + assert not inference.same_inference(base) + + trainer = resolved(recipe(miles_cfg__max_tokens_per_gpu=8192)) + assert trainer.trainer_app_name == base.trainer_app_name + assert not trainer.same_trainer(base) + assert trainer.same_inference(base) + + other = resolved(recipe("qwen35-9b-lora-64k")) + assert other.trainer_app_name != base.trainer_app_name + assert other.inference_app_name != base.inference_app_name + assert "implementation" not in other.model_dump() + + +def test_examples_only_contain_model_infrastructure(): + for path in Path(config_path("qwen35-9b-lora-16k")).parent.glob("qwen*.py"): + config = load(path) + assert not hasattr(config, "deployment") + assert not hasattr(config, "revision") + + +@pytest.mark.parametrize("value", ["invalid", 7]) +def test_backend_parser_validates_configured_types_and_choices(value): + parser = argparse.ArgumentParser() + parser.add_argument("--count", type=int, choices=[1, 2], required=True) + argv = ["--count", "1"] + apply_config_overrides(parser, {"count": value}, argv) + with pytest.raises(SystemExit): + parser.parse_args(argv) + + +@pytest.mark.parametrize( + "section,field", + [ + ("optimizer_overrides", "lr"), + ("optimizer_overrides", "adam_eps"), + ("provider_overrides", "calculate_per_token_loss"), + ("provider_overrides", "attention_backend"), + ("distributed_overrides", "overlap_grad_reduce"), + ], +) +def test_managed_backend_values_fail_before_record_creation(section, field): + base = recipe("qwen35-4b-fft-64k") + candidate = deepcopy(base) + candidate.megatron_cfg[section] = {field: 1} + with pytest.raises(ValueError, match=field): + DeploymentConfig.create(candidate) + + +def test_multinode_ownership_and_topology(): + config = load(config_path("qwen38-27b-lora-256k")) + row = DeploymentConfig.create(config) + miles = row.trainer_settings["miles"] + assert miles["actor_num_nodes"] == 2 + assert miles["actor_num_gpus_per_node"] == 8 + assert miles["tensor_model_parallel_size"] == 2 + assert miles["context_parallel_size"] == 8 + invalid = deepcopy(config) + invalid.miles_cfg["actor_num_nodes"] = 3 + with pytest.raises(ValueError, match="actor_num_nodes"): + DeploymentConfig.create(invalid) + + +def test_config_inheritance_and_constructor_overrides_copy_nested_options(): + class Child(Parent): + name = "child" + max_context_length = 8192 + overrides = {"miles_cfg.max_tokens_per_gpu": 8192} + + first, second = Child(), Child(name="second") + first.miles_cfg["target_modules"].append("extra") + first.miles_cfg["cli_options"]["recompute_num_layers"] = 2 + assert second.name == "second" + assert second.max_context_length == 8192 + assert second.trainer_gpu == "H100" + assert backend_config(second)["miles"]["max_tokens_per_gpu"] == 8192 + assert Parent().max_context_length == 16384 + for config in (second, Parent()): + assert "extra" not in config.miles_cfg["target_modules"] + assert config.miles_cfg["cli_options"]["recompute_num_layers"] == 1 + + +def test_backend_dictionary_assignment_replaces_inherited_values(): + class Child(Parent): + inference_gpu = "H100" + sglang_cfg = {"max_running_requests": 4} + + config = Child() + assert config.inference_gpu == "H100" + assert config.inference_max_replicas == 8 + assert config.sglang_cfg == {"max_running_requests": 4} + + +def test_overrides_compose_across_generations_and_constructor(): + class Child(Parent): + overrides = { + "trainer_gpu": "H200", + "trainer_env.FIRST": "1", + "miles_cfg.max_tokens_per_gpu": 8192, + "sglang_cfg.future_option.nested": [1, 2], + } + + class Grandchild(Child): + overrides = { + "miles_cfg.max_tokens_per_gpu": 4096, + "trainer_env.SECOND": "2", + } + + config = Grandchild( + name="custom", + overrides={"miles_cfg.max_tokens_per_gpu": 2048}, + ) + assert config.name == "custom" + assert config.trainer_gpu == "H200" + assert config.trainer_env["FIRST"] == "1" + assert config.trainer_env["SECOND"] == "2" + assert config.miles_cfg["max_tokens_per_gpu"] == 2048 + assert Grandchild().miles_cfg["max_tokens_per_gpu"] == 4096 + assert Child().miles_cfg["max_tokens_per_gpu"] == 8192 + config.sglang_cfg["future_option"]["nested"].append(3) + assert Child.overrides["sglang_cfg.future_option.nested"] == [1, 2] + assert Grandchild().sglang_cfg["future_option"]["nested"] == [1, 2] + assert "FIRST" not in Parent().trainer_env + + +def test_override_values_replace_dictionaries_and_lists(): + class Child(Parent): + overrides = {"trainer_env": {"FIRST": "1"}} + + class Grandchild(Child): + overrides = { + "trainer_env": {}, + "miles_cfg.target_modules": ["q_proj"], + } + + config = Grandchild() + assert config.trainer_env == {} + assert config.miles_cfg["target_modules"] == ["q_proj"] + assert Child().trainer_env == {"FIRST": "1"} + # Constructor fields replace the inherited section, then overrides apply. + config = Grandchild( + sglang_cfg={}, inference_gpu="H100", overrides={"inference_gpu": "H200"} + ) + assert config.inference_gpu == "H200" + assert config.sglang_cfg == {} diff --git a/tests/test_system.py b/tests/test_system.py index f81a190..1d2e2ce 100644 --- a/tests/test_system.py +++ b/tests/test_system.py @@ -5,21 +5,22 @@ import tinker from tinker import types -from tests.support import SingleEnginePlatform, TinkerStubExecutor, serve from lilo.control_plane import ControlPlane, create_control_plane_app from lilo.engine import Engine from lilo.engine.backend_http import HttpBackendClient, create_backend_app from lilo.engine.http import HttpEngineClient, create_engine_app from lilo.providers.local import InMemoryKeyValueStore +from tests.support import SingleEnginePlatform, TinkerStubExecutor, serve BASE_MODEL = "Qwen/Qwen3-8B" DEFINITION = "qwen3_8b" DEFINITIONS = ( SimpleNamespace( - DEFINITION_ID=DEFINITION, - MODEL_NAME=BASE_MODEL, - PARAMETERIZATION="lora", - CATALOG_VISIBLE=True, + definition_id=DEFINITION, + name=DEFINITION, + model=BASE_MODEL, + parameterization="lora", + max_context_length=16384, ), ) API_KEY = "tml-test" diff --git a/uv.lock b/uv.lock index 2f3031b..3d76c80 100644 --- a/uv.lock +++ b/uv.lock @@ -476,6 +476,7 @@ source = { editable = "." } dependencies = [ { name = "fastapi" }, { name = "httpx" }, + { name = "huggingface-hub" }, { name = "modal" }, { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, @@ -497,6 +498,7 @@ dev = [ requires-dist = [ { name = "fastapi", specifier = ">=0.141.1" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "huggingface-hub", specifier = ">=0.34" }, { name = "modal", specifier = ">=1.5.3" }, { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.39,<2" }, { name = "opentelemetry-sdk", specifier = ">=1.39,<2" },