| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Weight-transfer metrics scraped from a real two-API-server deployment.""" |
| 4 | |
| 5 | from pathlib import Path |
| 6 | |
| 7 | import requests |
| 8 | import torch |
| 9 | from prometheus_client.multiprocess import MultiProcessCollector |
| 10 | from prometheus_client.parser import text_string_to_metric_families |
| 11 | from transformers import AutoModelForCausalLM |
| 12 | |
| 13 | from tests.entrypoints.serve.dev.rlhf.conftest import MODEL_NAME |
| 14 | from tests.utils import RemoteOpenAIServer |
| 15 | from vllm.distributed.weight_transfer import ( |
| 16 | HTTPVLLMWeightSyncClient, |
| 17 | ModuleSource, |
| 18 | WeightTransferTrainerFactory, |
| 19 | ) |
| 20 | from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo |
| 21 | |
| 22 | DURATION = "vllm:rl_weight_update_operation_duration_seconds" |
| 23 | IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight" |
| 24 | # The API servers accept from one shared listening socket, so the kernel may |
| 25 | # hand one process most connections. Sync until both have recorded, within a cap. |
| 26 | MIN_SYNCS = 8 |
| 27 | MAX_SYNCS = 64 |
| 28 | |
| 29 | |
| 30 | def scrape(url: str) -> dict[tuple[str, str], float]: |
| 31 | text = requests.get(f"{url}/metrics", timeout=30).text |
| 32 | return { |
| 33 | (sample.name, sample.labels["operation"]): sample.value |
| 34 | for family in text_string_to_metric_families(text) |
| 35 | for sample in family.samples |
| 36 | if sample.name in (f"{DURATION}_count", IN_FLIGHT) |
| 37 | } |
| 38 | |
| 39 | |
| 40 | def pids_that_recorded(multiproc_dir: Path) -> set[str]: |
| 41 | return { |
| 42 | path.stem.rsplit("_", 1)[1] |
| 43 | for path in multiproc_dir.glob("histogram_*.db") |
| 44 | if any( |
| 45 | metric.name == DURATION |
| 46 | for metric in MultiProcessCollector.merge([str(path)]) |
| 47 | ) |
| 48 | } |
| 49 | |
| 50 | |
| 51 | def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path): |
| 52 | args = [ |
| 53 | "--load-format", |
| 54 | "dummy", |
| 55 | "--enforce-eager", |
| 56 | "--max-model-len", |
| 57 | "512", |
| 58 | "--gpu-memory-utilization", |
| 59 | "0.5", |
| 60 | "--api-server-count", |
| 61 | "2", |
| 62 | "--weight-transfer-config", |
| 63 | '{"backend": "ipc"}', |
| 64 | ] |
| 65 | env = { |
| 66 | "VLLM_SERVER_DEV_MODE": "1", |
| 67 | "VLLM_ALLOW_INSECURE_SERIALIZATION": "1", |
| 68 | "PROMETHEUS_MULTIPROC_DIR": str(tmp_path), |
| 69 | } |
| 70 | with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env) as server: |
| 71 | url = server.url_root |
| 72 | trainer = AutoModelForCausalLM.from_pretrained( |
| 73 | MODEL_NAME, torch_dtype=torch.bfloat16 |
| 74 | ).cuda() |
| 75 | engine = WeightTransferTrainerFactory.trainer_init( |
| 76 | IPCTrainerInitInfo(rank=0), |
| 77 | client=HTTPVLLMWeightSyncClient(url), |
| 78 | source=ModuleSource(trainer), |
| 79 | ) |
| 80 | syncs = 0 |
| 81 | while syncs < MIN_SYNCS or ( |
| 82 | len(pids_that_recorded(tmp_path)) < 2 and syncs < MAX_SYNCS |
| 83 | ): |
| 84 | engine.send_weights() |
| 85 | syncs += 1 |
| 86 | metrics = scrape(url) |
| 87 | |
| 88 | calls = {"init": 1, "start": syncs, "update": syncs, "finish": syncs} |
| 89 | for operation, count in calls.items(): |
| 90 | assert metrics[(f"{DURATION}_count", operation)] == count, operation |
| 91 | assert metrics[(IN_FLIGHT, operation)] == 0, operation |
| 92 | # The counts above are sums over processes only if both servers recorded. |
| 93 | assert len(pids_that_recorded(tmp_path)) == 2 |